fix: register supplier response waiter before reader can consume reply

gateway/internal/upstream/immediate_response_test.go

@@ -0,0 +1,103 @@
+package upstream
+
+import (
+	"context"
+	"encoding/binary"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"cmpp-platform/gateway/internal/queue"
+	cmpp "github.com/bigwhite/gocmpp"
+)
+
+func TestImmediateSupplierResponseIsNeverLost(t *testing.T) {
+	for _, version := range []string{"2.0", "3.0"} {
+		t.Run(version, func(t *testing.T) {
+			listener, err := net.Listen("tcp", "127.0.0.1:0")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer listener.Close()
+			var wire atomic.Int32
+			done := make(chan struct{})
+			go func() {
+				defer close(done)
+				socket, err := listener.Accept()
+				if err != nil {
+					return
+				}
+				defer socket.Close()
+				for {
+					header := make([]byte, 12)
+					if _, err = io.ReadFull(socket, header); err != nil {
+						return
+					}
+					n := binary.BigEndian.Uint32(header)
+					if n < 12 || n > 4096 {
+						return
+					}
+					body := make([]byte, int(n)-12)
+					if _, err = io.ReadFull(socket, body); err != nil {
+						return
+					}
+					kind := binary.BigEndian.Uint32(header[4:])
+					var response []byte
+					switch kind {
+					case 1:
+						size := 18
+						if version == "3.0" {
+							size = 21
+						}
+						response = make([]byte, size)
+						response[size-1] = byte(protocolVersion(version))
+					case 4:
+						wire.Add(1)
+						size := 9
+						if version == "3.0" {
+							size = 12
+						}
+						response = make([]byte, size)
+						binary.BigEndian.PutUint64(response, uint64(wire.Load()))
+					default:
+						return
+					}
+					binary.BigEndian.PutUint32(header, uint32(12+len(response)))
+					binary.BigEndian.PutUint32(header[4:], kind|0x80000000)
+					if _, err = socket.Write(append(header, response...)); err != nil {
+						return
+					}
+				}
+			}()
+			api := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(204) }))
+			defer api.Close()
+			client := cmpp.NewClient(protocolVersion(version))
+			if err = client.Connect(listener.Addr().String(), "qa", "qa", time.Second); err != nil {
+				t.Fatal(err)
+			}
+			c := &connection{client: client, retired: true, config: queue.UpstreamConfig{CMPPVersion: version}, pending: make(map[uint32]chan submitPartResponse), tracker: make(map[uint64]queue.SubmitCommand), apiBaseURL: api.URL, httpClient: api.Client()}
+			defer func() { c.close(); <-done }()
+			go c.readLoop()
+			cmd := submitCommandForPacketTest(version)
+			parts, err := splitSubmitContent(cmd.CMPP.MsgFmt, cmd.Content)
+			if err != nil {
+				t.Fatal(err)
+			}
+			for i := 0; i < 500; i++ {
+				ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
+				_, _, result, err := c.submitPart(ctx, cmd, parts[0])
+				cancel()
+				if err != nil || result.SubmitStatus != "accepted" {
+					t.Fatalf("response lost at wire submit %d: %+v / %v", i+1, result, err)
+				}
+			}
+			if wire.Load() != 500 {
+				t.Fatalf("wire submits = %d", wire.Load())
+			}
+		})
+	}
+}

gateway/internal/upstream/submit.go

@@ -138,8 +138,8 @@ func (c *connection) submitPart(ctx context.Context, cmd queue.SubmitCommand, pa
 	c.mu.Lock()
 	client := c.client
 	closed := c.closed
-	c.mu.Unlock()
 	if closed || client == nil {
+		c.mu.Unlock()
 		err := fmt.Errorf("supplier connection is not available")
 		result := submitResult(cmd, 0, "", "timeout", "CONNECTION_LOST", err.Error())
 		return 0, "", result, err
@@ -153,6 +153,12 @@ func (c *connection) submitPart(ctx context.Context, cmd queue.SubmitCommand, pa
 		wireSource = "gateway_write_complete"
 	}
 	c.sendMu.Unlock()
+	// The reader must not consume an immediate response before its waiter is
+	// registered. Keep the same mu -> sendMu lock order as the heartbeat path.
+	if err == nil {
+		c.pending[seq] = rspCh
+	}
+	c.mu.Unlock()
 	if err != nil {
 		c.emitProtocolLog(protocolLogEvent{
 			Protocol:      "cmpp",
@@ -195,9 +201,6 @@ func (c *connection) submitPart(ctx context.Context, cmd queue.SubmitCommand, pa
 		},
 	})
 
-	c.mu.Lock()
-	c.pending[seq] = rspCh
-	c.mu.Unlock()
 	defer func() {
 		c.mu.Lock()
 		delete(c.pending, seq)