MiFaZhan
128 lines
3.3 KiB
Go
128 lines
3.3 KiB
Go
package rabbitmq
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
amqp "github.com/rabbitmq/amqp091-go"
|
|
)
|
|
|
|
func PublishConfirmed(ctx context.Context, queue string, data any) error {
|
|
body, err := json.Marshal(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return publishConfirmedBody(ctx, queue, body, "", "", 0)
|
|
}
|
|
|
|
func PublishDelayedConfirmed(ctx context.Context, delayQueue, targetQueue string, delay time.Duration, data any) error {
|
|
body, err := json.Marshal(data)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return publishConfirmedBody(ctx, delayQueue, body, delayQueue, targetQueue, delay)
|
|
}
|
|
|
|
func publishConfirmedBody(ctx context.Context, queue string, body []byte, delayQueue, targetQueue string, delay time.Duration) error {
|
|
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
|
defer cancel()
|
|
cfg := GetConfig(ctx)
|
|
if !cfg.Enabled || cfg.Dsn == "" {
|
|
return gerror.New("rabbitmq is disabled or dsn is empty")
|
|
}
|
|
connection, cleanup, err := openConfirmedConnection(ctx, cfg.Dsn)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer cleanup()
|
|
channel, err := connection.Channel()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if delayQueue != "" {
|
|
if err = declareQueue(channel, targetQueue); err != nil {
|
|
return err
|
|
}
|
|
if err = declareDelayQueue(channel, delayQueue, targetQueue); err != nil {
|
|
return err
|
|
}
|
|
} else if err = declareQueue(channel, queue); err != nil {
|
|
return err
|
|
}
|
|
if err = channel.Confirm(false); err != nil {
|
|
return err
|
|
}
|
|
confirmations := channel.NotifyPublish(make(chan amqp.Confirmation, 1))
|
|
returns := channel.NotifyReturn(make(chan amqp.Return, 1))
|
|
if err = channel.PublishWithContext(ctx, "", queue, true, false, amqp.Publishing{
|
|
ContentType: "application/json", DeliveryMode: amqp.Persistent, Timestamp: time.Now(), Body: body, Expiration: confirmedExpiration(delay),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case returned := <-returns:
|
|
return gerror.Newf("rabbitmq message unroutable: %s", returned.ReplyText)
|
|
case confirmation, ok := <-confirmations:
|
|
if !ok || !confirmation.Ack {
|
|
return gerror.New("rabbitmq message was not confirmed")
|
|
}
|
|
select {
|
|
case returned := <-returns:
|
|
return gerror.Newf("rabbitmq message unroutable: %s", returned.ReplyText)
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
|
|
func openConfirmedConnection(ctx context.Context, dsn string) (*amqp.Connection, func(), error) {
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
var (
|
|
socket net.Conn
|
|
stop func() bool
|
|
)
|
|
cleanup := func() {
|
|
if stop != nil {
|
|
stop()
|
|
}
|
|
if socket != nil {
|
|
_ = socket.Close()
|
|
}
|
|
}
|
|
connection, err := amqp.DialConfig(dsn, amqp.Config{
|
|
Dial: func(network, address string) (net.Conn, error) {
|
|
var dialErr error
|
|
socket, dialErr = (&net.Dialer{}).DialContext(ctx, network, address)
|
|
if dialErr != nil {
|
|
return nil, dialErr
|
|
}
|
|
transport := socket
|
|
stop = context.AfterFunc(ctx, func() { _ = transport.Close() })
|
|
return socket, nil
|
|
},
|
|
})
|
|
if err != nil {
|
|
cleanup()
|
|
if ctx.Err() != nil {
|
|
return nil, nil, ctx.Err()
|
|
}
|
|
return nil, nil, gerror.Wrap(err, "rabbitmq confirmed publisher connect failed")
|
|
}
|
|
return connection, cleanup, nil
|
|
}
|
|
|
|
func confirmedExpiration(delay time.Duration) string {
|
|
if delay <= 0 {
|
|
return ""
|
|
}
|
|
return strconv.FormatInt(delay.Milliseconds(), 10)
|
|
}
|