sync: 同步 v1.0.4 并脱敏配置与镜像

server/internal/library/rabbitmq/confirmed.go

@@ -3,6 +3,8 @@ package rabbitmq
 import (
 	"context"
 	"encoding/json"
+	"net"
+	"strconv"
 	"time"
 
 	"github.com/gogf/gf/v2/errors/gerror"
@@ -14,12 +16,41 @@ func PublishConfirmed(ctx context.Context, queue string, data any) error {
 	if err != nil {
 		return err
 	}
-	channel, err := Channel(ctx)
+	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
 	}
-	defer channel.Close()
-	if err = declareQueue(channel, queue); err != nil {
+	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 {
@@ -27,10 +58,8 @@ func PublishConfirmed(ctx context.Context, queue string, data any) error {
 	}
 	confirmations := channel.NotifyPublish(make(chan amqp.Confirmation, 1))
 	returns := channel.NotifyReturn(make(chan amqp.Return, 1))
-	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
-	defer cancel()
 	if err = channel.PublishWithContext(ctx, "", queue, true, false, amqp.Publishing{
-		ContentType: "application/json", DeliveryMode: amqp.Persistent, Timestamp: time.Now(), Body: body,
+		ContentType: "application/json", DeliveryMode: amqp.Persistent, Timestamp: time.Now(), Body: body, Expiration: confirmedExpiration(delay),
 	}); err != nil {
 		return err
 	}
@@ -38,16 +67,61 @@ func PublishConfirmed(ctx context.Context, queue string, data any) error {
 	case <-ctx.Done():
 		return ctx.Err()
 	case returned := <-returns:
-		return gerror.Newf("traffic usage message unroutable: %s", returned.ReplyText)
+		return gerror.Newf("rabbitmq message unroutable: %s", returned.ReplyText)
 	case confirmation, ok := <-confirmations:
 		if !ok || !confirmation.Ack {
-			return gerror.New("traffic usage message was not confirmed")
+			return gerror.New("rabbitmq message was not confirmed")
 		}
 		select {
 		case returned := <-returns:
-			return gerror.Newf("traffic usage message unroutable: %s", returned.ReplyText)
+			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)
+}

server/internal/library/rabbitmq/confirmed_test.go

@@ -0,0 +1,83 @@
+package rabbitmq
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"testing"
+	"time"
+)
+
+func TestConfirmedExpiration(t *testing.T) {
+	for _, item := range []struct {
+		delay time.Duration
+		want  string
+	}{{0, ""}, {-time.Second, ""}, {time.Second, "1000"}, {1500 * time.Millisecond, "1500"}} {
+		if got := confirmedExpiration(item.delay); got != item.want {
+			t.Fatalf("expiration(%v)=%q, want %q", item.delay, got, item.want)
+		}
+	}
+}
+
+func TestConfirmedConnectionHandshakeTimeout(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer listener.Close()
+	finished := make(chan struct{})
+	go func() {
+		defer close(finished)
+		connection, acceptErr := listener.Accept()
+		if acceptErr != nil {
+			return
+		}
+		defer connection.Close()
+		_ = connection.SetDeadline(time.Now().Add(3 * time.Second))
+		_, _ = io.Copy(io.Discard, connection)
+	}()
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	started := time.Now()
+	connection, cleanup, err := openConfirmedConnection(ctx, "amqp://guest:guest@"+listener.Addr().String()+"/")
+	if cleanup != nil {
+		cleanup()
+	}
+	if connection != nil || !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("expected handshake timeout, connection=%v err=%v", connection, err)
+	}
+	if time.Since(started) > 2*time.Second {
+		t.Fatal("handshake exceeded the context deadline")
+	}
+	select {
+	case <-finished:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out socket was not closed")
+	}
+}
+
+func TestConfirmedConnectionCanceledContext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	connection, cleanup, err := openConfirmedConnection(ctx, "invalid-dsn")
+	if connection != nil || cleanup != nil || !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected cancellation before dialing, connection=%v err=%v", connection, err)
+	}
+}
+
+func TestConsumerCloserUnregister(t *testing.T) {
+	closerMu.Lock()
+	initial := len(closers)
+	closerMu.Unlock()
+	for attempt := 0; attempt < 1000; attempt++ {
+		unregister := registerCloser(func() error { return nil })
+		unregister()
+		unregister()
+	}
+	closerMu.Lock()
+	defer closerMu.Unlock()
+	if len(closers) != initial {
+		t.Fatalf("stale consumer closers retained: initial=%d actual=%d", initial, len(closers))
+	}
+}

server/internal/library/rabbitmq/rabbitmq.go

@@ -27,7 +27,8 @@ type JSONHandler[T any] func(ctx context.Context, message T) error
 var (
 	conn     *amqp.Connection
 	connMu   sync.Mutex
-	closers  []func() error
+	closers  = make(map[uint64]func() error)
+	closerID uint64
 	closerMu sync.Mutex
 )
 
@@ -217,7 +218,7 @@ func Close() error {
 	for _, closer := range closers {
 		_ = closer()
 	}
-	closers = nil
+	closers = make(map[uint64]func() error)
 	closerMu.Unlock()
 
 	connMu.Lock()
@@ -252,8 +253,107 @@ func declareDelayQueue(ch *amqp.Channel, delayQueue, targetQueue string) error {
 }
 
 // registerCloser 记录消费者 channel，服务退出时统一关闭。
-func registerCloser(closer func() error) {
+func registerCloser(closer func() error) func() {
 	closerMu.Lock()
 	defer closerMu.Unlock()
-	closers = append(closers, closer)
+	closerID++
+	id := closerID
+	closers[id] = closer
+	return func() {
+		closerMu.Lock()
+		defer closerMu.Unlock()
+		delete(closers, id)
+	}
+}
+
+func ConsumeJSONRequeue[T any](ctx context.Context, queue, consumer string, handler JSONHandler[T]) error {
+	ctx, cancel := context.WithCancel(ctx)
+	open := func() (*amqp.Channel, <-chan amqp.Delivery, error) {
+		channel, err := Channel(ctx)
+		if err != nil {
+			return nil, nil, err
+		}
+		if err = declareQueue(channel, queue); err == nil {
+			err = channel.Qos(1, 0, false)
+		}
+		if err != nil {
+			_ = channel.Close()
+			return nil, nil, err
+		}
+		deliveries, err := channel.Consume(queue, consumer, false, false, false, false, nil)
+		if err != nil {
+			_ = channel.Close()
+			return nil, nil, err
+		}
+		return channel, deliveries, nil
+	}
+	channel, deliveries, err := open()
+	if err != nil {
+		cancel()
+		return err
+	}
+	var channelMu sync.Mutex
+	unregister := registerCloser(func() error {
+		cancel()
+		channelMu.Lock()
+		defer channelMu.Unlock()
+		return channel.Close()
+	})
+	go func() {
+		defer func() {
+			cancel()
+			unregister()
+			_ = channel.Close()
+		}()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case delivery, ok := <-deliveries:
+				if !ok {
+					_ = channel.Close()
+					for {
+						select {
+						case <-ctx.Done():
+							return
+						case <-time.After(3 * time.Second):
+						}
+						next, stream, reconnectErr := open()
+						if reconnectErr != nil {
+							g.Log().Errorf(ctx, "role sync consumer reconnect failed, queue:%s, err:%v", queue, reconnectErr)
+							continue
+						}
+						channelMu.Lock()
+						if ctx.Err() != nil {
+							channelMu.Unlock()
+							_ = next.Close()
+							return
+						}
+						channel, deliveries = next, stream
+						channelMu.Unlock()
+						break
+					}
+					continue
+				}
+				var message T
+				if decodeErr := json.Unmarshal(delivery.Body, &message); decodeErr != nil {
+					g.Log().Errorf(ctx, "rabbitmq invalid role sync message, queue:%s, err:%v", queue, decodeErr)
+					_ = delivery.Ack(false)
+					continue
+				}
+				if handleErr := handler(ctx, message); handleErr != nil {
+					g.Log().Errorf(ctx, "role sync message retained, queue:%s, err:%v", queue, handleErr)
+					select {
+					case <-ctx.Done():
+						return
+					case <-time.After(time.Second):
+					}
+					_ = delivery.Nack(false, true)
+					continue
+				}
+				_ = delivery.Ack(false)
+			}
+		}
+	}()
+	return nil
 }