diff --git a/docs/agent_guides/streaming-system/message/message-semantic-collection.md b/docs/agent_guides/streaming-system/message/message-semantic-collection.md index d7f9a62009..6b93f0e491 100644 --- a/docs/agent_guides/streaming-system/message/message-semantic-collection.md +++ b/docs/agent_guides/streaming-system/message/message-semantic-collection.md @@ -18,6 +18,7 @@ All broadcast messages implicitly carry **SharedCluster** via the Broadcaster. | CreateSnapshot | Broadcast: CChannel | No | SharedDBName + ExclusiveCollectionName + ExclusiveSnapshotName | | DropSnapshot | Broadcast: CChannel | No | ExclusiveSnapshotName | | RestoreSnapshot | Broadcast: CChannel | No | SharedDBName + ExclusiveCollectionName + ExclusiveSnapshotName | +| DropSnapshotsByCollection | Broadcast: CChannel | No | SharedDBName + SharedCollectionName | | Import | Broadcast: VChannels (no CChannel) | No | — | | Insert | Single VChannel | No | — | | Delete | Single VChannel | No | — | @@ -27,6 +28,7 @@ All broadcast messages implicitly carry **SharedCluster** via the Broadcaster. | AlterLoadConfig | Broadcast: CChannel | No | SharedDBName + ExclusiveCollectionName | | DropLoadConfig | Broadcast: CChannel | No | SharedDBName + ExclusiveCollectionName (or ExclusiveCluster) | | BatchUpdateManifest | Broadcast: CChannel | No | SharedDBName + SharedCollectionName | +| RefreshExternalCollection | Broadcast: CChannel | No | — | ## Message Descriptions @@ -36,7 +38,7 @@ All broadcast messages implicitly carry **SharedCluster** via the Broadcaster. - **TruncateCollection**: Logically truncates by sealing and dropping all segments before the truncation timestamp. Implicitly flushes all growing segments. Uses AckSyncUp. - **CreatePartition** / **DropPartition**: Creates or drops a partition. DropPartition implicitly flushes the partition's growing segments. - **CreateIndex** / **AlterIndex** / **DropIndex**: Manages indexes on a collection's field. CChannel-only. -- **CreateSnapshot** / **DropSnapshot** / **RestoreSnapshot**: Manages collection snapshots. CChannel-only. +- **CreateSnapshot** / **DropSnapshot** / **RestoreSnapshot** / **DropSnapshotsByCollection**: Manages collection snapshots. CChannel-only. - **Import**: Initiates a bulk import job for a collection. - **Insert** / **Delete**: DML on a single VChannel. CipherEnabled. - **CreateSegment** / **Flush**: WAL-generated (SelfControlled). Allocates or seals a growing segment. @@ -44,6 +46,11 @@ All broadcast messages implicitly carry **SharedCluster** via the Broadcaster. - **AlterLoadConfig**: Modifies load configuration — partition set, replica count, load fields, etc. CChannel-only, consumed by QueryCoord. - **DropLoadConfig**: Removes load configuration, unloading/releasing from query nodes. Uses ExclusiveCluster when part of DropCollection flow. - **BatchUpdateManifest**: Updates segment manifest versions in batch. Used after compaction or index building. CChannel-only. +- **RefreshExternalCollection**: Submits an external collection refresh job using a pre-allocated job ID from the WAL message. CChannel-only. + +## Replication Compatibility + +Current producers explicitly mark these collection-scoped broadcast messages with `Unreplicable` (`_ur`): CreateSnapshot, DropSnapshot, RestoreSnapshot, BatchUpdateManifest, and RefreshExternalCollection. Replication skips the concrete marked messages instead of classifying the whole message type as permanently unsupported, so newly generated messages can become replicable later by no longer setting `_ur`. ## Data Lifecycle Ordering Invariants diff --git a/docs/agent_guides/streaming-system/message/message.md b/docs/agent_guides/streaming-system/message/message.md index 1f8cbeccbc..4f48d413b6 100644 --- a/docs/agent_guides/streaming-system/message/message.md +++ b/docs/agent_guides/streaming-system/message/message.md @@ -12,7 +12,7 @@ Messages transition through three stages: - **ImmutableMessage**: The post-append, read-only state. Carries a backend-assigned MessageID and LastConfirmedMessageID. Returned by WAL Read/Consume operations. For transactions, multiple ImmutableMessages are assembled into an **ImmutableTxnMessage** (Begin + body messages + Commit) by the consumer-side TxnBuffer. -- **ReplicateMutableMessage**: Created from a source cluster's ImmutableMessage for cross-cluster replication. Attaches a `ReplicateHeader` preserving the source cluster's original Message Properties, then re-enters the local WAL as a MutableMessage, +- **ReplicateMutableMessage**: Created from a source cluster's ImmutableMessage for cross-cluster replication. Attaches a `ReplicateHeader` preserving the source cluster's original Message Properties, then re-enters the local WAL as a MutableMessage. Replication honors message-level compatibility markers such as `Unreplicable`. ## Key Properties @@ -27,6 +27,7 @@ Messages transition through three stages: | **LastConfirmedMessageID** | Reading from this MessageID guarantees all subsequent messages have TimeTick greater than this message's TimeTick (including txn messages). | | **TxnContext** | Links the message to a transaction. Nil if non-transactional. | | **ReplicateHeader** | Source cluster's original message metadata for cross-cluster replication. | +| **Unreplicable** (`_ur`) | Marks this concrete message as unsafe for cross-cluster replication. Replication skips only messages carrying this property; absence means replicable, including old WAL messages written before the marker existed. | | **MessageID** | Backend-assigned unique identifier. | | **PChannel** | The PChannel this message belongs to. | diff --git a/docs/agent_guides/streaming-system/replication/replicate.md b/docs/agent_guides/streaming-system/replication/replicate.md index 8d98528a46..7c0d8f8ceb 100644 --- a/docs/agent_guides/streaming-system/replication/replicate.md +++ b/docs/agent_guides/streaming-system/replication/replicate.md @@ -13,10 +13,16 @@ Milvus supports multi-cluster WAL replication via a star topology: one PRIMARY c ## Data Flow -1. **Primary WAL** → **CDC ChannelReplicator** (per-PChannel, runs on primary StreamingNode): reads messages from the primary WAL starting at the secondary's `ReplicateCheckpoint`. Self-controlled messages (TimeTick, CreateSegment, Flush) are skipped. +1. **Primary WAL** → **CDC ChannelReplicator** (per-PChannel, runs on primary StreamingNode): reads messages from the primary WAL starting at the secondary's `ReplicateCheckpoint`. Self-controlled messages (TimeTick, CreateSegment, Flush) and messages carrying the `Unreplicable` (`_ur`) property are skipped. 2. **ChannelReplicator** → **Secondary Proxy** via `CreateReplicateStream` gRPC bidirectional stream: sends each message with its original `MessageID`, `Properties`, and `Payload`, along with the `SourceClusterID`. 3. **Secondary Proxy** → **Secondary WAL**: the Proxy remaps VChannel names and appends to the local WAL. The **Replicate Interceptor** validates the incoming message (cluster ID match, TimeTick deduplication) and tracks checkpoint. +## Message-Level Replication Skip + +Some DDL/control messages cannot be safely replayed on a SECONDARY until their replay contract is deterministic across clusters. Producers mark those concrete WAL messages with the `Unreplicable` (`_ur`) message property. The CDC sender treats them like ignored messages and advances replication progress without sending them. The SECONDARY replicate interceptor also ignores replicated messages that carry `_ur`, which protects mixed-version or already-forwarded traffic. + +This is a **message property**, not a static `MessageType` rule. Future support for one of these DDLs should stop setting `_ur` on newly generated messages; old WAL messages that already carry `_ur` remain skipped for rolling-upgrade compatibility. + ## Checkpoint & Consistency The secondary maintains a `ReplicateCheckpoint` per PChannel: `{ClusterID, PChannel, MessageID, TimeTick}`. diff --git a/internal/cdc/replication/replicatestream/replicate_stream_client_impl.go b/internal/cdc/replication/replicatestream/replicate_stream_client_impl.go index 0d2991cb95..065479eec9 100644 --- a/internal/cdc/replication/replicatestream/replicate_stream_client_impl.go +++ b/internal/cdc/replication/replicatestream/replicate_stream_client_impl.go @@ -169,7 +169,7 @@ func (r *replicateStreamClient) Replicate(msg message.ImmutableMessage) error { case <-r.ctx.Done(): return nil default: - if msg.MessageType().IsSelfControlled() { + if msg.MessageType().IsSelfControlled() || msg.IsUnreplicable() { // If no messages are being replicated, update the last replicated time tick. if r.pendingMessages.Len() == 0 { r.metrics.UpdateLastReplicatedTimeTick(msg.TimeTick()) diff --git a/internal/cdc/replication/replicatestream/replicate_stream_client_impl_test.go b/internal/cdc/replication/replicatestream/replicate_stream_client_impl_test.go index 6faa356732..fefcc31f8e 100644 --- a/internal/cdc/replication/replicatestream/replicate_stream_client_impl_test.go +++ b/internal/cdc/replication/replicatestream/replicate_stream_client_impl_test.go @@ -86,6 +86,7 @@ func TestReplicateStreamClient_Replicate(t *testing.T) { mockMsg.EXPECT().TimeTick().Return(tt) mockMsg.EXPECT().EstimateSize().Return(1024) mockMsg.EXPECT().MessageType().Return(message.MessageTypeInsert) + mockMsg.EXPECT().IsUnreplicable().Return(false) mockMsg.EXPECT().MessageID().Return(messageID) mockMsg.EXPECT().IntoImmutableMessageProto().Return(&commonpb.ImmutableMessage{ Id: messageID.IntoProto(), @@ -153,6 +154,31 @@ func TestReplicateStreamClient_Replicate_ContextCanceled(t *testing.T) { mockStreamClient.Close() } +func TestReplicateStreamClient_Replicate_UnreplicableMessageIgnored(t *testing.T) { + ctx := context.Background() + repMeta := &streamingpb.ReplicatePChannelMeta{ + SourceChannelName: "test-source-channel", + TargetChannelName: "test-target-channel", + TargetCluster: &commonpb.MilvusCluster{ + ClusterId: "test-cluster", + }, + } + client := &replicateStreamClient{ + ctx: ctx, + pendingMessages: NewMsgQueue(MsgQueueOptions{Capacity: 16, MaxSize: 1024}), + metrics: NewReplicateMetrics(repMeta), + } + + mockMsg := mock_message.NewMockImmutableMessage(t) + mockMsg.EXPECT().MessageType().Return(message.MessageTypeCreateSnapshot) + mockMsg.EXPECT().IsUnreplicable().Return(true) + mockMsg.EXPECT().TimeTick().Return(uint64(100)).Maybe() + + err := client.Replicate(mockMsg) + assert.ErrorIs(t, err, ErrReplicateIgnored) + assert.Equal(t, 0, client.pendingMessages.Len()) +} + func TestReplicateStreamClient_Reconnect(t *testing.T) { ctx := context.Background() targetCluster := &commonpb.MilvusCluster{ @@ -205,6 +231,7 @@ func TestReplicateStreamClient_Reconnect(t *testing.T) { mockMsg.EXPECT().MessageID().Return(messageID) mockMsg.EXPECT().EstimateSize().Return(1024) mockMsg.EXPECT().MessageType().Return(message.MessageTypeInsert) + mockMsg.EXPECT().IsUnreplicable().Return(false) mockMsg.EXPECT().IntoImmutableMessageProto().Return(&commonpb.ImmutableMessage{ Id: messageID.IntoProto(), Payload: []byte("test-payload"), diff --git a/internal/datacoord/services.go b/internal/datacoord/services.go index 401b96f4ca..d45f54bf62 100644 --- a/internal/datacoord/services.go +++ b/internal/datacoord/services.go @@ -2234,6 +2234,7 @@ func (s *Server) CreateSnapshot(ctx context.Context, req *datapb.CreateSnapshotR }). WithBody(&message.CreateSnapshotMessageBody{}). WithBroadcast([]string{streaming.WAL().ControlChannel()}). + WithUnreplicable(). MustBuildBroadcast(), ); err != nil { mlog.Error(context.TODO(), "CreateSnapshot broadcast failed", mlog.Err(err)) @@ -2284,6 +2285,7 @@ func (s *Server) BatchUpdateManifest(ctx context.Context, req *datapb.BatchUpdat Items: items, }). WithBroadcast([]string{streaming.WAL().ControlChannel()}). + WithUnreplicable(). MustBuildBroadcast(), ); err != nil { mlog.Error(context.TODO(), "BatchUpdateManifest broadcast failed", mlog.Err(err)) @@ -2387,6 +2389,7 @@ func (s *Server) DropSnapshot(ctx context.Context, req *datapb.DropSnapshotReque }). WithBody(&message.DropSnapshotMessageBody{}). WithBroadcast([]string{streaming.WAL().ControlChannel()}). + WithUnreplicable(). MustBuildBroadcast(), ); err != nil { mlog.Error(context.TODO(), "DropSnapshot broadcast failed", mlog.Err(err)) @@ -2744,6 +2747,7 @@ func (s *Server) RefreshExternalCollection(ctx context.Context, req *datapb.Refr }). WithBody(&message.RefreshExternalCollectionMessageBody{}). WithBroadcast([]string{streaming.WAL().ControlChannel()}). + WithUnreplicable(). MustBuildBroadcast() if _, err := b.Broadcast(ctx, msg); err != nil { diff --git a/internal/datacoord/services_commit_backfill.go b/internal/datacoord/services_commit_backfill.go index c3ce65110e..44fec0909a 100644 --- a/internal/datacoord/services_commit_backfill.go +++ b/internal/datacoord/services_commit_backfill.go @@ -155,6 +155,7 @@ func broadcastBackfillBatch( Items: items, }). WithBroadcast(channels). + WithUnreplicable(). MustBuildBroadcast(), ) return err diff --git a/internal/datacoord/snapshot_manager.go b/internal/datacoord/snapshot_manager.go index 6600372a2d..63917de966 100644 --- a/internal/datacoord/snapshot_manager.go +++ b/internal/datacoord/snapshot_manager.go @@ -761,6 +761,7 @@ func (sm *snapshotManager) RestoreSnapshot( }). WithBody(&message.RestoreSnapshotMessageBody{}). WithBroadcast([]string{streaming.WAL().ControlChannel()}). + WithUnreplicable(). MustBuildBroadcast() if _, bcErr := restoreBroadcaster.Broadcast(ctx, msg); bcErr != nil { diff --git a/internal/streamingnode/server/wal/interceptors/replicate/replicates/impl.go b/internal/streamingnode/server/wal/interceptors/replicate/replicates/impl.go index ab6a213c41..80f059fbc1 100644 --- a/internal/streamingnode/server/wal/interceptors/replicate/replicates/impl.go +++ b/internal/streamingnode/server/wal/interceptors/replicate/replicates/impl.go @@ -111,6 +111,12 @@ func (impl *replicatesManagerImpl) BeginReplicateMessage(ctx context.Context, ms } return nil, ErrNotHandledByReplicateManager } + if msg.IsUnreplicable() { + if rh != nil { + return nil, status.NewIgnoreOperation("wal unreplicable message cannot be replicated") + } + return nil, ErrNotHandledByReplicateManager + } impl.mu.Lock() defer func() { diff --git a/internal/streamingnode/server/wal/interceptors/replicate/replicates/manager_test.go b/internal/streamingnode/server/wal/interceptors/replicate/replicates/manager_test.go index f9f1c7c39c..c2d76bb097 100644 --- a/internal/streamingnode/server/wal/interceptors/replicate/replicates/manager_test.go +++ b/internal/streamingnode/server/wal/interceptors/replicate/replicates/manager_test.go @@ -96,6 +96,48 @@ func TestSecondaryReplicateManager(t *testing.T) { testMessageOnSecondary(t, rm) } +func TestReplicateManager_UnreplicableMessage(t *testing.T) { + 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("test2", "test1"), + ReplicateCheckpoint: &utility.ReplicateCheckpoint{ + ClusterID: "test2", + PChannel: "test2-rootcoord-dml_0", + MessageID: walimplstest.NewTestMessageID(1), + TimeTick: 1, + }, + }, + TxnBuffer: utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics()), + }, + }) + assert.NoError(t, err) + assert.Equal(t, replicateutil.RoleSecondary, rm.Role()) + + msg := newCreateSnapshotMutableMessage() + g, err := rm.BeginReplicateMessage(context.Background(), msg) + assert.ErrorIs(t, err, ErrNotHandledByReplicateManager) + assert.Nil(t, g) + + msg = newCreateSnapshotMutableMessage().WithReplicateHeader(&message.ReplicateHeader{ + ClusterID: "test2", + MessageID: walimplstest.NewTestMessageID(2), + LastConfirmedMessageID: walimplstest.NewTestMessageID(1), + TimeTick: 2, + VChannel: "test2-rootcoord-dml_0", + }) + g, err = rm.BeginReplicateMessage(context.Background(), msg) + assert.True(t, status.AsStreamingError(err).IsIgnoredOperation()) + assert.Nil(t, g) +} + func TestSalvageCheckpointCaptureOnForcePromote(t *testing.T) { // Setup: cluster starts as secondary with a checkpoint txnBuffer := utility.NewTxnBuffer(mlog.With(), metricsutil.NewScanMetrics(types.PChannelInfo{}).NewScannerMetrics()) @@ -602,6 +644,15 @@ func testMessageOnSecondary(t *testing.T, rm ReplicatesManager) { assert.Nil(t, g) } +func newCreateSnapshotMutableMessage() message.MutableMessage { + return message.NewCreateSnapshotMessageBuilderV2(). + WithHeader(&message.CreateSnapshotMessageHeader{}). + WithBody(&message.CreateSnapshotMessageBody{}). + WithVChannel("test1-rootcoord-dml_0"). + WithUnreplicable(). + MustBuildMutable() +} + // newReplicateConfiguration creates a valid replicate configuration for testing func newReplicateConfiguration(primaryClusterID string, secondaryClusterID ...string) *commonpb.ReplicateConfiguration { clusters := []*commonpb.MilvusCluster{ diff --git a/pkg/mocks/streaming/util/mock_message/mock_BroadcastMutableMessage.go b/pkg/mocks/streaming/util/mock_message/mock_BroadcastMutableMessage.go index 541037cb73..e16aa4dd49 100644 --- a/pkg/mocks/streaming/util/mock_message/mock_BroadcastMutableMessage.go +++ b/pkg/mocks/streaming/util/mock_message/mock_BroadcastMutableMessage.go @@ -297,6 +297,51 @@ func (_c *MockBroadcastMutableMessage_IsPersisted_Call) RunAndReturn(run func() return _c } +// IsUnreplicable provides a mock function with no fields +func (_m *MockBroadcastMutableMessage) IsUnreplicable() bool { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for IsUnreplicable") + } + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockBroadcastMutableMessage_IsUnreplicable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsUnreplicable' +type MockBroadcastMutableMessage_IsUnreplicable_Call struct { + *mock.Call +} + +// IsUnreplicable is a helper method to define mock.On call +func (_e *MockBroadcastMutableMessage_Expecter) IsUnreplicable() *MockBroadcastMutableMessage_IsUnreplicable_Call { + return &MockBroadcastMutableMessage_IsUnreplicable_Call{Call: _e.mock.On("IsUnreplicable")} +} + +func (_c *MockBroadcastMutableMessage_IsUnreplicable_Call) Run(run func()) *MockBroadcastMutableMessage_IsUnreplicable_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockBroadcastMutableMessage_IsUnreplicable_Call) Return(_a0 bool) *MockBroadcastMutableMessage_IsUnreplicable_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockBroadcastMutableMessage_IsUnreplicable_Call) RunAndReturn(run func() bool) *MockBroadcastMutableMessage_IsUnreplicable_Call { + _c.Call.Return(run) + return _c +} + // MarshalLogObject provides a mock function with given fields: _a0 func (_m *MockBroadcastMutableMessage) MarshalLogObject(_a0 mlog.ObjectEncoder) error { ret := _m.Called(_a0) diff --git a/pkg/mocks/streaming/util/mock_message/mock_ImmutableMessage.go b/pkg/mocks/streaming/util/mock_message/mock_ImmutableMessage.go index 3ba2506279..903b8ad95e 100644 --- a/pkg/mocks/streaming/util/mock_message/mock_ImmutableMessage.go +++ b/pkg/mocks/streaming/util/mock_message/mock_ImmutableMessage.go @@ -393,6 +393,51 @@ func (_c *MockImmutableMessage_IsPersisted_Call) RunAndReturn(run func() bool) * return _c } +// IsUnreplicable provides a mock function with no fields +func (_m *MockImmutableMessage) IsUnreplicable() bool { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for IsUnreplicable") + } + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockImmutableMessage_IsUnreplicable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsUnreplicable' +type MockImmutableMessage_IsUnreplicable_Call struct { + *mock.Call +} + +// IsUnreplicable is a helper method to define mock.On call +func (_e *MockImmutableMessage_Expecter) IsUnreplicable() *MockImmutableMessage_IsUnreplicable_Call { + return &MockImmutableMessage_IsUnreplicable_Call{Call: _e.mock.On("IsUnreplicable")} +} + +func (_c *MockImmutableMessage_IsUnreplicable_Call) Run(run func()) *MockImmutableMessage_IsUnreplicable_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockImmutableMessage_IsUnreplicable_Call) Return(_a0 bool) *MockImmutableMessage_IsUnreplicable_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockImmutableMessage_IsUnreplicable_Call) RunAndReturn(run func() bool) *MockImmutableMessage_IsUnreplicable_Call { + _c.Call.Return(run) + return _c +} + // LastConfirmedMessageID provides a mock function with no fields func (_m *MockImmutableMessage) LastConfirmedMessageID() message.MessageID { ret := _m.Called() diff --git a/pkg/mocks/streaming/util/mock_message/mock_ImmutableTxnMessage.go b/pkg/mocks/streaming/util/mock_message/mock_ImmutableTxnMessage.go index f3b2e50b10..15880a4c8a 100644 --- a/pkg/mocks/streaming/util/mock_message/mock_ImmutableTxnMessage.go +++ b/pkg/mocks/streaming/util/mock_message/mock_ImmutableTxnMessage.go @@ -487,6 +487,51 @@ func (_c *MockImmutableTxnMessage_IsPersisted_Call) RunAndReturn(run func() bool return _c } +// IsUnreplicable provides a mock function with no fields +func (_m *MockImmutableTxnMessage) IsUnreplicable() bool { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for IsUnreplicable") + } + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockImmutableTxnMessage_IsUnreplicable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsUnreplicable' +type MockImmutableTxnMessage_IsUnreplicable_Call struct { + *mock.Call +} + +// IsUnreplicable is a helper method to define mock.On call +func (_e *MockImmutableTxnMessage_Expecter) IsUnreplicable() *MockImmutableTxnMessage_IsUnreplicable_Call { + return &MockImmutableTxnMessage_IsUnreplicable_Call{Call: _e.mock.On("IsUnreplicable")} +} + +func (_c *MockImmutableTxnMessage_IsUnreplicable_Call) Run(run func()) *MockImmutableTxnMessage_IsUnreplicable_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockImmutableTxnMessage_IsUnreplicable_Call) Return(_a0 bool) *MockImmutableTxnMessage_IsUnreplicable_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockImmutableTxnMessage_IsUnreplicable_Call) RunAndReturn(run func() bool) *MockImmutableTxnMessage_IsUnreplicable_Call { + _c.Call.Return(run) + return _c +} + // LastConfirmedMessageID provides a mock function with no fields func (_m *MockImmutableTxnMessage) LastConfirmedMessageID() message.MessageID { ret := _m.Called() diff --git a/pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go b/pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go index 810ec3f84f..70118d407d 100644 --- a/pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go +++ b/pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go @@ -345,6 +345,51 @@ func (_c *MockMutableMessage_IsPersisted_Call) RunAndReturn(run func() bool) *Mo return _c } +// IsUnreplicable provides a mock function with no fields +func (_m *MockMutableMessage) IsUnreplicable() bool { + ret := _m.Called() + + if len(ret) == 0 { + panic("no return value specified for IsUnreplicable") + } + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockMutableMessage_IsUnreplicable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsUnreplicable' +type MockMutableMessage_IsUnreplicable_Call struct { + *mock.Call +} + +// IsUnreplicable is a helper method to define mock.On call +func (_e *MockMutableMessage_Expecter) IsUnreplicable() *MockMutableMessage_IsUnreplicable_Call { + return &MockMutableMessage_IsUnreplicable_Call{Call: _e.mock.On("IsUnreplicable")} +} + +func (_c *MockMutableMessage_IsUnreplicable_Call) Run(run func()) *MockMutableMessage_IsUnreplicable_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockMutableMessage_IsUnreplicable_Call) Return(_a0 bool) *MockMutableMessage_IsUnreplicable_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockMutableMessage_IsUnreplicable_Call) RunAndReturn(run func() bool) *MockMutableMessage_IsUnreplicable_Call { + _c.Call.Return(run) + return _c +} + // MarshalLogObject provides a mock function with given fields: _a0 func (_m *MockMutableMessage) MarshalLogObject(_a0 mlog.ObjectEncoder) error { ret := _m.Called(_a0) diff --git a/pkg/streaming/util/message/builder.go b/pkg/streaming/util/message/builder.go index 8aaf50ed6c..e1f628fde8 100644 --- a/pkg/streaming/util/message/builder.go +++ b/pkg/streaming/util/message/builder.go @@ -149,6 +149,12 @@ func (b *mutableMesasgeBuilder[H, B]) WithNotPersisted() *mutableMesasgeBuilder[ return b } +// WithUnreplicable marks this concrete message as unsafe for cross-cluster replication. +func (b *mutableMesasgeBuilder[H, B]) WithUnreplicable() *mutableMesasgeBuilder[H, B] { + b.WithProperty(messageUnreplicable, "") + return b +} + // WithBody creates a new builder with message body. func (b *mutableMesasgeBuilder[H, B]) WithBody(body B) *mutableMesasgeBuilder[H, B] { b.body = body diff --git a/pkg/streaming/util/message/message.go b/pkg/streaming/util/message/message.go index 8c9e7bada6..de4aaf45d9 100644 --- a/pkg/streaming/util/message/message.go +++ b/pkg/streaming/util/message/message.go @@ -43,6 +43,9 @@ type BasicMessage interface { // Should be used with read-only promise. Properties() RProperties + // IsUnreplicable returns true if the message cannot be replicated to a secondary cluster. + IsUnreplicable() bool + // TimeTick returns the time tick of current message. // Available only when the message's version greater than 0. // Otherwise, it will panic. diff --git a/pkg/streaming/util/message/message_impl.go b/pkg/streaming/util/message/message_impl.go index f0c2bb9b5e..bf0fde50c6 100644 --- a/pkg/streaming/util/message/message_impl.go +++ b/pkg/streaming/util/message/message_impl.go @@ -64,6 +64,11 @@ func (m *messageImpl) Properties() RProperties { return m.properties } +// IsUnreplicable returns true if the message cannot be replicated. +func (m *messageImpl) IsUnreplicable() bool { + return m.properties.Exist(messageUnreplicable) +} + // IsPersisted returns true if the message is persisted. func (m *messageImpl) IsPersisted() bool { return !m.properties.Exist(messageNotPersisteted) diff --git a/pkg/streaming/util/message/message_test.go b/pkg/streaming/util/message/message_test.go index 3cfce5edbe..7469b655fa 100644 --- a/pkg/streaming/util/message/message_test.go +++ b/pkg/streaming/util/message/message_test.go @@ -66,6 +66,69 @@ func TestMessageType(t *testing.T) { assert.False(t, MessageTypeTimeTick.IsDMLMessageType()) } +func TestMessageUnreplicableProperty(t *testing.T) { + insertMsg := NewInsertMessageBuilderV1(). + WithHeader(&InsertMessageHeader{}). + WithBody(&msgpb.InsertRequest{ShardName: "v1"}). + WithVChannel("v1"). + MustBuildMutable() + assert.False(t, insertMsg.IsUnreplicable()) + + createCollectionMsg := NewCreateCollectionMessageBuilderV1(). + WithHeader(&CreateCollectionMessageHeader{}). + WithBody(&msgpb.CreateCollectionRequest{}). + WithBroadcast([]string{"v1"}). + MustBuildBroadcast() + assert.False(t, createCollectionMsg.IsUnreplicable()) + + assert.False(t, NewCreateSnapshotMessageBuilderV2(). + WithHeader(&CreateSnapshotMessageHeader{}). + WithBody(&CreateSnapshotMessageBody{}). + WithBroadcast([]string{"v1"}). + MustBuildBroadcast(). + IsUnreplicable()) + assert.True(t, NewCreateSnapshotMessageBuilderV2(). + WithHeader(&CreateSnapshotMessageHeader{}). + WithBody(&CreateSnapshotMessageBody{}). + WithBroadcast([]string{"v1"}). + WithUnreplicable(). + MustBuildBroadcast(). + IsUnreplicable()) + assert.True(t, NewDropSnapshotMessageBuilderV2(). + WithHeader(&DropSnapshotMessageHeader{}). + WithBody(&DropSnapshotMessageBody{}). + WithBroadcast([]string{"v1"}). + WithUnreplicable(). + MustBuildBroadcast(). + IsUnreplicable()) + assert.True(t, NewRestoreSnapshotMessageBuilderV2(). + WithHeader(&RestoreSnapshotMessageHeader{}). + WithBody(&RestoreSnapshotMessageBody{}). + WithBroadcast([]string{"v1"}). + WithUnreplicable(). + MustBuildBroadcast(). + IsUnreplicable()) + assert.True(t, NewBatchUpdateManifestMessageBuilderV2(). + WithHeader(&BatchUpdateManifestMessageHeader{}). + WithBody(&BatchUpdateManifestMessageBody{}). + WithBroadcast([]string{"v1"}). + WithUnreplicable(). + MustBuildBroadcast(). + IsUnreplicable()) + assert.True(t, NewRefreshExternalCollectionMessageBuilderV2(). + WithHeader(&RefreshExternalCollectionMessageHeader{}). + WithBody(&RefreshExternalCollectionMessageBody{}). + WithBroadcast([]string{"v1"}). + WithUnreplicable(). + MustBuildBroadcast(). + IsUnreplicable()) + + legacySnapshotMsg := NewMutableMessageBeforeAppend(nil, map[string]string{ + messageTypeKey: MessageTypeCreateSnapshot.marshal(), + }) + assert.False(t, legacySnapshotMsg.IsUnreplicable()) +} + func TestVersion(t *testing.T) { v := newMessageVersionFromString("") assert.Equal(t, VersionOld, v) diff --git a/pkg/streaming/util/message/properties.go b/pkg/streaming/util/message/properties.go index c4d458c305..26e3d36b04 100644 --- a/pkg/streaming/util/message/properties.go +++ b/pkg/streaming/util/message/properties.go @@ -17,6 +17,7 @@ const ( messageNotPersisteted = "_np" // check if the message is unpersisted. messagePChannelLevel = "_pcl" // mark the message as pchannel level message. messageReplicateMesssageHeader = "_rh" // replicate message header. + messageUnreplicable = "_ur" // mark the message as unsafe to replicate. messageTraceContext = "_tc" // Trace context subset header. )