fix: stop retrying unrecoverable streaming writes (#50894)

issue: #50878

- streaming redo: convert dropped collection or partition waits into
unrecoverable errors so stale inserts stop retrying
- streaming write errors: return unrecoverable errors for pre-WAL schema
and function materialization failures
- streaming status: map terminal streaming errors to non-retryable gRPC
statuses
- streaming handler: keep RO channel assignment checks retryable while
waiting for RW promotion

---------

Signed-off-by: chyezh <chyezh@outlook.com>

internal/streamingnode/client/handler/handler_client_impl.go

@@ -77,7 +77,7 @@ func (hc *handlerClientImpl) GetReplicateCheckpoint(ctx context.Context, pchanne
 	logger := mlog.With(mlog.FieldPChannel(pchannel), mlog.String("handler", "replicate checkpoint"))
 	cp, err := hc.createHandlerAfterStreamingNodeReady(ctx, logger, pchannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) {
 		if assign.Channel.AccessMode != types.AccessModeRW {
-			return nil, status.NewInvalidArgument("replicate checkpoint can only be read for RW channel")
+			return nil, status.NewInner("replicate checkpoint can only be read for RW channel")
 		}
 		localWAL, err := registry.GetLocalAvailableWAL(assign.Channel)
 		if err == nil {
@@ -114,7 +114,7 @@ func (hc *handlerClientImpl) GetSalvageCheckpoint(ctx context.Context, pchannel
 	logger := mlog.With(mlog.FieldPChannel(pchannel), mlog.String("handler", "salvage checkpoint"))
 	cps, err := hc.createHandlerAfterStreamingNodeReady(ctx, logger, pchannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) {
 		if assign.Channel.AccessMode != types.AccessModeRW {
-			return nil, status.NewInvalidArgument("salvage checkpoint can only be read for RW channel")
+			return nil, status.NewInner("salvage checkpoint can only be read for RW channel")
 		}
 		localWAL, err := registry.GetLocalAvailableWAL(assign.Channel)
 		if err == nil {
@@ -217,7 +217,7 @@ func (hc *handlerClientImpl) CreateProducer(ctx context.Context, opts *ProducerO
 	logger := mlog.With(mlog.FieldPChannel(opts.PChannel), mlog.String("handler", "producer"))
 	p, err := hc.createHandlerAfterStreamingNodeReady(ctx, logger, opts.PChannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) {
 		if assign.Channel.AccessMode != types.AccessModeRW {
-			return nil, status.NewInvalidArgument("producer can only be created for RW channel")
+			return nil, status.NewInner("producer can only be created for RW channel")
 		}
 		// Check if the localWAL is assigned at local
 		localWAL, err := registry.GetLocalAvailableWAL(assign.Channel)

internal/streamingnode/client/handler/handler_client_test.go

@@ -194,6 +194,59 @@ func TestHandlerClient_GetSalvageCheckpoint(t *testing.T) {
 	assert.Nil(t, cps)
 }
 
+func TestHandlerClientReadOnlyAssignmentWaitsForNextAssignment(t *testing.T) {
+	assignment := &types.PChannelInfoAssigned{
+		Channel: types.PChannelInfo{Name: "pchannel", Term: 1, AccessMode: types.AccessModeRO},
+		Node:    types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
+	}
+
+	testcases := []struct {
+		name string
+		run  func(*handlerClientImpl) error
+	}{
+		{
+			name: "replicate checkpoint",
+			run: func(handler *handlerClientImpl) error {
+				cp, err := handler.GetReplicateCheckpoint(context.Background(), "pchannel")
+				assert.Nil(t, cp)
+				return err
+			},
+		},
+		{
+			name: "salvage checkpoint",
+			run: func(handler *handlerClientImpl) error {
+				cps, err := handler.GetSalvageCheckpoint(context.Background(), "pchannel")
+				assert.Nil(t, cps)
+				return err
+			},
+		},
+		{
+			name: "producer",
+			run: func(handler *handlerClientImpl) error {
+				producer, err := handler.CreateProducer(context.Background(), &ProducerOptions{PChannel: "pchannel"})
+				assert.Nil(t, producer)
+				return err
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			w := mock_assignment.NewMockWatcher(t)
+			w.EXPECT().Get(mock.Anything, "pchannel").Return(assignment)
+			w.EXPECT().Watch(mock.Anything, "pchannel", assignment).Return(context.Canceled)
+
+			handler := &handlerClientImpl{
+				lifetime: typeutil.NewLifetime(),
+				watcher:  w,
+			}
+
+			err := tc.run(handler)
+			assert.ErrorIs(t, err, context.Canceled)
+		})
+	}
+}
+
 func TestHandlerClient_PrepareReleaseManualFlush(t *testing.T) {
 	assignment := &types.PChannelInfoAssigned{
 		Channel: types.PChannelInfo{Name: "pchannel", Term: 1, AccessMode: types.AccessModeRO},

internal/streamingnode/server/wal/interceptors/redo/redo_interceptor.go

@@ -56,6 +56,9 @@ func (r *redoAppendInterceptor) waitUntilGrowingSegmentReady(ctx context.Context
 			uniqueKey := shards.PartitionUniqueKey{CollectionID: h.CollectionId, PartitionID: partition.PartitionId}
 			ready, err := r.shardManager.WaitUntilGrowingSegmentReady(uniqueKey)
 			if err != nil {
+				if errors.IsAny(err, shards.ErrCollectionNotFound, shards.ErrPartitionNotFound) {
+					return status.NewUnrecoverableError("fail to wait growing segment ready, %s", err.Error())
+				}
 				return err
 			}
 			select {

internal/streamingnode/server/wal/interceptors/redo/redo_interceptor_test.go

@@ -0,0 +1,106 @@
+package redo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/shards"
+	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/recovery"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
+	"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
+)
+
+type testTxnManager struct{}
+
+func (m *testTxnManager) RecoverDone() <-chan struct{} {
+	ch := make(chan struct{})
+	close(ch)
+	return ch
+}
+
+func TestRedoInterceptorWaitUntilGrowingSegmentReadyReturnsUnrecoverableForDroppedCollectionOrPartition(t *testing.T) {
+	paramtable.Init()
+	resource.InitForTest(t)
+
+	tests := []struct {
+		name       string
+		vchannels  map[string]*streamingpb.VChannelMeta
+		collection int64
+		partition  int64
+	}{
+		{
+			name:       "collection not found",
+			vchannels:  map[string]*streamingpb.VChannelMeta{},
+			collection: 1,
+			partition:  2,
+		},
+		{
+			name: "partition not found",
+			vchannels: map[string]*streamingpb.VChannelMeta{
+				"test-vchannel": {
+					Vchannel: "test-vchannel",
+					State:    streamingpb.VChannelState_VCHANNEL_STATE_NORMAL,
+					CollectionInfo: &streamingpb.CollectionInfoOfVChannel{
+						CollectionId: 1,
+						Partitions: []*streamingpb.PartitionInfoOfVChannel{
+							{PartitionId: 2},
+						},
+					},
+				},
+			},
+			collection: 1,
+			partition:  3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			shardManager := shards.RecoverShardManager(&shards.ShardManagerRecoverParam{
+				ChannelInfo: types.PChannelInfo{Name: "test-channel", Term: 1},
+				WAL:         syncutil.NewFuture[wal.WAL](),
+				InitialRecoverSnapshot: &recovery.RecoverySnapshot{
+					VChannels:          tt.vchannels,
+					SegmentAssignments: map[int64]*streamingpb.SegmentAssignmentMeta{},
+					Checkpoint:         &recovery.WALCheckpoint{TimeTick: 100},
+				},
+				TxnManager: &testTxnManager{},
+			})
+			defer shardManager.Close()
+
+			interceptor := NewInterceptorBuilder().Build(&interceptors.InterceptorBuildParam{
+				ShardManager: shardManager,
+			})
+			defer interceptor.Close()
+
+			msg := message.NewInsertMessageBuilderV1().
+				WithVChannel("test-vchannel").
+				WithHeader(&messagespb.InsertMessageHeader{
+					CollectionId: tt.collection,
+					Partitions: []*messagespb.PartitionSegmentAssignment{
+						{PartitionId: tt.partition, Rows: 1},
+					},
+				}).
+				WithBody(&msgpb.InsertRequest{}).
+				MustBuildMutable()
+
+			msgID, err := interceptor.DoAppend(context.Background(), msg, func(context.Context, message.MutableMessage) (message.MessageID, error) {
+				t.Fatal("append should not be called when collection or partition is already dropped")
+				return nil, nil
+			})
+			require.Nil(t, msgID)
+			require.Error(t, err)
+			require.True(t, status.AsStreamingError(err).IsUnrecoverable())
+		})
+	}
+}

internal/streamingnode/server/wal/interceptors/shard/shard_interceptor.go

@@ -174,7 +174,7 @@ func (impl *shardInterceptor) handleInsertMessage(ctx context.Context, msg messa
 			mlog.Bool("schemaVersionProvided", header.SchemaVersion != nil),
 			mlog.Int32("schemaVersion", schemaVersion),
 			mlog.Err(err))
-		return nil, errors.Wrap(err, "CheckIfCollectionSchemaVersionMatch")
+		return nil, status.NewUnrecoverableError("unexpected error from CheckIfCollectionSchemaVersionMatch: %s", err.Error())
 	}
 	schemaVersion = correctSchemaVersion
 	if err := impl.materializeFunctionFields(ctx, insertMsg, header.GetCollectionId(), schemaVersion); err != nil {
@@ -182,7 +182,7 @@ func (impl *shardInterceptor) handleInsertMessage(ctx context.Context, msg messa
 			mlog.Int64("collectionID", header.GetCollectionId()),
 			mlog.Int32("schemaVersion", schemaVersion),
 			mlog.Err(err))
-		return nil, status.NewInner("failed to materialize function fields before WAL append: %s", err.Error())
+		return nil, status.NewUnrecoverableError("failed to materialize function fields before WAL append: %s", err.Error())
 	}
 	for _, partition := range header.GetPartitions() {
 		if partition.BinarySize == 0 {

internal/streamingnode/server/wal/interceptors/shard/shard_interceptor_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/shard/shards"
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
 	"github.com/milvus-io/milvus/internal/util/function"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
 	"github.com/milvus-io/milvus/pkg/v3/common"
 	"github.com/milvus-io/milvus/pkg/v3/mlog"
 	"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
@@ -476,10 +477,11 @@ func TestShardInterceptor(t *testing.T) {
 	assert.Error(t, err)
 	assert.Nil(t, msgID)
 
-	// Unexpected error from the schema version check must be propagated as-is.
+	// Unexpected error from the schema version check must stop producer retry.
 	shardManager.EXPECT().CheckIfCollectionSchemaVersionMatch(insertHdrMatcher).Return(int32(-1), mockErr).Once()
 	msgID, err = i.DoAppend(ctx, msg, appender)
 	assert.Error(t, err)
+	assert.True(t, status.AsStreamingError(err).IsUnrecoverable())
 	assert.Nil(t, msgID)
 
 	msg = message.NewDeleteMessageBuilderV1().

internal/util/streamingutil/status/rpc_error.go

@@ -24,6 +24,12 @@ var streamingErrorToGRPCStatus = map[streamingpb.StreamingCode]codes.Code{
 	streamingpb.StreamingCode_STREAMING_CODE_INVAILD_ARGUMENT:          codes.InvalidArgument,
 	streamingpb.StreamingCode_STREAMING_CODE_TRANSACTION_EXPIRED:       codes.FailedPrecondition,
 	streamingpb.StreamingCode_STREAMING_CODE_INVALID_TRANSACTION_STATE: codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_UNRECOVERABLE:             codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_RESOURCE_ACQUIRED:         codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_REPLICATE_VIOLATION:       codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_WALNAME_MISMATCH:          codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_SCHEMA_VERSION_MISMATCH:   codes.FailedPrecondition,
+	streamingpb.StreamingCode_STREAMING_CODE_RATE_LIMIT_REJECTED:       codes.ResourceExhausted,
 	streamingpb.StreamingCode_STREAMING_CODE_UNKNOWN:                   codes.Unknown,
 }
 

internal/util/streamingutil/status/rpc_error_test.go

@@ -41,6 +41,21 @@ func TestNewGRPCStatusFromStreamingError(t *testing.T) {
 	)
 	assert.Equal(t, codes.FailedPrecondition, st.Code())
 
+	st = NewGRPCStatusFromStreamingError(
+		NewUnrecoverableError("test"),
+	)
+	assert.Equal(t, codes.FailedPrecondition, st.Code())
+
+	st = NewGRPCStatusFromStreamingError(
+		NewSchemaVersionMismatch("test"),
+	)
+	assert.Equal(t, codes.FailedPrecondition, st.Code())
+
+	st = NewGRPCStatusFromStreamingError(
+		NewRateLimitRejected("test"),
+	)
+	assert.Equal(t, codes.ResourceExhausted, st.Code())
+
 	st = NewGRPCStatusFromStreamingError(
 		New(10086, "test"),
 	)

internal/util/streamingutil/status/streaming_error.go

@@ -56,6 +56,7 @@ func (e *StreamingError) IsSkippedOperation() bool {
 // Stop resuming retry and report to user.
 func (e *StreamingError) IsUnrecoverable() bool {
 	return e.Code == streamingpb.StreamingCode_STREAMING_CODE_UNRECOVERABLE ||
+		e.IsInvalidArgument() ||
 		e.IsReplicateViolation() ||
 		e.IsTxnUnavilable() || e.IsSchemaVersionMismatch()
 }
@@ -91,6 +92,11 @@ func (e *StreamingError) IsSchemaVersionMismatch() bool {
 	return e.Code == streamingpb.StreamingCode_STREAMING_CODE_SCHEMA_VERSION_MISMATCH
 }
 
+// IsInvalidArgument returns true if the error is caused by invalid argument.
+func (e *StreamingError) IsInvalidArgument() bool {
+	return e.Code == streamingpb.StreamingCode_STREAMING_CODE_INVAILD_ARGUMENT
+}
+
 // IsOnShutdown returns true if the error is caused by on shutdown.
 func (e *StreamingError) IsOnShutdown() bool {
 	return e.Code == streamingpb.StreamingCode_STREAMING_CODE_ON_SHUTDOWN

internal/util/streamingutil/status/streaming_error_test.go

@@ -70,4 +70,10 @@ func TestStreamingError(t *testing.T) {
 	assert.True(t, streamingErr.IsUnrecoverable())
 	pbErr = streamingErr.AsPBError()
 	assert.Equal(t, streamingpb.StreamingCode_STREAMING_CODE_TRANSACTION_EXPIRED, pbErr.Code)
+
+	streamingErr = NewInvalidArgument("test")
+	assert.True(t, streamingErr.IsInvalidArgument())
+	assert.True(t, streamingErr.IsUnrecoverable())
+	pbErr = streamingErr.AsPBError()
+	assert.Equal(t, streamingpb.StreamingCode_STREAMING_CODE_INVAILD_ARGUMENT, pbErr.Code)
 }