mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
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:
co-authored by
Claude Opus 4.5
parent
695f27bacb
commit
4ba08692ce
@@ -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
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
+51
-4
@@ -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)
|
||||
|
||||
|
||||
@@ -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
@@ -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{
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
+202
-3
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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() {
|
||||
|
||||
Reference in New Issue
Block a user