fix: reuse CN pipeline connections after clean teardown

pkg/cnservice/server.go

@@ -905,9 +905,10 @@ func handleWaitingNextMsg(ctx context.Context, message morpc.Message, cs morpc.C
 		if cache, err = cs.CreateCache(ctx, message.GetID()); err != nil {
 			return err
 		}
-		cache.Add(message)
+		return cache.Add(message)
+	default:
+		return moerr.NewInvalidInputNoCtx("only pipeline messages may be fragmented")
 	}
-	return nil
 }
 
 func handleAssemblePipeline(ctx context.Context, message morpc.Message, cs morpc.ClientSession) error {
@@ -917,19 +918,27 @@ func handleAssemblePipeline(ctx context.Context, message morpc.Message, cs morpc
 	if err != nil {
 		return err
 	}
+	// CreateCache also returns a cache for an unfragmented message. Always
+	// remove the map entry on every terminal assembly path, not just close the
+	// queue object.
+	defer cs.DeleteCache(message.GetID())
+	finalMessage := message.(*pipeline.Message)
 	for {
-		msg, ok, err := cache.Pop()
+		cached, ok, err := cache.Pop()
 		if err != nil {
 			return err
 		}
 		if !ok {
-			cache.Close()
 			break
 		}
-		data = append(data, msg.(*pipeline.Message).GetData()...)
+		fragment, ok := cached.(*pipeline.Message)
+		if !ok || fragment.GetCmd() != finalMessage.GetCmd() ||
+			fragment.GetRequestedTeardownMode() != finalMessage.GetRequestedTeardownMode() {
+			return moerr.NewInvalidInputNoCtx("inconsistent pipeline message fragments")
+		}
+		data = append(data, fragment.GetData()...)
 	}
-	msg := message.(*pipeline.Message)
-	msg.SetData(append(data, msg.GetData()...))
+	finalMessage.SetData(append(data, finalMessage.GetData()...))
 	return nil
 }
 

pkg/cnservice/server_test.go

@@ -141,6 +141,7 @@ func Test_InitServer(t *testing.T) {
 	session := mock_morpc.NewMockClientSession(ctrl)
 	msg.Cmd = pipeline.Method_PipelineMessage
 	session.EXPECT().CreateCache(ctx, uint64(0)).Return(&testMessageCache{}, nil).Times(2)
+	session.EXPECT().DeleteCache(uint64(0)).Times(1)
 
 	msg.Sid = pipeline.Status_WaitingNext
 	err = srv.handleRequest(
@@ -173,6 +174,36 @@ type testMessageCache struct {
 	cache []morpc.Message
 }
 
+func TestHandleAssemblePipelineDeletesCacheAndRejectsMixedNegotiation(t *testing.T) {
+	t.Run("assembles and deletes", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		session := mock_morpc.NewMockClientSession(ctrl)
+		ctx := context.Background()
+		cache := &testMessageCache{cache: []morpc.Message{
+			&pipeline.Message{Id: 42, Cmd: pipeline.Method_PipelineMessage, Data: []byte("a"), RequestedTeardownMode: pipeline.StreamTeardownMode_FinishAck},
+			&pipeline.Message{Id: 42, Cmd: pipeline.Method_PipelineMessage, Data: []byte("b"), RequestedTeardownMode: pipeline.StreamTeardownMode_FinishAck},
+		}}
+		session.EXPECT().CreateCache(ctx, uint64(42)).Return(cache, nil)
+		session.EXPECT().DeleteCache(uint64(42))
+		final := &pipeline.Message{Id: 42, Cmd: pipeline.Method_PipelineMessage, Data: []byte("c"), RequestedTeardownMode: pipeline.StreamTeardownMode_FinishAck}
+		require.NoError(t, handleAssemblePipeline(ctx, final, session))
+		require.Equal(t, []byte("abc"), final.GetData())
+	})
+
+	t.Run("mixed teardown mode is protocol error", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		session := mock_morpc.NewMockClientSession(ctrl)
+		ctx := context.Background()
+		cache := &testMessageCache{cache: []morpc.Message{
+			&pipeline.Message{Id: 43, Cmd: pipeline.Method_PipelineMessage, RequestedTeardownMode: pipeline.StreamTeardownMode_LegacyClose},
+		}}
+		session.EXPECT().CreateCache(ctx, uint64(43)).Return(cache, nil)
+		session.EXPECT().DeleteCache(uint64(43))
+		final := &pipeline.Message{Id: 43, Cmd: pipeline.Method_PipelineMessage, RequestedTeardownMode: pipeline.StreamTeardownMode_FinishAck}
+		require.Error(t, handleAssemblePipeline(ctx, final, session))
+	})
+}
+
 func (c *testMessageCache) Add(val morpc.Message) error {
 	c.cache = append(c.cache, val)
 	return nil

pkg/cnservice/types.go

@@ -158,6 +158,9 @@ type Config struct {
 		BatchRows int64 `toml:"batch-rows"`
 		// BatchSize is the memory limit for one batch
 		BatchSize int64 `toml:"batch-size"`
+		// DisableStreamReuse is an emergency rollback gate for negotiated
+		// FIN/FIN_ACK pipeline teardown. It is enabled by default when false.
+		DisableStreamReuse bool `toml:"disable-stream-reuse"`
 	}
 
 	// Frontend parameters for the frontend
@@ -454,6 +457,10 @@ func (c *Config) Validate() error {
 		moruntime.EnableCheckInvalidRCErrors,
 		c.Txn.EnableCheckRCInvalidError,
 	)
+	moruntime.ServiceRuntime(c.UUID).SetGlobalVariables(
+		moruntime.EnablePipelineStreamReuse,
+		!c.Pipeline.DisableStreamReuse,
+	)
 	return nil
 }
 

pkg/common/morpc/metrics.go

@@ -130,6 +130,9 @@ type serverMetrics struct {
 	outputBytesCounter            prometheus.Counter
 	sendingQueueSizeGauge         prometheus.Gauge
 	sessionSizeGauge              prometheus.Gauge
+	receivedStreamStateGauge      prometheus.Gauge
+	sentStreamStateGauge          prometheus.Gauge
+	messageCacheStateGauge        prometheus.Gauge
 	sendingBatchSizeGauge         prometheus.Gauge
 	writeDurationHistogram        prometheus.Observer
 	writeLatencyDurationHistogram prometheus.Observer
@@ -144,6 +147,9 @@ func newServerMetrics(name string) *serverMetrics {
 		sendingQueueSizeGauge:         v2.NewRPCServerSendingQueueSizeGaugeByName(name),
 		writeLatencyDurationHistogram: v2.NewRPCServerWriteLatencyDurationHistogramByName(name),
 		sessionSizeGauge:              v2.NewRPCServerSessionSizeGaugeByName(name),
+		receivedStreamStateGauge:      v2.NewRPCServerStreamStateGaugeByName(name, "received_sequence"),
+		sentStreamStateGauge:          v2.NewRPCServerStreamStateGaugeByName(name, "sent_sequence"),
+		messageCacheStateGauge:        v2.NewRPCServerStreamStateGaugeByName(name, "message_cache"),
 		inputBytesCounter:             v2.NewRPCInputCounter(),
 		outputBytesCounter:            v2.NewRPCOutputCounter(),
 	}

pkg/common/morpc/server.go

@@ -244,12 +244,17 @@ func (s *server) onMessage(rs goetty.IOSession, value any, sequence uint64) erro
 	if request.stream &&
 		!cs.validateStreamRequest(requestID, request.streamSequence) {
 		s.logger.Error("failed to handle stream request",
-			zap.Uint32("last-sequence", cs.receivedStreamSequences[requestID]),
+			zap.Uint32("last-sequence", cs.lastReceivedStreamSequence(requestID)),
 			zap.Uint32("current-sequence", request.streamSequence),
 			zap.String("client", rs.RemoteAddress()))
 		cs.cancelWrite()
 		return moerr.NewStreamClosedNoCtx()
 	}
+	if request.stream {
+		request.Ctx = context.WithValue(request.Ctx, streamTerminalTokenContextKey{}, StreamTerminalToken{
+			owner: cs, streamID: requestID, sequence: request.streamSequence,
+		})
+	}
 
 	// handle internal message
 	if request.internal {
@@ -557,7 +562,10 @@ type clientSession struct {
 	conn          goetty.IOSession
 	c             chan *Future
 	newFutureFunc func() *Future
-	// streaming id -> last received sequence, no concurrent, access in io goroutine
+	// streaming id -> last received sequence. FinishStream also accesses this
+	// map from a handler goroutine, so streamStateMu is the synchronization
+	// boundary for validation and terminal retirement.
+	streamStateMu           sync.Mutex
 	receivedStreamSequences map[uint64]uint32
 	// streaming id -> last sent sequence, multi-stream access in multi-goroutines if
 	// the tcp connection is shared. But no concurrent in one stream.
@@ -604,6 +612,8 @@ func (cs *clientSession) RemoteAddress() string {
 }
 
 func (cs *clientSession) Close() error {
+	cs.streamStateMu.Lock()
+	defer cs.streamStateMu.Unlock()
 	cs.mu.Lock()
 	defer cs.mu.Unlock()
 	if cs.mu.closed {
@@ -613,6 +623,15 @@ func (cs *clientSession) Close() error {
 	cs.cleanSend()
 	close(cs.c)
 	cs.mu.closed = true
+	if cs.metrics != nil {
+		cs.metrics.receivedStreamStateGauge.Sub(float64(len(cs.receivedStreamSequences)))
+		sentCount := 0
+		cs.sentStreamSequences.Range(func(_, _ any) bool { sentCount++; return true })
+		cs.metrics.sentStreamStateGauge.Sub(float64(sentCount))
+		cs.metrics.messageCacheStateGauge.Sub(float64(len(cs.mu.caches)))
+	}
+	clear(cs.receivedStreamSequences)
+	cs.sentStreamSequences.Clear()
 	for _, c := range cs.mu.caches {
 		c.cache.Close()
 	}
@@ -749,7 +768,11 @@ func (cs *clientSession) checkCacheTimeout() {
 				cs.mu.Lock()
 				for k, c := range cs.mu.caches {
 					if c.closeIfTimeout() {
+						c.cache.Close()
 						delete(cs.mu.caches, k)
+						if cs.metrics != nil {
+							cs.metrics.messageCacheStateGauge.Dec()
+						}
 					}
 				}
 				cs.mu.Unlock()
@@ -766,17 +789,73 @@ func (cs *clientSession) cancelWrite() {
 func (cs *clientSession) validateStreamRequest(
 	id uint64,
 	sequence uint32) bool {
+	cs.streamStateMu.Lock()
+	defer cs.streamStateMu.Unlock()
 	expectSequence := cs.receivedStreamSequences[id] + 1
 	if sequence != expectSequence {
 		return false
 	}
 	cs.receivedStreamSequences[id] = sequence
 	if sequence == 1 {
 		cs.sentStreamSequences.Store(id, uint32(0))
+		if cs.metrics != nil {
+			cs.metrics.receivedStreamStateGauge.Inc()
+			cs.metrics.sentStreamStateGauge.Inc()
+		}
 	}
 	return true
 }
 
+func (cs *clientSession) lastReceivedStreamSequence(id uint64) uint32 {
+	cs.streamStateMu.Lock()
+	defer cs.streamStateMu.Unlock()
+	return cs.receivedStreamSequences[id]
+}
+
+// FinishStream synchronously flushes the final response before removing both
+// receive and send sequence entries. Holding streamStateMu prevents the IO loop
+// from validating a later request against half-retired state.
+func (cs *clientSession) FinishStream(
+	ctx context.Context,
+	token StreamTerminalToken,
+	response Message,
+) error {
+	cs.streamStateMu.Lock()
+	valid := token.owner == cs &&
+		cs.receivedStreamSequences[token.streamID] == token.sequence &&
+		response != nil && response.GetID() == token.streamID
+	if !valid {
+		cs.streamStateMu.Unlock()
+		_ = cs.Close()
+		return moerr.NewStreamClosedNoCtx()
+	}
+
+	cache, err := cs.GetCache(token.streamID)
+	if err != nil || cache != nil {
+		cs.streamStateMu.Unlock()
+		_ = cs.Close()
+		if err != nil {
+			return err
+		}
+		return moerr.NewInternalErrorNoCtx("cannot finish stream with pending message cache")
+	}
+
+	err = cs.Write(ctx, response)
+	if err == nil {
+		delete(cs.receivedStreamSequences, token.streamID)
+		cs.sentStreamSequences.Delete(token.streamID)
+		if cs.metrics != nil {
+			cs.metrics.receivedStreamStateGauge.Dec()
+			cs.metrics.sentStreamStateGauge.Dec()
+		}
+	}
+	cs.streamStateMu.Unlock()
+	if err != nil {
+		_ = cs.Close()
+	}
+	return err
+}
+
 func (cs *clientSession) CreateCache(
 	ctx context.Context,
 	cacheID uint64) (MessageCache, error) {
@@ -791,6 +870,9 @@ func (cs *clientSession) CreateCache(
 	if !ok {
 		v = cacheWithContext{ctx: ctx, cache: newCache()}
 		cs.mu.caches[cacheID] = v
+		if cs.metrics != nil {
+			cs.metrics.messageCacheStateGauge.Inc()
+		}
 		cs.startCheckCacheTimeout()
 	}
 	return v.cache, nil
@@ -806,6 +888,9 @@ func (cs *clientSession) DeleteCache(cacheID uint64) {
 	if c, ok := cs.mu.caches[cacheID]; ok {
 		c.cache.Close()
 		delete(cs.mu.caches, cacheID)
+		if cs.metrics != nil {
+			cs.metrics.messageCacheStateGauge.Dec()
+		}
 	}
 }