feat: implement data salvage for force failover (#47599)

issue: https://github.com/milvus-io/milvus/issues/47598

design doc:
https://github.com/milvus-io/milvus-design-docs/blob/main/design_docs/20260205-data_salvage_for_force_failover.md

## Summary

Implement data salvage feature that enables zero data loss after force
failover by capturing and exposing the salvage checkpoint, along with a
DumpMessages API to retrieve unsynchronized messages.

### SalvageCheckpoint Capture & Persistence
- Add `salvageCheckpoint` field to `ReplicateManager`
- Capture last synced position on force promote before clearing
secondary state
- Persist to separate etcd key via `StreamingNodeCataLog`
- Load salvage checkpoint on WAL recovery

### API Exposure
- Return `salvage_checkpoint` in `GetReplicateInfoResponse`
- Add `GetSalvageCheckpoint` to `HandlerClient` and `ReplicateService`

### DumpMessages Streaming API
- Implement `DumpMessages` RPC for streaming WAL messages
- Support `start_timetick` and `end_timetick` filters
- Filter out system messages (TimeTick, CreateSegment, Flush,
RollbackTxn)
- Use oneof response for status or message

### Tests
- Unit tests for salvage checkpoint capture
- Unit tests for `shouldDumpMessage` filter
- Integration tests for `GetReplicateInfo` and `DumpMessages`

## Test plan
- [ ] Unit tests pass
- [ ] Integration tests pass
- [ ] Manual test with force failover scenario

proto pr: https://github.com/milvus-io/milvus-proto/pull/568

🤖 Generated with [Claude Code](https://claude.ai/code)

---------

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
Signed-off-by: Yihao Dai <yihao.dai@zilliz.com>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
yihao.dai
2026-03-27 13:47:30 +08:00
committed by GitHub
co-authored by Claude Opus 4.5
parent 695f27bacb
commit 4ba08692ce
51 changed files with 2442 additions and 298 deletions
+1 -1
View File
@@ -19,7 +19,7 @@ require (
github.com/google/btree v1.1.2
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/klauspost/compress v1.18.0
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b
github.com/minio/minio-go/v7 v7.0.73
github.com/panjf2000/ants/v2 v2.11.3 // indirect
github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 // indirect
+2 -2
View File
@@ -814,8 +814,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6 h1:YHMFI6L
github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2 h1:OEqYnraL/Tf2nGMjRJji6xaaDe40zSQVBvruzEBZ1hU=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b h1:tUp6a72FFeosfjN+8jy/Hax7BqVbtgsaMMVrAfzKnEc=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
+5
View File
@@ -1192,6 +1192,11 @@ func (s *Server) CreateReplicateStream(stream milvuspb.MilvusService_CreateRepli
return s.proxy.CreateReplicateStream(stream)
}
// DumpMessages streams messages from a WAL range for data salvage.
func (s *Server) DumpMessages(req *milvuspb.DumpMessagesRequest, stream milvuspb.MilvusService_DumpMessagesServer) error {
return s.proxy.DumpMessages(req, stream)
}
// ComputePhraseMatchSlop computes the minimum slop required for phrase matching.
func (s *Server) ComputePhraseMatchSlop(ctx context.Context, req *milvuspb.ComputePhraseMatchSlopRequest) (*milvuspb.ComputePhraseMatchSlopResponse, error) {
return s.proxy.ComputePhraseMatchSlop(ctx, req)
@@ -101,6 +101,15 @@ func (s replicateService) shouldSkipReplicateMessageType(msgType message.Message
return ok
}
func (s replicateService) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error) {
if !s.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrWALAccesserClosed
}
defer s.lifetime.Done()
return s.handlerClient.GetSalvageCheckpoint(ctx, channelName)
}
// overwriteReplicateMessage overwrites the replicate message.
// because some message such as create collection message write vchannel in its body, so we need to overwrite the message.
func (s replicateService) overwriteReplicateMessage(ctx context.Context, msg message.ReplicateMutableMessage, rh *message.ReplicateHeader) (message.MutableMessage, error) {
@@ -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)
@@ -931,7 +999,8 @@ func TestReplicateServiceAlterReplicateConfigMessage(t *testing.T) {
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
}
replicateMsgs := createReplicateAlterConfigMessages(newConfig, []string{"primary-rootcoord-dml_0_1v0", "primary-rootcoord-dml_1_1v1"}, false)
replicateMsgs := createSimpleReplicateAlterConfigMessages(newConfig)
for _, msg := range replicateMsgs {
_, err := rs.Append(context.Background(), msg)
assert.NoError(t, err)
@@ -957,7 +1026,8 @@ func TestReplicateServiceAlterReplicateConfigMessage(t *testing.T) {
{ClusterId: "primary", Pchannels: []string{"primary-rootcoord-dml_0", "primary-rootcoord-dml_1"}},
},
}
replicateMsgs := createReplicateAlterConfigMessages(newConfig, []string{"primary-rootcoord-dml_0_1v0", "primary-rootcoord-dml_1_1v1"}, false)
replicateMsgs := createSimpleReplicateAlterConfigMessages(newConfig)
for _, msg := range replicateMsgs {
_, err := rs.Append(context.Background(), msg)
assert.NoError(t, err)
@@ -1055,7 +1125,8 @@ func TestReplicateServiceAlterReplicateConfigMessage(t *testing.T) {
{SourceClusterId: "by-dev", TargetClusterId: "nonexistent"},
},
}
replicateMsgs := createReplicateAlterConfigMessages(invalidConfig, []string{"primary-rootcoord-dml_0_1v0", "primary-rootcoord-dml_1_1v1"}, false)
replicateMsgs := createSimpleReplicateAlterConfigMessages(invalidConfig)
_, err := rs.Append(context.Background(), replicateMsgs[0])
assert.Error(t, err)
})
@@ -1208,6 +1279,25 @@ func TestReplicateServiceGetConfigError(t *testing.T) {
assert.Contains(t, err.Error(), "config unavailable")
}
func createSimpleReplicateAlterConfigMessages(newConfig *commonpb.ReplicateConfiguration) []message.ReplicateMutableMessage {
alterMsg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: newConfig,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"primary-rootcoord-dml_0_1v0", "primary-rootcoord-dml_1_1v1"}).
MustBuildBroadcast()
msgs := alterMsg.WithBroadcastID(200).SplitIntoMutableMessage()
replicateMsgs := make([]message.ReplicateMutableMessage, 0, len(msgs))
for _, msg := range msgs {
immutableMsg := msg.WithLastConfirmedUseMessageID().WithTimeTick(1).IntoImmutableMessage(pulsar2.NewPulsarID(
pulsar.NewMessageID(1, 2, 3, 4),
))
replicateMsgs = append(replicateMsgs, message.MustNewReplicateMessage("primary", immutableMsg.IntoImmutableMessageProto()))
}
return replicateMsgs
}
func createReplicateControlChannelMessages() []message.ReplicateMutableMessage {
// Create a non-broadcast message on the control channel (vchannel ends with "_vcchan")
msg := message.NewCreateDatabaseMessageBuilderV2().
@@ -97,6 +97,10 @@ type ReplicateService interface {
// GetReplicateCheckpoint returns the WAL checkpoint that will be used to create scanner
// from the correct position, ensuring no duplicate or missing messages.
GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error)
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
// Returns an empty slice if no force promote has occurred.
GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error)
}
// Balancer is the interface for managing the balancer of the wal.
@@ -73,6 +73,10 @@ func (n *noopReplicateService) GetReplicateCheckpoint(ctx context.Context, chann
return nil, nil
}
func (n *noopReplicateService) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error) {
return nil, nil
}
type noopBalancer struct{}
func (n *noopBalancer) ListStreamingNode(ctx context.Context) ([]types.StreamingNodeInfo, error) {
+9
View File
@@ -5,6 +5,7 @@ import (
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus/internal/metastore/model"
@@ -357,4 +358,12 @@ type StreamingNodeCataLog interface {
// SaveConsumeCheckpoint saves the consuming checkpoint of the wal.
SaveConsumeCheckpoint(ctx context.Context, pChannelName string, checkpoint *streamingpb.WALCheckpoint) error
// SaveSalvageCheckpoint saves the salvage checkpoint.
// The checkpoint is captured during force promote.
SaveSalvageCheckpoint(ctx context.Context, pChannelName string, checkpoint *commonpb.ReplicateCheckpoint) error
// GetSalvageCheckpoint gets all salvage checkpoints for a channel.
// Returns an empty slice if none exist. One checkpoint per source cluster.
GetSalvageCheckpoint(ctx context.Context, pChannelName string) ([]*commonpb.ReplicateCheckpoint, error)
}
@@ -9,4 +9,5 @@ const (
DirectorySchema = "schema"
KeyConsumeCheckpoint = "consume-checkpoint"
KeySalvageCheckpoint = "salvage-checkpoint"
)
@@ -10,6 +10,7 @@ import (
"github.com/cockroachdb/errors"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/pkg/v2/kv"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
@@ -271,6 +272,35 @@ func (c *catalog) SaveConsumeCheckpoint(ctx context.Context, pchannelName string
return c.metaKV.Save(ctx, key, string(value))
}
// SaveSalvageCheckpoint saves the salvage checkpoint, keyed by the source cluster ID.
// Multiple salvage checkpoints can exist for a channel (one per source cluster).
func (c *catalog) SaveSalvageCheckpoint(ctx context.Context, pchannelName string, checkpoint *commonpb.ReplicateCheckpoint) error {
key := buildSalvageCheckpointPath(pchannelName, checkpoint.GetClusterId())
value, err := proto.Marshal(checkpoint)
if err != nil {
return err
}
return c.metaKV.Save(ctx, key, string(value))
}
// GetSalvageCheckpoint gets all salvage checkpoints for a channel (one per source cluster).
func (c *catalog) GetSalvageCheckpoint(ctx context.Context, pchannelName string) ([]*commonpb.ReplicateCheckpoint, error) {
prefix := buildSalvageCheckpointPrefix(pchannelName)
_, values, err := c.metaKV.LoadWithPrefix(ctx, prefix)
if err != nil {
return nil, err
}
checkpoints := make([]*commonpb.ReplicateCheckpoint, 0, len(values))
for _, value := range values {
val := &commonpb.ReplicateCheckpoint{}
if err = proto.Unmarshal([]byte(value), val); err != nil {
return nil, err
}
checkpoints = append(checkpoints, val)
}
return checkpoints, nil
}
// Prefix functions: return paths ending with "/" for LoadWithPrefix queries.
// buildWALPrefix returns the prefix for all WAL metadata under a pchannel.
@@ -317,3 +347,13 @@ func removePrefix(prefix string, keys []string) []string {
}
return keys
}
// buildSalvageCheckpointPrefix builds the prefix for all salvage checkpoints under a pchannel.
func buildSalvageCheckpointPrefix(pchannelName string) string {
return buildWALPrefix(pchannelName) + KeySalvageCheckpoint + "/"
}
// buildSalvageCheckpointPath builds the path for salvage checkpoint for a specific source cluster.
func buildSalvageCheckpointPath(pchannelName, sourceClusterID string) string {
return buildSalvageCheckpointPrefix(pchannelName) + sourceClusterID
}
@@ -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/vchannel/", buildVChannelPrefix("p1"))
assert.Equal(t, "streamingnode-meta/wal/p2/vchannel/", buildVChannelPrefix("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/segment-assign/1", buildSegmentAssignmentKey("p1", 1))
assert.Equal(t, "streamingnode-meta/wal/p2/segment-assign/2", buildSegmentAssignmentKey("p2", 2))
@@ -219,4 +316,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"))
}
@@ -7,15 +7,15 @@ import (
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
milvuspb "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
milvuspb "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
mock "github.com/stretchr/testify/mock"
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
wal "github.com/milvus-io/milvus/internal/streamingnode/server/wal"
utility "github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
)
// MockReplicateService is an autogenerated mock type for the ReplicateService type
@@ -91,23 +91,23 @@ func (_c *MockReplicateService_Append_Call) RunAndReturn(run func(context.Contex
}
// GetReplicateCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockReplicateService) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
func (_m *MockReplicateService) GetReplicateCheckpoint(ctx context.Context, channelName string) (*utility.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetReplicateCheckpoint")
}
var r0 *wal.ReplicateCheckpoint
var r0 *utility.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) (*wal.ReplicateCheckpoint, error)); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) (*utility.ReplicateCheckpoint, error)); ok {
return rf(ctx, channelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) *wal.ReplicateCheckpoint); ok {
if rf, ok := ret.Get(0).(func(context.Context, string) *utility.ReplicateCheckpoint); ok {
r0 = rf(ctx, channelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*wal.ReplicateCheckpoint)
r0 = ret.Get(0).(*utility.ReplicateCheckpoint)
}
}
@@ -139,12 +139,12 @@ func (_c *MockReplicateService_GetReplicateCheckpoint_Call) Run(run func(ctx con
return _c
}
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) Return(_a0 *wal.ReplicateCheckpoint, _a1 error) *MockReplicateService_GetReplicateCheckpoint_Call {
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) Return(_a0 *utility.ReplicateCheckpoint, _a1 error) *MockReplicateService_GetReplicateCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) RunAndReturn(run func(context.Context, string) (*wal.ReplicateCheckpoint, error)) *MockReplicateService_GetReplicateCheckpoint_Call {
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) RunAndReturn(run func(context.Context, string) (*utility.ReplicateCheckpoint, error)) *MockReplicateService_GetReplicateCheckpoint_Call {
_c.Call.Return(run)
return _c
}
@@ -207,6 +207,65 @@ func (_c *MockReplicateService_GetReplicateConfiguration_Call) RunAndReturn(run
return _c
}
// GetSalvageCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockReplicateService) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*utility.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*utility.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) ([]*utility.ReplicateCheckpoint, error)); ok {
return rf(ctx, channelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) []*utility.ReplicateCheckpoint); ok {
r0 = rf(ctx, channelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*utility.ReplicateCheckpoint)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
r1 = rf(ctx, channelName)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockReplicateService_GetSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSalvageCheckpoint'
type MockReplicateService_GetSalvageCheckpoint_Call struct {
*mock.Call
}
// GetSalvageCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - channelName string
func (_e *MockReplicateService_Expecter) GetSalvageCheckpoint(ctx interface{}, channelName interface{}) *MockReplicateService_GetSalvageCheckpoint_Call {
return &MockReplicateService_GetSalvageCheckpoint_Call{Call: _e.mock.On("GetSalvageCheckpoint", ctx, channelName)}
}
func (_c *MockReplicateService_GetSalvageCheckpoint_Call) Run(run func(ctx context.Context, channelName string)) *MockReplicateService_GetSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *MockReplicateService_GetSalvageCheckpoint_Call) Return(_a0 []*utility.ReplicateCheckpoint, _a1 error) *MockReplicateService_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockReplicateService_GetSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string) ([]*utility.ReplicateCheckpoint, error)) *MockReplicateService_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// UpdateReplicateConfiguration provides a mock function with given fields: ctx, req
func (_m *MockReplicateService) UpdateReplicateConfiguration(ctx context.Context, req *milvuspb.UpdateReplicateConfigurationRequest) error {
ret := _m.Called(ctx, req)
+1 -2
View File
@@ -3869,8 +3869,7 @@ func (_c *MockDataCoord_WatchChannels_Call) RunAndReturn(run func(context.Contex
func NewMockDataCoord(t interface {
mock.TestingT
Cleanup(func())
},
) *MockDataCoord {
}) *MockDataCoord {
mock := &MockDataCoord{}
mock.Mock.Test(t)
+1 -2
View File
@@ -4451,8 +4451,7 @@ func (_c *MockDataCoordClient_WatchChannels_Call) RunAndReturn(run func(context.
func NewMockDataCoordClient(t interface {
mock.TestingT
Cleanup(func())
},
) *MockDataCoordClient {
}) *MockDataCoordClient {
mock := &MockDataCoordClient{}
mock.Mock.Test(t)
@@ -5,6 +5,8 @@ package mock_metastore
import (
context "context"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
mock "github.com/stretchr/testify/mock"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
@@ -82,6 +84,65 @@ func (_c *MockStreamingNodeCataLog_GetConsumeCheckpoint_Call) RunAndReturn(run f
return _c
}
// GetSalvageCheckpoint provides a mock function with given fields: ctx, pChannelName
func (_m *MockStreamingNodeCataLog) GetSalvageCheckpoint(ctx context.Context, pChannelName string) ([]*commonpb.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, pChannelName)
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*commonpb.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) ([]*commonpb.ReplicateCheckpoint, error)); ok {
return rf(ctx, pChannelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) []*commonpb.ReplicateCheckpoint); ok {
r0 = rf(ctx, pChannelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*commonpb.ReplicateCheckpoint)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
r1 = rf(ctx, pChannelName)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockStreamingNodeCataLog_GetSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSalvageCheckpoint'
type MockStreamingNodeCataLog_GetSalvageCheckpoint_Call struct {
*mock.Call
}
// GetSalvageCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - pChannelName string
func (_e *MockStreamingNodeCataLog_Expecter) GetSalvageCheckpoint(ctx interface{}, pChannelName interface{}) *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call {
return &MockStreamingNodeCataLog_GetSalvageCheckpoint_Call{Call: _e.mock.On("GetSalvageCheckpoint", ctx, pChannelName)}
}
func (_c *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call) Run(run func(ctx context.Context, pChannelName string)) *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call) Return(_a0 []*commonpb.ReplicateCheckpoint, _a1 error) *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string) ([]*commonpb.ReplicateCheckpoint, error)) *MockStreamingNodeCataLog_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// ListSegmentAssignment provides a mock function with given fields: ctx, pChannelName
func (_m *MockStreamingNodeCataLog) ListSegmentAssignment(ctx context.Context, pChannelName string) ([]*streamingpb.SegmentAssignmentMeta, error) {
ret := _m.Called(ctx, pChannelName)
@@ -248,6 +309,54 @@ func (_c *MockStreamingNodeCataLog_SaveConsumeCheckpoint_Call) RunAndReturn(run
return _c
}
// SaveSalvageCheckpoint provides a mock function with given fields: ctx, pChannelName, checkpoint
func (_m *MockStreamingNodeCataLog) SaveSalvageCheckpoint(ctx context.Context, pChannelName string, checkpoint *commonpb.ReplicateCheckpoint) error {
ret := _m.Called(ctx, pChannelName, checkpoint)
if len(ret) == 0 {
panic("no return value specified for SaveSalvageCheckpoint")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, *commonpb.ReplicateCheckpoint) error); ok {
r0 = rf(ctx, pChannelName, checkpoint)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveSalvageCheckpoint'
type MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call struct {
*mock.Call
}
// SaveSalvageCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - pChannelName string
// - checkpoint *commonpb.ReplicateCheckpoint
func (_e *MockStreamingNodeCataLog_Expecter) SaveSalvageCheckpoint(ctx interface{}, pChannelName interface{}, checkpoint interface{}) *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call {
return &MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call{Call: _e.mock.On("SaveSalvageCheckpoint", ctx, pChannelName, checkpoint)}
}
func (_c *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call) Run(run func(ctx context.Context, pChannelName string, checkpoint *commonpb.ReplicateCheckpoint)) *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(*commonpb.ReplicateCheckpoint))
})
return _c
}
func (_c *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call) Return(_a0 error) *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string, *commonpb.ReplicateCheckpoint) error) *MockStreamingNodeCataLog_SaveSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// SaveSegmentAssignments provides a mock function with given fields: ctx, pChannelName, infos
func (_m *MockStreamingNodeCataLog) SaveSegmentAssignments(ctx context.Context, pChannelName string, infos map[int64]*streamingpb.SegmentAssignmentMeta) error {
ret := _m.Called(ctx, pChannelName, infos)
+1 -2
View File
@@ -10518,8 +10518,7 @@ func (_c *MixCoord_WatchChannels_Call) RunAndReturn(run func(context.Context, *d
func NewMixCoord(t interface {
mock.TestingT
Cleanup(func())
},
) *MixCoord {
}) *MixCoord {
mock := &MixCoord{}
mock.Mock.Test(t)
+1 -2
View File
@@ -12227,8 +12227,7 @@ func (_c *MockMixCoordClient_WatchChannels_Call) RunAndReturn(run func(context.C
func NewMockMixCoordClient(t interface {
mock.TestingT
Cleanup(func())
},
) *MockMixCoordClient {
}) *MockMixCoordClient {
mock := &MockMixCoordClient{}
mock.Mock.Test(t)
+48 -2
View File
@@ -3265,6 +3265,53 @@ func (_c *MockProxy_Dummy_Call) RunAndReturn(run func(context.Context, *milvuspb
return _c
}
// DumpMessages provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) DumpMessages(_a0 *milvuspb.DumpMessagesRequest, _a1 milvuspb.MilvusService_DumpMessagesServer) error {
ret := _m.Called(_a0, _a1)
if len(ret) == 0 {
panic("no return value specified for DumpMessages")
}
var r0 error
if rf, ok := ret.Get(0).(func(*milvuspb.DumpMessagesRequest, milvuspb.MilvusService_DumpMessagesServer) error); ok {
r0 = rf(_a0, _a1)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockProxy_DumpMessages_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DumpMessages'
type MockProxy_DumpMessages_Call struct {
*mock.Call
}
// DumpMessages is a helper method to define mock.On call
// - _a0 *milvuspb.DumpMessagesRequest
// - _a1 milvuspb.MilvusService_DumpMessagesServer
func (_e *MockProxy_Expecter) DumpMessages(_a0 interface{}, _a1 interface{}) *MockProxy_DumpMessages_Call {
return &MockProxy_DumpMessages_Call{Call: _e.mock.On("DumpMessages", _a0, _a1)}
}
func (_c *MockProxy_DumpMessages_Call) Run(run func(_a0 *milvuspb.DumpMessagesRequest, _a1 milvuspb.MilvusService_DumpMessagesServer)) *MockProxy_DumpMessages_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*milvuspb.DumpMessagesRequest), args[1].(milvuspb.MilvusService_DumpMessagesServer))
})
return _c
}
func (_c *MockProxy_DumpMessages_Call) Return(_a0 error) *MockProxy_DumpMessages_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockProxy_DumpMessages_Call) RunAndReturn(run func(*milvuspb.DumpMessagesRequest, milvuspb.MilvusService_DumpMessagesServer) error) *MockProxy_DumpMessages_Call {
_c.Call.Return(run)
return _c
}
// Flush provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) Flush(_a0 context.Context, _a1 *milvuspb.FlushRequest) (*milvuspb.FlushResponse, error) {
ret := _m.Called(_a0, _a1)
@@ -9112,8 +9159,7 @@ func (_c *MockProxy_Upsert_Call) RunAndReturn(run func(context.Context, *milvusp
func NewMockProxy(t interface {
mock.TestingT
Cleanup(func())
},
) *MockProxy {
}) *MockProxy {
mock := &MockProxy{}
mock.Mock.Test(t)
@@ -350,6 +350,65 @@ func (_c *MockHandlerClient_GetWALMetricsIfLocal_Call) RunAndReturn(run func(con
return _c
}
// GetSalvageCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockHandlerClient) GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*wal.ReplicateCheckpoint
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) ([]*wal.ReplicateCheckpoint, error)); ok {
return rf(ctx, channelName)
}
if rf, ok := ret.Get(0).(func(context.Context, string) []*wal.ReplicateCheckpoint); ok {
r0 = rf(ctx, channelName)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*wal.ReplicateCheckpoint)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
r1 = rf(ctx, channelName)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockHandlerClient_GetSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSalvageCheckpoint'
type MockHandlerClient_GetSalvageCheckpoint_Call struct {
*mock.Call
}
// GetSalvageCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - channelName string
func (_e *MockHandlerClient_Expecter) GetSalvageCheckpoint(ctx interface{}, channelName interface{}) *MockHandlerClient_GetSalvageCheckpoint_Call {
return &MockHandlerClient_GetSalvageCheckpoint_Call{Call: _e.mock.On("GetSalvageCheckpoint", ctx, channelName)}
}
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) Run(run func(ctx context.Context, channelName string)) *MockHandlerClient_GetSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) Return(_a0 []*wal.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_GetSalvageCheckpoint_Call) RunAndReturn(run func(context.Context, string) ([]*wal.ReplicateCheckpoint, error)) *MockHandlerClient_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// NewMockHandlerClient creates a new instance of MockHandlerClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewMockHandlerClient(t interface {
@@ -360,6 +360,53 @@ func (_c *MockWAL_GetReplicateCheckpoint_Call) RunAndReturn(run func() (*wal.Rep
return _c
}
// GetSalvageCheckpoint provides a mock function with no fields
func (_m *MockWAL) GetSalvageCheckpoint() []*wal.ReplicateCheckpoint {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*wal.ReplicateCheckpoint
if rf, ok := ret.Get(0).(func() []*wal.ReplicateCheckpoint); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*wal.ReplicateCheckpoint)
}
}
return r0
}
// MockWAL_GetSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSalvageCheckpoint'
type MockWAL_GetSalvageCheckpoint_Call struct {
*mock.Call
}
// GetSalvageCheckpoint is a helper method to define mock.On call
func (_e *MockWAL_Expecter) GetSalvageCheckpoint() *MockWAL_GetSalvageCheckpoint_Call {
return &MockWAL_GetSalvageCheckpoint_Call{Call: _e.mock.On("GetSalvageCheckpoint")}
}
func (_c *MockWAL_GetSalvageCheckpoint_Call) Run(run func()) *MockWAL_GetSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockWAL_GetSalvageCheckpoint_Call) Return(_a0 []*wal.ReplicateCheckpoint) *MockWAL_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWAL_GetSalvageCheckpoint_Call) RunAndReturn(run func() []*wal.ReplicateCheckpoint) *MockWAL_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// IsAvailable provides a mock function with no fields
func (_m *MockWAL) IsAvailable() bool {
ret := _m.Called()
@@ -6,6 +6,7 @@ import (
context "context"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
replicates "github.com/milvus-io/milvus/internal/streamingnode/server/wal/interceptors/replicate/replicates"
@@ -15,7 +16,7 @@ import (
utility "github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
)
// MockReplicatesManager is an autogenerated mock type for the ReplicatesManager type
// ReplicateManager is an autogenerated mock type for the ReplicateManager type
type MockReplicatesManager struct {
mock.Mock
}
@@ -144,6 +145,53 @@ func (_c *MockReplicatesManager_GetReplicateCheckpoint_Call) RunAndReturn(run fu
return _c
}
// GetSalvageCheckpoint provides a mock function with no fields
func (_m *MockReplicatesManager) GetSalvageCheckpoint() []*utility.ReplicateCheckpoint {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetSalvageCheckpoint")
}
var r0 []*utility.ReplicateCheckpoint
if rf, ok := ret.Get(0).(func() []*utility.ReplicateCheckpoint); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*utility.ReplicateCheckpoint)
}
}
return r0
}
// MockReplicatesManager_GetSalvageCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetSalvageCheckpoint'
type MockReplicatesManager_GetSalvageCheckpoint_Call struct {
*mock.Call
}
// GetSalvageCheckpoint is a helper method to define mock.On call
func (_e *MockReplicatesManager_Expecter) GetSalvageCheckpoint() *MockReplicatesManager_GetSalvageCheckpoint_Call {
return &MockReplicatesManager_GetSalvageCheckpoint_Call{Call: _e.mock.On("GetSalvageCheckpoint")}
}
func (_c *MockReplicatesManager_GetSalvageCheckpoint_Call) Run(run func()) *MockReplicatesManager_GetSalvageCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockReplicatesManager_GetSalvageCheckpoint_Call) Return(_a0 []*utility.ReplicateCheckpoint) *MockReplicatesManager_GetSalvageCheckpoint_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockReplicatesManager_GetSalvageCheckpoint_Call) RunAndReturn(run func() []*utility.ReplicateCheckpoint) *MockReplicatesManager_GetSalvageCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// Role provides a mock function with no fields
func (_m *MockReplicatesManager) Role() replicateutil.Role {
ret := _m.Called()
@@ -214,7 +262,7 @@ type MockReplicatesManager_SwitchReplicateMode_Call struct {
// SwitchReplicateMode is a helper method to define mock.On call
// - ctx context.Context
// - msg message.MutableAlterReplicateConfigMessageV2
// - msg message.specializedMutableMessage[*messagespb.AlterReplicateConfigMessageHeader,*messagespb.AlterReplicateConfigMessageBody]
func (_e *MockReplicatesManager_Expecter) SwitchReplicateMode(ctx interface{}, msg interface{}) *MockReplicatesManager_SwitchReplicateMode_Call {
return &MockReplicatesManager_SwitchReplicateMode_Call{Call: _e.mock.On("SwitchReplicateMode", ctx, msg)}
}
@@ -241,8 +289,7 @@ func (_c *MockReplicatesManager_SwitchReplicateMode_Call) RunAndReturn(run func(
func NewMockReplicatesManager(t interface {
mock.TestingT
Cleanup(func())
},
) *MockReplicatesManager {
}) *MockReplicatesManager {
mock := &MockReplicatesManager{}
mock.Mock.Test(t)
+438
View File
@@ -0,0 +1,438 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
import (
"context"
"sync"
"testing"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message/adaptor"
pulsar2 "github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/pulsar"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
)
func TestShouldDumpMessage(t *testing.T) {
// Messages that SHOULD be dumped (replicable data)
assert.True(t, shouldDumpMessage(message.MessageTypeInsert), "Insert should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeDelete), "Delete should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeCreateCollection), "CreateCollection should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeDropCollection), "DropCollection should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeCreatePartition), "CreatePartition should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeDropPartition), "DropPartition should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeBeginTxn), "BeginTxn should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeCommitTxn), "CommitTxn should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeTxn), "Txn should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeImport), "Import should be dumped")
assert.True(t, shouldDumpMessage(message.MessageTypeManualFlush), "ManualFlush should be dumped")
// Messages that should NOT be dumped (system messages)
assert.False(t, shouldDumpMessage(message.MessageTypeTimeTick), "TimeTick should NOT be dumped")
assert.False(t, shouldDumpMessage(message.MessageTypeCreateSegment), "CreateSegment should NOT be dumped")
assert.False(t, shouldDumpMessage(message.MessageTypeFlush), "Flush should NOT be dumped")
assert.False(t, shouldDumpMessage(message.MessageTypeRollbackTxn), "RollbackTxn should NOT be dumped")
}
// mockDumpMessagesServer is a minimal implementation of milvuspb.MilvusService_DumpMessagesServer for testing.
type mockDumpMessagesServer struct {
ctx context.Context
mu sync.Mutex
sent []*milvuspb.DumpMessagesResponse
sendErr error
}
func (m *mockDumpMessagesServer) Send(resp *milvuspb.DumpMessagesResponse) error {
if m.sendErr != nil {
return m.sendErr
}
m.mu.Lock()
defer m.mu.Unlock()
m.sent = append(m.sent, resp)
return nil
}
func (m *mockDumpMessagesServer) getSent() []*milvuspb.DumpMessagesResponse {
m.mu.Lock()
defer m.mu.Unlock()
return m.sent
}
func (m *mockDumpMessagesServer) Context() context.Context { return m.ctx }
func (m *mockDumpMessagesServer) SetHeader(metadata.MD) error { return nil }
func (m *mockDumpMessagesServer) SendHeader(metadata.MD) error { return nil }
func (m *mockDumpMessagesServer) SetTrailer(metadata.MD) {}
func (m *mockDumpMessagesServer) SendMsg(interface{}) error { return nil }
func (m *mockDumpMessagesServer) RecvMsg(interface{}) error { return nil }
// Ensure mockDumpMessagesServer satisfies the grpc.ServerStream interface.
var _ grpc.ServerStream = (*mockDumpMessagesServer)(nil)
// testStartMessageID returns a valid *commonpb.MessageID for testing.
func testStartMessageID() *commonpb.MessageID {
return pulsar2.NewPulsarID(pulsar.EarliestMessageID()).IntoProto()
}
// buildTestImmutableMessage creates an ImmutableMessage with the given timetick for testing.
func buildTestImmutableMessage(timetick uint64) message.ImmutableMessage {
msgID := pulsar2.NewPulsarID(pulsar.EarliestMessageID())
return message.NewCreateDatabaseMessageBuilderV2().
WithHeader(&message.CreateDatabaseMessageHeader{}).
WithBody(&message.CreateDatabaseMessageBody{}).
WithVChannel("test-channel").
MustBuildMutable().
WithTimeTick(timetick).
WithLastConfirmed(msgID).
IntoImmutableMessage(msgID)
}
// buildRollbackTxnMessage creates a RollbackTxn ImmutableMessage (which shouldDumpMessage returns false for).
func buildRollbackTxnMessage(timetick uint64) message.ImmutableMessage {
msgID := pulsar2.NewPulsarID(pulsar.EarliestMessageID())
return message.NewRollbackTxnMessageBuilderV2().
WithHeader(&message.RollbackTxnMessageHeader{}).
WithBody(&message.RollbackTxnMessageBody{}).
WithVChannel("test-channel").
MustBuildMutable().
WithTimeTick(timetick).
WithLastConfirmed(msgID).
IntoImmutableMessage(msgID)
}
func TestDumpMessages_NodeUnhealthy(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Abnormal)
stream := &mockDumpMessagesServer{ctx: context.Background()}
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{}, stream)
assert.Error(t, err)
assert.True(t, errors.Is(err, merr.ErrServiceNotReady))
}
func TestDumpMessages_MissingPChannel(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
stream := &mockDumpMessagesServer{ctx: context.Background()}
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
StartMessageId: testStartMessageID(),
}, stream)
assert.Error(t, err)
}
func TestDumpMessages_MissingStartMessageId(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
stream := &mockDumpMessagesServer{ctx: context.Background()}
// nil StartMessageId
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
}, stream)
assert.Error(t, err)
}
func TestDumpMessages_EmptyStartMessageIdBytes(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
stream := &mockDumpMessagesServer{ctx: context.Background()}
// StartMessageId present but with empty Id string
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: &commonpb.MessageID{Id: ""},
}, stream)
assert.Error(t, err)
}
func TestDumpMessages_ContextCanceled(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{})) // never fires
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).Return(mockScanner)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
// Cancel before DumpMessages is called so ctx.Done() fires immediately in the select
cancel()
stream := &mockDumpMessagesServer{ctx: ctx}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
}, stream)
assert.ErrorIs(t, err, context.Canceled)
}
func TestDumpMessages_ScannerError(t *testing.T) {
scannerDone := make(chan struct{})
close(scannerDone) // immediately done
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(scannerDone)
mockScanner.EXPECT().Error().Return(errors.New("scanner error"))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).Return(mockScanner)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
}, stream)
assert.EqualError(t, err, "scanner error")
}
func TestDumpMessages_ScannerDoneSuccessfully(t *testing.T) {
scannerDone := make(chan struct{})
close(scannerDone)
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(scannerDone)
mockScanner.EXPECT().Error().Return(nil)
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).Return(mockScanner)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
}, stream)
assert.NoError(t, err)
assert.Empty(t, stream.getSent())
}
func TestDumpMessages_MessageSentSuccessfully(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
// Write a data message within the [0, 100] timetick range, then one past end to stop loop.
ch <- buildTestImmutableMessage(50)
ch <- buildTestImmutableMessage(200) // > endTimetick (100), triggers return nil
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
EndTimetick: 100,
}, stream)
assert.NoError(t, err)
assert.Len(t, stream.getSent(), 1)
}
func TestDumpMessages_FilterByStartTimetick(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
ch <- buildTestImmutableMessage(50) // < startTimetick (100), skipped
ch <- buildTestImmutableMessage(150) // within [100, 200], sent
ch <- buildTestImmutableMessage(250) // > endTimetick (200), stops loop
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
StartTimetick: 100,
EndTimetick: 200,
}, stream)
assert.NoError(t, err)
assert.Len(t, stream.getSent(), 1)
}
func TestDumpMessages_StopAtEndTimetick(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
ch <- buildTestImmutableMessage(200) // > endTimetick (100), stops immediately
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
EndTimetick: 100,
}, stream)
assert.NoError(t, err)
assert.Empty(t, stream.getSent())
}
func TestDumpMessages_FilterSystemMessages(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
ch <- buildRollbackTxnMessage(50) // system message, should be filtered
ch <- buildTestImmutableMessage(50) // data message, should be sent
ch <- buildTestImmutableMessage(200) // > endTimetick, stops loop
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
EndTimetick: 100,
}, stream)
assert.NoError(t, err)
// Only the data message should be sent, not the RollbackTxn
assert.Len(t, stream.getSent(), 1)
}
func TestDumpMessages_SendError(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
ch <- buildTestImmutableMessage(50)
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
sendErr := errors.New("send error")
stream := &mockDumpMessagesServer{ctx: context.Background(), sendErr: sendErr}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
}, stream)
assert.EqualError(t, err, "send error")
}
func TestDumpMessages_ChannelClosed(t *testing.T) {
mockScanner := mock_streaming.NewMockScanner(t)
mockScanner.EXPECT().Done().Return(make(chan struct{}))
mockScanner.EXPECT().Close()
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Read(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, opts streaming.ReadOption) streaming.Scanner {
ch, ok := opts.MessageHandler.(adaptor.ChanMessageHandler)
require.True(t, ok)
// Closing the channel triggers the !ok case in the select loop
close(ch)
return mockScanner
})
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
stream := &mockDumpMessagesServer{ctx: context.Background()}
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
err := node.DumpMessages(&milvuspb.DumpMessagesRequest{
Pchannel: "test-channel",
StartMessageId: testStartMessageID(),
}, stream)
assert.NoError(t, err)
}
+133 -1
View File
@@ -60,6 +60,8 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/proxypb"
"github.com/milvus-io/milvus/pkg/v2/proto/querypb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message/adaptor"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/options"
"github.com/milvus-io/milvus/pkg/v2/util"
"github.com/milvus-io/milvus/pkg/v2/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v2/util/crypto"
@@ -7169,8 +7171,24 @@ func (node *Proxy) GetReplicateInfo(ctx context.Context, req *milvuspb.GetReplic
if err != nil {
return nil, err
}
// Get the salvage checkpoint for the specified source cluster.
// Returns nil if source_cluster_id is not provided or no checkpoint exists for that cluster.
var salvageCheckpointProto *commonpb.ReplicateCheckpoint
salvageCheckpoints, err := streaming.WAL().Replicate().GetSalvageCheckpoint(ctx, req.GetTargetPchannel())
if err != nil {
return nil, err
}
for _, cp := range salvageCheckpoints {
if cp.ClusterID == req.GetSourceClusterId() {
salvageCheckpointProto = cp.IntoProto()
break
}
}
return &milvuspb.GetReplicateInfoResponse{
Checkpoint: checkpoint.IntoProto(),
Checkpoint: checkpoint.IntoProto(),
SalvageCheckpoint: salvageCheckpointProto,
}, nil
}
@@ -7200,6 +7218,120 @@ func (node *Proxy) CreateReplicateStream(stream milvuspb.MilvusService_CreateRep
return s.Execute()
}
// shouldDumpMessage returns true if the message should be included in dump output.
// Filters out system messages that are not useful for data salvage.
func shouldDumpMessage(msgType message.MessageType) bool {
// Self-controlled messages (TimeTick, CreateSegment, Flush) are internal system messages
if msgType.IsSelfControlled() {
return false
}
// RollbackTxn is also not useful for data salvage
if msgType == message.MessageTypeRollbackTxn {
return false
}
return true
}
// DumpMessages streams messages from a WAL range for data salvage.
func (node *Proxy) DumpMessages(req *milvuspb.DumpMessagesRequest, stream milvuspb.MilvusService_DumpMessagesServer) error {
ctx := stream.Context()
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-DumpMessages")
defer sp.End()
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return err
}
logger := log.Ctx(ctx).With(
zap.String("pchannel", req.GetPchannel()),
zap.Uint64("startTimetick", req.GetStartTimetick()),
zap.Uint64("endTimetick", req.GetEndTimetick()),
)
logger.Info("DumpMessages received")
// Validate request
if req.GetPchannel() == "" {
return merr.WrapErrParameterMissing("pchannel")
}
if req.GetStartMessageId() == nil || len(req.GetStartMessageId().GetId()) == 0 {
return merr.WrapErrParameterMissing("start_message_id")
}
startMsgID := message.MustUnmarshalMessageID(req.GetStartMessageId())
// Use exclusive start position (dump messages AFTER start_message_id)
// This is appropriate for salvage scenarios where start_message_id is the last synced message
deliverPolicy := options.DeliverPolicyStartAfter(startMsgID)
// Create a channel-based message handler
msgCh := make(adaptor.ChanMessageHandler, 16)
// Open scanner
scanner := streaming.WAL().Read(ctx, streaming.ReadOption{
PChannel: req.GetPchannel(),
DeliverPolicy: deliverPolicy,
MessageHandler: msgCh,
})
defer scanner.Close()
// Get timetick filters
startTimetick := req.GetStartTimetick()
endTimetick := req.GetEndTimetick()
msgCount := 0
// Stream messages
for {
select {
case <-ctx.Done():
logger.Info("DumpMessages context canceled", zap.Int("messageCount", msgCount))
return ctx.Err()
case <-scanner.Done():
// Scanner closed
if err := scanner.Error(); err != nil {
logger.Warn("DumpMessages scanner error", zap.Error(err), zap.Int("messageCount", msgCount))
return err
}
logger.Info("DumpMessages completed", zap.Int("messageCount", msgCount))
return nil
case msg, ok := <-msgCh:
if !ok {
// Channel closed
logger.Info("DumpMessages channel closed", zap.Int("messageCount", msgCount))
return nil
}
msgTimetick := msg.TimeTick()
// Check start timetick filter
if startTimetick > 0 && msgTimetick < startTimetick {
continue
}
// Check end timetick condition
if endTimetick > 0 && msgTimetick > endTimetick {
logger.Info("DumpMessages reached end timetick", zap.Int("messageCount", msgCount))
return nil
}
// Filter system messages
if !shouldDumpMessage(msg.MessageType()) {
continue
}
// Send message to stream (using oneof - only message, no status)
if err := stream.Send(&milvuspb.DumpMessagesResponse{
Response: &milvuspb.DumpMessagesResponse_Message{
Message: msg.IntoImmutableMessageProto(),
},
}); err != nil {
logger.Warn("DumpMessages send failed", zap.Error(err))
return err
}
msgCount++
}
}
}
func (node *Proxy) ComputePhraseMatchSlop(ctx context.Context, req *milvuspb.ComputePhraseMatchSlopRequest) (*milvuspb.ComputePhraseMatchSlopResponse, error) {
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return &milvuspb.ComputePhraseMatchSlopResponse{
+209
View File
@@ -0,0 +1,209 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package proxy
import (
"context"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
)
func TestProxy_GetReplicateInfo_NodeUnhealthy(t *testing.T) {
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Abnormal)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
})
assert.Error(t, err)
assert.True(t, errors.Is(err, merr.ErrServiceNotReady))
assert.Nil(t, resp)
}
func TestProxy_GetReplicateInfo_GetCheckpointError(t *testing.T) {
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(nil, errors.New("checkpoint error"))
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
})
assert.EqualError(t, err, "checkpoint error")
assert.Nil(t, resp)
}
func TestProxy_GetReplicateInfo_GetSalvageCheckpointError(t *testing.T) {
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(&utility.ReplicateCheckpoint{ClusterID: "cluster-a", TimeTick: 100}, nil)
replicateService.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-pchannel").
Return(nil, errors.New("salvage error"))
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
SourceClusterId: "source-cluster",
})
assert.EqualError(t, err, "salvage error")
assert.Nil(t, resp)
}
func TestProxy_GetReplicateInfo_Success_NoSalvageCheckpoints(t *testing.T) {
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(&utility.ReplicateCheckpoint{ClusterID: "cluster-a", PChannel: "test-pchannel", TimeTick: 100}, nil)
replicateService.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-pchannel").
Return(nil, nil)
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
SourceClusterId: "source-cluster",
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "cluster-a", resp.GetCheckpoint().GetClusterId())
assert.Equal(t, uint64(100), resp.GetCheckpoint().GetTimeTick())
assert.Nil(t, resp.GetSalvageCheckpoint())
}
func TestProxy_GetReplicateInfo_Success_MatchingSourceCluster(t *testing.T) {
salvageCPs := []*utility.ReplicateCheckpoint{
{ClusterID: "other-cluster", PChannel: "other-pchannel", TimeTick: 50},
{ClusterID: "source-cluster", PChannel: "source-pchannel", TimeTick: 200},
}
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(&utility.ReplicateCheckpoint{ClusterID: "cluster-a", PChannel: "test-pchannel", TimeTick: 100}, nil)
replicateService.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-pchannel").
Return(salvageCPs, nil)
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
SourceClusterId: "source-cluster",
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "cluster-a", resp.GetCheckpoint().GetClusterId())
assert.NotNil(t, resp.GetSalvageCheckpoint())
assert.Equal(t, "source-cluster", resp.GetSalvageCheckpoint().GetClusterId())
assert.Equal(t, "source-pchannel", resp.GetSalvageCheckpoint().GetPchannel())
assert.Equal(t, uint64(200), resp.GetSalvageCheckpoint().GetTimeTick())
}
func TestProxy_GetReplicateInfo_Success_NoMatchingSourceCluster(t *testing.T) {
salvageCPs := []*utility.ReplicateCheckpoint{
{ClusterID: "other-cluster", PChannel: "other-pchannel", TimeTick: 50},
}
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(&utility.ReplicateCheckpoint{ClusterID: "cluster-a", PChannel: "test-pchannel", TimeTick: 100}, nil)
replicateService.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-pchannel").
Return(salvageCPs, nil)
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
SourceClusterId: "source-cluster", // not in salvage list
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "cluster-a", resp.GetCheckpoint().GetClusterId())
assert.Nil(t, resp.GetSalvageCheckpoint()) // no match
}
func TestProxy_GetReplicateInfo_Success_NoSourceClusterIDFilter(t *testing.T) {
// When SourceClusterId is empty, no salvage checkpoint is returned
salvageCPs := []*utility.ReplicateCheckpoint{
{ClusterID: "some-cluster", PChannel: "some-pchannel", TimeTick: 75},
}
replicateService := mock_streaming.NewMockReplicateService(t)
replicateService.EXPECT().GetReplicateCheckpoint(mock.Anything, "test-pchannel").
Return(&utility.ReplicateCheckpoint{ClusterID: "cluster-a", PChannel: "test-pchannel", TimeTick: 100}, nil)
replicateService.EXPECT().GetSalvageCheckpoint(mock.Anything, "test-pchannel").
Return(salvageCPs, nil)
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
prevWAL := streaming.WAL()
streaming.SetWALForTest(mockWAL)
defer streaming.SetWALForTest(prevWAL)
node := &Proxy{}
node.UpdateStateCode(commonpb.StateCode_Healthy)
resp, err := node.GetReplicateInfo(context.Background(), &milvuspb.GetReplicateInfoRequest{
TargetPchannel: "test-pchannel",
// SourceClusterId intentionally empty
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "cluster-a", resp.GetCheckpoint().GetClusterId())
assert.Nil(t, resp.GetSalvageCheckpoint()) // no source cluster filter → no match
}
@@ -134,6 +134,12 @@ func (s *ackCallbackScheduler) triggerAckCallback() {
pendingTasks := make([]*broadcastTask, 0, len(s.pendingAckedTasks))
for _, task := range s.pendingAckedTasks {
if task.IsForcePromoteMessage() {
// Force promote: fix incomplete broadcasts in background (BlockUntilAllAck → fix).
// The task still goes through normal FastLock → doAckCallback below.
go s.doForcePromoteFixIncompleteBroadcasts(task)
}
g, err := s.rkLocker.FastLock(task.Header().ResourceKeys.Collect()...)
if err != nil {
s.Logger().Warn("lock is occupied, delay the ack callback", zap.Uint64("broadcastID", task.Header().BroadcastID), zap.Error(err))
@@ -36,6 +36,7 @@ func TestStartBroadcastWithSecondaryClusterResourceKey(t *testing.T) {
snmanager.ResetStreamingNodeManager()
mb := newMockBalancerForTest(t)
mb.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
balance.Register(mb)
@@ -59,6 +60,7 @@ func TestStartBroadcastWithSecondaryClusterResourceKey(t *testing.T) {
snmanager.ResetStreamingNodeManager()
mb := newMockBalancerForTest(t)
mb.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
balance.Register(mb)
@@ -79,6 +81,7 @@ func TestStartBroadcastWithSecondaryClusterResourceKey(t *testing.T) {
snmanager.ResetStreamingNodeManager()
mb := newMockBalancerForTest(t)
mb.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(errors.New("balance error")).Maybe()
balance.Register(mb)
@@ -100,6 +103,7 @@ func TestStartBroadcastWithResourceKeys(t *testing.T) {
snmanager.ResetStreamingNodeManager()
mb := newMockBalancerForTest(t)
mb.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
balance.Register(mb)
@@ -123,6 +127,7 @@ func TestStartBroadcastWithResourceKeys(t *testing.T) {
snmanager.ResetStreamingNodeManager()
mb := newMockBalancerForTest(t)
mb.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
balance.Register(mb)
@@ -420,8 +420,8 @@ func TestWithSecondaryClusterResourceKey(t *testing.T) {
mb := mock_balancer.NewMockBalancer(t)
mb.EXPECT().ReplicateRole().Return(replicateutil.RoleSecondary).Maybe()
mb.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
time.Sleep(100 * time.Second)
return nil
}).Maybe()
balance.Register(mb)
@@ -456,8 +456,8 @@ func TestWithSecondaryClusterResourceKey(t *testing.T) {
mb := mock_balancer.NewMockBalancer(t)
mb.EXPECT().ReplicateRole().Return(replicateutil.RolePrimary).Maybe()
mb.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
time.Sleep(100 * time.Second)
return nil
}).Maybe()
balance.Register(mb)
@@ -492,8 +492,8 @@ func TestWithSecondaryClusterResourceKey(t *testing.T) {
mb := mock_balancer.NewMockBalancer(t)
mb.EXPECT().ReplicateRole().Return(replicateutil.RoleSecondary).Maybe()
mb.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
time.Sleep(100 * time.Second)
return nil
}).Maybe()
balance.Register(mb)
@@ -183,19 +183,52 @@ func (s *assignmentServiceImpl) validateReplicateConfiguration(ctx context.Conte
return b, nil
}
// validateForcePromoteConfiguration validates that the force promote configuration is safe.
// Requirements:
// 1. Must contain ONLY the current cluster
// 2. Must have NO topology (no replication relationships)
func validateForcePromoteConfiguration(config *commonpb.ReplicateConfiguration, currentClusterID string) error {
// Use config helper to validate the configuration structure
helper, err := replicateutil.NewConfigHelper(currentClusterID, config)
if err != nil {
return status.NewInvaildArgument("invalid replicate configuration for force promote: %v", err)
}
// Check that configuration contains exactly one cluster (the current cluster)
if len(config.Clusters) != 1 {
return status.NewInvaildArgument(
"force promote requires configuration with exactly one cluster (current cluster only), got %d clusters",
len(config.Clusters))
}
// Check that the single cluster is the current cluster
if config.Clusters[0].ClusterId != currentClusterID {
return status.NewInvaildArgument(
"force promote requires configuration with only current cluster %s, got cluster %s",
currentClusterID,
config.Clusters[0].ClusterId)
}
// Check that there is NO topology (no replication)
if len(config.CrossClusterTopology) > 0 {
return status.NewInvaildArgument(
"force promote requires configuration with no topology (single primary cluster), got %d topology edges",
len(config.CrossClusterTopology))
}
// Verify the cluster role is primary (should be true for single cluster with no topology)
if helper.GetCurrentCluster().Role() != replicateutil.RolePrimary {
return status.NewInvaildArgument("force promote configuration must result in current cluster being primary")
}
return nil
}
// handleForcePromote handles force promote logic for replicate configuration.
// It promotes a secondary cluster to standalone primary immediately without waiting for CDC replication.
func (s *assignmentServiceImpl) handleForcePromote(ctx context.Context, config *commonpb.ReplicateConfiguration) (*streamingpb.UpdateReplicateConfigurationResponse, error) {
log.Ctx(ctx).Warn("Force promote replicate configuration requested")
// VALIDATION 1: Force promote requires empty cluster and topology fields
// The configuration will be constructed from the current cluster's existing meta
// Use len() instead of != nil because empty slices (from client builder) are not nil
if len(config.GetClusters()) > 0 || len(config.GetCrossClusterTopology()) > 0 {
return nil, status.NewInvaildArgument(
"force promote requires empty cluster and topology fields; it promotes the cluster to primary automatically")
}
// Use WithSecondaryClusterResourceKey to:
// 1. Acquire exclusive cluster-level resource key
// 2. Verify the cluster is secondary (not primary)
@@ -332,6 +365,7 @@ func (s *assignmentServiceImpl) alterReplicateConfiguration(ctx context.Context,
// Update the configuration
// For force promote, incomplete broadcasts are already fixed by ackCallbackScheduler
// before this callback is invoked.
return balancer.UpdateReplicateConfiguration(ctx, result)
}
@@ -198,97 +198,54 @@ func TestAssignmentService(t *testing.T) {
}
func TestForcePromoteValidation(t *testing.T) {
// Test force promote validation - force promote requires empty cluster and topology fields
// The configuration is auto-constructed from the current cluster's existing meta
t.Run("invalid_non_empty_clusters", func(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with non-empty clusters should fail
t.Run("valid_single_cluster_no_topology", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote requires empty cluster and topology fields")
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.NoError(t, err)
})
t.Run("invalid_non_empty_topology", func(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with non-empty topology should fail
// Test validateForcePromoteConfiguration standalone function directly
t.Run("valid_single_cluster_no_topology_direct", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.NoError(t, err)
})
t.Run("invalid_multiple_clusters", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.Error(t, err)
// Config helper validates this as "primary count is not 1"
assert.Contains(t, err.Error(), "primary count is not 1")
})
t.Run("invalid_with_topology_direct", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
ForcePromote: true,
})
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote requires empty cluster and topology fields")
})
}
@@ -459,57 +416,6 @@ func TestForcePromoteIdempotent(t *testing.T) {
assert.NotNil(t, resp)
}
func TestForcePromoteInvalidConfig(t *testing.T) {
// Test that force promote rejects non-empty cluster/topology fields
// This test is now covered by TestForcePromoteValidation
// Keep this for additional coverage
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with multiple clusters (invalid for force promote - non-empty clusters)
invalidCfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "other", Pchannels: []string{"other-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://other:19530", Token: "other"}},
},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: invalidCfg,
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote requires empty cluster and topology fields")
}
func TestAlterReplicateConfigCallbackForcePromote(t *testing.T) {
resource.InitForTest()
@@ -81,6 +81,10 @@ type HandlerClient interface {
// GetReplicateCheckpoint returns the WAL checkpoint that will be used to create scanner.
GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error)
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
// Returns an empty slice if no force promote has occurred.
GetSalvageCheckpoint(ctx context.Context, channelName string) ([]*wal.ReplicateCheckpoint, error)
// GetWALMetricsIfLocal gets the metrics of the local wal.
// It will only return the metrics of the local wal but not the remote wal.
GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error)
@@ -105,6 +105,40 @@ func (hc *handlerClientImpl) GetReplicateCheckpoint(ctx context.Context, pchanne
return cp.(*wal.ReplicateCheckpoint), nil
}
// GetSalvageCheckpoint gets all salvage checkpoints of the wal.
// Note: Currently only supports local WAL. Remote WAL support requires streaming proto update.
func (hc *handlerClientImpl) GetSalvageCheckpoint(ctx context.Context, pchannel string) ([]*wal.ReplicateCheckpoint, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrClientClosed
}
defer hc.lifetime.Done()
logger := log.With(zap.String("pchannel", pchannel), zap.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, errors.New("salvage checkpoint can only be read for RW channel")
}
localWAL, err := registry.GetLocalAvailableWAL(assign.Channel)
if err == nil {
// Local WAL - get salvage checkpoints directly
return localWAL.GetSalvageCheckpoint(), nil
}
if !shouldUseRemoteWAL(err) {
return nil, err
}
// The salvage checkpoint is only meaningful on the streaming node that performed the force promote,
// so this path should not be reached in practice.
return nil, errors.New("GetSalvageCheckpoint is not implemented for remote WAL")
})
if err != nil {
return nil, err
}
if cps == nil {
return nil, nil
}
return cps.([]*wal.ReplicateCheckpoint), nil
}
// GetWALMetricsIfLocal gets the metrics of the local wal.
func (hc *handlerClientImpl) GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {
@@ -153,6 +153,45 @@ func TestHandlerClient(t *testing.T) {
handler.GetLatestMVCCTimestampIfLocal(ctx, "pchannel")
}
func TestHandlerClient_GetSalvageCheckpoint(t *testing.T) {
assignment := &types.PChannelInfoAssigned{
Channel: types.PChannelInfo{Name: "pchannel", Term: 1},
Node: types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
}
service := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeHandlerServiceClient](t)
rb := mock_resolver.NewMockBuilder(t)
rb.EXPECT().Close().Run(func() {})
w := mock_assignment.NewMockWatcher(t)
w.EXPECT().Close().Run(func() {})
// Always return the assignment so the create func is invoked.
w.EXPECT().Get(mock.Anything, mock.Anything).Return(assignment)
// Watch returns context.Canceled to break the backoff retry loop.
w.EXPECT().Watch(mock.Anything, mock.Anything, mock.Anything).Return(context.Canceled)
rebalanceTrigger := mock_types.NewMockAssignmentRebalanceTrigger(t)
handler := &handlerClientImpl{
lifetime: typeutil.NewLifetime(),
service: service,
rb: rb,
watcher: w,
rebalanceTrigger: rebalanceTrigger,
}
ctx := context.Background()
// Remote WAL returns "not implemented"; Watch returns Canceled to exit the loop.
cps, err := handler.GetSalvageCheckpoint(ctx, "pchannel")
assert.Error(t, err)
assert.Nil(t, cps)
// After close, GetSalvageCheckpoint returns ErrClientClosed immediately.
service.EXPECT().Close().Return()
handler.Close()
cps, err = handler.GetSalvageCheckpoint(ctx, "pchannel")
assert.ErrorIs(t, err, ErrClientClosed)
assert.Nil(t, cps)
}
func TestDial(t *testing.T) {
paramtable.Init()
@@ -225,11 +225,24 @@ func (o *openerAdaptorImpl) openRWWAL(ctx context.Context, l walimpls.WALImpls,
InitialRecoverSnapshot: snapshot,
TxnManager: param.TxnManager,
})
if param.ReplicateManager, err = replicates.RecoverReplicateManager(&replicates.ReplicateManagerRecoverParam{
ChannelInfo: param.ChannelInfo,
CurrentClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
InitialRecoverSnapshot: snapshot,
}); err != nil {
// Load salvage checkpoints from etcd (one per source cluster that was force-promoted from).
var salvageCheckpoints []*utility.ReplicateCheckpoint
if salvageCPProtos, err := resource.Resource().StreamingNodeCatalog().GetSalvageCheckpoint(ctx, param.ChannelInfo.Name); err != nil {
log.Ctx(ctx).Info("failed to load salvage checkpoints", zap.Error(err))
} else {
for _, proto := range salvageCPProtos {
salvageCheckpoints = append(salvageCheckpoints, utility.NewReplicateCheckpointFromProto(proto))
}
}
if param.ReplicateManager, err = replicates.RecoverReplicateManager(
&replicates.ReplicateManagerRecoverParam{
ChannelInfo: param.ChannelInfo,
CurrentClusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
InitialRecoverSnapshot: snapshot,
SalvageCheckpoints: salvageCheckpoints,
},
); err != nil {
return nil, err
}
@@ -61,6 +61,10 @@ func (w *roWALAdaptorImpl) GetReplicateCheckpoint() (*wal.ReplicateCheckpoint, e
panic("we cannot get replicate checkpoint from a read only wal")
}
func (w *roWALAdaptorImpl) GetSalvageCheckpoint() []*wal.ReplicateCheckpoint {
panic("we cannot get salvage checkpoint from a read only wal")
}
// Append writes a record to the log.
func (w *roWALAdaptorImpl) Append(ctx context.Context, msg message.MutableMessage) (*wal.AppendResult, error) {
panic("we cannot append message into a read only wal")
@@ -140,6 +140,16 @@ func (w *walAdaptorImpl) GetReplicateCheckpoint() (*utility.ReplicateCheckpoint,
return w.param.ReplicateManager.GetReplicateCheckpoint()
}
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
func (w *walAdaptorImpl) GetSalvageCheckpoint() []*utility.ReplicateCheckpoint {
if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil
}
defer w.lifetime.Done()
return w.param.ReplicateManager.GetSalvageCheckpoint()
}
// Append writes a record to the log.
func (w *walAdaptorImpl) Append(ctx context.Context, msg message.MutableMessage) (*wal.AppendResult, error) {
if !w.lifetime.Add(typeutil.LifetimeStateWorking) {
@@ -86,6 +86,8 @@ func initResourceForTest(t *testing.T) {
catalog.EXPECT().SaveSegmentAssignments(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.EXPECT().ListVChannel(mock.Anything, mock.Anything).Return(nil, nil)
catalog.EXPECT().SaveVChannels(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.EXPECT().GetSalvageCheckpoint(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
catalog.EXPECT().SaveSalvageCheckpoint(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
fMixCoordClient := syncutil.NewFuture[internaltypes.MixCoordClient]()
fMixCoordClient.Set(rc)
resource.InitForTest(
@@ -187,6 +189,10 @@ func (f *testOneWALFramework) testReadAndWrite(ctx context.Context, rwWAL wal.WA
require.True(f.t, status.AsStreamingError(err).IsReplicateViolation())
require.Nil(f.t, cp)
// No force promote has occurred, so salvage checkpoints should be empty.
salvageCPs := rwWAL.GetSalvageCheckpoint()
require.Nil(f.t, salvageCPs)
f.testSendCreateCollection(ctx, rwWAL)
defer f.testSendDropCollection(ctx, rwWAL)
@@ -23,15 +23,15 @@ type (
// InterceptorBuildParam is the parameter to build a interceptor.
type InterceptorBuildParam struct {
ChannelInfo types.PChannelInfo
WAL *syncutil.Future[wal.WAL] // The wal final object, can be used after interceptor is ready.
LastTimeTickMessage message.ImmutableMessage // The last time tick message in wal.
LastConfirmedMessageID message.MessageID // The last confirmed message id in wal.
WriteAheadBuffer *wab.WriteAheadBuffer // The write ahead buffer for the wal, used to erase the subscription of underlying wal.
MVCCManager *mvcc.MVCCManager // The MVCC manager for the wal, can be used to get the latest mvcc timetick.
InitialRecoverSnapshot *recovery.RecoverySnapshot // The initial recover snapshot for the wal, used to recover the wal state.
TxnManager *txn.TxnManager // The transaction manager for the wal, used to manage the transactions.
ShardManager shards.ShardManager // The shard manager for the wal, used to manage the shards, segment assignment, partition.
ReplicateManager replicates.ReplicateManager // The replicates manager for the wal, used to manage the replicates.
WAL *syncutil.Future[wal.WAL] // The wal final object, can be used after interceptor is ready.
LastTimeTickMessage message.ImmutableMessage // The last time tick message in wal.
LastConfirmedMessageID message.MessageID // The last confirmed message id in wal.
WriteAheadBuffer *wab.WriteAheadBuffer // The write ahead buffer for the wal, used to erase the subscription of underlying wal.
MVCCManager *mvcc.MVCCManager // The MVCC manager for the wal, can be used to get the latest mvcc timetick.
InitialRecoverSnapshot *recovery.RecoverySnapshot // The initial recover snapshot for the wal, used to recover the wal state.
TxnManager *txn.TxnManager // The transaction manager for the wal, used to manage the transactions.
ShardManager shards.ShardManager // The shard manager for the wal, used to manage the shards, segment assignment, partition.
ReplicateManager replicates.ReplicatesManager // The replicates manager for the wal, used to manage the replicates.
}
// Clear release the resources in the interceptor build param.
@@ -15,7 +15,7 @@ import (
const interceptorName = "replicate"
type replicateInterceptor struct {
replicateManager replicates.ReplicateManager
replicateManager replicates.ReplicatesManager
txnManager *txn.TxnManager
}
@@ -23,22 +23,28 @@ var ErrNotHandledByReplicateManager = errors.New("not handled by replicate manag
type ReplicateManagerRecoverParam struct {
ChannelInfo types.PChannelInfo
CurrentClusterID string
InitialRecoverSnapshot *recovery.RecoverySnapshot // the initial recover snapshot of the replicate manager.
InitialRecoverSnapshot *recovery.RecoverySnapshot // the initial recover snapshot of the replicate manager.
SalvageCheckpoints []*utility.ReplicateCheckpoint // loaded from etcd, one per source cluster
}
// RecoverReplicateManager recovers the replicate manager from the initial recover snapshot.
// It will recover the replicate manager from the initial recover snapshot.
// If the wal is on replicating mode, it will recover the replicate state.
func RecoverReplicateManager(param *ReplicateManagerRecoverParam) (ReplicateManager, error) {
func RecoverReplicateManager(param *ReplicateManagerRecoverParam) (ReplicatesManager, error) {
replicateConfigHelper, err := replicateutil.NewConfigHelper(param.CurrentClusterID, param.InitialRecoverSnapshot.Checkpoint.ReplicateConfig)
if err != nil {
return nil, newReplicateViolationErrorForConfig(param.InitialRecoverSnapshot.Checkpoint.ReplicateConfig, err)
}
salvageCheckpoints := make(map[string]*utility.ReplicateCheckpoint, len(param.SalvageCheckpoints))
for _, cp := range param.SalvageCheckpoints {
salvageCheckpoints[cp.ClusterID] = cp
}
rm := &replicatesManagerImpl{
mu: sync.Mutex{},
currentClusterID: param.CurrentClusterID,
pchannel: param.ChannelInfo,
replicateConfigHelper: replicateConfigHelper,
salvageCheckpoints: salvageCheckpoints,
}
if !rm.isPrimaryRole() {
// if current cluster is not the primary role,
@@ -56,7 +62,8 @@ type replicatesManagerImpl struct {
pchannel types.PChannelInfo
currentClusterID string
replicateConfigHelper *replicateutil.ConfigHelper
secondaryState *secondaryState // if the current cluster is not the primary role, it will have secondaryState.
secondaryState *secondaryState // if the current cluster is not the primary role, it will have secondaryState.
salvageCheckpoints map[string]*utility.ReplicateCheckpoint // captured on force promote, keyed by source clusterID
}
// SwitchReplicateMode switches the replicates manager between replicating mode and non-replicating mode.
@@ -72,6 +79,12 @@ func (impl *replicatesManagerImpl) SwitchReplicateMode(_ context.Context, msg me
incomingCurrentClusterConfig := newGraph.GetCurrentCluster()
switch incomingCurrentClusterConfig.Role() {
case replicateutil.RolePrimary:
// Capture salvage checkpoint before dropping secondary state on force promote.
// Store keyed by source clusterID so multiple force promotes don't overwrite each other.
if msg.Header().ForcePromote && impl.secondaryState != nil {
cp := impl.secondaryState.GetCheckpoint().Clone()
impl.salvageCheckpoints[cp.ClusterID] = cp
}
// drop the replicating state if the current cluster is switched to primary.
impl.secondaryState = nil
case replicateutil.RoleSecondary:
@@ -133,6 +146,20 @@ func (impl *replicatesManagerImpl) GetReplicateCheckpoint() (*utility.ReplicateC
return impl.secondaryState.GetCheckpoint(), nil
}
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
func (impl *replicatesManagerImpl) GetSalvageCheckpoint() []*utility.ReplicateCheckpoint {
impl.mu.Lock()
defer impl.mu.Unlock()
if len(impl.salvageCheckpoints) == 0 {
return nil
}
result := make([]*utility.ReplicateCheckpoint, 0, len(impl.salvageCheckpoints))
for _, cp := range impl.salvageCheckpoints {
result = append(result, cp)
}
return result
}
// beginReplicateMessage begins the replicate message operation.
func (impl *replicatesManagerImpl) beginReplicateMessage(ctx context.Context, msg message.MutableMessage) (ReplicateAcker, error) {
rh := msg.ReplicateHeader()
@@ -21,11 +21,11 @@ type ReplicateAcker interface {
Ack(err error)
}
// ReplicateManager manages the replicate operation on one wal.
// ReplicatesManager manages the replicate operation on one wal.
// There are two states:
// 1. primary: wal will only receive the non-replicate message.
// 2. secondary: wal will only receive the replicate message.
type ReplicateManager interface {
type ReplicatesManager interface {
// Role returns the role of the replicate manager.
Role() replicateutil.Role
@@ -45,4 +45,8 @@ type ReplicateManager interface {
// GetReplicateCheckpoint gets current replicate checkpoint.
// return ReplicateViolationError if the replicate mode is not replicating.
GetReplicateCheckpoint() (*utility.ReplicateCheckpoint, error)
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
// Returns an empty slice if no force promote has occurred.
GetSalvageCheckpoint() []*utility.ReplicateCheckpoint
}
@@ -96,6 +96,200 @@ func TestSecondaryReplicateManager(t *testing.T) {
testMessageOnSecondary(t, rm)
}
func TestSalvageCheckpointCaptureOnForcePromote(t *testing.T) {
// Setup: cluster starts as secondary with a checkpoint
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: "test2",
PChannel: "test2-rootcoord-dml_0",
MessageID: walimplstest.NewTestMessageID(100),
TimeTick: 1000,
},
ReplicateConfig: newReplicateConfiguration("test2", "test1"),
},
TxnBuffer: txnBuffer,
},
})
assert.NoError(t, err)
assert.Equal(t, replicateutil.RoleSecondary, rm.Role())
// Initially no salvage checkpoint
assert.Empty(t, rm.GetSalvageCheckpoint())
// Force promote to primary
err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "test2"))
assert.NoError(t, err)
assert.Equal(t, replicateutil.RolePrimary, rm.Role())
// Salvage checkpoint should be captured
salvageCPs := rm.GetSalvageCheckpoint()
assert.Len(t, salvageCPs, 1)
salvageCP := salvageCPs[0]
assert.Equal(t, "test2", salvageCP.ClusterID)
assert.Equal(t, "test2-rootcoord-dml_0", salvageCP.PChannel)
assert.Equal(t, uint64(1000), salvageCP.TimeTick)
}
func TestSalvageCheckpointNotCapturedOnNormalPromote(t *testing.T) {
// Setup: cluster starts as secondary
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: "test2",
PChannel: "test2-rootcoord-dml_0",
MessageID: walimplstest.NewTestMessageID(100),
TimeTick: 1000,
},
ReplicateConfig: newReplicateConfiguration("test2", "test1"),
},
TxnBuffer: txnBuffer,
},
})
assert.NoError(t, err)
assert.Equal(t, replicateutil.RoleSecondary, rm.Role())
// Normal promote (not force)
err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", false, "test2"))
assert.NoError(t, err)
assert.Equal(t, replicateutil.RolePrimary, rm.Role())
// Salvage checkpoint should NOT be captured
assert.Empty(t, rm.GetSalvageCheckpoint())
}
func TestSalvageCheckpointNotCapturedWhenAlreadyPrimary(t *testing.T) {
// Setup: cluster starts as primary (no secondary state)
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: nil,
ReplicateConfig: newReplicateConfiguration("test1", "test2"),
},
},
})
assert.NoError(t, err)
assert.Equal(t, replicateutil.RolePrimary, rm.Role())
// Force promote when already primary (no secondary state)
err = rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessageWithForcePromote("test1", true, "test2"))
assert.NoError(t, err)
// Salvage checkpoint should NOT be captured (no secondary state to capture)
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)
@@ -144,7 +338,7 @@ func TestSecondaryReplicateManagerWithTxn(t *testing.T) {
}
}
func testSwitchReplicateMode(t *testing.T, rm ReplicateManager, primaryClusterID, secondaryClusterID string) {
func testSwitchReplicateMode(t *testing.T, rm ReplicatesManager, primaryClusterID, secondaryClusterID string) {
ctx := context.Background()
// switch to primary
@@ -284,7 +478,7 @@ func testSwitchReplicateMode(t *testing.T, rm ReplicateManager, primaryClusterID
assert.Equal(t, cp.TimeTick, uint64(0))
}
func testMessageOnPrimary(t *testing.T, rm ReplicateManager) {
func testMessageOnPrimary(t *testing.T, rm ReplicatesManager) {
// switch to primary
err := rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("test1", "test2"))
assert.NoError(t, err)
@@ -311,7 +505,7 @@ func testMessageOnPrimary(t *testing.T, rm ReplicateManager) {
assert.Nil(t, g)
}
func testMessageOnSecondary(t *testing.T, rm ReplicateManager) {
func testMessageOnSecondary(t *testing.T, rm ReplicatesManager) {
// switch to secondary
err := rm.SwitchReplicateMode(context.Background(), newAlterReplicateConfigMessage("test2", "test1"))
assert.NoError(t, err)
@@ -389,9 +583,14 @@ func newReplicateConfiguration(primaryClusterID string, secondaryClusterID ...st
}
func newAlterReplicateConfigMessage(primaryClusterID string, secondaryClusterID ...string) message.MutableAlterReplicateConfigMessageV2 {
return newAlterReplicateConfigMessageWithForcePromote(primaryClusterID, false, secondaryClusterID...)
}
func newAlterReplicateConfigMessageWithForcePromote(primaryClusterID string, forcePromote bool, secondaryClusterID ...string) message.MutableAlterReplicateConfigMessageV2 {
return message.MustAsMutableAlterReplicateConfigMessageV2(message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: newReplicateConfiguration(primaryClusterID, secondaryClusterID...),
ForcePromote: forcePromote,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithVChannel(primaryClusterID + "-rootcoord-dml_0").
@@ -212,6 +212,11 @@ func (m *TxnManager) RollbackAllInFlightTransactions() {
m.Logger().Info("Rolled back in-flight transactions",
zap.Int64s("txnIDs", ids))
// Signal GracefulClose if it's already waiting and all sessions are now cleared.
if len(m.sessions) == 0 && m.closed != nil {
m.closed.Close()
}
m.notifyRecoverDone()
}
@@ -27,7 +27,7 @@ func (rs *recoveryStorageImpl) isDirty() bool {
rs.mu.Lock()
defer rs.mu.Unlock()
return rs.dirtyCounter > 0
return rs.dirtyCounter > 0 || rs.pendingSalvageCheckpoint != nil
}
// TODO: !!! all recovery persist operation should be a compare-and-swap operation to
@@ -133,6 +133,16 @@ func (rs *recoveryStorageImpl) persistDirtySnapshot(ctx context.Context, lvl zap
return err
}
// Salvage checkpoint must be persisted before the consume checkpoint to guarantee ordering:
// if the node crashes between these two writes, the next snapshot retry will re-persist both.
if snapshot.SalvageCheckpoint != nil {
if err := rs.retryOperationWithBackoff(ctx, rs.Logger().With(zap.String("op", "persistSalvageCheckpoint")), func(ctx context.Context) error {
return resource.Resource().StreamingNodeCatalog().SaveSalvageCheckpoint(ctx, rs.channel.Name, snapshot.SalvageCheckpoint.IntoProto())
}); err != nil {
return err
}
}
// checkpoint updates should always be persisted after other updates success.
if err := rs.retryOperationWithBackoff(ctx, rs.Logger().With(zap.String("op", "persistCheckpoint")), func(ctx context.Context) error {
return resource.Resource().StreamingNodeCatalog().
@@ -24,6 +24,9 @@ type RecoverySnapshot struct {
TxnBuffer *utility.TxnBuffer
// Used during WAL alteration process
AlterWALInfo *AlterWALInfo
// SalvageCheckpoint captures the replicate checkpoint at force-promote time.
// It must be persisted before the consume checkpoint so that the ordering guarantee holds.
SalvageCheckpoint *utility.ReplicateCheckpoint
}
// AlterWALInfo contains information about WAL alteration process.
@@ -99,6 +99,9 @@ type recoveryStorageImpl struct {
pendingPersistSnapshot *RecoverySnapshot
// used to mark switch MQ msg found
alterWALInfo *AlterWALInfo
// pendingSalvageCheckpoint holds the salvage checkpoint captured during force promote.
// Set under r.mu; consumed and persisted by the background task to avoid holding the lock.
pendingSalvageCheckpoint *utility.ReplicateCheckpoint
}
// Metrics gets the metrics of the wal.
@@ -181,7 +184,7 @@ func (r *recoveryStorageImpl) notifyPersist() {
func (r *recoveryStorageImpl) consumeDirtySnapshot() *RecoverySnapshot {
r.mu.Lock()
defer r.mu.Unlock()
if r.dirtyCounter == 0 {
if r.dirtyCounter == 0 && r.pendingSalvageCheckpoint == nil {
return nil
}
@@ -205,12 +208,17 @@ func (r *recoveryStorageImpl) consumeDirtySnapshot() *RecoverySnapshot {
vchannels[vchannel.meta.Vchannel] = dirtySnapshot
}
}
// Atomically capture the salvage checkpoint alongside other dirty state.
// Clearing it here (under r.mu) ensures it is only consumed once.
salvageCP := r.pendingSalvageCheckpoint
r.pendingSalvageCheckpoint = nil
// clear the dirty counter.
r.dirtyCounter = 0
return &RecoverySnapshot{
VChannels: vchannels,
SegmentAssignments: segments,
Checkpoint: r.checkpoint.Clone(),
SalvageCheckpoint: salvageCP,
}
}
@@ -257,6 +265,11 @@ func (r *recoveryStorageImpl) updateCheckpoint(msg message.ImmutableMessage) {
clusterRole := replicateutil.MustNewConfigHelper(r.currentClusterID, header.ReplicateConfiguration).GetCurrentCluster()
switch clusterRole.Role() {
case replicateutil.RolePrimary:
if header.GetForcePromote() && r.checkpoint.ReplicateCheckpoint != nil {
// Store for background task to persist; never call etcd while holding r.mu.
r.pendingSalvageCheckpoint = r.checkpoint.ReplicateCheckpoint
r.notifyPersist()
}
r.checkpoint.ReplicateCheckpoint = nil
case replicateutil.RoleSecondary:
// Update the replicate checkpoint if the cluster role is secondary.
@@ -0,0 +1,243 @@
package recovery
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/mocks/mock_metastore"
"github.com/milvus-io/milvus/internal/streamingnode/server/resource"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal/utility"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
)
// newAlterReplicateConfigMessageWithForcePromote creates an AlterReplicateConfig message with ForcePromote set.
func newAlterReplicateConfigMessageWithForcePromote(primaryClusterID string, secondaryClusterID []string, timetick uint64, messageID message.MessageID) message.ImmutableMessage {
return message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: newReplicateConfiguration(primaryClusterID, secondaryClusterID...),
ForcePromote: true,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithVChannel("test1-rootcoord-dml_0").
MustBuildMutable().
WithTimeTick(timetick).
WithLastConfirmed(messageID).
IntoImmutableMessage(walimplstest.NewTestMessageID(10086))
}
func TestUpdateCheckpointForcePromote(t *testing.T) {
t.Run("force_promote_captures_salvage_checkpoint", func(t *testing.T) {
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(1),
TimeTick: 1,
ReplicateCheckpoint: &utility.ReplicateCheckpoint{
ClusterID: "test2",
PChannel: "test2-rootcoord-dml_0",
MessageID: walimplstest.NewTestMessageID(50),
TimeTick: 500,
},
},
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
// Start as secondary of test2
rs.updateCheckpoint(newAlterReplicateConfigMessage("test2", []string{"test1"}, 2, walimplstest.NewTestMessageID(2)))
assert.NotNil(t, rs.checkpoint.ReplicateCheckpoint)
assert.Nil(t, rs.pendingSalvageCheckpoint)
// Force promote to primary — should capture the salvage checkpoint
rs.updateCheckpoint(newAlterReplicateConfigMessageWithForcePromote("test1", []string{"test2"}, 3, walimplstest.NewTestMessageID(3)))
assert.Nil(t, rs.checkpoint.ReplicateCheckpoint)
assert.NotNil(t, rs.pendingSalvageCheckpoint)
assert.Equal(t, "test2", rs.pendingSalvageCheckpoint.ClusterID)
})
t.Run("normal_promote_does_not_capture_salvage_checkpoint", func(t *testing.T) {
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(1),
TimeTick: 1,
ReplicateCheckpoint: &utility.ReplicateCheckpoint{
ClusterID: "test2",
PChannel: "test2-rootcoord-dml_0",
MessageID: walimplstest.NewTestMessageID(50),
TimeTick: 500,
},
},
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
rs.updateCheckpoint(newAlterReplicateConfigMessage("test2", []string{"test1"}, 2, walimplstest.NewTestMessageID(2)))
// Normal promote (no ForcePromote flag)
rs.updateCheckpoint(newAlterReplicateConfigMessage("test1", []string{"test2"}, 3, walimplstest.NewTestMessageID(3)))
assert.Nil(t, rs.checkpoint.ReplicateCheckpoint)
assert.Nil(t, rs.pendingSalvageCheckpoint)
})
t.Run("force_promote_no_replicate_checkpoint_noop", func(t *testing.T) {
// Force promote when there's no existing replicate checkpoint should not set pendingSalvageCheckpoint
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(1),
TimeTick: 1,
ReplicateCheckpoint: nil,
},
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
rs.updateCheckpoint(newAlterReplicateConfigMessageWithForcePromote("test1", []string{"test2"}, 2, walimplstest.NewTestMessageID(2)))
assert.Nil(t, rs.checkpoint.ReplicateCheckpoint)
assert.Nil(t, rs.pendingSalvageCheckpoint)
})
}
func TestConsumeDirtySnapshotWithSalvageCheckpoint(t *testing.T) {
t.Run("salvage_checkpoint_included_and_cleared", func(t *testing.T) {
cp := &utility.ReplicateCheckpoint{
ClusterID: "cluster-x",
PChannel: "cluster-x-rootcoord-dml_0",
TimeTick: 999,
}
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(10),
TimeTick: 10,
},
segments: map[int64]*segmentRecoveryInfo{},
vchannels: map[string]*vchannelRecoveryInfo{},
pendingSalvageCheckpoint: cp,
dirtyCounter: 0,
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
// consumeDirtySnapshot should pick up the pending salvage checkpoint even when dirtyCounter==0
snapshot := rs.consumeDirtySnapshot()
assert.NotNil(t, snapshot)
assert.Equal(t, cp, snapshot.SalvageCheckpoint)
// After consuming, pendingSalvageCheckpoint should be cleared
assert.Nil(t, rs.pendingSalvageCheckpoint)
// A second call with nothing dirty should return nil
snapshot2 := rs.consumeDirtySnapshot()
assert.Nil(t, snapshot2)
})
t.Run("dirty_messages_without_salvage_checkpoint", func(t *testing.T) {
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(10),
TimeTick: 10,
},
segments: map[int64]*segmentRecoveryInfo{},
vchannels: map[string]*vchannelRecoveryInfo{},
pendingSalvageCheckpoint: nil,
dirtyCounter: 3,
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
snapshot := rs.consumeDirtySnapshot()
assert.NotNil(t, snapshot)
assert.Nil(t, snapshot.SalvageCheckpoint)
assert.Equal(t, 0, rs.dirtyCounter)
})
t.Run("both_dirty_messages_and_salvage_checkpoint", func(t *testing.T) {
cp := &utility.ReplicateCheckpoint{
ClusterID: "cluster-y",
PChannel: "cluster-y-rootcoord-dml_0",
TimeTick: 777,
}
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test1-rootcoord-dml_0"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(10),
TimeTick: 10,
},
segments: map[int64]*segmentRecoveryInfo{},
vchannels: map[string]*vchannelRecoveryInfo{},
pendingSalvageCheckpoint: cp,
dirtyCounter: 5,
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test1-rootcoord-dml_0"}),
}
snapshot := rs.consumeDirtySnapshot()
assert.NotNil(t, snapshot)
assert.Equal(t, cp, snapshot.SalvageCheckpoint)
assert.Equal(t, 0, rs.dirtyCounter)
assert.Nil(t, rs.pendingSalvageCheckpoint)
})
}
func TestIsDirtyWithSalvageCheckpoint(t *testing.T) {
// dirtyCounter==0 but pendingSalvageCheckpoint != nil → isDirty() should be true.
cp := &utility.ReplicateCheckpoint{
ClusterID: "cluster-x",
PChannel: "test-pchannel",
TimeTick: 500,
}
rs := &recoveryStorageImpl{
currentClusterID: "test1",
channel: types.PChannelInfo{Name: "test-pchannel"},
checkpoint: &WALCheckpoint{MessageID: walimplstest.NewTestMessageID(10), TimeTick: 10},
segments: map[int64]*segmentRecoveryInfo{},
vchannels: map[string]*vchannelRecoveryInfo{},
pendingSalvageCheckpoint: cp,
dirtyCounter: 0,
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test-pchannel"}),
}
assert.True(t, rs.isDirty())
// Consuming the snapshot clears pendingSalvageCheckpoint.
snapshot := rs.consumeDirtySnapshot()
assert.NotNil(t, snapshot)
assert.False(t, rs.isDirty())
}
func TestPersistDirtySnapshotWithSalvageCheckpoint(t *testing.T) {
snCatalog := mock_metastore.NewMockStreamingNodeCataLog(t)
snCatalog.EXPECT().SaveSalvageCheckpoint(mock.Anything, "test-pchannel", mock.Anything).Return(nil)
snCatalog.EXPECT().SaveConsumeCheckpoint(mock.Anything, "test-pchannel", mock.Anything).Return(nil)
resource.InitForTest(t, resource.OptStreamingNodeCatalog(snCatalog))
cp := &utility.ReplicateCheckpoint{
ClusterID: "cluster-x",
PChannel: "test-pchannel",
MessageID: walimplstest.NewTestMessageID(50),
TimeTick: 500,
}
rs := &recoveryStorageImpl{
cfg: newConfig(),
channel: types.PChannelInfo{Name: "test-pchannel"},
checkpoint: &WALCheckpoint{
MessageID: walimplstest.NewTestMessageID(10),
TimeTick: 10,
},
segments: map[int64]*segmentRecoveryInfo{},
vchannels: map[string]*vchannelRecoveryInfo{},
pendingSalvageCheckpoint: cp,
metrics: newRecoveryStorageMetrics(types.PChannelInfo{Name: "test-pchannel"}),
}
err := rs.persistDirtySnapshot(context.Background(), zap.InfoLevel)
assert.NoError(t, err)
assert.Nil(t, rs.pendingPersistSnapshot)
}
+4
View File
@@ -31,6 +31,10 @@ type WAL interface {
// the message id of the replicate checkpoint will be 0.
GetReplicateCheckpoint() (*ReplicateCheckpoint, error)
// GetSalvageCheckpoint returns all salvage checkpoints captured during force promote.
// Returns an empty slice if no force promote has occurred.
GetSalvageCheckpoint() []*ReplicateCheckpoint
// Append writes a record to the log.
Append(ctx context.Context, msg message.MutableMessage) (*AppendResult, error)
+1 -1
View File
@@ -22,7 +22,7 @@ require (
github.com/jolestar/go-commons-pool/v2 v2.1.2
github.com/json-iterator/go v1.1.13-0.20220915233716-71ac16282d12
github.com/klauspost/compress v1.18.0
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b
github.com/minio/minio-go/v7 v7.0.73
github.com/panjf2000/ants/v2 v2.11.3
github.com/prometheus/client_golang v1.20.5
+2 -2
View File
@@ -517,8 +517,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6 h1:YHMFI6L
github.com/milvus-io/cgosymbolizer v0.0.0-20250318084424-114f4050c3a6/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2 h1:OEqYnraL/Tf2nGMjRJji6xaaDe40zSQVBvruzEBZ1hU=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260304062516-e7e54d966ef2/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b h1:tUp6a72FFeosfjN+8jy/Hax7BqVbtgsaMMVrAfzKnEc=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.6-0.20260309063517-8b5776d31f2b/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34=
github.com/minio/md5-simd v1.1.2/go.mod h1:MzdKDxYpY2BT9XQFocsiZf/NKVtR7nkE4RoEpN+20RM=
github.com/minio/minio-go/v7 v7.0.73 h1:qr2vi96Qm7kZ4v7LLebjte+MQh621fFWnv93p12htEo=
@@ -0,0 +1,267 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package replication
import (
"context"
"io"
"testing"
"github.com/stretchr/testify/suite"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/tests/integration"
)
type DataSalvageSuite struct {
integration.MiniClusterSuite
}
func TestDataSalvage(t *testing.T) {
suite.Run(t, new(DataSalvageSuite))
}
// TestGetReplicateInfoOnPrimaryCluster verifies that GetReplicateInfo
// returns checkpoint info on a primary cluster. The salvage checkpoint
// should be nil since no force promote has occurred.
func (s *DataSalvageSuite) TestGetReplicateInfoOnPrimaryCluster() {
ctx := context.Background()
clusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
pchannel := clusterID + "-pchan0"
// First set up replication config to make the cluster a primary
config := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: clusterID,
Pchannels: []string{pchannel},
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
updateResp, err := s.Cluster.MilvusClient.UpdateReplicateConfiguration(ctx, &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
ForcePromote: false,
})
s.NoError(err)
s.NoError(merr.Error(updateResp))
// Get replicate info
resp, err := s.Cluster.MilvusClient.GetReplicateInfo(ctx, &milvuspb.GetReplicateInfoRequest{
TargetPchannel: pchannel,
})
s.NoError(err)
// On a primary cluster, checkpoint should exist but salvage checkpoint should be nil
// (no force promote has occurred)
log.Info("GetReplicateInfo response",
zap.Any("checkpoint", resp.GetCheckpoint()),
zap.Any("salvageCheckpoint", resp.GetSalvageCheckpoint()))
// Salvage checkpoint should be nil on primary cluster
s.Nil(resp.GetSalvageCheckpoint(), "salvage checkpoint should be nil on primary cluster")
}
// TestDumpMessagesBasic verifies that DumpMessages can stream messages
// from a WAL channel after inserting some data.
func (s *DataSalvageSuite) TestDumpMessagesBasic() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
const (
dim = 128
dbName = ""
rowNum = 100
)
collectionName := "TestDumpMessages" + funcutil.GenRandomStr()
// Create collection
schema := integration.ConstructSchemaOfVecDataType(collectionName, dim, true, schemapb.DataType_FloatVector)
marshaledSchema, err := proto.Marshal(schema)
s.NoError(err)
createResp, err := s.Cluster.MilvusClient.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName,
CollectionName: collectionName,
Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum,
})
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, createResp.GetErrorCode())
// Insert some data
fVecColumn := integration.NewFloatVectorFieldData(integration.FloatVecField, rowNum, dim)
hashKeys := integration.GenerateHashKeys(rowNum)
insertResp, err := s.Cluster.MilvusClient.Insert(ctx, &milvuspb.InsertRequest{
DbName: dbName,
CollectionName: collectionName,
FieldsData: []*schemapb.FieldData{fVecColumn},
HashKeys: hashKeys,
NumRows: uint32(rowNum),
})
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, insertResp.GetStatus().GetErrorCode())
// Get pchannel for the collection
descResp, err := s.Cluster.MilvusClient.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
CollectionName: collectionName,
})
s.NoError(err)
s.NotEmpty(descResp.GetVirtualChannelNames())
// Get pchannel from vchannel
vchannel := descResp.GetVirtualChannelNames()[0]
pchannel := funcutil.ToPhysicalChannel(vchannel)
log.Info("Testing DumpMessages",
zap.String("pchannel", pchannel),
zap.String("vchannel", vchannel))
// Set up replication config first
clusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
config := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: clusterID,
Pchannels: []string{pchannel},
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
updateResp, err := s.Cluster.MilvusClient.UpdateReplicateConfiguration(ctx, &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
ForcePromote: false,
})
s.NoError(err)
s.NoError(merr.Error(updateResp))
// Get replicate checkpoint as start position
infoResp, err := s.Cluster.MilvusClient.GetReplicateInfo(ctx, &milvuspb.GetReplicateInfoRequest{
TargetPchannel: pchannel,
})
s.NoError(err)
s.NotNil(infoResp.GetCheckpoint())
s.NotNil(infoResp.GetCheckpoint().GetMessageId())
// Dump messages from start position (exclusive - messages after the checkpoint)
stream, err := s.Cluster.MilvusClient.DumpMessages(ctx, &milvuspb.DumpMessagesRequest{
Pchannel: pchannel,
StartMessageId: infoResp.GetCheckpoint().GetMessageId(),
})
s.NoError(err)
// Read some messages (with timeout via context cancellation)
messages := make([]*milvuspb.DumpMessagesResponse, 0)
readCtx, readCancel := context.WithCancel(ctx)
// Read in a goroutine with limit
go func() {
for i := 0; i < 10; i++ { // Read up to 10 messages
resp, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
log.Warn("error receiving message", zap.Error(err))
break
}
// Check if response contains a message (not a status)
if msg := resp.GetMessage(); msg != nil {
messages = append(messages, resp)
log.Info("received message", zap.String("messageId", msg.GetId().GetId()))
} else if status := resp.GetStatus(); status != nil {
log.Warn("received status response", zap.Any("status", status))
break
}
}
readCancel()
}()
<-readCtx.Done()
log.Info("DumpMessages test completed", zap.Int("messageCount", len(messages)))
// We should have received at least some messages
// Note: The exact count depends on timing and what messages are in the WAL
// For now, just verify the API works without error
// Clean up
dropResp, err := s.Cluster.MilvusClient.DropCollection(ctx, &milvuspb.DropCollectionRequest{
CollectionName: collectionName,
})
s.NoError(err)
s.Equal(commonpb.ErrorCode_Success, dropResp.GetErrorCode())
}
// TestDumpMessagesWithMissingPchannel verifies that DumpMessages returns
// an error when pchannel is not provided.
func (s *DataSalvageSuite) TestDumpMessagesWithMissingPchannel() {
ctx := context.Background()
stream, err := s.Cluster.MilvusClient.DumpMessages(ctx, &milvuspb.DumpMessagesRequest{
Pchannel: "", // Missing pchannel
StartMessageId: &commonpb.MessageID{
Id: "test",
},
})
// The error might come during stream creation or first Recv()
if err == nil {
_, err = stream.Recv()
}
s.Error(err)
}
// TestDumpMessagesWithMissingStartMessageId verifies that DumpMessages returns
// an error when start_message_id is not provided.
func (s *DataSalvageSuite) TestDumpMessagesWithMissingStartMessageId() {
ctx := context.Background()
stream, err := s.Cluster.MilvusClient.DumpMessages(ctx, &milvuspb.DumpMessagesRequest{
Pchannel: "test-pchannel",
StartMessageId: nil, // Missing start message ID
})
// The error might come during stream creation or first Recv()
if err == nil {
_, err = stream.Recv()
}
s.Error(err)
}
@@ -50,102 +50,6 @@ func (s *ForcePromoteSuite) getPChannelNames() []string {
return pchannels
}
// TestForcePromoteOnPrimaryClusterShouldFail verifies that force promote
// returns an error when called on a primary cluster (no replication configured).
// Force promote is only intended for secondary clusters during failover.
func (s *ForcePromoteSuite) TestForcePromoteOnPrimaryClusterShouldFail() {
ctx := context.Background()
// Force promote requires empty cluster and topology fields
// The configuration is auto-constructed from the cluster's existing meta
config := &commonpb.ReplicateConfiguration{}
// Call UpdateReplicateConfiguration with force_promote=true on primary cluster
req := &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
ForcePromote: true,
}
resp, err := s.Cluster.MilvusClient.UpdateReplicateConfiguration(ctx, req)
// Should return an error because we're on a primary cluster
s.NoError(err) // RPC should succeed
s.NotNil(resp)
err = merr.Error(resp)
s.Error(err)
s.Contains(err.Error(), "force promote can only be used on secondary clusters")
}
// TestForcePromoteWithNonEmptyClustersShouldFail verifies that force promote
// rejects configurations with non-empty clusters field.
// Force promote auto-constructs the configuration from existing meta.
func (s *ForcePromoteSuite) TestForcePromoteWithNonEmptyClustersShouldFail() {
ctx := context.Background()
clusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
pchannels := s.getPChannelNames()
// Create config with clusters (invalid for force promote - must be empty)
config := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: clusterID,
Pchannels: pchannels,
ConnectionParam: &commonpb.ConnectionParam{
Uri: "http://localhost:19530",
Token: "test-token",
},
},
},
}
req := &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
ForcePromote: true,
}
resp, err := s.Cluster.MilvusClient.UpdateReplicateConfiguration(ctx, req)
// Should return an error about non-empty cluster/topology fields
s.NoError(err)
s.NotNil(resp)
err = merr.Error(resp)
s.Error(err)
s.Contains(err.Error(), "force promote requires empty cluster and topology fields")
}
// TestForcePromoteWithTopologyShouldFail verifies that force promote
// rejects configurations that contain cross-cluster topology.
func (s *ForcePromoteSuite) TestForcePromoteWithTopologyShouldFail() {
ctx := context.Background()
clusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
// Create config with only topology (no clusters) - still invalid for force promote
config := &commonpb.ReplicateConfiguration{
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{
SourceClusterId: clusterID,
TargetClusterId: "other-cluster",
},
},
}
req := &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
ForcePromote: true,
}
resp, err := s.Cluster.MilvusClient.UpdateReplicateConfiguration(ctx, req)
// Should return an error about non-empty topology field
s.NoError(err)
s.NotNil(resp)
err = merr.Error(resp)
s.Error(err)
s.Contains(err.Error(), "force promote requires empty cluster and topology fields")
}
// TestNormalUpdateReplicateConfiguration verifies that normal (non-force) updates
// work correctly on a primary cluster.
func (s *ForcePromoteSuite) TestNormalUpdateReplicateConfiguration() {