test: add unit tests for salvage checkpoint in data salvage feature

Add unit test coverage for salvage checkpoint functionality introduced in PR #47599:

- kv_catalog_test.go: TestCatalogSalvageCheckpoint (save/get success, errors,
  multiple clusters) and extend TestBuildPrefixAndKey with salvage key builders
- replicate_service_test.go: TestReplicateServiceGetSalvageCheckpoint
  (success, error, closed_lifetime)
- manager_test.go: TestSalvageCheckpointLoadedFromEtcd and
  TestSalvageCheckpointMultipleForcePromotes
- salvage_checkpoint_test.go (new): TestUpdateCheckpointForcePromote and
  TestConsumeDirtySnapshotWithSalvageCheckpoint covering the force-promote
  path in recovery_storage_impl

Signed-off-by: Yihao Dai <yihao.dai@zilliz.com>

Author: bigsheeper

internal/distributed/streaming/replicate_service_test.go

@@ -876,6 +876,74 @@ func TestReplicateServiceGetCheckpoint(t *testing.T) {
 	})
 }
 
+func TestReplicateServiceGetSalvageCheckpoint(t *testing.T) {
+	t.Run("success", func(t *testing.T) {
+		c := mock_client.NewMockClient(t)
+		h := mock_handler.NewMockHandlerClient(t)
+
+		expected := []*wal.ReplicateCheckpoint{
+			{ClusterID: "primary", PChannel: "primary-rootcoord-dml_0", TimeTick: 500},
+		}
+		h.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-channel").Return(expected, nil)
+
+		rs := &replicateService{
+			walAccesserImpl: &walAccesserImpl{
+				lifetime:             typeutil.NewLifetime(),
+				clusterID:            "by-dev",
+				streamingCoordClient: c,
+				handlerClient:        h,
+				producers:            make(map[string]*producer.ResumableProducer),
+			},
+		}
+
+		checkpoints, err := rs.GetSalvageCheckpoint(context.Background(), "test-channel")
+		assert.NoError(t, err)
+		assert.Equal(t, expected, checkpoints)
+	})
+
+	t.Run("error", func(t *testing.T) {
+		c := mock_client.NewMockClient(t)
+		h := mock_handler.NewMockHandlerClient(t)
+
+		h.EXPECT().GetSalvageCheckpoint(mock.Anything, "bad-channel").Return(nil, errors.New("not found"))
+
+		rs := &replicateService{
+			walAccesserImpl: &walAccesserImpl{
+				lifetime:             typeutil.NewLifetime(),
+				clusterID:            "by-dev",
+				streamingCoordClient: c,
+				handlerClient:        h,
+				producers:            make(map[string]*producer.ResumableProducer),
+			},
+		}
+
+		checkpoints, err := rs.GetSalvageCheckpoint(context.Background(), "bad-channel")
+		assert.Error(t, err)
+		assert.Nil(t, checkpoints)
+	})
+
+	t.Run("closed_lifetime", func(t *testing.T) {
+		c := mock_client.NewMockClient(t)
+		h := mock_handler.NewMockHandlerClient(t)
+
+		rs := &replicateService{
+			walAccesserImpl: &walAccesserImpl{
+				lifetime:             typeutil.NewLifetime(),
+				clusterID:            "by-dev",
+				streamingCoordClient: c,
+				handlerClient:        h,
+				producers:            make(map[string]*producer.ResumableProducer),
+			},
+		}
+		rs.lifetime.SetState(typeutil.LifetimeStateStopped)
+		rs.lifetime.Wait()
+
+		checkpoints, err := rs.GetSalvageCheckpoint(context.Background(), "test-channel")
+		assert.ErrorIs(t, err, ErrWALAccesserClosed)
+		assert.Nil(t, checkpoints)
+	})
+}
+
 func TestReplicateServiceAppendClosed(t *testing.T) {
 	c := mock_client.NewMockClient(t)
 	h := mock_handler.NewMockHandlerClient(t)

internal/metastore/kv/streamingnode/kv_catalog_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/stretchr/testify/mock"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
 	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
 	"github.com/milvus-io/milvus/internal/kv/mocks"
@@ -197,6 +198,99 @@ func TestCatalogVChannel(t *testing.T) {
 	}
 }
 
+func TestCatalogSalvageCheckpoint(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("save_and_get_success", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		cp := &commonpb.ReplicateCheckpoint{
+			ClusterId: "source-cluster",
+			Pchannel:  "source-cluster-rootcoord-dml_0",
+		}
+		cpBytes, err := proto.Marshal(cp)
+		assert.NoError(t, err)
+
+		kv.EXPECT().Save(mock.Anything, mock.Anything, string(cpBytes)).Return(nil)
+		err = catalog.SaveSalvageCheckpoint(ctx, "p1", cp)
+		assert.NoError(t, err)
+
+		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(
+			[]string{"streamingnode-meta/wal/p1/salvage-checkpoint/source-cluster"},
+			[]string{string(cpBytes)},
+			nil,
+		)
+		checkpoints, err := catalog.GetSalvageCheckpoint(ctx, "p1")
+		assert.NoError(t, err)
+		assert.Len(t, checkpoints, 1)
+		assert.Equal(t, "source-cluster", checkpoints[0].ClusterId)
+		assert.Equal(t, "source-cluster-rootcoord-dml_0", checkpoints[0].Pchannel)
+	})
+
+	t.Run("save_error", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		kv.EXPECT().Save(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("etcd error"))
+		err := catalog.SaveSalvageCheckpoint(ctx, "p1", &commonpb.ReplicateCheckpoint{ClusterId: "c1"})
+		assert.Error(t, err)
+	})
+
+	t.Run("get_load_error", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, errors.New("etcd error"))
+		checkpoints, err := catalog.GetSalvageCheckpoint(ctx, "p1")
+		assert.Error(t, err)
+		assert.Nil(t, checkpoints)
+	})
+
+	t.Run("get_unmarshal_error", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(
+			[]string{"key"},
+			[]string{"invalid-proto-bytes"},
+			nil,
+		)
+		checkpoints, err := catalog.GetSalvageCheckpoint(ctx, "p1")
+		assert.Error(t, err)
+		assert.Nil(t, checkpoints)
+	})
+
+	t.Run("get_empty", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(nil, nil, nil)
+		checkpoints, err := catalog.GetSalvageCheckpoint(ctx, "p1")
+		assert.NoError(t, err)
+		assert.Empty(t, checkpoints)
+	})
+
+	t.Run("get_multiple_clusters", func(t *testing.T) {
+		kv := mocks.NewMetaKv(t)
+		catalog := NewCataLog(kv)
+
+		cp1 := &commonpb.ReplicateCheckpoint{ClusterId: "cluster-a"}
+		cp2 := &commonpb.ReplicateCheckpoint{ClusterId: "cluster-b"}
+		bytes1, _ := proto.Marshal(cp1)
+		bytes2, _ := proto.Marshal(cp2)
+
+		kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).Return(
+			[]string{"key1", "key2"},
+			[]string{string(bytes1), string(bytes2)},
+			nil,
+		)
+		checkpoints, err := catalog.GetSalvageCheckpoint(ctx, "p1")
+		assert.NoError(t, err)
+		assert.Len(t, checkpoints, 2)
+	})
+}
+
 func TestBuildPrefixAndKey(t *testing.T) {
 	// Prefix functions
 	assert.Equal(t, "streamingnode-meta/wal/p1/", buildWALPrefix("p1"))
@@ -208,6 +302,9 @@ func TestBuildPrefixAndKey(t *testing.T) {
 	assert.Equal(t, "streamingnode-meta/wal/p1/segment-assign/", buildSegmentAssignmentPrefix("p1"))
 	assert.Equal(t, "streamingnode-meta/wal/p2/segment-assign/", buildSegmentAssignmentPrefix("p2"))
 
+	assert.Equal(t, "streamingnode-meta/wal/p1/salvage-checkpoint/", buildSalvageCheckpointPrefix("p1"))
+	assert.Equal(t, "streamingnode-meta/wal/p2/salvage-checkpoint/", buildSalvageCheckpointPrefix("p2"))
+
 	// Key functions
 	assert.Equal(t, "streamingnode-meta/wal/p1/vchannel/v1", buildVChannelKey("p1", "v1"))
 	assert.Equal(t, "streamingnode-meta/wal/p2/vchannel/v2", buildVChannelKey("p2", "v2"))
@@ -220,4 +317,7 @@ func TestBuildPrefixAndKey(t *testing.T) {
 
 	assert.Equal(t, "streamingnode-meta/wal/p1/consume-checkpoint", buildConsumeCheckpointKey("p1"))
 	assert.Equal(t, "streamingnode-meta/wal/p2/consume-checkpoint", buildConsumeCheckpointKey("p2"))
+
+	assert.Equal(t, "streamingnode-meta/wal/p1/salvage-checkpoint/cluster-a", buildSalvageCheckpointPath("p1", "cluster-a"))
+	assert.Equal(t, "streamingnode-meta/wal/p2/salvage-checkpoint/cluster-b", buildSalvageCheckpointPath("p2", "cluster-b"))
 }

internal/streamingnode/server/wal/interceptors/replicate/replicates/manager_test.go

@@ -195,6 +195,101 @@ func TestSalvageCheckpointNotCapturedWhenAlreadyPrimary(t *testing.T) {
 	assert.Empty(t, rm.GetSalvageCheckpoint())
 }
 
+func TestSalvageCheckpointLoadedFromEtcd(t *testing.T) {
+	// Simulate WAL recovery where salvage checkpoints were previously persisted to etcd
+	// and loaded back into the ReplicateManagerRecoverParam.
+	preLoaded := []*utility.ReplicateCheckpoint{
+		{ClusterID: "cluster-a", PChannel: "cluster-a-rootcoord-dml_0", TimeTick: 100},
+		{ClusterID: "cluster-b", PChannel: "cluster-b-rootcoord-dml_0", TimeTick: 200},
+	}
+
+	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
+		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
+		CurrentClusterID: "test1",
+		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
+			Checkpoint: &utility.WALCheckpoint{
+				MessageID:       walimplstest.NewTestMessageID(1),
+				TimeTick:        1,
+				ReplicateConfig: newReplicateConfiguration("test1", "test2"),
+			},
+		},
+		SalvageCheckpoints: preLoaded,
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, replicateutil.RolePrimary, rm.Role())
+
+	salvageCPs := rm.GetSalvageCheckpoint()
+	assert.Len(t, salvageCPs, 2)
+
+	byCluster := make(map[string]*utility.ReplicateCheckpoint, 2)
+	for _, cp := range salvageCPs {
+		byCluster[cp.ClusterID] = cp
+	}
+	assert.Equal(t, uint64(100), byCluster["cluster-a"].TimeTick)
+	assert.Equal(t, "cluster-a-rootcoord-dml_0", byCluster["cluster-a"].PChannel)
+	assert.Equal(t, uint64(200), byCluster["cluster-b"].TimeTick)
+	assert.Equal(t, "cluster-b-rootcoord-dml_0", byCluster["cluster-b"].PChannel)
+}
+
+func TestSalvageCheckpointMultipleForcePromotes(t *testing.T) {
+	// Start as secondary of cluster-a, force promote to primary,
+	// then become secondary of cluster-b, force promote again.
+	// Both salvage checkpoints should accumulate (keyed by source cluster).
+	txnBuffer := utility.NewTxnBuffer(log.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
+	rm, err := RecoverReplicateManager(&ReplicateManagerRecoverParam{
+		ChannelInfo:      types.PChannelInfo{Name: "test1-rootcoord-dml_0", Term: 1},
+		CurrentClusterID: "test1",
+		InitialRecoverSnapshot: &recovery.RecoverySnapshot{
+			Checkpoint: &utility.WALCheckpoint{
+				MessageID: walimplstest.NewTestMessageID(1),
+				TimeTick:  1,
+				ReplicateCheckpoint: &utility.ReplicateCheckpoint{
+					ClusterID: "cluster-a",
+					PChannel:  "cluster-a-rootcoord-dml_0",
+					MessageID: walimplstest.NewTestMessageID(10),
+					TimeTick:  100,
+				},
+				ReplicateConfig: newReplicateConfiguration("cluster-a", "test1"),
+			},
+			TxnBuffer: txnBuffer,
+		},
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())
+	assert.Empty(t, rm.GetSalvageCheckpoint())
+
+	// First force promote: secondary(cluster-a) -> primary
+	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "cluster-a"))
+	assert.NoError(t, err)
+	assert.Equal(t, replicateutil.RolePrimary, rm.Role())
+	cps := rm.GetSalvageCheckpoint()
+	assert.Len(t, cps, 1)
+	assert.Equal(t, "cluster-a", cps[0].ClusterID)
+
+	// Now become secondary of cluster-b
+	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("cluster-b", "test1"))
+	assert.NoError(t, err)
+	assert.Equal(t, replicateutil.RoleSecondary, rm.Role())
+
+	// Still one salvage checkpoint from first promote
+	assert.Len(t, rm.GetSalvageCheckpoint(), 1)
+
+	// Second force promote: secondary(cluster-b) -> primary
+	err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "cluster-b"))
+	assert.NoError(t, err)
+	assert.Equal(t, replicateutil.RolePrimary, rm.Role())
+
+	// Both salvage checkpoints should now be present
+	cps = rm.GetSalvageCheckpoint()
+	assert.Len(t, cps, 2)
+	byCluster := make(map[string]*utility.ReplicateCheckpoint, 2)
+	for _, cp := range cps {
+		byCluster[cp.ClusterID] = cp
+	}
+	assert.Contains(t, byCluster, "cluster-a")
+	assert.Contains(t, byCluster, "cluster-b")
+}
+
 func TestSecondaryReplicateManagerWithTxn(t *testing.T) {
 	txnBuffer := utility.NewTxnBuffer(log.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics())
 	txnMsgs := newReplicateTxnMessage("test1", "test2", 2)