feat: Add CDC support (#44124)

This PR implements a new CDC service for Milvus 2.6, providing log-based
cross-cluster replication.

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

---------

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
Signed-off-by: chyezh <chyezh@outlook.com>
Co-authored-by: chyezh <chyezh@outlook.com>
This commit is contained in:
yihao.dai 2025-09-16 16:32:01 +08:00 committed by GitHub
parent 98d23de36c
commit 51f69f32d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
210 changed files with 12883 additions and 2566 deletions

View File

@ -361,6 +361,10 @@ test-mixcoord:
@echo "Running go unittests..."
@(env bash $(PWD)/scripts/run_go_unittest.sh -t mixcoord)
test-cdc:
@echo "Running cdc unittests..."
@(env bash $(PWD)/scripts/run_go_unittest.sh -t cdc)
test-go: build-cpp-with-unittest
@echo "Running go unittests..."
@(env bash $(PWD)/scripts/run_go_unittest.sh)
@ -536,6 +540,9 @@ generate-mockery-pkg:
generate-mockery-internal: getdeps
$(INSTALL_PATH)/mockery --config $(PWD)/internal/.mockery.yaml
generate-mockery-cdc: getdeps
$(INSTALL_PATH)/mockery --config $(PWD)/internal/cdc/.mockery.yaml
generate-mockery: generate-mockery-types generate-mockery-kv generate-mockery-rootcoord generate-mockery-proxy generate-mockery-querycoord generate-mockery-querynode generate-mockery-datacoord generate-mockery-pkg generate-mockery-internal
generate-yaml: milvus-tools

View File

@ -6,15 +6,13 @@ require (
github.com/blang/semver/v4 v4.0.0
github.com/cockroachdb/errors v1.9.1
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250901025546-3cddc9bb1357
github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e
github.com/quasilyte/go-ruleguard/dsl v0.3.22
github.com/samber/lo v1.27.0
github.com/stretchr/testify v1.10.0
github.com/tidwall/gjson v1.17.1
go.opentelemetry.io/otel v1.28.0
go.uber.org/atomic v1.11.0
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
google.golang.org/grpc v1.65.0
google.golang.org/protobuf v1.34.2
)
@ -89,6 +87,7 @@ require (
go.etcd.io/etcd/raft/v3 v3.5.5 // indirect
go.etcd.io/etcd/server/v3 v3.5.5 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
go.opentelemetry.io/otel v1.28.0 // indirect
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.20.0 // indirect
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.20.0 // indirect
go.opentelemetry.io/otel/metric v1.28.0 // indirect
@ -99,6 +98,7 @@ require (
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.27.0 // indirect
golang.org/x/crypto v0.36.0 // indirect
golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842 // indirect
golang.org/x/net v0.38.0 // indirect
golang.org/x/sync v0.12.0 // indirect
golang.org/x/sys v0.31.0 // indirect

View File

@ -318,8 +318,8 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr
github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1 h1:NLoSWXvlJD8t91G3CUsooXqYnm5nfsBngztQYYT58V0=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250901025546-3cddc9bb1357 h1:OYM9ylL42FTVL5kAHOZtsOPqzXq9Pn0/H1YLfXcS/e4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250901025546-3cddc9bb1357/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e h1:VCr43pG4efacDbM4au70fh8/5hNTftoWzm1iEumvDWM=
github.com/milvus-io/milvus/pkg/v2 v2.0.0-20250319085209-5a6b4e56d59e/go.mod h1:37AWzxVs2NS4QUJrkcbeLUwi+4Av0h5mEdjLI62EANU=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=

View File

@ -1266,6 +1266,52 @@ func (_c *MilvusServiceServer_CreatePrivilegeGroup_Call) RunAndReturn(run func(c
return _c
}
// CreateReplicateStream provides a mock function with given fields: _a0
func (_m *MilvusServiceServer) CreateReplicateStream(_a0 milvuspb.MilvusService_CreateReplicateStreamServer) error {
ret := _m.Called(_a0)
if len(ret) == 0 {
panic("no return value specified for CreateReplicateStream")
}
var r0 error
if rf, ok := ret.Get(0).(func(milvuspb.MilvusService_CreateReplicateStreamServer) error); ok {
r0 = rf(_a0)
} else {
r0 = ret.Error(0)
}
return r0
}
// MilvusServiceServer_CreateReplicateStream_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateReplicateStream'
type MilvusServiceServer_CreateReplicateStream_Call struct {
*mock.Call
}
// CreateReplicateStream is a helper method to define mock.On call
// - _a0 milvuspb.MilvusService_CreateReplicateStreamServer
func (_e *MilvusServiceServer_Expecter) CreateReplicateStream(_a0 interface{}) *MilvusServiceServer_CreateReplicateStream_Call {
return &MilvusServiceServer_CreateReplicateStream_Call{Call: _e.mock.On("CreateReplicateStream", _a0)}
}
func (_c *MilvusServiceServer_CreateReplicateStream_Call) Run(run func(_a0 milvuspb.MilvusService_CreateReplicateStreamServer)) *MilvusServiceServer_CreateReplicateStream_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(milvuspb.MilvusService_CreateReplicateStreamServer))
})
return _c
}
func (_c *MilvusServiceServer_CreateReplicateStream_Call) Return(_a0 error) *MilvusServiceServer_CreateReplicateStream_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MilvusServiceServer_CreateReplicateStream_Call) RunAndReturn(run func(milvuspb.MilvusService_CreateReplicateStreamServer) error) *MilvusServiceServer_CreateReplicateStream_Call {
_c.Call.Return(run)
return _c
}
// CreateResourceGroup provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) CreateResourceGroup(_a0 context.Context, _a1 *milvuspb.CreateResourceGroupRequest) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)
@ -3685,6 +3731,65 @@ func (_c *MilvusServiceServer_GetReplicas_Call) RunAndReturn(run func(context.Co
return _c
}
// GetReplicateInfo provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) GetReplicateInfo(_a0 context.Context, _a1 *milvuspb.GetReplicateInfoRequest) (*milvuspb.GetReplicateInfoResponse, error) {
ret := _m.Called(_a0, _a1)
if len(ret) == 0 {
panic("no return value specified for GetReplicateInfo")
}
var r0 *milvuspb.GetReplicateInfoResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetReplicateInfoRequest) (*milvuspb.GetReplicateInfoResponse, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.GetReplicateInfoRequest) *milvuspb.GetReplicateInfoResponse); ok {
r0 = rf(_a0, _a1)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*milvuspb.GetReplicateInfoResponse)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.GetReplicateInfoRequest) error); ok {
r1 = rf(_a0, _a1)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MilvusServiceServer_GetReplicateInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateInfo'
type MilvusServiceServer_GetReplicateInfo_Call struct {
*mock.Call
}
// GetReplicateInfo is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.GetReplicateInfoRequest
func (_e *MilvusServiceServer_Expecter) GetReplicateInfo(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_GetReplicateInfo_Call {
return &MilvusServiceServer_GetReplicateInfo_Call{Call: _e.mock.On("GetReplicateInfo", _a0, _a1)}
}
func (_c *MilvusServiceServer_GetReplicateInfo_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.GetReplicateInfoRequest)) *MilvusServiceServer_GetReplicateInfo_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.GetReplicateInfoRequest))
})
return _c
}
func (_c *MilvusServiceServer_GetReplicateInfo_Call) Return(_a0 *milvuspb.GetReplicateInfoResponse, _a1 error) *MilvusServiceServer_GetReplicateInfo_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MilvusServiceServer_GetReplicateInfo_Call) RunAndReturn(run func(context.Context, *milvuspb.GetReplicateInfoRequest) (*milvuspb.GetReplicateInfoResponse, error)) *MilvusServiceServer_GetReplicateInfo_Call {
_c.Call.Return(run)
return _c
}
// GetUserTags provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) GetUserTags(_a0 context.Context, _a1 *milvuspb.GetUserTagsRequest) (*milvuspb.GetUserTagsResponse, error) {
ret := _m.Called(_a0, _a1)
@ -6222,6 +6327,65 @@ func (_c *MilvusServiceServer_UpdateCredential_Call) RunAndReturn(run func(conte
return _c
}
// UpdateReplicateConfiguration provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) UpdateReplicateConfiguration(_a0 context.Context, _a1 *milvuspb.UpdateReplicateConfigurationRequest) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)
if len(ret) == 0 {
panic("no return value specified for UpdateReplicateConfiguration")
}
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.UpdateReplicateConfigurationRequest) (*commonpb.Status, error)); ok {
return rf(_a0, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.UpdateReplicateConfigurationRequest) *commonpb.Status); ok {
r0 = rf(_a0, _a1)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*commonpb.Status)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.UpdateReplicateConfigurationRequest) error); ok {
r1 = rf(_a0, _a1)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MilvusServiceServer_UpdateReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateReplicateConfiguration'
type MilvusServiceServer_UpdateReplicateConfiguration_Call struct {
*mock.Call
}
// UpdateReplicateConfiguration is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.UpdateReplicateConfigurationRequest
func (_e *MilvusServiceServer_Expecter) UpdateReplicateConfiguration(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_UpdateReplicateConfiguration_Call {
return &MilvusServiceServer_UpdateReplicateConfiguration_Call{Call: _e.mock.On("UpdateReplicateConfiguration", _a0, _a1)}
}
func (_c *MilvusServiceServer_UpdateReplicateConfiguration_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.UpdateReplicateConfigurationRequest)) *MilvusServiceServer_UpdateReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.UpdateReplicateConfigurationRequest))
})
return _c
}
func (_c *MilvusServiceServer_UpdateReplicateConfiguration_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_UpdateReplicateConfiguration_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MilvusServiceServer_UpdateReplicateConfiguration_Call) RunAndReturn(run func(context.Context, *milvuspb.UpdateReplicateConfigurationRequest) (*commonpb.Status, error)) *MilvusServiceServer_UpdateReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// UpdateResourceGroups provides a mock function with given fields: _a0, _a1
func (_m *MilvusServiceServer) UpdateResourceGroups(_a0 context.Context, _a1 *milvuspb.UpdateResourceGroupsRequest) (*commonpb.Status, error) {
ret := _m.Called(_a0, _a1)

View File

@ -0,0 +1,57 @@
package milvusclient
import (
"context"
"github.com/cockroachdb/errors"
"google.golang.org/grpc"
"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/pkg/v2/util/merr"
)
// UpdateReplicateConfiguration updates the replicate configuration to the Milvus cluster.
// Use ReplicateConfigurationBuilder to build the configuration.
func (c *Client) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration, opts ...grpc.CallOption) error {
req := &milvuspb.UpdateReplicateConfigurationRequest{
ReplicateConfiguration: config,
}
err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.UpdateReplicateConfiguration(ctx, req, opts...)
return merr.CheckRPCCall(resp, err)
})
return err
}
// GetReplicateInfo gets replicate information from the Milvus cluster
func (c *Client) GetReplicateInfo(ctx context.Context, sourceClusterID string, opts ...grpc.CallOption) (*milvuspb.GetReplicateInfoResponse, error) {
req := &milvuspb.GetReplicateInfoRequest{
SourceClusterId: sourceClusterID,
}
var resp *milvuspb.GetReplicateInfoResponse
err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
var err error
resp, err = milvusService.GetReplicateInfo(ctx, req, opts...)
return merr.CheckRPCCall(resp, err)
})
return resp, err
}
// CreateReplicateStream creates a replicate stream
func (c *Client) CreateReplicateStream(ctx context.Context, opts ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error) {
var streamClient milvuspb.MilvusService_CreateReplicateStreamClient
err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
var err error
streamClient, err = milvusService.CreateReplicateStream(ctx, opts...)
if err != nil {
return err
}
if streamClient == nil {
return errors.New("stream client is nil")
}
return nil
})
return streamClient, err
}

View File

@ -0,0 +1,113 @@
package milvusclient
import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
// MilvusClusterBuilder defines the interface for building Milvus cluster configuration
type MilvusClusterBuilder interface {
// WithPchannels adds physical channels
WithPchannels(pchannels ...string) MilvusClusterBuilder
// WithToken sets the authentication token
WithToken(token string) MilvusClusterBuilder
// WithURI sets the connection URI
WithURI(uri string) MilvusClusterBuilder
// Build constructs and returns the MilvusCluster object
Build() *commonpb.MilvusCluster
}
// ReplicateConfigurationBuilder defines the interface for building replicate configuration
type ReplicateConfigurationBuilder interface {
// WithCluster adds a cluster configuration, you can use MilvusClusterBuilder to create a cluster
WithCluster(cluster *commonpb.MilvusCluster) ReplicateConfigurationBuilder
// WithTopology adds a cross-cluster topology configuration
WithTopology(sourceClusterID, targetClusterID string) ReplicateConfigurationBuilder
// Build constructs and returns the configuration object
Build() *commonpb.ReplicateConfiguration
}
// NewMilvusClusterBuilder creates a new Milvus cluster builder
func NewMilvusClusterBuilder(clusterID string) MilvusClusterBuilder {
return newMilvusClusterBuilder(clusterID)
}
// NewReplicateConfigurationBuilder creates a new replicate configuration builder
func NewReplicateConfigurationBuilder() ReplicateConfigurationBuilder {
return newReplicateConfigurationBuilder()
}
type milvusClusterBuilder struct {
clusterID string
uri string
token string
pchannels []string
}
func newMilvusClusterBuilder(clusterID string) MilvusClusterBuilder {
return &milvusClusterBuilder{
clusterID: clusterID,
uri: "localhost:19530", // default URI
token: "", // default empty token
pchannels: []string{},
}
}
func (b *milvusClusterBuilder) WithPchannels(pchannels ...string) MilvusClusterBuilder {
b.pchannels = append(b.pchannels, pchannels...)
return b
}
func (b *milvusClusterBuilder) WithToken(token string) MilvusClusterBuilder {
b.token = token
return b
}
func (b *milvusClusterBuilder) WithURI(uri string) MilvusClusterBuilder {
b.uri = uri
return b
}
func (b *milvusClusterBuilder) Build() *commonpb.MilvusCluster {
return &commonpb.MilvusCluster{
ClusterId: b.clusterID,
ConnectionParam: &commonpb.ConnectionParam{
Uri: b.uri,
Token: b.token,
},
Pchannels: b.pchannels,
}
}
type replicateConfigurationBuilder struct {
config *commonpb.ReplicateConfiguration
}
func newReplicateConfigurationBuilder() ReplicateConfigurationBuilder {
return &replicateConfigurationBuilder{
config: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
},
}
}
func (b *replicateConfigurationBuilder) WithCluster(cluster *commonpb.MilvusCluster) ReplicateConfigurationBuilder {
b.config.Clusters = append(b.config.Clusters, cluster)
return b
}
func (b *replicateConfigurationBuilder) WithTopology(sourceClusterID, targetClusterID string) ReplicateConfigurationBuilder {
topology := &commonpb.CrossClusterTopology{
SourceClusterId: sourceClusterID,
TargetClusterId: targetClusterID,
}
b.config.CrossClusterTopology = append(b.config.CrossClusterTopology, topology)
return b
}
func (b *replicateConfigurationBuilder) Build() *commonpb.ReplicateConfiguration {
return b.config
}

View File

@ -0,0 +1,116 @@
// 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.
// nolint
package milvusclient_test
import (
"context"
"log"
"github.com/milvus-io/milvus/client/v2/milvusclient"
)
const (
exampleMilvusAddr = `127.0.0.1:19530`
)
func ExampleClient_UpdateReplicateConfiguration() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: exampleMilvusAddr,
})
if err != nil {
log.Fatal(err)
}
// Use builder pattern to create cluster configuration (chained calls)
sourceCluster := milvusclient.NewMilvusClusterBuilder("source-cluster").
WithURI("localhost:19530").
WithToken("source-token").
WithPchannels("source-channel-1", "source-channel-2").
Build()
targetCluster := milvusclient.NewMilvusClusterBuilder("target-cluster").
WithURI("localhost:19531").
WithToken("target-token").
WithPchannels("target-channel-1", "target-channel-2").
Build()
// Use builder pattern to build replicate configuration
config := milvusclient.NewReplicateConfigurationBuilder().
WithCluster(sourceCluster).
WithCluster(targetCluster).
WithTopology("source-cluster", "target-cluster").
Build()
// Update replicate configuration
err = cli.UpdateReplicateConfiguration(ctx, config)
if err != nil {
log.Printf("Failed to update replicate configuration: %v", err)
return
}
log.Println("Replicate configuration updated successfully")
}
func ExampleClient_GetReplicateInfo() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: exampleMilvusAddr,
})
if err != nil {
log.Fatal(err)
}
// Get replicate information for a specific source cluster
resp, err := cli.GetReplicateInfo(ctx, "source-cluster")
if err != nil {
log.Printf("Failed to get replicate information: %v", err)
return
}
log.Printf("Replicate information retrieved successfully, checkpoint count: %d", len(resp.GetCheckpoints()))
}
func ExampleClient_CreateReplicateStream() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: exampleMilvusAddr,
})
if err != nil {
log.Fatal(err)
}
// Create replicate stream
stream, err := cli.CreateReplicateStream(ctx)
if err != nil {
log.Printf("Failed to create replicate stream: %v", err)
return
}
defer stream.CloseSend()
log.Println("Replicate stream created successfully")
// Here you can continue to use the stream for data transmission
// For example: stream.Send(&milvuspb.ReplicateMessage{...})
}

44
cmd/components/cdc.go Normal file
View File

@ -0,0 +1,44 @@
// 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 components
import (
"context"
"github.com/milvus-io/milvus/internal/distributed/cdc"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
type CDC struct {
*cdc.Server
}
// NewCDC creates a new CDC
func NewCDC(ctx context.Context, factory dependency.Factory) (*CDC, error) {
svr, err := cdc.NewServer(ctx)
if err != nil {
return nil, err
}
return &CDC{
Server: svr,
}, nil
}
func (c *CDC) GetName() string {
return typeutil.CDCRole
}

View File

@ -147,6 +147,7 @@ func GetMilvusRoles(args []string, flags *flag.FlagSet) *roles.MilvusRoles {
role.EnableQueryNode = true
role.EnableDataNode = true
role.EnableStreamingNode = true
role.EnableCDC = true
role.Local = true
role.Embedded = serverType == typeutil.EmbeddedRole
case typeutil.MixCoordRole:
@ -164,6 +165,9 @@ func GetMilvusRoles(args []string, flags *flag.FlagSet) *roles.MilvusRoles {
role.EnableQueryNode = true
streamingutil.EnableEmbededQueryNode()
}
case typeutil.CDCRole:
role.EnableCDC = true
default:
fmt.Fprintf(os.Stderr, "Unknown server type = %s\n%s", serverType, getHelp())
os.Exit(-1)

View File

@ -48,6 +48,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/metrics"
rocksmqimpl "github.com/milvus-io/milvus/pkg/v2/mq/mqimpl/rocksmq/server"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/tracer"
"github.com/milvus-io/milvus/pkg/v2/util/etcd"
"github.com/milvus-io/milvus/pkg/v2/util/expr"
@ -73,7 +74,7 @@ func init() {
// stopRocksmqIfUsed closes the RocksMQ if it is used.
func stopRocksmqIfUsed() {
if name := util.MustSelectWALName(); name == util.WALTypeRocksmq {
if name := util.MustSelectWALName(); name == message.WALNameRocksmq {
rocksmqimpl.CloseRocksMQ()
}
}
@ -149,6 +150,7 @@ type MilvusRoles struct {
EnableRootCoord bool `env:"ENABLE_ROOT_COORD"`
EnableQueryCoord bool `env:"ENABLE_QUERY_COORD"`
EnableDataCoord bool `env:"ENABLE_DATA_COORD"`
EnableCDC bool `env:"ENABLE_CDC"`
Local bool
Alias string
Embedded bool
@ -207,6 +209,11 @@ func (mr *MilvusRoles) runDataNode(ctx context.Context, localMsg bool, wg *sync.
return runComponent(ctx, localMsg, wg, components.NewDataNode, metrics.RegisterDataNode)
}
func (mr *MilvusRoles) runCDC(ctx context.Context, localMsg bool, wg *sync.WaitGroup) component {
wg.Add(1)
return runComponent(ctx, localMsg, wg, components.NewCDC, metrics.RegisterCDC)
}
func (mr *MilvusRoles) setupLogger() {
params := paramtable.Get()
logConfig := log.Config{
@ -371,6 +378,7 @@ func (mr *MilvusRoles) Run() {
mr.EnableRootCoord,
mr.EnableQueryCoord,
mr.EnableDataCoord,
mr.EnableCDC,
}
enableComponents = lo.Filter(enableComponents, func(v bool, _ int) bool {
return v
@ -400,7 +408,7 @@ func (mr *MilvusRoles) Run() {
componentMap := make(map[string]component)
var mixCoord component
var proxy, dataNode, queryNode, streamingNode component
var proxy, dataNode, queryNode, streamingNode, cdc component
if (mr.EnableRootCoord && mr.EnableDataCoord && mr.EnableQueryCoord) || mr.EnableMixCoord {
paramtable.SetLocalComponentEnabled(typeutil.MixCoordRole)
@ -433,6 +441,12 @@ func (mr *MilvusRoles) Run() {
componentMap[typeutil.StreamingNodeRole] = streamingNode
}
if mr.EnableCDC {
paramtable.SetLocalComponentEnabled(typeutil.CDCRole)
cdc = mr.runCDC(ctx, local, &wg)
componentMap[typeutil.CDCRole] = cdc
}
wg.Wait()
http.RegisterStopComponent(func(role string) error {
@ -503,7 +517,7 @@ func (mr *MilvusRoles) Run() {
log.Info("All coordinators have stopped")
// stop nodes
nodes := []component{streamingNode, queryNode, dataNode}
nodes := []component{streamingNode, queryNode, dataNode, cdc}
stopNodeWG := &sync.WaitGroup{}
for _, node := range nodes {
if node != nil {
@ -546,5 +560,8 @@ func (mr *MilvusRoles) GetRoles() []string {
if mr.EnableDataNode {
roles = append(roles, typeutil.DataNodeRole)
}
if mr.EnableCDC {
roles = append(roles, typeutil.CDCRole)
}
return roles
}

7
go.mod
View File

@ -58,6 +58,7 @@ require (
require (
cloud.google.com/go/storage v1.50.0
github.com/antlr4-go/antlr/v4 v4.13.1
github.com/apache/pulsar-client-go v0.15.1
github.com/aws/aws-sdk-go-v2 v1.36.3
github.com/aws/aws-sdk-go-v2/config v1.29.9
github.com/aws/aws-sdk-go-v2/credentials v1.17.62
@ -72,7 +73,8 @@ require (
github.com/hashicorp/golang-lru/v2 v2.0.7
github.com/jolestar/go-commons-pool/v2 v2.1.2
github.com/magiconair/properties v1.8.7
github.com/milvus-io/milvus/pkg/v2 v2.0.0-00010101000000-000000000000
github.com/milvus-io/milvus/client/v2 v2.6.0
github.com/milvus-io/milvus/pkg/v2 v2.5.7
github.com/pkg/errors v0.9.1
github.com/remeh/sizedwaitgroup v1.0.0
github.com/shirou/gopsutil/v4 v4.24.10
@ -93,6 +95,8 @@ require (
github.com/smartystreets/goconvey v1.7.2 // indirect
)
replace github.com/milvus-io/milvus/client/v2 => ./client
require (
cel.dev/expr v0.19.1 // indirect
cloud.google.com/go v0.118.3 // indirect
@ -116,7 +120,6 @@ require (
github.com/alibabacloud-go/debug v0.0.0-20190504072949-9472017b5c68 // indirect
github.com/alibabacloud-go/tea v1.1.8 // indirect
github.com/andybalholm/brotli v1.1.0 // indirect
github.com/apache/pulsar-client-go v0.15.1 // indirect
github.com/apache/thrift v0.20.0 // indirect
github.com/ardielle/ardielle-go v1.5.2 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 // indirect

4
go.sum
View File

@ -786,10 +786,6 @@ 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.1 h1:NLoSWXvlJD8t91G3CUsooXqYnm5nfsBngztQYYT58V0=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.1/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250903074146-fe93017b6822 h1:aLZGg4Hfy37RDTqMXsVAcqv+CqzpRtN13tLfgcGK1nQ=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250903074146-fe93017b6822/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250903080546-f1a74984d9e4 h1:9hoFUhw6wVRGBZFPDNjEaUihvZy5DEOV+SxRIWLEbRM=
github.com/milvus-io/milvus-proto/go-api/v2 v2.6.2-0.20250903080546-f1a74984d9e4/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=

View File

@ -9,6 +9,7 @@ packages:
github.com/milvus-io/milvus/internal/distributed/streaming:
interfaces:
WALAccesser:
ReplicateService:
Utility:
Broadcast:
Local:
@ -18,6 +19,7 @@ packages:
Balancer:
github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster:
interfaces:
Broadcaster:
AppendOperator:
Watcher:
github.com/milvus-io/milvus/internal/streamingcoord/client:
@ -78,6 +80,7 @@ packages:
Manager:
github.com/milvus-io/milvus/internal/metastore:
interfaces:
ReplicationCatalog:
StreamingCoordCataLog:
StreamingNodeCataLog:
github.com/milvus-io/milvus/internal/util/segcore:

View File

@ -0,0 +1,24 @@
quiet: False
with-expecter: True
inpackage: True
filename: "mock_{{.InterfaceNameSnake}}.go"
mockname: "Mock{{.InterfaceName}}"
outpkg: "{{.PackageName}}"
dir: "{{.InterfaceDir}}"
packages:
github.com/milvus-io/milvus/internal/cdc/cluster:
interfaces:
ClusterClient:
MilvusClient:
github.com/milvus-io/milvus/internal/cdc/replication:
interfaces:
ReplicateManagerClient:
github.com/milvus-io/milvus/internal/cdc/replication/replicatestream:
interfaces:
ReplicateStreamClient:
github.com/milvus-io/milvus/internal/cdc/replication/replicatemanager:
interfaces:
ChannelReplicator:
github.com/milvus-io/milvus/internal/cdc/controller:
interfaces:
Controller:

View File

@ -0,0 +1,51 @@
// 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 cluster
import (
"context"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/client/v2/milvusclient"
"github.com/milvus-io/milvus/pkg/v2/log"
)
type ClusterClient interface {
CreateMilvusClient(ctx context.Context, cluster *commonpb.MilvusCluster) (MilvusClient, error)
}
var _ ClusterClient = (*clusterClient)(nil)
type clusterClient struct{}
func NewClusterClient() ClusterClient {
return &clusterClient{}
}
func (c *clusterClient) CreateMilvusClient(ctx context.Context, cluster *commonpb.MilvusCluster) (MilvusClient, error) {
cli, err := milvusclient.New(ctx, &milvusclient.ClientConfig{
Address: cluster.GetConnectionParam().GetUri(),
APIKey: cluster.GetConnectionParam().GetToken(),
})
if err != nil {
log.Warn("failed to create milvus client", zap.Error(err))
return nil, err
}
return cli, nil
}

View File

@ -0,0 +1,34 @@
// 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 cluster
import (
"context"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
)
type MilvusClient interface {
// GetReplicateInfo gets the replicate information from the milvus cluster.
GetReplicateInfo(ctx context.Context, sourceClusterID string, opts ...grpc.CallOption) (*milvuspb.GetReplicateInfoResponse, error)
// CreateReplicateStream creates a replicate stream to the milvus cluster.
CreateReplicateStream(ctx context.Context, opts ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error)
// Close closes the milvus client.
Close(ctx context.Context) error
}

View File

@ -0,0 +1,97 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package cluster
import (
context "context"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
mock "github.com/stretchr/testify/mock"
)
// MockClusterClient is an autogenerated mock type for the ClusterClient type
type MockClusterClient struct {
mock.Mock
}
type MockClusterClient_Expecter struct {
mock *mock.Mock
}
func (_m *MockClusterClient) EXPECT() *MockClusterClient_Expecter {
return &MockClusterClient_Expecter{mock: &_m.Mock}
}
// CreateMilvusClient provides a mock function with given fields: ctx, _a1
func (_m *MockClusterClient) CreateMilvusClient(ctx context.Context, _a1 *commonpb.MilvusCluster) (MilvusClient, error) {
ret := _m.Called(ctx, _a1)
if len(ret) == 0 {
panic("no return value specified for CreateMilvusClient")
}
var r0 MilvusClient
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *commonpb.MilvusCluster) (MilvusClient, error)); ok {
return rf(ctx, _a1)
}
if rf, ok := ret.Get(0).(func(context.Context, *commonpb.MilvusCluster) MilvusClient); ok {
r0 = rf(ctx, _a1)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(MilvusClient)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *commonpb.MilvusCluster) error); ok {
r1 = rf(ctx, _a1)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockClusterClient_CreateMilvusClient_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateMilvusClient'
type MockClusterClient_CreateMilvusClient_Call struct {
*mock.Call
}
// CreateMilvusClient is a helper method to define mock.On call
// - ctx context.Context
// - _a1 *commonpb.MilvusCluster
func (_e *MockClusterClient_Expecter) CreateMilvusClient(ctx interface{}, _a1 interface{}) *MockClusterClient_CreateMilvusClient_Call {
return &MockClusterClient_CreateMilvusClient_Call{Call: _e.mock.On("CreateMilvusClient", ctx, _a1)}
}
func (_c *MockClusterClient_CreateMilvusClient_Call) Run(run func(ctx context.Context, _a1 *commonpb.MilvusCluster)) *MockClusterClient_CreateMilvusClient_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*commonpb.MilvusCluster))
})
return _c
}
func (_c *MockClusterClient_CreateMilvusClient_Call) Return(_a0 MilvusClient, _a1 error) *MockClusterClient_CreateMilvusClient_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockClusterClient_CreateMilvusClient_Call) RunAndReturn(run func(context.Context, *commonpb.MilvusCluster) (MilvusClient, error)) *MockClusterClient_CreateMilvusClient_Call {
_c.Call.Return(run)
return _c
}
// NewMockClusterClient creates a new instance of MockClusterClient. 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 NewMockClusterClient(t interface {
mock.TestingT
Cleanup(func())
}) *MockClusterClient {
mock := &MockClusterClient{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,233 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package cluster
import (
context "context"
grpc "google.golang.org/grpc"
milvuspb "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
mock "github.com/stretchr/testify/mock"
)
// MockMilvusClient is an autogenerated mock type for the MilvusClient type
type MockMilvusClient struct {
mock.Mock
}
type MockMilvusClient_Expecter struct {
mock *mock.Mock
}
func (_m *MockMilvusClient) EXPECT() *MockMilvusClient_Expecter {
return &MockMilvusClient_Expecter{mock: &_m.Mock}
}
// Close provides a mock function with given fields: ctx
func (_m *MockMilvusClient) Close(ctx context.Context) error {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for Close")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context) error); ok {
r0 = rf(ctx)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockMilvusClient_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockMilvusClient_Close_Call struct {
*mock.Call
}
// Close is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockMilvusClient_Expecter) Close(ctx interface{}) *MockMilvusClient_Close_Call {
return &MockMilvusClient_Close_Call{Call: _e.mock.On("Close", ctx)}
}
func (_c *MockMilvusClient_Close_Call) Run(run func(ctx context.Context)) *MockMilvusClient_Close_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockMilvusClient_Close_Call) Return(_a0 error) *MockMilvusClient_Close_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockMilvusClient_Close_Call) RunAndReturn(run func(context.Context) error) *MockMilvusClient_Close_Call {
_c.Call.Return(run)
return _c
}
// CreateReplicateStream provides a mock function with given fields: ctx, opts
func (_m *MockMilvusClient) CreateReplicateStream(ctx context.Context, opts ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error) {
_va := make([]interface{}, len(opts))
for _i := range opts {
_va[_i] = opts[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for CreateReplicateStream")
}
var r0 milvuspb.MilvusService_CreateReplicateStreamClient
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error)); ok {
return rf(ctx, opts...)
}
if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) milvuspb.MilvusService_CreateReplicateStreamClient); ok {
r0 = rf(ctx, opts...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(milvuspb.MilvusService_CreateReplicateStreamClient)
}
}
if rf, ok := ret.Get(1).(func(context.Context, ...grpc.CallOption) error); ok {
r1 = rf(ctx, opts...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockMilvusClient_CreateReplicateStream_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateReplicateStream'
type MockMilvusClient_CreateReplicateStream_Call struct {
*mock.Call
}
// CreateReplicateStream is a helper method to define mock.On call
// - ctx context.Context
// - opts ...grpc.CallOption
func (_e *MockMilvusClient_Expecter) CreateReplicateStream(ctx interface{}, opts ...interface{}) *MockMilvusClient_CreateReplicateStream_Call {
return &MockMilvusClient_CreateReplicateStream_Call{Call: _e.mock.On("CreateReplicateStream",
append([]interface{}{ctx}, opts...)...)}
}
func (_c *MockMilvusClient_CreateReplicateStream_Call) Run(run func(ctx context.Context, opts ...grpc.CallOption)) *MockMilvusClient_CreateReplicateStream_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]grpc.CallOption, len(args)-1)
for i, a := range args[1:] {
if a != nil {
variadicArgs[i] = a.(grpc.CallOption)
}
}
run(args[0].(context.Context), variadicArgs...)
})
return _c
}
func (_c *MockMilvusClient_CreateReplicateStream_Call) Return(_a0 milvuspb.MilvusService_CreateReplicateStreamClient, _a1 error) *MockMilvusClient_CreateReplicateStream_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockMilvusClient_CreateReplicateStream_Call) RunAndReturn(run func(context.Context, ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error)) *MockMilvusClient_CreateReplicateStream_Call {
_c.Call.Return(run)
return _c
}
// GetReplicateInfo provides a mock function with given fields: ctx, sourceClusterID, opts
func (_m *MockMilvusClient) GetReplicateInfo(ctx context.Context, sourceClusterID string, opts ...grpc.CallOption) (*milvuspb.GetReplicateInfoResponse, error) {
_va := make([]interface{}, len(opts))
for _i := range opts {
_va[_i] = opts[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, sourceClusterID)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for GetReplicateInfo")
}
var r0 *milvuspb.GetReplicateInfoResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string, ...grpc.CallOption) (*milvuspb.GetReplicateInfoResponse, error)); ok {
return rf(ctx, sourceClusterID, opts...)
}
if rf, ok := ret.Get(0).(func(context.Context, string, ...grpc.CallOption) *milvuspb.GetReplicateInfoResponse); ok {
r0 = rf(ctx, sourceClusterID, opts...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*milvuspb.GetReplicateInfoResponse)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string, ...grpc.CallOption) error); ok {
r1 = rf(ctx, sourceClusterID, opts...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockMilvusClient_GetReplicateInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateInfo'
type MockMilvusClient_GetReplicateInfo_Call struct {
*mock.Call
}
// GetReplicateInfo is a helper method to define mock.On call
// - ctx context.Context
// - sourceClusterID string
// - opts ...grpc.CallOption
func (_e *MockMilvusClient_Expecter) GetReplicateInfo(ctx interface{}, sourceClusterID interface{}, opts ...interface{}) *MockMilvusClient_GetReplicateInfo_Call {
return &MockMilvusClient_GetReplicateInfo_Call{Call: _e.mock.On("GetReplicateInfo",
append([]interface{}{ctx, sourceClusterID}, opts...)...)}
}
func (_c *MockMilvusClient_GetReplicateInfo_Call) Run(run func(ctx context.Context, sourceClusterID string, opts ...grpc.CallOption)) *MockMilvusClient_GetReplicateInfo_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]grpc.CallOption, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(grpc.CallOption)
}
}
run(args[0].(context.Context), args[1].(string), variadicArgs...)
})
return _c
}
func (_c *MockMilvusClient_GetReplicateInfo_Call) Return(_a0 *milvuspb.GetReplicateInfoResponse, _a1 error) *MockMilvusClient_GetReplicateInfo_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockMilvusClient_GetReplicateInfo_Call) RunAndReturn(run func(context.Context, string, ...grpc.CallOption) (*milvuspb.GetReplicateInfoResponse, error)) *MockMilvusClient_GetReplicateInfo_Call {
_c.Call.Return(run)
return _c
}
// NewMockMilvusClient creates a new instance of MockMilvusClient. 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 NewMockMilvusClient(t interface {
mock.TestingT
Cleanup(func())
}) *MockMilvusClient {
mock := &MockMilvusClient{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,24 @@
// 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 controller
// Controller controls and schedules the CDC process.
// It will periodically update the replications by the replicate configuration.
type Controller interface {
Start()
Stop()
}

View File

@ -0,0 +1,83 @@
// 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 controllerimpl
import (
"context"
"sync"
"time"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/pkg/v2/log"
)
const checkInterval = 10 * time.Second
type controller struct {
ctx context.Context
wg sync.WaitGroup
stopOnce sync.Once
stopChan chan struct{}
}
func NewController() *controller {
return &controller{
ctx: context.Background(),
stopChan: make(chan struct{}),
}
}
func (c *controller) Start() {
log.Ctx(c.ctx).Info("CDC controller started")
c.wg.Add(1)
go func() {
defer c.wg.Done()
timer := time.NewTicker(checkInterval)
defer timer.Stop()
for {
select {
case <-c.stopChan:
return
case <-timer.C:
c.run()
}
}
}()
}
func (c *controller) Stop() {
c.stopOnce.Do(func() {
log.Ctx(c.ctx).Info("CDC controller stopping...")
// TODO: sheep, gracefully stop the replicators
close(c.stopChan)
c.wg.Wait()
log.Ctx(c.ctx).Info("CDC controller stopped")
})
}
func (c *controller) run() {
replicatePChannels, err := resource.Resource().ReplicationCatalog().ListReplicatePChannels(c.ctx)
if err != nil {
log.Ctx(c.ctx).Error("failed to get replicate pchannels", zap.Error(err))
return
}
for _, replicatePChannel := range replicatePChannels {
resource.Resource().ReplicateManagerClient().CreateReplicator(replicatePChannel)
}
}

View File

@ -0,0 +1,83 @@
// 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 controllerimpl
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/cdc/replication"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/internal/mocks/mock_metastore"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
)
func TestController_StartAndStop(t *testing.T) {
mockReplicateManagerClient := replication.NewMockReplicateManagerClient(t)
resource.InitForTest(t,
resource.OptReplicateManagerClient(mockReplicateManagerClient),
)
ctrl := NewController()
assert.NotPanics(t, func() {
ctrl.Start()
})
assert.NotPanics(t, func() {
ctrl.Stop()
})
}
func TestController_Run(t *testing.T) {
mockReplicateManagerClient := replication.NewMockReplicateManagerClient(t)
replicatePChannels := []*streamingpb.ReplicatePChannelMeta{
{
SourceChannelName: "test-source-channel-1",
TargetChannelName: "test-target-channel-1",
},
}
mockReplicationCatalog := mock_metastore.NewMockReplicationCatalog(t)
mockReplicationCatalog.EXPECT().ListReplicatePChannels(mock.Anything).Return(replicatePChannels, nil)
mockReplicateManagerClient.EXPECT().CreateReplicator(replicatePChannels[0]).Return()
resource.InitForTest(t,
resource.OptReplicateManagerClient(mockReplicateManagerClient),
resource.OptReplicationCatalog(mockReplicationCatalog),
)
ctrl := NewController()
ctrl.Start()
defer ctrl.Stop()
ctrl.run()
}
func TestController_RunError(t *testing.T) {
mockReplicateManagerClient := replication.NewMockReplicateManagerClient(t)
mockReplicationCatalog := mock_metastore.NewMockReplicationCatalog(t)
mockReplicationCatalog.EXPECT().ListReplicatePChannels(mock.Anything).Return(nil, assert.AnError)
resource.InitForTest(t,
resource.OptReplicateManagerClient(mockReplicateManagerClient),
resource.OptReplicationCatalog(mockReplicationCatalog),
)
ctrl := NewController()
ctrl.Start()
defer ctrl.Stop()
ctrl.run()
}

View File

@ -0,0 +1,96 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package controller
import mock "github.com/stretchr/testify/mock"
// MockController is an autogenerated mock type for the Controller type
type MockController struct {
mock.Mock
}
type MockController_Expecter struct {
mock *mock.Mock
}
func (_m *MockController) EXPECT() *MockController_Expecter {
return &MockController_Expecter{mock: &_m.Mock}
}
// Start provides a mock function with no fields
func (_m *MockController) Start() {
_m.Called()
}
// MockController_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
type MockController_Start_Call struct {
*mock.Call
}
// Start is a helper method to define mock.On call
func (_e *MockController_Expecter) Start() *MockController_Start_Call {
return &MockController_Start_Call{Call: _e.mock.On("Start")}
}
func (_c *MockController_Start_Call) Run(run func()) *MockController_Start_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockController_Start_Call) Return() *MockController_Start_Call {
_c.Call.Return()
return _c
}
func (_c *MockController_Start_Call) RunAndReturn(run func()) *MockController_Start_Call {
_c.Run(run)
return _c
}
// Stop provides a mock function with no fields
func (_m *MockController) Stop() {
_m.Called()
}
// MockController_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
type MockController_Stop_Call struct {
*mock.Call
}
// Stop is a helper method to define mock.On call
func (_e *MockController_Expecter) Stop() *MockController_Stop_Call {
return &MockController_Stop_Call{Call: _e.mock.On("Stop")}
}
func (_c *MockController_Stop_Call) Run(run func()) *MockController_Stop_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockController_Stop_Call) Return() *MockController_Stop_Call {
_c.Call.Return()
return _c
}
func (_c *MockController_Stop_Call) RunAndReturn(run func()) *MockController_Stop_Call {
_c.Run(run)
return _c
}
// NewMockController creates a new instance of MockController. 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 NewMockController(t interface {
mock.TestingT
Cleanup(func())
}) *MockController {
mock := &MockController{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,68 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package replication
import (
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
mock "github.com/stretchr/testify/mock"
)
// MockReplicateManagerClient is an autogenerated mock type for the ReplicateManagerClient type
type MockReplicateManagerClient struct {
mock.Mock
}
type MockReplicateManagerClient_Expecter struct {
mock *mock.Mock
}
func (_m *MockReplicateManagerClient) EXPECT() *MockReplicateManagerClient_Expecter {
return &MockReplicateManagerClient_Expecter{mock: &_m.Mock}
}
// CreateReplicator provides a mock function with given fields: replicateInfo
func (_m *MockReplicateManagerClient) CreateReplicator(replicateInfo *streamingpb.ReplicatePChannelMeta) {
_m.Called(replicateInfo)
}
// MockReplicateManagerClient_CreateReplicator_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateReplicator'
type MockReplicateManagerClient_CreateReplicator_Call struct {
*mock.Call
}
// CreateReplicator is a helper method to define mock.On call
// - replicateInfo *streamingpb.ReplicatePChannelMeta
func (_e *MockReplicateManagerClient_Expecter) CreateReplicator(replicateInfo interface{}) *MockReplicateManagerClient_CreateReplicator_Call {
return &MockReplicateManagerClient_CreateReplicator_Call{Call: _e.mock.On("CreateReplicator", replicateInfo)}
}
func (_c *MockReplicateManagerClient_CreateReplicator_Call) Run(run func(replicateInfo *streamingpb.ReplicatePChannelMeta)) *MockReplicateManagerClient_CreateReplicator_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(*streamingpb.ReplicatePChannelMeta))
})
return _c
}
func (_c *MockReplicateManagerClient_CreateReplicator_Call) Return() *MockReplicateManagerClient_CreateReplicator_Call {
_c.Call.Return()
return _c
}
func (_c *MockReplicateManagerClient_CreateReplicator_Call) RunAndReturn(run func(*streamingpb.ReplicatePChannelMeta)) *MockReplicateManagerClient_CreateReplicator_Call {
_c.Run(run)
return _c
}
// NewMockReplicateManagerClient creates a new instance of MockReplicateManagerClient. 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 NewMockReplicateManagerClient(t interface {
mock.TestingT
Cleanup(func())
}) *MockReplicateManagerClient {
mock := &MockReplicateManagerClient{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,25 @@
// 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 "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
// ReplicateManagerClient is the client that manages the replicate configuration.
type ReplicateManagerClient interface {
// CreateReplicator creates a new replicator for the replicate pchannel.
CreateReplicator(replicateInfo *streamingpb.ReplicatePChannelMeta)
}

View File

@ -0,0 +1,237 @@
// 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 replicatemanager
import (
"context"
"fmt"
"time"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/cdc/replication/replicatestream"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"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/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
const scannerHandlerChanSize = 64
// Replicator is the client that replicates the message to the channel in the target cluster.
type Replicator interface {
// StartReplicate starts the replicate for the channel.
StartReplicate()
// StopReplicate stops the replicate loop
// and wait for the loop to exit.
StopReplicate()
// GetState returns the current state of the replicator.
GetState() typeutil.LifetimeState
}
var _ Replicator = (*channelReplicator)(nil)
// channelReplicator is the implementation of ChannelReplicator.
type channelReplicator struct {
replicateInfo *streamingpb.ReplicatePChannelMeta
createRscFunc replicatestream.CreateReplicateStreamClientFunc
ctx context.Context
cancel context.CancelFunc
lifetime *typeutil.Lifetime
}
// NewChannelReplicator creates a new ChannelReplicator.
func NewChannelReplicator(replicateMeta *streamingpb.ReplicatePChannelMeta) Replicator {
ctx, cancel := context.WithCancel(context.Background())
createRscFunc := replicatestream.NewReplicateStreamClient
return &channelReplicator{
replicateInfo: replicateMeta,
createRscFunc: createRscFunc,
ctx: ctx,
cancel: cancel,
lifetime: typeutil.NewLifetime(),
}
}
func (r *channelReplicator) StartReplicate() {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
if !r.lifetime.Add(typeutil.LifetimeStateWorking) {
logger.Warn("replicate channel already started")
return
}
logger.Info("start replicate channel")
go func() {
defer r.lifetime.Done()
for {
err := r.replicateLoop()
if err != nil {
logger.Warn("replicate channel failed", zap.Error(err))
time.Sleep(10 * time.Second)
continue
}
break
}
logger.Info("stop replicate channel")
}()
}
// replicateLoop starts the replicate loop.
func (r *channelReplicator) replicateLoop() error {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
startFrom, err := r.getReplicateStartMessageID()
if err != nil {
return err
}
ch := make(adaptor.ChanMessageHandler, scannerHandlerChanSize)
var deliverPolicy options.DeliverPolicy
if startFrom == nil {
// No checkpoint found, seek from the earliest position
deliverPolicy = options.DeliverPolicyAll()
} else {
// Seek from the checkpoint
deliverPolicy = options.DeliverPolicyStartFrom(startFrom)
}
scanner := streaming.WAL().Read(r.ctx, streaming.ReadOption{
PChannel: r.replicateInfo.GetSourceChannelName(),
DeliverPolicy: deliverPolicy,
DeliverFilters: []options.DeliverFilter{},
MessageHandler: ch,
})
defer scanner.Close()
rsc := r.createRscFunc(r.ctx, r.replicateInfo)
defer rsc.Close()
logger.Info("start replicate channel loop", zap.Any("startFrom", startFrom))
for {
select {
case <-r.ctx.Done():
logger.Info("replicate channel stopped")
return nil
case msg := <-ch:
// TODO: Should be done at streamingnode.
if msg.MessageType().IsSelfControlled() {
logger.Debug("skip self-controlled message", log.FieldMessage(msg))
continue
}
err := rsc.Replicate(msg)
if err != nil {
panic(fmt.Sprintf("replicate message failed due to unrecoverable error: %v", err))
}
logger.Debug("replicate message success", log.FieldMessage(msg))
if msg.MessageType() == message.MessageTypePutReplicateConfig {
roleChanged := r.handlePutReplicateConfigMessage(msg)
if roleChanged {
// Role changed, return and stop replicate.
return nil
}
}
}
}
}
func (r *channelReplicator) getReplicateStartMessageID() (message.MessageID, error) {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
ctx, cancel := context.WithTimeout(r.ctx, 30*time.Second)
defer cancel()
milvusClient, err := resource.Resource().ClusterClient().CreateMilvusClient(ctx, r.replicateInfo.GetTargetCluster())
if err != nil {
return nil, err
}
defer milvusClient.Close(ctx)
sourceClusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
replicateInfo, err := milvusClient.GetReplicateInfo(ctx, sourceClusterID)
if err != nil {
return nil, err
}
var checkpoint *commonpb.ReplicateCheckpoint
for _, cp := range replicateInfo.GetCheckpoints() {
if cp.GetPchannel() == r.replicateInfo.GetSourceChannelName() {
checkpoint = cp
break
}
}
if checkpoint == nil || checkpoint.MessageId == nil {
logger.Info("channel not found in replicate info, will start from the beginning")
return nil, nil
}
startFrom := message.MustUnmarshalMessageID(checkpoint.GetMessageId())
logger.Info("replicate messages from position",
zap.Any("checkpoint", checkpoint),
zap.Any("startFromMessageID", startFrom),
)
return startFrom, nil
}
func (r *channelReplicator) handlePutReplicateConfigMessage(msg message.ImmutableMessage) (roleChanged bool) {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
logger.Info("handle PutReplicateConfigMessage", log.FieldMessage(msg))
prcMsg := message.MustAsImmutablePutReplicateConfigMessageV2(msg)
replicateConfig := prcMsg.Header().ReplicateConfiguration
currentClusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
currentCluster := replicateutil.MustNewConfigHelper(currentClusterID, replicateConfig).GetCurrentCluster()
if currentCluster.Role() == replicateutil.RolePrimary {
logger.Info("primary cluster, skip handle PutReplicateConfigMessage")
return false
}
// Current cluster role changed, not primary cluster,
// we need to remove the replicate pchannel.
err := resource.Resource().ReplicationCatalog().RemoveReplicatePChannel(r.ctx,
r.replicateInfo.GetSourceChannelName(), r.replicateInfo.GetTargetChannelName())
if err != nil {
panic(fmt.Sprintf("failed to remove replicate pchannel: %v", err))
}
logger.Info("handle PutReplicateConfigMessage done, replicate pchannel removed")
return true
}
func (r *channelReplicator) StopReplicate() {
r.lifetime.SetState(typeutil.LifetimeStateStopped)
r.cancel()
r.lifetime.Wait()
}
func (r *channelReplicator) GetState() typeutil.LifetimeState {
return r.lifetime.GetState()
}

View File

@ -0,0 +1,96 @@
// 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 replicatemanager
import (
"context"
"testing"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/stretchr/testify/assert"
mock "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/cdc/cluster"
"github.com/milvus-io/milvus/internal/cdc/replication/replicatestream"
"github.com/milvus-io/milvus/internal/cdc/resource"
"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/proto/streamingpb"
pulsar2 "github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/pulsar"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
func newMockPulsarMessageID() *commonpb.MessageID {
pulsarID := pulsar.EarliestMessageID()
msgID := pulsar2.NewPulsarID(pulsarID).Marshal()
return &commonpb.MessageID{
Id: msgID,
WALName: commonpb.WALName_Pulsar,
}
}
func TestChannelReplicator_StartReplicateChannel(t *testing.T) {
mockMilvusClient := cluster.NewMockMilvusClient(t)
mockMilvusClient.EXPECT().GetReplicateInfo(mock.Anything, mock.Anything).
Return(&milvuspb.GetReplicateInfoResponse{
Checkpoints: []*commonpb.ReplicateCheckpoint{
{
Pchannel: "test-source-channel",
MessageId: newMockPulsarMessageID(),
},
},
}, nil)
mockMilvusClient.EXPECT().Close(mock.Anything).Return(nil)
mockClusterClient := cluster.NewMockClusterClient(t)
mockClusterClient.EXPECT().CreateMilvusClient(mock.Anything, mock.Anything).
Return(mockMilvusClient, nil)
resource.InitForTest(t,
resource.OptClusterClient(mockClusterClient),
)
scanner := mock_streaming.NewMockScanner(t)
scanner.EXPECT().Close().Return()
wal := mock_streaming.NewMockWALAccesser(t)
wal.EXPECT().Read(mock.Anything, mock.Anything).Return(scanner)
streaming.SetWALForTest(wal)
rs := replicatestream.NewMockReplicateStreamClient(t)
rs.EXPECT().Close().Return()
cluster := &commonpb.MilvusCluster{ClusterId: "test-cluster"}
replicateInfo := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel",
TargetChannelName: "test-target-channel",
TargetCluster: cluster,
}
replicator := NewChannelReplicator(replicateInfo)
replicator.(*channelReplicator).createRscFunc = func(ctx context.Context,
replicateInfo *streamingpb.ReplicatePChannelMeta,
) replicatestream.ReplicateStreamClient {
return rs
}
assert.NotNil(t, replicator)
replicator.StartReplicate()
replicator.StopReplicate()
state := replicator.GetState()
assert.Equal(t, typeutil.LifetimeStateStopped, state)
}

View File

@ -0,0 +1,144 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package replicatemanager
import (
typeutil "github.com/milvus-io/milvus/pkg/v2/util/typeutil"
mock "github.com/stretchr/testify/mock"
)
// MockChannelReplicator is an autogenerated mock type for the ChannelReplicator type
type MockChannelReplicator struct {
mock.Mock
}
type MockChannelReplicator_Expecter struct {
mock *mock.Mock
}
func (_m *MockChannelReplicator) EXPECT() *MockChannelReplicator_Expecter {
return &MockChannelReplicator_Expecter{mock: &_m.Mock}
}
// GetState provides a mock function with no fields
func (_m *MockChannelReplicator) GetState() typeutil.LifetimeState {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetState")
}
var r0 typeutil.LifetimeState
if rf, ok := ret.Get(0).(func() typeutil.LifetimeState); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(typeutil.LifetimeState)
}
return r0
}
// MockChannelReplicator_GetState_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetState'
type MockChannelReplicator_GetState_Call struct {
*mock.Call
}
// GetState is a helper method to define mock.On call
func (_e *MockChannelReplicator_Expecter) GetState() *MockChannelReplicator_GetState_Call {
return &MockChannelReplicator_GetState_Call{Call: _e.mock.On("GetState")}
}
func (_c *MockChannelReplicator_GetState_Call) Run(run func()) *MockChannelReplicator_GetState_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockChannelReplicator_GetState_Call) Return(_a0 typeutil.LifetimeState) *MockChannelReplicator_GetState_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockChannelReplicator_GetState_Call) RunAndReturn(run func() typeutil.LifetimeState) *MockChannelReplicator_GetState_Call {
_c.Call.Return(run)
return _c
}
// StartReplicateChannel provides a mock function with no fields
func (_m *MockChannelReplicator) StartReplicateChannel() {
_m.Called()
}
// MockChannelReplicator_StartReplicateChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StartReplicateChannel'
type MockChannelReplicator_StartReplicateChannel_Call struct {
*mock.Call
}
// StartReplicateChannel is a helper method to define mock.On call
func (_e *MockChannelReplicator_Expecter) StartReplicateChannel() *MockChannelReplicator_StartReplicateChannel_Call {
return &MockChannelReplicator_StartReplicateChannel_Call{Call: _e.mock.On("StartReplicateChannel")}
}
func (_c *MockChannelReplicator_StartReplicateChannel_Call) Run(run func()) *MockChannelReplicator_StartReplicateChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockChannelReplicator_StartReplicateChannel_Call) Return() *MockChannelReplicator_StartReplicateChannel_Call {
_c.Call.Return()
return _c
}
func (_c *MockChannelReplicator_StartReplicateChannel_Call) RunAndReturn(run func()) *MockChannelReplicator_StartReplicateChannel_Call {
_c.Run(run)
return _c
}
// StopReplicateChannel provides a mock function with no fields
func (_m *MockChannelReplicator) StopReplicateChannel() {
_m.Called()
}
// MockChannelReplicator_StopReplicateChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StopReplicateChannel'
type MockChannelReplicator_StopReplicateChannel_Call struct {
*mock.Call
}
// StopReplicateChannel is a helper method to define mock.On call
func (_e *MockChannelReplicator_Expecter) StopReplicateChannel() *MockChannelReplicator_StopReplicateChannel_Call {
return &MockChannelReplicator_StopReplicateChannel_Call{Call: _e.mock.On("StopReplicateChannel")}
}
func (_c *MockChannelReplicator_StopReplicateChannel_Call) Run(run func()) *MockChannelReplicator_StopReplicateChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockChannelReplicator_StopReplicateChannel_Call) Return() *MockChannelReplicator_StopReplicateChannel_Call {
_c.Call.Return()
return _c
}
func (_c *MockChannelReplicator_StopReplicateChannel_Call) RunAndReturn(run func()) *MockChannelReplicator_StopReplicateChannel_Call {
_c.Run(run)
return _c
}
// NewMockChannelReplicator creates a new instance of MockChannelReplicator. 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 NewMockChannelReplicator(t interface {
mock.TestingT
Cleanup(func())
}) *MockChannelReplicator {
mock := &MockChannelReplicator{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,64 @@
// 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 replicatemanager
import (
"context"
"strings"
"go.uber.org/zap"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
// replicateManager is the implementation of ReplicateManagerClient.
type replicateManager struct {
ctx context.Context
// replicators is a map of replicate pchannel name to ChannelReplicator.
replicators map[string]Replicator
}
func NewReplicateManager() *replicateManager {
return &replicateManager{
ctx: context.Background(),
replicators: make(map[string]Replicator),
}
}
func (r *replicateManager) CreateReplicator(replicateInfo *streamingpb.ReplicatePChannelMeta) {
logger := log.With(
zap.String("sourceChannel", replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", replicateInfo.GetTargetChannelName()),
)
currentClusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
if !strings.Contains(replicateInfo.GetSourceChannelName(), currentClusterID) {
// current cluster is not source cluster, skip create replicator
return
}
_, ok := r.replicators[replicateInfo.GetSourceChannelName()]
if ok {
logger.Debug("replicator already exists, skip create replicator")
return
}
replicator := NewChannelReplicator(replicateInfo)
replicator.StartReplicate()
r.replicators[replicateInfo.GetSourceChannelName()] = replicator
logger.Info("created replicator for replicate pchannel")
}

View File

@ -0,0 +1,88 @@
// 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 replicatemanager
import (
"testing"
"github.com/stretchr/testify/assert"
mock "github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/cdc/cluster"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
func TestReplicateManager_CreateReplicator(t *testing.T) {
paramtable.Get().Save(paramtable.Get().CommonCfg.ClusterPrefix.Key, "test-source")
defer paramtable.Get().Reset(paramtable.Get().CommonCfg.ClusterPrefix.Key)
mockMilvusClient := cluster.NewMockMilvusClient(t)
mockMilvusClient.EXPECT().GetReplicateInfo(mock.Anything, mock.Anything).
Return(nil, assert.AnError).Maybe()
mockMilvusClient.EXPECT().Close(mock.Anything).Return(nil).Maybe()
mockClusterClient := cluster.NewMockClusterClient(t)
mockClusterClient.EXPECT().CreateMilvusClient(mock.Anything, mock.Anything).
Return(mockMilvusClient, nil).Maybe()
resource.InitForTest(t,
resource.OptClusterClient(mockClusterClient),
)
manager := NewReplicateManager()
// Test creating first replicator
replicateInfo := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel-1",
TargetChannelName: "test-target-channel-1",
TargetCluster: &commonpb.MilvusCluster{
ClusterId: "test-cluster-1",
},
}
manager.CreateReplicator(replicateInfo)
// Verify replicator was created
assert.Equal(t, 1, len(manager.replicators))
replicator, exists := manager.replicators["test-source-channel-1"]
assert.True(t, exists)
assert.NotNil(t, replicator)
// Test creating second replicator
replicateInfo2 := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel-2",
TargetChannelName: "test-target-channel-2",
TargetCluster: &commonpb.MilvusCluster{
ClusterId: "test-cluster-2",
},
}
manager.CreateReplicator(replicateInfo2)
// Verify second replicator was created
assert.Equal(t, 2, len(manager.replicators))
replicator2, exists := manager.replicators["test-source-channel-2"]
assert.True(t, exists)
assert.NotNil(t, replicator2)
// Verify first replicator still exists
replicator1, exists := manager.replicators["test-source-channel-1"]
assert.True(t, exists)
assert.NotNil(t, replicator1)
}

View File

@ -0,0 +1,121 @@
// 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 replicatestream
import (
"github.com/milvus-io/milvus/pkg/v2/metrics"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/util/timerecord"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
type ReplicateMetrics interface {
StartReplicate(msg message.ImmutableMessage)
OnSent(msg message.ImmutableMessage)
OnConfirmed(msg message.ImmutableMessage)
OnConnect()
OnDisconnect()
OnReconnect()
}
type msgMetrics struct {
tr *timerecord.TimeRecorder
}
type replicateMetrics struct {
replicateInfo *streamingpb.ReplicatePChannelMeta
msgsMetrics *typeutil.ConcurrentMap[string, msgMetrics] // message id -> msgMetrics
}
func NewReplicateMetrics(replicateInfo *streamingpb.ReplicatePChannelMeta) ReplicateMetrics {
return &replicateMetrics{
replicateInfo: replicateInfo,
msgsMetrics: typeutil.NewConcurrentMap[string, msgMetrics](),
}
}
func (m *replicateMetrics) StartReplicate(msg message.ImmutableMessage) {
msgID := msg.MessageID().String()
m.msgsMetrics.Insert(msgID, msgMetrics{
tr: timerecord.NewTimeRecorder("replicate_msg"),
})
}
func (m *replicateMetrics) OnSent(msg message.ImmutableMessage) {
sourceChannel := m.replicateInfo.GetSourceChannelName()
targetChannel := m.replicateInfo.GetTargetChannelName()
msgType := msg.MessageType().String()
metrics.CDCReplicatedMessagesTotal.WithLabelValues(
sourceChannel,
targetChannel,
msgType,
).Inc()
metrics.CDCReplicatedBytesTotal.WithLabelValues(
sourceChannel,
targetChannel,
msgType,
).Add(float64(msg.EstimateSize()))
}
func (m *replicateMetrics) OnConfirmed(msg message.ImmutableMessage) {
msgMetrics, ok := m.msgsMetrics.GetAndRemove(msg.MessageID().String())
if !ok {
return
}
replicateDuration := msgMetrics.tr.RecordSpan()
metrics.CDCReplicateEndToEndLatency.WithLabelValues(
m.replicateInfo.GetSourceChannelName(),
m.replicateInfo.GetTargetChannelName(),
).Observe(float64(replicateDuration.Milliseconds()))
}
func (m *replicateMetrics) OnConnect() {
metrics.CDCStreamRPCConnections.WithLabelValues(
m.replicateInfo.GetTargetCluster().GetClusterId(),
metrics.CDCStatusConnected,
).Inc()
}
func (m *replicateMetrics) OnDisconnect() {
clusterID := m.replicateInfo.GetTargetCluster().GetClusterId()
metrics.CDCStreamRPCConnections.WithLabelValues(
clusterID,
metrics.CDCStatusConnected,
).Dec()
metrics.CDCStreamRPCConnections.WithLabelValues(
clusterID,
metrics.CDCStatusDisconnected,
).Inc()
}
func (m *replicateMetrics) OnReconnect() {
clusterID := m.replicateInfo.GetTargetCluster().GetClusterId()
metrics.CDCStreamRPCConnections.WithLabelValues(
clusterID,
metrics.CDCStatusDisconnected,
).Dec()
metrics.CDCStreamRPCConnections.WithLabelValues(
clusterID,
metrics.CDCStatusConnected,
).Inc()
metrics.CDCStreamRPCReconnectTimes.WithLabelValues(
clusterID,
).Inc()
}

View File

@ -0,0 +1,113 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package replicatestream
import (
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
)
// MockReplicateStreamClient is an autogenerated mock type for the ReplicateStreamClient type
type MockReplicateStreamClient struct {
mock.Mock
}
type MockReplicateStreamClient_Expecter struct {
mock *mock.Mock
}
func (_m *MockReplicateStreamClient) EXPECT() *MockReplicateStreamClient_Expecter {
return &MockReplicateStreamClient_Expecter{mock: &_m.Mock}
}
// Close provides a mock function with no fields
func (_m *MockReplicateStreamClient) Close() {
_m.Called()
}
// MockReplicateStreamClient_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockReplicateStreamClient_Close_Call struct {
*mock.Call
}
// Close is a helper method to define mock.On call
func (_e *MockReplicateStreamClient_Expecter) Close() *MockReplicateStreamClient_Close_Call {
return &MockReplicateStreamClient_Close_Call{Call: _e.mock.On("Close")}
}
func (_c *MockReplicateStreamClient_Close_Call) Run(run func()) *MockReplicateStreamClient_Close_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockReplicateStreamClient_Close_Call) Return() *MockReplicateStreamClient_Close_Call {
_c.Call.Return()
return _c
}
func (_c *MockReplicateStreamClient_Close_Call) RunAndReturn(run func()) *MockReplicateStreamClient_Close_Call {
_c.Run(run)
return _c
}
// Replicate provides a mock function with given fields: msg
func (_m *MockReplicateStreamClient) Replicate(msg message.ImmutableMessage) error {
ret := _m.Called(msg)
if len(ret) == 0 {
panic("no return value specified for Replicate")
}
var r0 error
if rf, ok := ret.Get(0).(func(message.ImmutableMessage) error); ok {
r0 = rf(msg)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockReplicateStreamClient_Replicate_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Replicate'
type MockReplicateStreamClient_Replicate_Call struct {
*mock.Call
}
// Replicate is a helper method to define mock.On call
// - msg message.ImmutableMessage
func (_e *MockReplicateStreamClient_Expecter) Replicate(msg interface{}) *MockReplicateStreamClient_Replicate_Call {
return &MockReplicateStreamClient_Replicate_Call{Call: _e.mock.On("Replicate", msg)}
}
func (_c *MockReplicateStreamClient_Replicate_Call) Run(run func(msg message.ImmutableMessage)) *MockReplicateStreamClient_Replicate_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(message.ImmutableMessage))
})
return _c
}
func (_c *MockReplicateStreamClient_Replicate_Call) Return(_a0 error) *MockReplicateStreamClient_Replicate_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockReplicateStreamClient_Replicate_Call) RunAndReturn(run func(message.ImmutableMessage) error) *MockReplicateStreamClient_Replicate_Call {
_c.Call.Return(run)
return _c
}
// NewMockReplicateStreamClient creates a new instance of MockReplicateStreamClient. 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 NewMockReplicateStreamClient(t interface {
mock.TestingT
Cleanup(func())
}) *MockReplicateStreamClient {
mock := &MockReplicateStreamClient{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -0,0 +1,211 @@
// 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 replicatestream
import (
"context"
"sync"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
)
// MsgQueue exposes the required operations. We include context support to meet the
// blocking + cancel requirements.
// The names mirror the user's suggestion with small adjustments for clarity.
type MsgQueue interface {
// Enqueue appends a message. Blocks if capacity is full until space appears
// (via CleanupConfirmedMessages) or ctx is canceled.
Enqueue(ctx context.Context, msg message.ImmutableMessage) error
// Dequeue returns the next message from the current read cursor and advances
// the cursor by one. It does NOT delete the message from the queue storage.
// Blocks when there are no readable messages (i.e., cursor is at tail) until
// a new message is Enqueued or ctx is canceled.
Dequeue(ctx context.Context) (message.ImmutableMessage, error)
// SeekToHead moves the read cursor to the first not-yet-deleted message.
SeekToHead()
// CleanupConfirmedMessages permanently removes all messages whose Timetick()
// <= lastConfirmedTimeTick. This may free capacity and wake Enqueue waiters.
// Returns the messages that were cleaned up.
CleanupConfirmedMessages(lastConfirmedTimeTick uint64) []message.ImmutableMessage
// Len returns the number of stored (not yet deleted) messages.
Len() int
// Cap returns the maximum capacity of the queue.
Cap() int
}
var _ MsgQueue = (*msgQueue)(nil)
// msgQueue is a bounded, in-memory implementation of MsgQueue.
// It uses a simple slice for storage and an integer read cursor that is always
// within [0, len(buf)].
type msgQueue struct {
mu sync.Mutex
notEmpty *sync.Cond
notFull *sync.Cond
buf []message.ImmutableMessage
readIdx int
cap int
closed bool
}
// NewMsgQueue creates a queue with a fixed capacity (>0).
func NewMsgQueue(capacity int) *msgQueue {
if capacity <= 0 {
panic("capacity must be > 0")
}
q := &msgQueue{cap: capacity}
q.notEmpty = sync.NewCond(&q.mu)
q.notFull = sync.NewCond(&q.mu)
return q
}
func (q *msgQueue) Len() int {
q.mu.Lock()
defer q.mu.Unlock()
return len(q.buf)
}
func (q *msgQueue) Cap() int { return q.cap }
// Enqueue implements a blocking producer. It respects ctx cancellation.
func (q *msgQueue) Enqueue(ctx context.Context, msg message.ImmutableMessage) error {
q.mu.Lock()
defer q.mu.Unlock()
for len(q.buf) >= q.cap {
if ctx.Err() != nil {
return context.Canceled
}
q.waitWithContext(ctx, q.notFull)
if ctx.Err() != nil {
return context.Canceled
}
}
// Optional runtime check: enforce non-decreasing timetick
// if n := len(q.buf); n > 0 {
// if q.buf[n-1].Timetick() > msg.Timetick() {
// return fmt.Errorf("enqueue timetick order violation: last=%d new=%d", q.buf[n-1].Timetick(), msg.Timetick())
// }
// }
q.buf = append(q.buf, msg)
// New data is available for readers.
q.notEmpty.Signal()
return nil
}
// Dequeue returns the next message at the read cursor. Does not delete it.
func (q *msgQueue) Dequeue(ctx context.Context) (message.ImmutableMessage, error) {
q.mu.Lock()
defer q.mu.Unlock()
for q.readIdx >= len(q.buf) { // no readable messages
if ctx.Err() != nil {
return nil, context.Canceled
}
q.waitWithContext(ctx, q.notEmpty)
if ctx.Err() != nil {
return nil, context.Canceled
}
}
m := q.buf[q.readIdx]
q.readIdx++ // advance read cursor only; storage remains intact
// Readers advancing may not directly free capacity, but producers may still
// be waiting due to full buffer; signal in case cleanup made space earlier.
q.notFull.Signal()
return m, nil
}
// SeekToHead resets the read cursor to the first existing element.
func (q *msgQueue) SeekToHead() {
q.mu.Lock()
q.readIdx = 0
// Let potential consumers know there's readable data (if any exists).
if len(q.buf) > 0 {
q.notEmpty.Broadcast()
}
q.mu.Unlock()
}
// CleanupConfirmedMessages permanently drops messages with Timetick <= watermark.
// This frees capacity and may move the read cursor backward proportionally to the
// number of deleted messages (but clamped at 0).
// Returns the messages that were cleaned up.
func (q *msgQueue) CleanupConfirmedMessages(lastConfirmedTimeTick uint64) []message.ImmutableMessage {
q.mu.Lock()
defer q.mu.Unlock()
if len(q.buf) == 0 {
return nil
}
// Find first index whose Timetick() > lastConfirmedTimeTick
cut := 0
for cut < len(q.buf) && q.buf[cut].TimeTick() <= lastConfirmedTimeTick {
cut++
}
var cleanedMessages []message.ImmutableMessage
if cut > 0 {
// Collect the messages that will be cleaned up
cleanedMessages = make([]message.ImmutableMessage, cut)
copy(cleanedMessages, q.buf[:cut])
// Drop the prefix [0:cut)
q.buf = q.buf[cut:]
// Adjust read cursor relative to the new slice
q.readIdx -= cut
if q.readIdx < 0 {
q.readIdx = 0
}
// Free space became available; wake blocked producers.
q.notFull.Broadcast()
}
return cleanedMessages
}
// waitWithContext waits on cond while allowing ctx cancellation to wake it up.
// Implementation detail:
// - cond.Wait() requires holding q.mu; Wait atomically unlocks and re-locks.
// - We spawn a short-lived goroutine that Broadcasts on ctx.Done(). This is
// safe and bounded: after Wait returns, we close a local channel to let the
// goroutine exit immediately (even if ctx wasn't canceled).
func (q *msgQueue) waitWithContext(ctx context.Context, cond *sync.Cond) {
done := make(chan struct{})
go func() {
select {
case <-ctx.Done():
q.mu.Lock()
cond.Broadcast()
q.mu.Unlock()
case <-done:
}
}()
cond.Wait() // releases q.mu while waiting; re-locks before returning
close(done)
}

View File

@ -0,0 +1,287 @@
// 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 replicatestream
import (
"context"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/pkg/v2/mocks/streaming/util/mock_message"
)
func TestMsgQueue_BasicOperations(t *testing.T) {
// Test basic queue operations
queue := NewMsgQueue(3)
assert.Equal(t, 3, queue.Cap())
assert.Equal(t, 0, queue.Len())
// Test enqueue and dequeue
ctx := context.Background()
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
assert.Equal(t, 1, queue.Len())
// Test dequeue
dequeuedMsg, err := queue.Dequeue(ctx)
assert.NoError(t, err)
assert.Equal(t, msg1, dequeuedMsg)
assert.Equal(t, 1, queue.Len()) // Length doesn't change after dequeue
}
func TestMsgQueue_EnqueueBlocking(t *testing.T) {
// Test enqueue blocking when queue is full
queue := NewMsgQueue(2)
ctx := context.Background()
// Fill the queue
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
msg2 := mock_message.NewMockImmutableMessage(t)
msg2.EXPECT().TimeTick().Return(uint64(200)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg2)
assert.NoError(t, err)
assert.Equal(t, 2, queue.Len())
// Try to enqueue when full - should block
msg3 := mock_message.NewMockImmutableMessage(t)
msg3.EXPECT().TimeTick().Return(uint64(300)).Maybe()
// Use a context with timeout to test blocking
ctxWithTimeout, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer cancel()
err = queue.Enqueue(ctxWithTimeout, msg3)
assert.Error(t, err)
// Context timeout will cause context.Canceled error, not DeadlineExceeded
assert.Equal(t, context.Canceled, err)
}
func TestMsgQueue_DequeueBlocking(t *testing.T) {
// Test dequeue blocking when queue is empty
queue := NewMsgQueue(2)
ctx := context.Background()
// Try to dequeue from empty queue - should block
ctxWithTimeout, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
defer cancel()
_, err := queue.Dequeue(ctxWithTimeout)
assert.Error(t, err)
// Context timeout will cause context.Canceled error, not DeadlineExceeded
assert.Equal(t, context.Canceled, err)
}
func TestMsgQueue_SeekToHead(t *testing.T) {
// Test seek to head functionality
queue := NewMsgQueue(3)
ctx := context.Background()
// Add messages
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
msg2 := mock_message.NewMockImmutableMessage(t)
msg2.EXPECT().TimeTick().Return(uint64(200)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg2)
assert.NoError(t, err)
// Dequeue first message
dequeuedMsg, err := queue.Dequeue(ctx)
assert.NoError(t, err)
assert.Equal(t, msg1, dequeuedMsg)
// Seek to head
queue.SeekToHead()
// Should be able to dequeue first message again
dequeuedMsg, err = queue.Dequeue(ctx)
assert.NoError(t, err)
assert.Equal(t, msg1, dequeuedMsg)
}
func TestMsgQueue_CleanupConfirmedMessages(t *testing.T) {
// Test cleanup functionality
queue := NewMsgQueue(5)
ctx := context.Background()
// Add messages with different timeticks
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
msg2 := mock_message.NewMockImmutableMessage(t)
msg2.EXPECT().TimeTick().Return(uint64(200)).Maybe()
msg3 := mock_message.NewMockImmutableMessage(t)
msg3.EXPECT().TimeTick().Return(uint64(300)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg2)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg3)
assert.NoError(t, err)
assert.Equal(t, 3, queue.Len())
// Cleanup messages with timetick <= 200
cleanedMessages := queue.CleanupConfirmedMessages(200)
assert.Equal(t, 1, queue.Len())
assert.Equal(t, 2, len(cleanedMessages))
assert.Equal(t, msg1, cleanedMessages[0])
assert.Equal(t, msg2, cleanedMessages[1])
// First two messages should be removed
dequeuedMsg, err := queue.Dequeue(ctx)
assert.NoError(t, err)
assert.Equal(t, msg3, dequeuedMsg) // Only msg3 remains
}
func TestMsgQueue_CleanupWithReadCursor(t *testing.T) {
// Test cleanup when read cursor is advanced
queue := NewMsgQueue(5)
ctx := context.Background()
// Add messages
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
msg2 := mock_message.NewMockImmutableMessage(t)
msg2.EXPECT().TimeTick().Return(uint64(200)).Maybe()
msg3 := mock_message.NewMockImmutableMessage(t)
msg3.EXPECT().TimeTick().Return(uint64(300)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg2)
assert.NoError(t, err)
err = queue.Enqueue(ctx, msg3)
assert.NoError(t, err)
// Dequeue first message (advance read cursor)
dequeuedMsg, err := queue.Dequeue(ctx)
assert.NoError(t, err)
assert.Equal(t, msg1, dequeuedMsg)
assert.Equal(t, 1, queue.readIdx)
// Cleanup messages with timetick <= 150
cleanedMessages := queue.CleanupConfirmedMessages(150)
assert.Equal(t, 2, queue.Len()) // msg1 removed, msg2 and msg3 remain
assert.Equal(t, 0, queue.readIdx) // read cursor adjusted
assert.Equal(t, 1, len(cleanedMessages))
assert.Equal(t, msg1, cleanedMessages[0])
}
func TestMsgQueue_ContextCancellation(t *testing.T) {
// Test context cancellation
queue := NewMsgQueue(1)
ctx, cancel := context.WithCancel(context.Background())
// Fill the queue
msg1 := mock_message.NewMockImmutableMessage(t)
msg1.EXPECT().TimeTick().Return(uint64(100)).Maybe()
err := queue.Enqueue(ctx, msg1)
assert.NoError(t, err)
// Try to enqueue when full
msg2 := mock_message.NewMockImmutableMessage(t)
msg2.EXPECT().TimeTick().Return(uint64(200)).Maybe()
// Cancel context before enqueue
cancel()
err = queue.Enqueue(ctx, msg2)
assert.Error(t, err)
assert.Equal(t, context.Canceled, err)
}
func TestMsgQueue_NewMsgQueueValidation(t *testing.T) {
// Test constructor validation
assert.Panics(t, func() {
NewMsgQueue(0)
})
assert.Panics(t, func() {
NewMsgQueue(-1)
})
// Valid capacity
queue := NewMsgQueue(1)
assert.NotNil(t, queue)
assert.Equal(t, 1, queue.Cap())
}
func TestMsgQueue_ConcurrentOperations(t *testing.T) {
// Test concurrent enqueue and dequeue operations
queue := NewMsgQueue(10)
ctx := context.Background()
numMessages := 100
wg := sync.WaitGroup{}
// Start producer goroutine
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < numMessages; i++ {
msg := mock_message.NewMockImmutableMessage(t)
msg.EXPECT().TimeTick().Return(uint64(i)).Maybe()
err := queue.Enqueue(ctx, msg)
assert.NoError(t, err)
}
}()
// Start consumer goroutine
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < numMessages; i++ {
dequeuedMsg, err := queue.Dequeue(ctx)
assert.NoError(t, err)
cleanedMessages := queue.CleanupConfirmedMessages(dequeuedMsg.TimeTick())
assert.Equal(t, 1, len(cleanedMessages))
assert.Equal(t, dequeuedMsg, cleanedMessages[0])
}
}()
wg.Wait()
assert.Equal(t, 0, queue.Len())
}
func TestMsgQueue_EdgeCases(t *testing.T) {
// Test edge cases
queue := NewMsgQueue(1)
// Test with nil message (if allowed by interface)
// This depends on the actual message.ImmutableMessage interface implementation
// For now, we'll test with valid messages
// Test cleanup on empty queue
cleanedMessages := queue.CleanupConfirmedMessages(100)
assert.Equal(t, 0, queue.Len())
assert.Nil(t, cleanedMessages)
// Test seek to head on empty queue
queue.SeekToHead()
assert.Equal(t, 0, queue.readIdx)
}

View File

@ -0,0 +1,38 @@
// 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 replicatestream
import (
context "context"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
)
// ReplicateStreamClient is the client that replicates the message to the given cluster.
type ReplicateStreamClient interface {
// Replicate replicates the message to the target cluster.
// Replicate opeartion doesn't promise the message is delivered to the target cluster.
// It will cache the message in memory and retry until the message is delivered to the target cluster or the client is closed.
// Once the error is returned, the replicate operation will be unrecoverable.
Replicate(msg message.ImmutableMessage) error
// Stop stops the replicate operation.
Close()
}
type CreateReplicateStreamClientFunc func(ctx context.Context, replicateInfo *streamingpb.ReplicatePChannelMeta) ReplicateStreamClient

View File

@ -0,0 +1,293 @@
// 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 replicatestream
import (
"context"
"sync"
"time"
"github.com/cenkalti/backoff/v4"
"go.uber.org/zap"
"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/cdc/resource"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
const pendingMessageQueueLength = 128
// replicateStreamClient is the implementation of ReplicateStreamClient.
type replicateStreamClient struct {
replicateInfo *streamingpb.ReplicatePChannelMeta
clusterID string
client milvuspb.MilvusService_CreateReplicateStreamClient
pendingMessages MsgQueue
metrics ReplicateMetrics
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
// NewReplicateStreamClient creates a new ReplicateStreamClient.
func NewReplicateStreamClient(ctx context.Context, replicateInfo *streamingpb.ReplicatePChannelMeta) ReplicateStreamClient {
ctx1, cancel := context.WithCancel(ctx)
ctx1 = contextutil.WithClusterID(ctx1, replicateInfo.GetTargetCluster().GetClusterId())
rs := &replicateStreamClient{
clusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
replicateInfo: replicateInfo,
pendingMessages: NewMsgQueue(pendingMessageQueueLength),
metrics: NewReplicateMetrics(replicateInfo),
ctx: ctx1,
cancel: cancel,
}
rs.metrics.OnConnect()
go rs.startInternal()
return rs
}
func (r *replicateStreamClient) startInternal() {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
defer func() {
r.metrics.OnDisconnect()
logger.Info("replicate stream client closed")
}()
backoff := backoff.NewExponentialBackOff()
backoff.InitialInterval = 100 * time.Millisecond
backoff.MaxInterval = 10 * time.Second
backoff.MaxElapsedTime = 0
backoff.Reset()
disconnect := func(stopCh chan struct{}, err error) {
r.metrics.OnDisconnect()
close(stopCh)
r.client.CloseSend()
r.wg.Wait()
time.Sleep(backoff.NextBackOff())
log.Warn("restart replicate stream client", zap.Error(err))
}
for {
select {
case <-r.ctx.Done():
return
default:
milvusClient, err := resource.Resource().ClusterClient().CreateMilvusClient(r.ctx, r.replicateInfo.GetTargetCluster())
if err != nil {
logger.Warn("create milvus client failed, retry...", zap.Error(err))
time.Sleep(backoff.NextBackOff())
continue
}
client, err := milvusClient.CreateReplicateStream(r.ctx)
if err != nil {
logger.Warn("create milvus replicate stream failed, retry...", zap.Error(err))
time.Sleep(backoff.NextBackOff())
continue
}
logger.Info("replicate stream client service started")
// reset client and pending messages
if oldClient := r.client; oldClient != nil {
r.metrics.OnReconnect()
}
r.client = client
r.pendingMessages.SeekToHead()
stopCh := make(chan struct{})
sendErrCh := r.startSendLoop(stopCh)
recvErrCh := r.startRecvLoop(stopCh)
select {
case <-r.ctx.Done():
r.client.CloseSend()
r.wg.Wait()
return
case err := <-sendErrCh:
disconnect(stopCh, err)
case err := <-recvErrCh:
disconnect(stopCh, err)
}
}
}
}
// Replicate replicates the message to the target cluster.
func (r *replicateStreamClient) Replicate(msg message.ImmutableMessage) error {
select {
case <-r.ctx.Done():
return nil
default:
r.metrics.StartReplicate(msg)
r.pendingMessages.Enqueue(r.ctx, msg)
return nil
}
}
func (r *replicateStreamClient) startSendLoop(stopCh <-chan struct{}) <-chan error {
errCh := make(chan error, 1)
r.wg.Add(1)
go func() {
defer r.wg.Done()
errCh <- r.sendLoop(stopCh)
}()
return errCh
}
func (r *replicateStreamClient) startRecvLoop(stopCh <-chan struct{}) <-chan error {
errCh := make(chan error, 1)
r.wg.Add(1)
go func() {
defer r.wg.Done()
errCh <- r.recvLoop(stopCh)
}()
return errCh
}
func (r *replicateStreamClient) sendLoop(stopCh <-chan struct{}) error {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
for {
select {
case <-r.ctx.Done():
logger.Info("send loop closed by ctx done")
return nil
case <-stopCh:
logger.Info("send loop closed by stopCh")
return nil
default:
msg, err := r.pendingMessages.Dequeue(r.ctx)
if err != nil {
// context canceled, return nil
return nil
}
if msg.MessageType() == message.MessageTypeTxn {
txnMsg := message.AsImmutableTxnMessage(msg)
// send txn begin message
beginMsg := txnMsg.Begin()
err := r.sendMessage(beginMsg)
if err != nil {
return err
}
// send txn messages
err = txnMsg.RangeOver(func(msg message.ImmutableMessage) error {
err = r.sendMessage(msg)
if err != nil {
return err
}
return nil
})
if err != nil {
return err
}
// send txn commit message
commitMsg := txnMsg.Commit()
err = r.sendMessage(commitMsg)
if err != nil {
return err
}
continue
}
err = r.sendMessage(msg)
if err != nil {
return err
}
}
}
}
func (r *replicateStreamClient) sendMessage(msg message.ImmutableMessage) (err error) {
defer func() {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
if err != nil {
logger.Warn("send message failed", zap.Error(err), log.FieldMessage(msg))
} else {
r.metrics.OnSent(msg)
logger.Debug("send message success", log.FieldMessage(msg))
}
}()
immutableMessage := msg.IntoImmutableMessageProto()
req := &milvuspb.ReplicateRequest{
Request: &milvuspb.ReplicateRequest_ReplicateMessage{
ReplicateMessage: &milvuspb.ReplicateMessage{
SourceClusterId: r.clusterID,
Message: &commonpb.ImmutableMessage{
Id: msg.MessageID().IntoProto(),
Payload: immutableMessage.GetPayload(),
Properties: immutableMessage.GetProperties(),
},
},
},
}
return r.client.Send(req)
}
func (r *replicateStreamClient) recvLoop(stopCh <-chan struct{}) error {
logger := log.With(
zap.String("sourceChannel", r.replicateInfo.GetSourceChannelName()),
zap.String("targetChannel", r.replicateInfo.GetTargetChannelName()),
)
for {
select {
case <-r.ctx.Done():
logger.Info("recv loop closed by ctx done")
return nil
case <-stopCh:
logger.Info("recv loop closed by stopCh")
return nil
default:
resp, err := r.client.Recv()
if err != nil {
logger.Warn("replicate stream recv failed", zap.Error(err))
return err
}
lastConfirmedMessageInfo := resp.GetReplicateConfirmedMessageInfo()
if lastConfirmedMessageInfo != nil {
messages := r.pendingMessages.CleanupConfirmedMessages(lastConfirmedMessageInfo.GetConfirmedTimeTick())
for _, msg := range messages {
r.metrics.OnConfirmed(msg)
}
}
}
}
}
func (r *replicateStreamClient) Close() {
r.cancel()
r.wg.Wait()
}

View File

@ -0,0 +1,322 @@
// 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 replicatestream
import (
"context"
"strconv"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
milvuspb "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/cdc/cluster"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
mock_message "github.com/milvus-io/milvus/pkg/v2/mocks/streaming/util/mock_message"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
)
func TestReplicateStreamClient_Replicate(t *testing.T) {
ctx := context.Background()
targetCluster := &commonpb.MilvusCluster{
ClusterId: "test-cluster",
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
}
// Setup mocks
mockStreamClient := newMockReplicateStreamClient(t)
mockMilvusClient := cluster.NewMockMilvusClient(t)
mockMilvusClient.EXPECT().CreateReplicateStream(mock.Anything).Return(mockStreamClient, nil)
mockClusterClient := cluster.NewMockClusterClient(t)
mockClusterClient.EXPECT().CreateMilvusClient(mock.Anything, mock.Anything).
Return(mockMilvusClient, nil)
resource.InitForTest(t,
resource.OptClusterClient(mockClusterClient),
)
wal := mock_streaming.NewMockWALAccesser(t)
streaming.SetWALForTest(wal)
replicateInfo := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel",
TargetChannelName: "test-target-channel",
TargetCluster: targetCluster,
}
replicateClient := NewReplicateStreamClient(ctx, replicateInfo)
// Test Replicate method
const msgCount = pendingMessageQueueLength * 10
go func() {
for i := 0; i < msgCount; i++ {
mockMsg := mock_message.NewMockImmutableMessage(t)
tt := uint64(i + 1)
messageID := walimplstest.NewTestMessageID(int64(tt))
mockMsg.EXPECT().TimeTick().Return(tt)
mockMsg.EXPECT().EstimateSize().Return(1024)
mockMsg.EXPECT().MessageType().Return(message.MessageTypeInsert)
mockMsg.EXPECT().MessageID().Return(messageID)
mockMsg.EXPECT().IntoImmutableMessageProto().Return(&commonpb.ImmutableMessage{
Id: messageID.IntoProto(),
Payload: []byte("test-payload"),
Properties: map[string]string{"key": "value"},
})
mockMsg.EXPECT().MarshalLogObject(mock.Anything).Return(nil)
err := replicateClient.Replicate(mockMsg)
assert.NoError(t, err)
}
}()
// recv the confirm message
for i := 0; i < msgCount; i++ {
mockStreamClient.ExpectRecv()
}
assert.Equal(t, msgCount, mockStreamClient.GetRecvCount())
assert.Equal(t, 0, replicateClient.(*replicateStreamClient).pendingMessages.Len())
replicateClient.Close()
}
func TestReplicateStreamClient_Replicate_ContextCanceled(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
targetCluster := &commonpb.MilvusCluster{
ClusterId: "test-cluster",
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
}
// Setup mocks
mockStreamClient := newMockReplicateStreamClient(t)
mockMilvusClient := cluster.NewMockMilvusClient(t)
mockMilvusClient.EXPECT().CreateReplicateStream(mock.Anything).Return(mockStreamClient, nil).Maybe()
mockMilvusClient.EXPECT().Close(mock.Anything).Return(nil).Maybe()
mockClusterClient := cluster.NewMockClusterClient(t)
mockClusterClient.EXPECT().CreateMilvusClient(mock.Anything, mock.Anything).
Return(mockMilvusClient, nil).Maybe()
resource.InitForTest(t,
resource.OptClusterClient(mockClusterClient),
)
wal := mock_streaming.NewMockWALAccesser(t)
streaming.SetWALForTest(wal)
replicateInfo := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel",
TargetChannelName: "test-target-channel",
TargetCluster: targetCluster,
}
client := NewReplicateStreamClient(ctx, replicateInfo)
defer client.Close()
// Cancel context
cancel()
// Test Replicate method with canceled context
mockMsg := mock_message.NewMockImmutableMessage(t)
err := client.Replicate(mockMsg)
assert.NoError(t, err) // Should return nil when context is canceled
}
func TestReplicateStreamClient_Reconnect(t *testing.T) {
ctx := context.Background()
targetCluster := &commonpb.MilvusCluster{
ClusterId: "test-cluster",
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
}
const reconnectTimes = 3
reconnectCount := 0
// Setup mocks with error to trigger retry logic
mockStreamClient := newMockReplicateStreamClient(t)
mockMilvusClient := cluster.NewMockMilvusClient(t)
mockMilvusClient.EXPECT().CreateReplicateStream(mock.Anything).RunAndReturn(func(ctx context.Context, opts ...grpc.CallOption) (milvuspb.MilvusService_CreateReplicateStreamClient, error) {
reconnectCount++
if reconnectCount < reconnectTimes {
return nil, assert.AnError
}
return mockStreamClient, nil
}).Times(reconnectTimes) // expect to be called reconnectTimes times
mockClusterClient := cluster.NewMockClusterClient(t)
mockClusterClient.EXPECT().CreateMilvusClient(mock.Anything, mock.Anything).
Return(mockMilvusClient, nil)
resource.InitForTest(t,
resource.OptClusterClient(mockClusterClient),
)
wal := mock_streaming.NewMockWALAccesser(t)
streaming.SetWALForTest(wal)
// Create client which will start internal retry loop
replicateInfo := &streamingpb.ReplicatePChannelMeta{
SourceChannelName: "test-source-channel",
TargetChannelName: "test-target-channel",
TargetCluster: targetCluster,
}
replicateClient := NewReplicateStreamClient(ctx, replicateInfo)
// Replicate after reconnected
const msgCount = 100
go func() {
for i := 0; i < msgCount; i++ {
tt := uint64(i + 1)
mockMsg := mock_message.NewMockImmutableMessage(t)
mockMsg.EXPECT().TimeTick().Return(tt)
messageID := walimplstest.NewTestMessageID(int64(tt))
mockMsg.EXPECT().MessageID().Return(messageID)
mockMsg.EXPECT().EstimateSize().Return(1024)
mockMsg.EXPECT().MessageType().Return(message.MessageTypeInsert)
mockMsg.EXPECT().IntoImmutableMessageProto().Return(&commonpb.ImmutableMessage{
Id: messageID.IntoProto(),
Payload: []byte("test-payload"),
Properties: map[string]string{"key": "value"},
})
mockMsg.EXPECT().MarshalLogObject(mock.Anything).Return(nil)
err := replicateClient.Replicate(mockMsg)
assert.NoError(t, err)
}
}()
for i := 0; i < msgCount; i++ {
mockStreamClient.ExpectRecv()
}
assert.Equal(t, msgCount, mockStreamClient.GetRecvCount())
assert.Equal(t, 0, replicateClient.(*replicateStreamClient).pendingMessages.Len())
replicateClient.Close()
}
// mockReplicateStreamClient implements the milvuspb.MilvusService_CreateReplicateStreamClient interface
type mockReplicateStreamClient struct {
sendError error
recvError error
ch chan *milvuspb.ReplicateRequest
expectRecvCh chan struct{}
recvCount int
t *testing.T
timeout time.Duration
closeCh chan struct{}
}
func newMockReplicateStreamClient(t *testing.T) *mockReplicateStreamClient {
return &mockReplicateStreamClient{
ch: make(chan *milvuspb.ReplicateRequest, 128),
expectRecvCh: make(chan struct{}, 128),
t: t,
timeout: 10 * time.Second,
closeCh: make(chan struct{}, 1),
}
}
func (m *mockReplicateStreamClient) Send(req *milvuspb.ReplicateRequest) error {
if m.sendError != nil {
return m.sendError
}
m.ch <- req
return m.sendError
}
func (m *mockReplicateStreamClient) Recv() (*milvuspb.ReplicateResponse, error) {
if m.recvError != nil {
return nil, m.recvError
}
select {
case <-m.closeCh:
return nil, nil
case req := <-m.ch:
// use id as time tick in mock
timeTick, err := strconv.ParseUint(req.GetReplicateMessage().GetMessage().GetId().GetId(), 10, 64)
if err != nil {
panic(err)
}
m.expectRecvCh <- struct{}{}
return &milvuspb.ReplicateResponse{
Response: &milvuspb.ReplicateResponse_ReplicateConfirmedMessageInfo{
ReplicateConfirmedMessageInfo: &milvuspb.ReplicateConfirmedMessageInfo{
ConfirmedTimeTick: timeTick,
},
},
}, nil
case <-time.After(m.timeout):
assert.Fail(m.t, "recv timeout")
return nil, assert.AnError
}
}
func (m *mockReplicateStreamClient) ExpectRecv() {
select {
case <-m.expectRecvCh:
m.recvCount++
return
case <-time.After(m.timeout):
assert.Fail(m.t, "expect recv timeout")
}
}
func (m *mockReplicateStreamClient) GetRecvCount() int {
return m.recvCount
}
func (m *mockReplicateStreamClient) RecvMsg(msg interface{}) error {
return nil
}
func (m *mockReplicateStreamClient) SendMsg(msg interface{}) error {
return nil
}
func (m *mockReplicateStreamClient) Header() (metadata.MD, error) {
return nil, nil
}
func (m *mockReplicateStreamClient) Trailer() metadata.MD {
return nil
}
func (m *mockReplicateStreamClient) CloseSend() error {
close(m.closeCh)
return nil
}
func (m *mockReplicateStreamClient) Context() context.Context {
return context.Background()
}

View File

@ -0,0 +1,143 @@
// 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 resource
import (
"reflect"
"github.com/milvus-io/milvus/internal/cdc/cluster"
"github.com/milvus-io/milvus/internal/cdc/controller"
"github.com/milvus-io/milvus/internal/cdc/replication"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/kv/streamingcoord"
"github.com/milvus-io/milvus/pkg/v2/kv"
)
var r *resourceImpl // singleton resource instance
// optResourceInit is the option to initialize the resource.
type optResourceInit func(r *resourceImpl)
// OptMetaKV provides the meta kv to the resource.
func OptMetaKV(metaKV kv.MetaKv) optResourceInit {
return func(r *resourceImpl) {
r.metaKV = metaKV
}
}
// OptReplicateManagerClient provides the replicate manager client to the resource.
func OptReplicateManagerClient(replicateManagerClient replication.ReplicateManagerClient) optResourceInit {
return func(r *resourceImpl) {
r.replicateManagerClient = replicateManagerClient
}
}
// OptReplicationCatalog provides the replication catalog to the resource.
func OptReplicationCatalog(catalog metastore.ReplicationCatalog) optResourceInit {
return func(r *resourceImpl) {
r.catalog = catalog
}
}
// OptClusterClient provides the cluster client to the resource.
func OptClusterClient(clusterClient cluster.ClusterClient) optResourceInit {
return func(r *resourceImpl) {
r.clusterClient = clusterClient
}
}
// OptController provides the controller to the resource.
func OptController(controller controller.Controller) optResourceInit {
return func(r *resourceImpl) {
r.controller = controller
}
}
// Done finish all initialization of resources.
func Init(opts ...optResourceInit) {
newR := &resourceImpl{}
for _, opt := range opts {
opt(newR)
}
newR.catalog = streamingcoord.NewReplicationCatalog(newR.MetaKV())
newR.clusterClient = cluster.NewClusterClient()
assertNotNil(newR.MetaKV())
assertNotNil(newR.ReplicationCatalog())
assertNotNil(newR.ClusterClient())
assertNotNil(newR.ReplicateManagerClient())
assertNotNil(newR.Controller())
r = newR
}
// Release releases the singleton of resources.
func Release() {}
// Resource access the underlying singleton of resources.
func Resource() *resourceImpl {
return r
}
// resourceImpl is a basic resource dependency for streamingnode server.
// All utility on it is concurrent-safe and singleton.
type resourceImpl struct {
metaKV kv.MetaKv
catalog metastore.ReplicationCatalog
clusterClient cluster.ClusterClient
replicateManagerClient replication.ReplicateManagerClient
controller controller.Controller
}
// MetaKV returns the meta kv.
func (r *resourceImpl) MetaKV() kv.MetaKv {
return r.metaKV
}
// ReplicationCatalog returns the replication catalog.
func (r *resourceImpl) ReplicationCatalog() metastore.ReplicationCatalog {
return r.catalog
}
// ClusterClient returns the cluster client.
func (r *resourceImpl) ClusterClient() cluster.ClusterClient {
return r.clusterClient
}
// ReplicateManagerClient returns the replicate manager client.
func (r *resourceImpl) ReplicateManagerClient() replication.ReplicateManagerClient {
return r.replicateManagerClient
}
// Controller returns the controller.
func (r *resourceImpl) Controller() controller.Controller {
return r.controller
}
// assertNotNil panics if the resource is nil.
func assertNotNil(v interface{}) {
iv := reflect.ValueOf(v)
if !iv.IsValid() {
panic("nil resource")
}
switch iv.Kind() {
case reflect.Ptr, reflect.Slice, reflect.Map, reflect.Func, reflect.Interface:
if iv.IsNil() {
panic("nil resource")
}
}
}

View File

@ -0,0 +1,37 @@
//go:build test
// +build test
package resource
import (
"testing"
"github.com/milvus-io/milvus/internal/cdc/cluster"
"github.com/milvus-io/milvus/internal/cdc/controller"
"github.com/milvus-io/milvus/internal/cdc/replication"
"github.com/milvus-io/milvus/internal/mocks/mock_metastore"
"github.com/milvus-io/milvus/pkg/v2/mocks/mock_kv"
)
// InitForTest initializes the singleton of resources for test.
func InitForTest(t *testing.T, opts ...optResourceInit) {
r = &resourceImpl{}
for _, opt := range opts {
opt(r)
}
if r.metaKV == nil {
r.metaKV = mock_kv.NewMockMetaKv(t)
}
if r.catalog == nil {
r.catalog = mock_metastore.NewMockReplicationCatalog(t)
}
if r.clusterClient == nil {
r.clusterClient = cluster.NewMockClusterClient(t)
}
if r.replicateManagerClient == nil {
r.replicateManagerClient = replication.NewMockReplicateManagerClient(t)
}
if r.controller == nil {
r.controller = controller.NewMockController(t)
}
}

55
internal/cdc/server.go Normal file
View File

@ -0,0 +1,55 @@
// 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 cdc
import (
"context"
"github.com/milvus-io/milvus/internal/cdc/resource"
"github.com/milvus-io/milvus/pkg/v2/log"
)
type CDCServer struct {
ctx context.Context
}
// NewCDCServer will return a CDCServer.
func NewCDCServer(ctx context.Context) *CDCServer {
return &CDCServer{
ctx: ctx,
}
}
// Init initializes the CDCServer.
func (svr *CDCServer) Init() error {
log.Ctx(svr.ctx).Info("CDCServer init successfully")
return nil
}
// Start starts CDCServer.
func (svr *CDCServer) Start() error {
resource.Resource().Controller().Start()
log.Ctx(svr.ctx).Info("CDCServer start successfully")
return nil
}
// Stop stops CDCServer.
func (svr *CDCServer) Stop() error {
resource.Resource().Controller().Stop()
log.Ctx(svr.ctx).Info("CDCServer stop successfully")
return nil
}

View File

@ -259,8 +259,7 @@ func (_c *MockNodeManager_Startup_Call) RunAndReturn(run func(context.Context, [
func NewMockNodeManager(t interface {
mock.TestingT
Cleanup(func())
},
) *MockNodeManager {
}) *MockNodeManager {
mock := &MockNodeManager{}
mock.Mock.Test(t)

View File

@ -0,0 +1,194 @@
// 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 cdc
import (
"context"
"sync"
"time"
"github.com/tikv/client-go/v2/txnkv"
clientv3 "go.etcd.io/etcd/client/v3"
"go.uber.org/zap"
"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/cdc"
"github.com/milvus-io/milvus/internal/cdc/controller/controllerimpl"
"github.com/milvus-io/milvus/internal/cdc/replication/replicatemanager"
"github.com/milvus-io/milvus/internal/cdc/resource"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
tikvkv "github.com/milvus-io/milvus/internal/kv/tikv"
"github.com/milvus-io/milvus/internal/util/componentutil"
kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
"github.com/milvus-io/milvus/pkg/v2/kv"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/util"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/tikv"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
// Server is the server of cdc.
type Server struct {
ctx context.Context
cancel context.CancelFunc
metaKV kv.MetaKv
cdcServer *cdc.CDCServer
etcdCli *clientv3.Client
tikvCli *txnkv.Client
componentState *componentutil.ComponentStateService
stopOnce sync.Once
}
// NewServer create a new CDC server.
func NewServer(ctx context.Context) (*Server, error) {
ctx1, cancel := context.WithCancel(ctx)
return &Server{
ctx: ctx1,
cancel: cancel,
componentState: componentutil.NewComponentStateService(typeutil.CDCRole),
stopOnce: sync.Once{},
}, nil
}
func (s *Server) Prepare() error {
return nil
}
// Run runs the server.
func (s *Server) Run() error {
if err := s.init(); err != nil {
return err
}
log.Ctx(s.ctx).Info("cdc init done")
if err := s.start(); err != nil {
return err
}
log.Ctx(s.ctx).Info("cdc start done")
return nil
}
// Stop stops the server, should be call after Run returned.
func (s *Server) Stop() (err error) {
s.stopOnce.Do(s.stop)
return nil
}
// stop stops the server.
func (s *Server) stop() {
s.componentState.OnStopping()
log := log.Ctx(s.ctx)
log.Info("stopping cdc...")
// Stop CDC service.
s.cdcServer.Stop()
// Stop etcd
if s.etcdCli != nil {
if err := s.etcdCli.Close(); err != nil {
log.Warn("cdc stop etcd client failed", zap.Error(err))
}
}
// Stop tikv
if s.tikvCli != nil {
if err := s.tikvCli.Close(); err != nil {
log.Warn("cdc stop tikv client failed", zap.Error(err))
}
}
log.Info("cdc stop done")
}
// Health check the health status of cdc.
func (s *Server) Health(ctx context.Context) commonpb.StateCode {
resp, _ := s.componentState.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{})
return resp.GetState().StateCode
}
func (s *Server) init() (err error) {
log := log.Ctx(s.ctx)
defer func() {
if err != nil {
log.Error("cdc init failed", zap.Error(err))
return
}
log.Info("init cdc server finished")
}()
// Create etcd client.
s.etcdCli, _ = kvfactory.GetEtcdAndPath()
if err := s.initMeta(); err != nil {
return err
}
// Create CDC service.
s.cdcServer = cdc.NewCDCServer(s.ctx)
resource.Init(
resource.OptMetaKV(s.metaKV),
resource.OptReplicateManagerClient(replicatemanager.NewReplicateManager()),
resource.OptController(controllerimpl.NewController()),
)
return nil
}
func (s *Server) start() (err error) {
log := log.Ctx(s.ctx)
defer func() {
if err != nil {
log.Error("CDC start failed", zap.Error(err))
return
}
log.Info("start CDC server finished")
}()
s.cdcServer.Start()
s.componentState.OnInitialized(0)
return nil
}
func (s *Server) initMeta() error {
params := paramtable.Get()
metaType := params.MetaStoreCfg.MetaStoreType.GetValue()
log := log.Ctx(s.ctx)
log.Info("cdc connecting to metadata store", zap.String("metaType", metaType))
metaRootPath := ""
if metaType == util.MetaStoreTypeTiKV {
var err error
s.tikvCli, err = tikv.GetTiKVClient(&paramtable.Get().TiKVCfg)
if err != nil {
log.Warn("cdc init tikv client failed", zap.Error(err))
return err
}
metaRootPath = params.TiKVCfg.MetaRootPath.GetValue()
s.metaKV = tikvkv.NewTiKV(s.tikvCli, metaRootPath,
tikvkv.WithRequestTimeout(paramtable.Get().ServiceParam.TiKVCfg.RequestTimeout.GetAsDuration(time.Millisecond)))
} else if metaType == util.MetaStoreTypeEtcd {
metaRootPath = params.EtcdCfg.MetaRootPath.GetValue()
s.metaKV = etcdkv.NewEtcdKV(s.etcdCli, metaRootPath,
etcdkv.WithRequestTimeout(paramtable.Get().ServiceParam.EtcdCfg.RequestTimeout.GetAsDuration(time.Millisecond)))
}
return nil
}

View File

@ -1149,3 +1149,18 @@ func (s *Server) RemoveFileResource(ctx context.Context, req *milvuspb.RemoveFil
func (s *Server) ListFileResources(ctx context.Context, req *milvuspb.ListFileResourcesRequest) (*milvuspb.ListFileResourcesResponse, error) {
return s.proxy.ListFileResources(ctx, req)
}
// UpdateReplicateConfiguration applies a full replacement of the current replication configuration across Milvus clusters.
func (s *Server) UpdateReplicateConfiguration(ctx context.Context, req *milvuspb.UpdateReplicateConfigurationRequest) (*commonpb.Status, error) {
return s.proxy.UpdateReplicateConfiguration(ctx, req)
}
// GetReplicateInfo retrieves replication-related metadata from a target Milvus cluster.
func (s *Server) GetReplicateInfo(ctx context.Context, req *milvuspb.GetReplicateInfoRequest) (*milvuspb.GetReplicateInfoResponse, error) {
return s.proxy.GetReplicateInfo(ctx, req)
}
// CreateReplicateStream establishes a replication stream on the target Milvus cluster.
func (s *Server) CreateReplicateStream(stream milvuspb.MilvusService_CreateReplicateStreamServer) error {
return s.proxy.CreateReplicateStream(stream)
}

View File

@ -39,6 +39,9 @@ func assertValidMessage(msgs ...message.MutableMessage) {
if msg.MessageType().IsSystem() {
panic("system message is not allowed to append from client")
}
if msg.MessageType().IsSelfControlled() {
panic("self controlled message is not allowed to append from client")
}
if msg.VChannel() == "" {
panic("we don't support sent all vchannel message at client now")
}
@ -50,6 +53,9 @@ func assertValidBroadcastMessage(msg message.BroadcastMutableMessage) {
if msg.MessageType().IsSystem() {
panic("system message is not allowed to broadcast append from client")
}
if msg.MessageType().IsSelfControlled() {
panic("self controlled message is not allowed to broadcast append from client")
}
}
// We only support delete and insert message for txn now.

View File

@ -10,4 +10,5 @@ var (
ErrCanceledOrDeadlineExceed = errors.New("canceled or deadline exceed")
ErrUnrecoverable = errors.New("unrecoverable")
ErrFenced = errors.New("fenced")
ErrIgnoredOperation = errors.New("ignored operation")
)

View File

@ -105,6 +105,9 @@ func (p *ResumableProducer) Produce(ctx context.Context, msg message.MutableMess
if sErr.IsUnrecoverable() {
return nil, errors.Mark(err, errs.ErrUnrecoverable)
}
if sErr.IsIgnoredOperation() {
return nil, errors.Mark(err, errs.ErrIgnoredOperation)
}
}
}
}

View File

@ -95,7 +95,7 @@ func (m *delegatorMsgstreamAdaptor) Seek(ctx context.Context, msgPositions []*ms
panic("should never be called if len(msgPositions) is not 1")
}
position := msgPositions[0]
startFrom := adaptor.MustGetMessageIDFromMQWrapperIDBytes(WAL().WALName(), position.MsgID)
startFrom := adaptor.MustGetMessageIDFromMQWrapperIDBytes(position.MsgID)
log.Info(
"delegator msgstream adaptor seeks from position with scanner",
zap.String("channel", position.GetChannelName()),

View File

@ -0,0 +1,149 @@
package streaming
import (
"context"
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"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/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
var _ ReplicateService = replicateService{}
type replicateService struct {
*walAccesserImpl
}
// Append appends the message into current cluster.
func (s replicateService) Append(ctx context.Context, rmsg message.ReplicateMutableMessage) (*types.AppendResult, error) {
rh := rmsg.ReplicateHeader()
if rh == nil {
panic("message is not a replicate message")
}
if !s.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrWALAccesserClosed
}
defer s.lifetime.Done()
msg, err := s.overwriteReplicateMessage(ctx, rmsg, rh)
if err != nil {
return nil, err
}
return s.appendToWAL(ctx, msg)
}
func (s replicateService) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
if !s.lifetime.Add(typeutil.LifetimeStateWorking) {
return ErrWALAccesserClosed
}
defer s.lifetime.Done()
return s.streamingCoordClient.Assignment().UpdateReplicateConfiguration(ctx, config)
}
func (s replicateService) GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
if !s.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrWALAccesserClosed
}
defer s.lifetime.Done()
config, err := s.streamingCoordClient.Assignment().GetReplicateConfiguration(ctx)
if err != nil {
return nil, err
}
return config, nil
}
func (s replicateService) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
if !s.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrWALAccesserClosed
}
defer s.lifetime.Done()
checkpoint, err := s.handlerClient.GetReplicateCheckpoint(ctx, channelName)
if err != nil {
return nil, err
}
return checkpoint, nil
}
// 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) {
cfg, err := s.streamingCoordClient.Assignment().GetReplicateConfiguration(ctx)
if err != nil {
return nil, err
}
// Get target vchannel on current cluster that should be written to
currentCluster := cfg.GetCluster(s.clusterID)
if currentCluster.Role() == replicateutil.RolePrimary {
return nil, status.NewReplicateViolation("primary cluster cannot receive replicate message")
}
sourceCluster := cfg.GetCluster(rh.ClusterID)
if sourceCluster == nil {
return nil, status.NewReplicateViolation("source cluster %s not found in replicate configuration", rh.ClusterID)
}
targetVChannel, err := s.getTargetVChannel(sourceCluster, msg.VChannel())
if err != nil {
return nil, err
}
// Get target broadcast vchannels on current cluster that should be written to
if bh := msg.BroadcastHeader(); bh != nil {
// broadcast header have vchannels, so we need to overwrite it.
targetBroadcastVChannels := make([]string, 0, len(bh.VChannels))
for _, vchannel := range bh.VChannels {
targetBroadcastVChannel, err := s.getTargetVChannel(sourceCluster, vchannel)
if err != nil {
return nil, status.NewReplicateViolation("failed to get target channel, %s", err.Error())
}
targetBroadcastVChannels = append(targetBroadcastVChannels, targetBroadcastVChannel)
}
msg.OverwriteReplicateVChannel(targetVChannel, targetBroadcastVChannels)
} else {
msg.OverwriteReplicateVChannel(targetVChannel)
}
// create collection message will set the vchannel in its body, so we need to overwrite it.
if msg.MessageType() == message.MessageTypeCreateCollection {
if err := s.overwriteCreateCollectionMessage(sourceCluster, msg); err != nil {
return nil, err
}
}
return msg, nil
}
// getTargetVChannel gets the target vchannel of the source vchannel.
func (s replicateService) getTargetVChannel(sourceCluster *replicateutil.MilvusCluster, sourceVChannel string) (string, error) {
sourcePChannel := funcutil.ToPhysicalChannel(sourceVChannel)
targetPChannel, err := sourceCluster.GetTargetChannel(sourcePChannel, s.clusterID)
if err != nil {
return "", status.NewReplicateViolation("failed to get target channel, %s", err.Error())
}
return strings.Replace(sourceVChannel, sourcePChannel, targetPChannel, 1), nil
}
// overwriteCreateCollectionMessage overwrites the create collection message.
func (s replicateService) overwriteCreateCollectionMessage(sourceCluster *replicateutil.MilvusCluster, msg message.ReplicateMutableMessage) error {
createCollectionMsg := message.MustAsMutableCreateCollectionMessageV1(msg)
body := createCollectionMsg.MustBody()
for idx, sourcePChannel := range body.PhysicalChannelNames {
targetPChannel, err := sourceCluster.GetTargetChannel(sourcePChannel, s.clusterID)
if err != nil {
return status.NewReplicateViolation("failed to get target channel, %s", err.Error())
}
body.PhysicalChannelNames[idx] = targetPChannel
body.VirtualChannelNames[idx] = strings.Replace(body.VirtualChannelNames[idx], sourcePChannel, targetPChannel, 1)
}
createCollectionMsg.OverwriteBody(body)
return nil
}

View File

@ -0,0 +1,124 @@
package streaming
import (
"context"
"strings"
"testing"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/stretchr/testify/assert"
"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/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/distributed/streaming/internal/producer"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/mock_client"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_producer"
"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/mock_handler"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
pulsar2 "github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/pulsar"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
func TestReplicateService(t *testing.T) {
c := mock_client.NewMockClient(t)
as := mock_client.NewMockAssignmentService(t)
c.EXPECT().Assignment().Return(as).Maybe()
h := mock_handler.NewMockHandlerClient(t)
p := mock_producer.NewMockProducer(t)
p.EXPECT().Append(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, mm message.MutableMessage) (*types.AppendResult, error) {
msg := message.MustAsMutableCreateCollectionMessageV1(mm)
assert.True(t, strings.HasPrefix(msg.VChannel(), "by-dev"))
for _, vchannel := range msg.BroadcastHeader().VChannels {
assert.True(t, strings.HasPrefix(vchannel, "by-dev"))
}
b := msg.MustBody()
for _, vchannel := range b.VirtualChannelNames {
assert.True(t, strings.HasPrefix(vchannel, "by-dev"))
}
for _, pchannel := range b.PhysicalChannelNames {
assert.True(t, strings.HasPrefix(pchannel, "by-dev"))
}
return &types.AppendResult{
MessageID: walimplstest.NewTestMessageID(1),
TimeTick: 1,
}, nil
}).Maybe()
p.EXPECT().IsAvailable().Return(true).Maybe()
p.EXPECT().Available().Return(make(chan struct{})).Maybe()
h.EXPECT().CreateProducer(mock.Anything, mock.Anything).Return(p, nil).Maybe()
as.EXPECT().GetReplicateConfiguration(mock.Anything).Return(replicateutil.MustNewConfigHelper(
"by-dev",
&commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "primary", Pchannels: []string{"primary-rootcoord-dml_0", "primary-rootcoord-dml_1"}},
{ClusterId: "by-dev", Pchannels: []string{"by-dev-rootcoord-dml_0", "by-dev-rootcoord-dml_1"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
},
), nil)
rs := &replicateService{
walAccesserImpl: &walAccesserImpl{
lifetime: typeutil.NewLifetime(),
clusterID: "by-dev",
streamingCoordClient: c,
handlerClient: h,
producers: make(map[string]*producer.ResumableProducer),
},
}
replicateMsgs := createReplicateCreateCollectionMessages()
for _, msg := range replicateMsgs {
_, err := rs.Append(context.Background(), msg)
assert.NoError(t, err)
}
}
func createReplicateCreateCollectionMessages() []message.ReplicateMutableMessage {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "ID", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "Vector", DataType: schemapb.DataType_FloatVector},
},
}
schemaBytes, _ := proto.Marshal(schema)
msg := message.NewCreateCollectionMessageBuilderV1().
WithHeader(&message.CreateCollectionMessageHeader{
CollectionId: 1,
PartitionIds: []int64{2},
}).
WithBody(&msgpb.CreateCollectionRequest{
CollectionID: 1,
CollectionName: "collection",
PartitionName: "partition",
PhysicalChannelNames: []string{
"primary-rootcoord-dml_0",
"primary-rootcoord-dml_1",
},
VirtualChannelNames: []string{
"primary-rootcoord-dml_0_1v0",
"primary-rootcoord-dml_1_1v1",
},
Schema: schemaBytes,
}).
WithBroadcast([]string{"primary-rootcoord-dml_0_1v0", "primary-rootcoord-dml_1_1v1"}).
MustBuildBroadcast()
msgs := msg.WithBroadcastID(100).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.NewReplicateMessage("primary", immutableMsg.IntoImmutableMessageProto()))
}
return replicateMsgs
}

View File

@ -4,10 +4,15 @@ import (
"context"
"time"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
"github.com/milvus-io/milvus/internal/util/hookutil"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/options"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
)
var singleton WALAccesser = nil
@ -16,6 +21,12 @@ var singleton WALAccesser = nil
// should be called before any other operations.
func Init() {
c, _ := kvfactory.GetEtcdAndPath()
// init and select wal name
util.InitAndSelectWALName()
// register cipher for cipher message
if hookutil.IsClusterEncyptionEnabled() {
message.RegisterCipher(hookutil.GetCipher())
}
singleton = newWALAccesser(c)
}
@ -81,6 +92,22 @@ type Scanner interface {
Close()
}
// ReplicateService is the interface for the replicate service.
type ReplicateService interface {
// Append appends the message into current cluster.
Append(ctx context.Context, msg message.ReplicateMutableMessage) (*types.AppendResult, error)
// UpdateReplicateConfiguration updates the replicate configuration to the milvus cluster.
UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error
// GetReplicateConfiguration returns the replicate configuration of the milvus cluster.
GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error)
// 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)
}
// Balancer is the interface for managing the balancer of the wal.
type Balancer interface {
// ListStreamingNode lists the streaming node.
@ -108,6 +135,9 @@ type Balancer interface {
// WALAccesser is the interfaces to interact with the milvus write ahead log.
type WALAccesser interface {
// Replicate returns the replicate service of the wal.
Replicate() ReplicateService
// ControlChannel returns the control channel name of the wal.
// It will return the channel name of the control channel of the wal.
ControlChannel() string
@ -115,9 +145,6 @@ type WALAccesser interface {
// Balancer returns the balancer management of the wal.
Balancer() Balancer
// WALName returns the name of the wal.
WALName() string
// Local returns the local services.
Local() Local

View File

@ -16,7 +16,6 @@ import (
"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/streaming/walimpls/impls/pulsar"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
@ -33,6 +32,56 @@ func TestMain(m *testing.M) {
m.Run()
}
func TestReplicate(t *testing.T) {
t.Skip("cat not running without streaming service at background")
streaming.Init()
defer streaming.Release()
pchannels1 := make([]string, 0, len(vChannels))
pchannels2 := make([]string, 0, len(vChannels))
for idx := 0; idx < 16; idx++ {
pchannels1 = append(pchannels1, fmt.Sprintf("primary-rootcoord-dml_%d", idx))
pchannels2 = append(pchannels2, fmt.Sprintf("by-dev-rootcoord-dml_%d", idx))
}
ctx := context.Background()
err := streaming.WAL().Replicate().UpdateReplicateConfiguration(ctx, &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: "primary",
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19530",
Token: "test-token",
},
Pchannels: pchannels1,
},
{
ClusterId: "by-dev",
ConnectionParam: &commonpb.ConnectionParam{
Uri: "localhost:19531",
Token: "test-token",
},
Pchannels: pchannels2,
},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{
SourceClusterId: "by-dev",
TargetClusterId: "primary",
},
},
})
if err != nil {
panic(err)
}
cfg, err := streaming.WAL().Replicate().GetReplicateConfiguration(ctx)
if err != nil {
panic(err)
}
t.Logf("cfg: %+v\n", cfg)
}
func TestStreamingBroadcast(t *testing.T) {
t.Skip("cat not running without streaming service at background")
streamingutil.SetStreamingServiceEnabled()

View File

@ -24,12 +24,15 @@ import (
"github.com/cockroachdb/errors"
"google.golang.org/protobuf/types/known/anypb"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
"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/rmq"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
)
var expectErr = make(chan error, 10)
@ -52,6 +55,24 @@ func SetupNoopWALForTest() {
singleton = &noopWALAccesser{}
}
type noopReplicateService struct{}
func (n *noopReplicateService) Append(ctx context.Context, msg message.ReplicateMutableMessage) (*types.AppendResult, error) {
return nil, nil
}
func (n *noopReplicateService) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
return nil
}
func (n *noopReplicateService) GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
return nil, nil
}
func (n *noopReplicateService) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
return nil, nil
}
type noopBalancer struct{}
func (n *noopBalancer) ListStreamingNode(ctx context.Context) ([]types.StreamingNodeInfo, error) {
@ -139,6 +160,10 @@ func (n *noopTxn) Rollback(ctx context.Context) error {
type noopWALAccesser struct{}
func (n *noopWALAccesser) Replicate() ReplicateService {
return &noopReplicateService{}
}
func (n *noopWALAccesser) ControlChannel() string {
return funcutil.GetControlChannel("noop")
}
@ -202,6 +227,18 @@ func (n *noopWALAccesser) AppendMessagesWithOption(ctx context.Context, opts App
return AppendResponses{}
}
func (n *noopWALAccesser) GetReplicateConfiguration(ctx context.Context) (*commonpb.ReplicateConfiguration, error) {
return nil, nil
}
func (n *noopWALAccesser) GetReplicateCheckpoint(ctx context.Context, channelName string) (*commonpb.ReplicateCheckpoint, error) {
return nil, nil
}
func (n *noopWALAccesser) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
return nil
}
type noopScanner struct{}
func (n *noopScanner) Done() <-chan struct{} {

View File

@ -12,12 +12,12 @@ import (
"github.com/milvus-io/milvus/internal/streamingcoord/client"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/pkg/v2/log"
"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/util/conc"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -31,6 +31,7 @@ func newWALAccesser(c *clientv3.Client) *walAccesserImpl {
handlerClient := handler.NewHandlerClient(streamingCoordClient.Assignment())
w := &walAccesserImpl{
lifetime: typeutil.NewLifetime(),
clusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
streamingCoordClient: streamingCoordClient,
handlerClient: handlerClient,
producerMutex: sync.Mutex{},
@ -47,7 +48,8 @@ func newWALAccesser(c *clientv3.Client) *walAccesserImpl {
// walAccesserImpl is the implementation of WALAccesser.
type walAccesserImpl struct {
log.Binder
lifetime *typeutil.Lifetime
lifetime *typeutil.Lifetime
clusterID string
// All services
streamingCoordClient client.Client
@ -59,12 +61,12 @@ type walAccesserImpl struct {
dispatchExecutionPool *conc.Pool[struct{}]
}
func (w *walAccesserImpl) Balancer() Balancer {
return balancerImpl{w}
func (w *walAccesserImpl) Replicate() ReplicateService {
return replicateService{w}
}
func (w *walAccesserImpl) WALName() string {
return util.MustSelectWALName()
func (w *walAccesserImpl) Balancer() Balancer {
return balancerImpl{w}
}
func (w *walAccesserImpl) Local() Local {

View File

@ -207,8 +207,23 @@ type QueryCoordCatalog interface {
GetCollectionTargets(ctx context.Context) (map[int64]*querypb.CollectionTarget, error)
}
// ReplicationCatalog is the interface for replication catalog
// it's used by CDC component.
type ReplicationCatalog interface {
// RemoveReplicatePChannel removes the replicate pchannel from metastore.
// Remove the task of CDC replication task of current cluster, should be called when a CDC replication task is finished.
RemoveReplicatePChannel(ctx context.Context, sourceChannelName, targetChannelName string) error
// ListReplicatePChannels lists all replicate pchannels from metastore.
// every ReplicatePChannelMeta is a task of CDC replication task of current cluster which is a source cluster in replication topology.
// the task is written by streaming coord, SaveReplicateConfiguration operation.
ListReplicatePChannels(ctx context.Context) ([]*streamingpb.ReplicatePChannelMeta, error)
}
// StreamingCoordCataLog is the interface for streamingcoord catalog
type StreamingCoordCataLog interface {
ReplicationCatalog
// GetCChannel get the control channel from metastore.
GetCChannel(ctx context.Context) (*streamingpb.CChannelMeta, error)
@ -237,6 +252,12 @@ type StreamingCoordCataLog interface {
// Make the task recoverable after restart.
// When broadcast task is done, it will be removed from metastore.
SaveBroadcastTask(ctx context.Context, broadcastID uint64, task *streamingpb.BroadcastTask) error
// SaveReplicateConfiguration saves the replicate configuration to metastore.
SaveReplicateConfiguration(ctx context.Context, config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta) error
// GetReplicateConfiguration gets the replicate configuration from metastore.
GetReplicateConfiguration(ctx context.Context) (*streamingpb.ReplicateConfigurationMeta, error)
}
// StreamingNodeCataLog is the interface for streamingnode catalog

View File

@ -6,4 +6,8 @@ const (
BroadcastTaskPrefix = MetaPrefix + "broadcast-task/"
VersionPrefix = MetaPrefix + "version/"
CChannelMetaPrefix = MetaPrefix + "cchannel/"
// Replicate
ReplicatePChannelMetaPrefix = MetaPrefix + "replicating-pchannel/"
ReplicateConfigurationKey = MetaPrefix + "replicate-configuration"
)

View File

@ -2,6 +2,7 @@ package streamingcoord
import (
"context"
"fmt"
"strconv"
"github.com/cockroachdb/errors"
@ -26,12 +27,25 @@ import (
//
// ├── pchannel-1
// └── pchannel-2
//
// └── replicate-configuration
// └── replicating-pchannel
// │   ├── cluster-1-pchannel-1
// │   └── cluster-1-pchannel-2
// │   ├── cluster-2-pchannel-1
// │   └── cluster-2-pchannel-2
func NewCataLog(metaKV kv.MetaKv) metastore.StreamingCoordCataLog {
return &catalog{
metaKV: metaKV,
}
}
func NewReplicationCatalog(metaKV kv.MetaKv) metastore.ReplicationCatalog {
return &catalog{
metaKV: metaKV,
}
}
// catalog is a kv based catalog.
type catalog struct {
metaKV kv.MetaKv
@ -163,3 +177,67 @@ func buildPChannelInfoPath(name string) string {
func buildBroadcastTaskPath(id uint64) string {
return BroadcastTaskPrefix + strconv.FormatUint(id, 10)
}
func (c *catalog) SaveReplicateConfiguration(ctx context.Context, config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta) error {
v, err := proto.Marshal(config)
if err != nil {
return errors.Wrapf(err, "marshal replicate configuration failed")
}
kvs := make(map[string]string, len(replicatingTasks)+1)
kvs[ReplicateConfigurationKey] = string(v)
for _, task := range replicatingTasks {
key := buildReplicatePChannelPath(task.GetTargetCluster().GetClusterId(), task.GetSourceChannelName())
v, err := proto.Marshal(task)
if err != nil {
return errors.Wrapf(err, "marshal replicate pchannel meta failed")
}
kvs[key] = string(v)
}
return etcd.SaveByBatchWithLimit(kvs, util.MaxEtcdTxnNum, func(partialKvs map[string]string) error {
return c.metaKV.MultiSave(ctx, partialKvs)
})
}
func (c *catalog) GetReplicateConfiguration(ctx context.Context) (*streamingpb.ReplicateConfigurationMeta, error) {
key := ReplicateConfigurationKey
value, err := c.metaKV.Load(ctx, key)
if err != nil {
if errors.Is(err, merr.ErrIoKeyNotFound) {
return nil, nil
}
return nil, err
}
config := &streamingpb.ReplicateConfigurationMeta{}
if err = proto.Unmarshal([]byte(value), config); err != nil {
return nil, errors.Wrapf(err, "unmarshal replicate configuration failed")
}
return config, nil
}
func (c *catalog) RemoveReplicatePChannel(ctx context.Context, targetClusterID, sourceChannelName string) error {
key := buildReplicatePChannelPath(targetClusterID, sourceChannelName)
return c.metaKV.Remove(ctx, key)
}
func (c *catalog) ListReplicatePChannels(ctx context.Context) ([]*streamingpb.ReplicatePChannelMeta, error) {
keys, values, err := c.metaKV.LoadWithPrefix(ctx, ReplicatePChannelMetaPrefix)
if err != nil {
return nil, err
}
infos := make([]*streamingpb.ReplicatePChannelMeta, 0, len(values))
for k, value := range values {
info := &streamingpb.ReplicatePChannelMeta{}
err = proto.Unmarshal([]byte(value), info)
if err != nil {
return nil, errors.Wrapf(err, "unmarshal replicate pchannel meta %s failed", keys[k])
}
infos = append(infos, info)
}
return infos, nil
}
func buildReplicatePChannelPath(targetClusterID, sourceChannelName string) string {
return fmt.Sprintf("%s%s-%s", ReplicatePChannelMetaPrefix, targetClusterID, sourceChannelName)
}

View File

@ -2,6 +2,7 @@ package streamingcoord
import (
"context"
"sort"
"strings"
"testing"
@ -9,6 +10,7 @@ import (
"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/pkg/v2/mocks/mock_kv"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
)
@ -139,3 +141,114 @@ func TestCatalog(t *testing.T) {
err = catalog.SaveBroadcastTask(context.Background(), 1, &streamingpb.BroadcastTask{})
assert.Error(t, err)
}
func TestCatalog_ReplicationCatalog(t *testing.T) {
kv := mock_kv.NewMockMetaKv(t)
kvStorage := make(map[string]string)
kv.EXPECT().Load(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) (string, error) {
return kvStorage[s], nil
})
kv.EXPECT().LoadWithPrefix(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) ([]string, []string, error) {
keys := make([]string, 0, len(kvStorage))
vals := make([]string, 0, len(kvStorage))
for k, v := range kvStorage {
if strings.HasPrefix(k, s) {
keys = append(keys, k)
vals = append(vals, v)
}
}
return keys, vals, nil
})
kv.EXPECT().MultiSave(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, kvs map[string]string) error {
for k, v := range kvs {
kvStorage[k] = v
}
return nil
})
kv.EXPECT().Remove(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, key string) error {
delete(kvStorage, key)
return nil
})
catalog := NewCataLog(kv)
// ReplicateConfiguration test
config := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: "source-cluster",
Pchannels: []string{"source-channel-1", "source-channel-2"},
},
{
ClusterId: "target-cluster-a",
Pchannels: []string{"target-channel-a-1", "target-channel-a-2"},
},
{
ClusterId: "target-cluster-b",
Pchannels: []string{"target-channel-b-1", "target-channel-b-2"},
},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{
SourceClusterId: "source-cluster",
TargetClusterId: "target-cluster-a",
},
{
SourceClusterId: "source-cluster",
TargetClusterId: "target-cluster-b",
},
},
}
err := catalog.SaveReplicateConfiguration(context.Background(), &streamingpb.ReplicateConfigurationMeta{ReplicateConfiguration: config}, nil)
assert.NoError(t, err)
cfg, err := catalog.GetReplicateConfiguration(context.Background())
assert.NoError(t, err)
assert.Equal(t, cfg.ReplicateConfiguration.GetClusters()[0].GetClusterId(), "source-cluster")
assert.Equal(t, cfg.ReplicateConfiguration.GetClusters()[1].GetClusterId(), "target-cluster-a")
assert.Equal(t, cfg.ReplicateConfiguration.GetClusters()[2].GetClusterId(), "target-cluster-b")
assert.Equal(t, cfg.ReplicateConfiguration.GetCrossClusterTopology()[0].GetSourceClusterId(), "source-cluster")
assert.Equal(t, cfg.ReplicateConfiguration.GetCrossClusterTopology()[0].GetTargetClusterId(), "target-cluster-a")
assert.Equal(t, cfg.ReplicateConfiguration.GetCrossClusterTopology()[1].GetSourceClusterId(), "source-cluster")
assert.Equal(t, cfg.ReplicateConfiguration.GetCrossClusterTopology()[1].GetTargetClusterId(), "target-cluster-b")
// ReplicatePChannel test
err = catalog.SaveReplicateConfiguration(context.Background(),
&streamingpb.ReplicateConfigurationMeta{ReplicateConfiguration: config},
[]*streamingpb.ReplicatePChannelMeta{
{
SourceChannelName: "source-channel-1",
TargetChannelName: "target-channel-1",
TargetCluster: &commonpb.MilvusCluster{ClusterId: "target-cluster"},
},
{
SourceChannelName: "source-channel-2",
TargetChannelName: "target-channel-2",
TargetCluster: &commonpb.MilvusCluster{ClusterId: "target-cluster"},
},
})
assert.NoError(t, err)
infos, err := catalog.ListReplicatePChannels(context.Background())
assert.NoError(t, err)
assert.Len(t, infos, 2)
sort.Slice(infos, func(i, j int) bool {
return infos[i].GetTargetChannelName() < infos[j].GetTargetChannelName()
})
assert.Equal(t, infos[0].GetSourceChannelName(), "source-channel-1")
assert.Equal(t, infos[0].GetTargetChannelName(), "target-channel-1")
assert.Equal(t, infos[0].GetTargetCluster().GetClusterId(), "target-cluster")
assert.Equal(t, infos[1].GetSourceChannelName(), "source-channel-2")
assert.Equal(t, infos[1].GetTargetChannelName(), "target-channel-2")
assert.Equal(t, infos[1].GetTargetCluster().GetClusterId(), "target-cluster")
err = catalog.RemoveReplicatePChannel(context.Background(), "target-cluster", "source-channel-1")
assert.NoError(t, err)
infos, err = catalog.ListReplicatePChannels(context.Background())
assert.NoError(t, err)
assert.Len(t, infos, 1)
assert.Equal(t, infos[0].GetSourceChannelName(), "source-channel-2")
assert.Equal(t, infos[0].GetTargetChannelName(), "target-channel-2")
assert.Equal(t, infos[0].GetTargetCluster().GetClusterId(), "target-cluster")
}

View File

@ -1767,64 +1767,6 @@ func (_c *RootCoordCatalog_ListDatabases_Call) RunAndReturn(run func(context.Con
return _c
}
// ListFileResource provides a mock function with given fields: ctx
func (_m *RootCoordCatalog) ListFileResource(ctx context.Context) ([]*model.FileResource, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for ListFileResource")
}
var r0 []*model.FileResource
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) ([]*model.FileResource, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) []*model.FileResource); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*model.FileResource)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// RootCoordCatalog_ListFileResource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListFileResource'
type RootCoordCatalog_ListFileResource_Call struct {
*mock.Call
}
// ListFileResource is a helper method to define mock.On call
// - ctx context.Context
func (_e *RootCoordCatalog_Expecter) ListFileResource(ctx interface{}) *RootCoordCatalog_ListFileResource_Call {
return &RootCoordCatalog_ListFileResource_Call{Call: _e.mock.On("ListFileResource", ctx)}
}
func (_c *RootCoordCatalog_ListFileResource_Call) Run(run func(ctx context.Context)) *RootCoordCatalog_ListFileResource_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *RootCoordCatalog_ListFileResource_Call) Return(_a0 []*model.FileResource, _a1 error) *RootCoordCatalog_ListFileResource_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *RootCoordCatalog_ListFileResource_Call) RunAndReturn(run func(context.Context) ([]*model.FileResource, error)) *RootCoordCatalog_ListFileResource_Call {
_c.Call.Return(run)
return _c
}
// ListGrant provides a mock function with given fields: ctx, tenant, entity
func (_m *RootCoordCatalog) ListGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
ret := _m.Called(ctx, tenant, entity)
@ -2183,53 +2125,6 @@ func (_c *RootCoordCatalog_ListUserRole_Call) RunAndReturn(run func(context.Cont
return _c
}
// RemoveFileResource provides a mock function with given fields: ctx, resourceID
func (_m *RootCoordCatalog) RemoveFileResource(ctx context.Context, resourceID int64) error {
ret := _m.Called(ctx, resourceID)
if len(ret) == 0 {
panic("no return value specified for RemoveFileResource")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, int64) error); ok {
r0 = rf(ctx, resourceID)
} else {
r0 = ret.Error(0)
}
return r0
}
// RootCoordCatalog_RemoveFileResource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveFileResource'
type RootCoordCatalog_RemoveFileResource_Call struct {
*mock.Call
}
// RemoveFileResource is a helper method to define mock.On call
// - ctx context.Context
// - resourceID int64
func (_e *RootCoordCatalog_Expecter) RemoveFileResource(ctx interface{}, resourceID interface{}) *RootCoordCatalog_RemoveFileResource_Call {
return &RootCoordCatalog_RemoveFileResource_Call{Call: _e.mock.On("RemoveFileResource", ctx, resourceID)}
}
func (_c *RootCoordCatalog_RemoveFileResource_Call) Run(run func(ctx context.Context, resourceID int64)) *RootCoordCatalog_RemoveFileResource_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(int64))
})
return _c
}
func (_c *RootCoordCatalog_RemoveFileResource_Call) Return(_a0 error) *RootCoordCatalog_RemoveFileResource_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *RootCoordCatalog_RemoveFileResource_Call) RunAndReturn(run func(context.Context, int64) error) *RootCoordCatalog_RemoveFileResource_Call {
_c.Call.Return(run)
return _c
}
// RestoreRBAC provides a mock function with given fields: ctx, tenant, meta
func (_m *RootCoordCatalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
ret := _m.Called(ctx, tenant, meta)
@ -2278,53 +2173,6 @@ func (_c *RootCoordCatalog_RestoreRBAC_Call) RunAndReturn(run func(context.Conte
return _c
}
// SaveFileResource provides a mock function with given fields: ctx, resource
func (_m *RootCoordCatalog) SaveFileResource(ctx context.Context, resource *model.FileResource) error {
ret := _m.Called(ctx, resource)
if len(ret) == 0 {
panic("no return value specified for SaveFileResource")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *model.FileResource) error); ok {
r0 = rf(ctx, resource)
} else {
r0 = ret.Error(0)
}
return r0
}
// RootCoordCatalog_SaveFileResource_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveFileResource'
type RootCoordCatalog_SaveFileResource_Call struct {
*mock.Call
}
// SaveFileResource is a helper method to define mock.On call
// - ctx context.Context
// - resource *model.FileResource
func (_e *RootCoordCatalog_Expecter) SaveFileResource(ctx interface{}, resource interface{}) *RootCoordCatalog_SaveFileResource_Call {
return &RootCoordCatalog_SaveFileResource_Call{Call: _e.mock.On("SaveFileResource", ctx, resource)}
}
func (_c *RootCoordCatalog_SaveFileResource_Call) Run(run func(ctx context.Context, resource *model.FileResource)) *RootCoordCatalog_SaveFileResource_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*model.FileResource))
})
return _c
}
func (_c *RootCoordCatalog_SaveFileResource_Call) Return(_a0 error) *RootCoordCatalog_SaveFileResource_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *RootCoordCatalog_SaveFileResource_Call) RunAndReturn(run func(context.Context, *model.FileResource) error) *RootCoordCatalog_SaveFileResource_Call {
_c.Call.Return(run)
return _c
}
// SavePrivilegeGroup provides a mock function with given fields: ctx, data
func (_m *RootCoordCatalog) SavePrivilegeGroup(ctx context.Context, data *milvuspb.PrivilegeGroupInfo) error {
ret := _m.Called(ctx, data)

View File

@ -0,0 +1,269 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mock_streaming
import (
context "context"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
replicateutil "github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
wal "github.com/milvus-io/milvus/internal/streamingnode/server/wal"
)
// MockReplicateService is an autogenerated mock type for the ReplicateService type
type MockReplicateService struct {
mock.Mock
}
type MockReplicateService_Expecter struct {
mock *mock.Mock
}
func (_m *MockReplicateService) EXPECT() *MockReplicateService_Expecter {
return &MockReplicateService_Expecter{mock: &_m.Mock}
}
// Append provides a mock function with given fields: ctx, msg
func (_m *MockReplicateService) Append(ctx context.Context, msg message.ReplicateMutableMessage) (*types.AppendResult, error) {
ret := _m.Called(ctx, msg)
if len(ret) == 0 {
panic("no return value specified for Append")
}
var r0 *types.AppendResult
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, message.ReplicateMutableMessage) (*types.AppendResult, error)); ok {
return rf(ctx, msg)
}
if rf, ok := ret.Get(0).(func(context.Context, message.ReplicateMutableMessage) *types.AppendResult); ok {
r0 = rf(ctx, msg)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*types.AppendResult)
}
}
if rf, ok := ret.Get(1).(func(context.Context, message.ReplicateMutableMessage) error); ok {
r1 = rf(ctx, msg)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockReplicateService_Append_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Append'
type MockReplicateService_Append_Call struct {
*mock.Call
}
// Append is a helper method to define mock.On call
// - ctx context.Context
// - msg message.ReplicateMutableMessage
func (_e *MockReplicateService_Expecter) Append(ctx interface{}, msg interface{}) *MockReplicateService_Append_Call {
return &MockReplicateService_Append_Call{Call: _e.mock.On("Append", ctx, msg)}
}
func (_c *MockReplicateService_Append_Call) Run(run func(ctx context.Context, msg message.ReplicateMutableMessage)) *MockReplicateService_Append_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(message.ReplicateMutableMessage))
})
return _c
}
func (_c *MockReplicateService_Append_Call) Return(_a0 *types.AppendResult, _a1 error) *MockReplicateService_Append_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockReplicateService_Append_Call) RunAndReturn(run func(context.Context, message.ReplicateMutableMessage) (*types.AppendResult, error)) *MockReplicateService_Append_Call {
_c.Call.Return(run)
return _c
}
// GetReplicateCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockReplicateService) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetReplicateCheckpoint")
}
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
}
// MockReplicateService_GetReplicateCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateCheckpoint'
type MockReplicateService_GetReplicateCheckpoint_Call struct {
*mock.Call
}
// GetReplicateCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - channelName string
func (_e *MockReplicateService_Expecter) GetReplicateCheckpoint(ctx interface{}, channelName interface{}) *MockReplicateService_GetReplicateCheckpoint_Call {
return &MockReplicateService_GetReplicateCheckpoint_Call{Call: _e.mock.On("GetReplicateCheckpoint", ctx, channelName)}
}
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) Run(run func(ctx context.Context, channelName string)) *MockReplicateService_GetReplicateCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *MockReplicateService_GetReplicateCheckpoint_Call) Return(_a0 *wal.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 {
_c.Call.Return(run)
return _c
}
// GetReplicateConfiguration provides a mock function with given fields: ctx
func (_m *MockReplicateService) GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for GetReplicateConfiguration")
}
var r0 *replicateutil.ConfigHelper
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*replicateutil.ConfigHelper, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *replicateutil.ConfigHelper); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*replicateutil.ConfigHelper)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockReplicateService_GetReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateConfiguration'
type MockReplicateService_GetReplicateConfiguration_Call struct {
*mock.Call
}
// GetReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockReplicateService_Expecter) GetReplicateConfiguration(ctx interface{}) *MockReplicateService_GetReplicateConfiguration_Call {
return &MockReplicateService_GetReplicateConfiguration_Call{Call: _e.mock.On("GetReplicateConfiguration", ctx)}
}
func (_c *MockReplicateService_GetReplicateConfiguration_Call) Run(run func(ctx context.Context)) *MockReplicateService_GetReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockReplicateService_GetReplicateConfiguration_Call) Return(_a0 *replicateutil.ConfigHelper, _a1 error) *MockReplicateService_GetReplicateConfiguration_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockReplicateService_GetReplicateConfiguration_Call) RunAndReturn(run func(context.Context) (*replicateutil.ConfigHelper, error)) *MockReplicateService_GetReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// UpdateReplicateConfiguration provides a mock function with given fields: ctx, config
func (_m *MockReplicateService) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
ret := _m.Called(ctx, config)
if len(ret) == 0 {
panic("no return value specified for UpdateReplicateConfiguration")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *commonpb.ReplicateConfiguration) error); ok {
r0 = rf(ctx, config)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockReplicateService_UpdateReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateReplicateConfiguration'
type MockReplicateService_UpdateReplicateConfiguration_Call struct {
*mock.Call
}
// UpdateReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
// - config *commonpb.ReplicateConfiguration
func (_e *MockReplicateService_Expecter) UpdateReplicateConfiguration(ctx interface{}, config interface{}) *MockReplicateService_UpdateReplicateConfiguration_Call {
return &MockReplicateService_UpdateReplicateConfiguration_Call{Call: _e.mock.On("UpdateReplicateConfiguration", ctx, config)}
}
func (_c *MockReplicateService_UpdateReplicateConfiguration_Call) Run(run func(ctx context.Context, config *commonpb.ReplicateConfiguration)) *MockReplicateService_UpdateReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*commonpb.ReplicateConfiguration))
})
return _c
}
func (_c *MockReplicateService_UpdateReplicateConfiguration_Call) Return(_a0 error) *MockReplicateService_UpdateReplicateConfiguration_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockReplicateService_UpdateReplicateConfiguration_Call) RunAndReturn(run func(context.Context, *commonpb.ReplicateConfiguration) error) *MockReplicateService_UpdateReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// NewMockReplicateService creates a new instance of MockReplicateService. 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 NewMockReplicateService(t interface {
mock.TestingT
Cleanup(func())
}) *MockReplicateService {
mock := &MockReplicateService{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -458,6 +458,53 @@ func (_c *MockWALAccesser_Read_Call) RunAndReturn(run func(context.Context, stre
return _c
}
// Replicate provides a mock function with no fields
func (_m *MockWALAccesser) Replicate() streaming.ReplicateService {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Replicate")
}
var r0 streaming.ReplicateService
if rf, ok := ret.Get(0).(func() streaming.ReplicateService); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(streaming.ReplicateService)
}
}
return r0
}
// MockWALAccesser_Replicate_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Replicate'
type MockWALAccesser_Replicate_Call struct {
*mock.Call
}
// Replicate is a helper method to define mock.On call
func (_e *MockWALAccesser_Expecter) Replicate() *MockWALAccesser_Replicate_Call {
return &MockWALAccesser_Replicate_Call{Call: _e.mock.On("Replicate")}
}
func (_c *MockWALAccesser_Replicate_Call) Run(run func()) *MockWALAccesser_Replicate_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockWALAccesser_Replicate_Call) Return(_a0 streaming.ReplicateService) *MockWALAccesser_Replicate_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWALAccesser_Replicate_Call) RunAndReturn(run func() streaming.ReplicateService) *MockWALAccesser_Replicate_Call {
_c.Call.Return(run)
return _c
}
// Txn provides a mock function with given fields: ctx, opts
func (_m *MockWALAccesser) Txn(ctx context.Context, opts streaming.TxnOption) (streaming.Txn, error) {
ret := _m.Called(ctx, opts)
@ -517,51 +564,6 @@ func (_c *MockWALAccesser_Txn_Call) RunAndReturn(run func(context.Context, strea
return _c
}
// WALName provides a mock function with no fields
func (_m *MockWALAccesser) WALName() string {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for WALName")
}
var r0 string
if rf, ok := ret.Get(0).(func() string); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(string)
}
return r0
}
// MockWALAccesser_WALName_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'WALName'
type MockWALAccesser_WALName_Call struct {
*mock.Call
}
// WALName is a helper method to define mock.On call
func (_e *MockWALAccesser_Expecter) WALName() *MockWALAccesser_WALName_Call {
return &MockWALAccesser_WALName_Call{Call: _e.mock.On("WALName")}
}
func (_c *MockWALAccesser_WALName_Call) Run(run func()) *MockWALAccesser_WALName_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockWALAccesser_WALName_Call) Return(_a0 string) *MockWALAccesser_WALName_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWALAccesser_WALName_Call) RunAndReturn(run func() string) *MockWALAccesser_WALName_Call {
_c.Call.Return(run)
return _c
}
// NewMockWALAccesser creates a new instance of MockWALAccesser. 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 NewMockWALAccesser(t interface {

View File

@ -0,0 +1,144 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mock_metastore
import (
context "context"
mock "github.com/stretchr/testify/mock"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
)
// MockReplicationCatalog is an autogenerated mock type for the ReplicationCatalog type
type MockReplicationCatalog struct {
mock.Mock
}
type MockReplicationCatalog_Expecter struct {
mock *mock.Mock
}
func (_m *MockReplicationCatalog) EXPECT() *MockReplicationCatalog_Expecter {
return &MockReplicationCatalog_Expecter{mock: &_m.Mock}
}
// ListReplicatePChannels provides a mock function with given fields: ctx
func (_m *MockReplicationCatalog) ListReplicatePChannels(ctx context.Context) ([]*streamingpb.ReplicatePChannelMeta, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for ListReplicatePChannels")
}
var r0 []*streamingpb.ReplicatePChannelMeta
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) ([]*streamingpb.ReplicatePChannelMeta, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) []*streamingpb.ReplicatePChannelMeta); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*streamingpb.ReplicatePChannelMeta)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockReplicationCatalog_ListReplicatePChannels_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListReplicatePChannels'
type MockReplicationCatalog_ListReplicatePChannels_Call struct {
*mock.Call
}
// ListReplicatePChannels is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockReplicationCatalog_Expecter) ListReplicatePChannels(ctx interface{}) *MockReplicationCatalog_ListReplicatePChannels_Call {
return &MockReplicationCatalog_ListReplicatePChannels_Call{Call: _e.mock.On("ListReplicatePChannels", ctx)}
}
func (_c *MockReplicationCatalog_ListReplicatePChannels_Call) Run(run func(ctx context.Context)) *MockReplicationCatalog_ListReplicatePChannels_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockReplicationCatalog_ListReplicatePChannels_Call) Return(_a0 []*streamingpb.ReplicatePChannelMeta, _a1 error) *MockReplicationCatalog_ListReplicatePChannels_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockReplicationCatalog_ListReplicatePChannels_Call) RunAndReturn(run func(context.Context) ([]*streamingpb.ReplicatePChannelMeta, error)) *MockReplicationCatalog_ListReplicatePChannels_Call {
_c.Call.Return(run)
return _c
}
// RemoveReplicatePChannel provides a mock function with given fields: ctx, sourceChannelName, targetChannelName
func (_m *MockReplicationCatalog) RemoveReplicatePChannel(ctx context.Context, sourceChannelName string, targetChannelName string) error {
ret := _m.Called(ctx, sourceChannelName, targetChannelName)
if len(ret) == 0 {
panic("no return value specified for RemoveReplicatePChannel")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
r0 = rf(ctx, sourceChannelName, targetChannelName)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockReplicationCatalog_RemoveReplicatePChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveReplicatePChannel'
type MockReplicationCatalog_RemoveReplicatePChannel_Call struct {
*mock.Call
}
// RemoveReplicatePChannel is a helper method to define mock.On call
// - ctx context.Context
// - sourceChannelName string
// - targetChannelName string
func (_e *MockReplicationCatalog_Expecter) RemoveReplicatePChannel(ctx interface{}, sourceChannelName interface{}, targetChannelName interface{}) *MockReplicationCatalog_RemoveReplicatePChannel_Call {
return &MockReplicationCatalog_RemoveReplicatePChannel_Call{Call: _e.mock.On("RemoveReplicatePChannel", ctx, sourceChannelName, targetChannelName)}
}
func (_c *MockReplicationCatalog_RemoveReplicatePChannel_Call) Run(run func(ctx context.Context, sourceChannelName string, targetChannelName string)) *MockReplicationCatalog_RemoveReplicatePChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string))
})
return _c
}
func (_c *MockReplicationCatalog_RemoveReplicatePChannel_Call) Return(_a0 error) *MockReplicationCatalog_RemoveReplicatePChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockReplicationCatalog_RemoveReplicatePChannel_Call) RunAndReturn(run func(context.Context, string, string) error) *MockReplicationCatalog_RemoveReplicatePChannel_Call {
_c.Call.Return(run)
return _c
}
// NewMockReplicationCatalog creates a new instance of MockReplicationCatalog. 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 NewMockReplicationCatalog(t interface {
mock.TestingT
Cleanup(func())
}) *MockReplicationCatalog {
mock := &MockReplicationCatalog{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -81,6 +81,64 @@ func (_c *MockStreamingCoordCataLog_GetCChannel_Call) RunAndReturn(run func(cont
return _c
}
// GetReplicateConfiguration provides a mock function with given fields: ctx
func (_m *MockStreamingCoordCataLog) GetReplicateConfiguration(ctx context.Context) (*streamingpb.ReplicateConfigurationMeta, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for GetReplicateConfiguration")
}
var r0 *streamingpb.ReplicateConfigurationMeta
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*streamingpb.ReplicateConfigurationMeta, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *streamingpb.ReplicateConfigurationMeta); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*streamingpb.ReplicateConfigurationMeta)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockStreamingCoordCataLog_GetReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateConfiguration'
type MockStreamingCoordCataLog_GetReplicateConfiguration_Call struct {
*mock.Call
}
// GetReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockStreamingCoordCataLog_Expecter) GetReplicateConfiguration(ctx interface{}) *MockStreamingCoordCataLog_GetReplicateConfiguration_Call {
return &MockStreamingCoordCataLog_GetReplicateConfiguration_Call{Call: _e.mock.On("GetReplicateConfiguration", ctx)}
}
func (_c *MockStreamingCoordCataLog_GetReplicateConfiguration_Call) Run(run func(ctx context.Context)) *MockStreamingCoordCataLog_GetReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockStreamingCoordCataLog_GetReplicateConfiguration_Call) Return(_a0 *streamingpb.ReplicateConfigurationMeta, _a1 error) *MockStreamingCoordCataLog_GetReplicateConfiguration_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockStreamingCoordCataLog_GetReplicateConfiguration_Call) RunAndReturn(run func(context.Context) (*streamingpb.ReplicateConfigurationMeta, error)) *MockStreamingCoordCataLog_GetReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// GetVersion provides a mock function with given fields: ctx
func (_m *MockStreamingCoordCataLog) GetVersion(ctx context.Context) (*streamingpb.StreamingVersion, error) {
ret := _m.Called(ctx)
@ -255,6 +313,112 @@ func (_c *MockStreamingCoordCataLog_ListPChannel_Call) RunAndReturn(run func(con
return _c
}
// ListReplicatePChannels provides a mock function with given fields: ctx
func (_m *MockStreamingCoordCataLog) ListReplicatePChannels(ctx context.Context) ([]*streamingpb.ReplicatePChannelMeta, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for ListReplicatePChannels")
}
var r0 []*streamingpb.ReplicatePChannelMeta
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) ([]*streamingpb.ReplicatePChannelMeta, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) []*streamingpb.ReplicatePChannelMeta); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).([]*streamingpb.ReplicatePChannelMeta)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockStreamingCoordCataLog_ListReplicatePChannels_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListReplicatePChannels'
type MockStreamingCoordCataLog_ListReplicatePChannels_Call struct {
*mock.Call
}
// ListReplicatePChannels is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockStreamingCoordCataLog_Expecter) ListReplicatePChannels(ctx interface{}) *MockStreamingCoordCataLog_ListReplicatePChannels_Call {
return &MockStreamingCoordCataLog_ListReplicatePChannels_Call{Call: _e.mock.On("ListReplicatePChannels", ctx)}
}
func (_c *MockStreamingCoordCataLog_ListReplicatePChannels_Call) Run(run func(ctx context.Context)) *MockStreamingCoordCataLog_ListReplicatePChannels_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockStreamingCoordCataLog_ListReplicatePChannels_Call) Return(_a0 []*streamingpb.ReplicatePChannelMeta, _a1 error) *MockStreamingCoordCataLog_ListReplicatePChannels_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockStreamingCoordCataLog_ListReplicatePChannels_Call) RunAndReturn(run func(context.Context) ([]*streamingpb.ReplicatePChannelMeta, error)) *MockStreamingCoordCataLog_ListReplicatePChannels_Call {
_c.Call.Return(run)
return _c
}
// RemoveReplicatePChannel provides a mock function with given fields: ctx, sourceChannelName, targetChannelName
func (_m *MockStreamingCoordCataLog) RemoveReplicatePChannel(ctx context.Context, sourceChannelName string, targetChannelName string) error {
ret := _m.Called(ctx, sourceChannelName, targetChannelName)
if len(ret) == 0 {
panic("no return value specified for RemoveReplicatePChannel")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, string) error); ok {
r0 = rf(ctx, sourceChannelName, targetChannelName)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockStreamingCoordCataLog_RemoveReplicatePChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveReplicatePChannel'
type MockStreamingCoordCataLog_RemoveReplicatePChannel_Call struct {
*mock.Call
}
// RemoveReplicatePChannel is a helper method to define mock.On call
// - ctx context.Context
// - sourceChannelName string
// - targetChannelName string
func (_e *MockStreamingCoordCataLog_Expecter) RemoveReplicatePChannel(ctx interface{}, sourceChannelName interface{}, targetChannelName interface{}) *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call {
return &MockStreamingCoordCataLog_RemoveReplicatePChannel_Call{Call: _e.mock.On("RemoveReplicatePChannel", ctx, sourceChannelName, targetChannelName)}
}
func (_c *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call) Run(run func(ctx context.Context, sourceChannelName string, targetChannelName string)) *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(string))
})
return _c
}
func (_c *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call) Return(_a0 error) *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call) RunAndReturn(run func(context.Context, string, string) error) *MockStreamingCoordCataLog_RemoveReplicatePChannel_Call {
_c.Call.Return(run)
return _c
}
// SaveBroadcastTask provides a mock function with given fields: ctx, broadcastID, task
func (_m *MockStreamingCoordCataLog) SaveBroadcastTask(ctx context.Context, broadcastID uint64, task *streamingpb.BroadcastTask) error {
ret := _m.Called(ctx, broadcastID, task)
@ -397,6 +561,54 @@ func (_c *MockStreamingCoordCataLog_SavePChannels_Call) RunAndReturn(run func(co
return _c
}
// SaveReplicateConfiguration provides a mock function with given fields: ctx, config, replicatingTasks
func (_m *MockStreamingCoordCataLog) SaveReplicateConfiguration(ctx context.Context, config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta) error {
ret := _m.Called(ctx, config, replicatingTasks)
if len(ret) == 0 {
panic("no return value specified for SaveReplicateConfiguration")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.ReplicateConfigurationMeta, []*streamingpb.ReplicatePChannelMeta) error); ok {
r0 = rf(ctx, config, replicatingTasks)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockStreamingCoordCataLog_SaveReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SaveReplicateConfiguration'
type MockStreamingCoordCataLog_SaveReplicateConfiguration_Call struct {
*mock.Call
}
// SaveReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
// - config *streamingpb.ReplicateConfigurationMeta
// - replicatingTasks []*streamingpb.ReplicatePChannelMeta
func (_e *MockStreamingCoordCataLog_Expecter) SaveReplicateConfiguration(ctx interface{}, config interface{}, replicatingTasks interface{}) *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call {
return &MockStreamingCoordCataLog_SaveReplicateConfiguration_Call{Call: _e.mock.On("SaveReplicateConfiguration", ctx, config, replicatingTasks)}
}
func (_c *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call) Run(run func(ctx context.Context, config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta)) *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*streamingpb.ReplicateConfigurationMeta), args[2].([]*streamingpb.ReplicatePChannelMeta))
})
return _c
}
func (_c *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call) Return(_a0 error) *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call) RunAndReturn(run func(context.Context, *streamingpb.ReplicateConfigurationMeta, []*streamingpb.ReplicatePChannelMeta) error) *MockStreamingCoordCataLog_SaveReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// SaveVersion provides a mock function with given fields: ctx, version
func (_m *MockStreamingCoordCataLog) SaveVersion(ctx context.Context, version *streamingpb.StreamingVersion) error {
ret := _m.Called(ctx, version)

View File

@ -5,8 +5,13 @@ package mock_client
import (
context "context"
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
commonpb "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
mock "github.com/stretchr/testify/mock"
replicateutil "github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
)
// MockAssignmentService is an autogenerated mock type for the AssignmentService type
@ -127,6 +132,64 @@ func (_c *MockAssignmentService_GetLatestAssignments_Call) RunAndReturn(run func
return _c
}
// GetReplicateConfiguration provides a mock function with given fields: ctx
func (_m *MockAssignmentService) GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
ret := _m.Called(ctx)
if len(ret) == 0 {
panic("no return value specified for GetReplicateConfiguration")
}
var r0 *replicateutil.ConfigHelper
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (*replicateutil.ConfigHelper, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) *replicateutil.ConfigHelper); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*replicateutil.ConfigHelper)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockAssignmentService_GetReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateConfiguration'
type MockAssignmentService_GetReplicateConfiguration_Call struct {
*mock.Call
}
// GetReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
func (_e *MockAssignmentService_Expecter) GetReplicateConfiguration(ctx interface{}) *MockAssignmentService_GetReplicateConfiguration_Call {
return &MockAssignmentService_GetReplicateConfiguration_Call{Call: _e.mock.On("GetReplicateConfiguration", ctx)}
}
func (_c *MockAssignmentService_GetReplicateConfiguration_Call) Run(run func(ctx context.Context)) *MockAssignmentService_GetReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *MockAssignmentService_GetReplicateConfiguration_Call) Return(_a0 *replicateutil.ConfigHelper, _a1 error) *MockAssignmentService_GetReplicateConfiguration_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockAssignmentService_GetReplicateConfiguration_Call) RunAndReturn(run func(context.Context) (*replicateutil.ConfigHelper, error)) *MockAssignmentService_GetReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// ReportAssignmentError provides a mock function with given fields: ctx, pchannel, err
func (_m *MockAssignmentService) ReportAssignmentError(ctx context.Context, pchannel types.PChannelInfo, err error) error {
ret := _m.Called(ctx, pchannel, err)
@ -175,6 +238,53 @@ func (_c *MockAssignmentService_ReportAssignmentError_Call) RunAndReturn(run fun
return _c
}
// UpdateReplicateConfiguration provides a mock function with given fields: ctx, config
func (_m *MockAssignmentService) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
ret := _m.Called(ctx, config)
if len(ret) == 0 {
panic("no return value specified for UpdateReplicateConfiguration")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, *commonpb.ReplicateConfiguration) error); ok {
r0 = rf(ctx, config)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockAssignmentService_UpdateReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateReplicateConfiguration'
type MockAssignmentService_UpdateReplicateConfiguration_Call struct {
*mock.Call
}
// UpdateReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
// - config *commonpb.ReplicateConfiguration
func (_e *MockAssignmentService_Expecter) UpdateReplicateConfiguration(ctx interface{}, config interface{}) *MockAssignmentService_UpdateReplicateConfiguration_Call {
return &MockAssignmentService_UpdateReplicateConfiguration_Call{Call: _e.mock.On("UpdateReplicateConfiguration", ctx, config)}
}
func (_c *MockAssignmentService_UpdateReplicateConfiguration_Call) Run(run func(ctx context.Context, config *commonpb.ReplicateConfiguration)) *MockAssignmentService_UpdateReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*commonpb.ReplicateConfiguration))
})
return _c
}
func (_c *MockAssignmentService_UpdateReplicateConfiguration_Call) Return(_a0 error) *MockAssignmentService_UpdateReplicateConfiguration_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockAssignmentService_UpdateReplicateConfiguration_Call) RunAndReturn(run func(context.Context, *commonpb.ReplicateConfiguration) error) *MockAssignmentService_UpdateReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// UpdateWALBalancePolicy provides a mock function with given fields: ctx, req
func (_m *MockAssignmentService) UpdateWALBalancePolicy(ctx context.Context, req *types.UpdateWALBalancePolicyRequest) (*types.UpdateWALBalancePolicyResponse, error) {
ret := _m.Called(ctx, req)

View File

@ -7,6 +7,8 @@ import (
balancer "github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
streamingpb "github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
@ -119,6 +121,63 @@ func (_c *MockBalancer_GetAllStreamingNodes_Call) RunAndReturn(run func(context.
return _c
}
// GetLatestChannelAssignment provides a mock function with no fields
func (_m *MockBalancer) GetLatestChannelAssignment() (*balancer.WatchChannelAssignmentsCallbackParam, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetLatestChannelAssignment")
}
var r0 *balancer.WatchChannelAssignmentsCallbackParam
var r1 error
if rf, ok := ret.Get(0).(func() (*balancer.WatchChannelAssignmentsCallbackParam, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() *balancer.WatchChannelAssignmentsCallbackParam); ok {
r0 = rf()
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*balancer.WatchChannelAssignmentsCallbackParam)
}
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockBalancer_GetLatestChannelAssignment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetLatestChannelAssignment'
type MockBalancer_GetLatestChannelAssignment_Call struct {
*mock.Call
}
// GetLatestChannelAssignment is a helper method to define mock.On call
func (_e *MockBalancer_Expecter) GetLatestChannelAssignment() *MockBalancer_GetLatestChannelAssignment_Call {
return &MockBalancer_GetLatestChannelAssignment_Call{Call: _e.mock.On("GetLatestChannelAssignment")}
}
func (_c *MockBalancer_GetLatestChannelAssignment_Call) Run(run func()) *MockBalancer_GetLatestChannelAssignment_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockBalancer_GetLatestChannelAssignment_Call) Return(_a0 *balancer.WatchChannelAssignmentsCallbackParam, _a1 error) *MockBalancer_GetLatestChannelAssignment_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockBalancer_GetLatestChannelAssignment_Call) RunAndReturn(run func() (*balancer.WatchChannelAssignmentsCallbackParam, error)) *MockBalancer_GetLatestChannelAssignment_Call {
_c.Call.Return(run)
return _c
}
// GetLatestWALLocated provides a mock function with given fields: ctx, pchannel
func (_m *MockBalancer) GetLatestWALLocated(ctx context.Context, pchannel string) (int64, bool) {
ret := _m.Called(ctx, pchannel)
@ -361,6 +420,67 @@ func (_c *MockBalancer_UpdateBalancePolicy_Call) RunAndReturn(run func(context.C
return _c
}
// UpdateReplicateConfiguration provides a mock function with given fields: ctx, msgs
func (_m *MockBalancer) UpdateReplicateConfiguration(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2) error {
_va := make([]interface{}, len(msgs))
for _i := range msgs {
_va[_i] = msgs[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
if len(ret) == 0 {
panic("no return value specified for UpdateReplicateConfiguration")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, ...message.ImmutablePutReplicateConfigMessageV2) error); ok {
r0 = rf(ctx, msgs...)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockBalancer_UpdateReplicateConfiguration_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'UpdateReplicateConfiguration'
type MockBalancer_UpdateReplicateConfiguration_Call struct {
*mock.Call
}
// UpdateReplicateConfiguration is a helper method to define mock.On call
// - ctx context.Context
// - msgs ...message.ImmutablePutReplicateConfigMessageV2
func (_e *MockBalancer_Expecter) UpdateReplicateConfiguration(ctx interface{}, msgs ...interface{}) *MockBalancer_UpdateReplicateConfiguration_Call {
return &MockBalancer_UpdateReplicateConfiguration_Call{Call: _e.mock.On("UpdateReplicateConfiguration",
append([]interface{}{ctx}, msgs...)...)}
}
func (_c *MockBalancer_UpdateReplicateConfiguration_Call) Run(run func(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2)) *MockBalancer_UpdateReplicateConfiguration_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]message.ImmutablePutReplicateConfigMessageV2, len(args)-1)
for i, a := range args[1:] {
if a != nil {
variadicArgs[i] = a.(message.ImmutablePutReplicateConfigMessageV2)
}
}
run(args[0].(context.Context), variadicArgs...)
})
return _c
}
func (_c *MockBalancer_UpdateReplicateConfiguration_Call) Return(_a0 error) *MockBalancer_UpdateReplicateConfiguration_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockBalancer_UpdateReplicateConfiguration_Call) RunAndReturn(run func(context.Context, ...message.ImmutablePutReplicateConfigMessageV2) error) *MockBalancer_UpdateReplicateConfiguration_Call {
_c.Call.Return(run)
return _c
}
// WatchChannelAssignments provides a mock function with given fields: ctx, cb
func (_m *MockBalancer) WatchChannelAssignments(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
ret := _m.Called(ctx, cb)

View File

@ -0,0 +1,225 @@
// Code generated by mockery v2.53.3. DO NOT EDIT.
package mock_broadcaster
import (
context "context"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
types "github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
)
// MockBroadcaster is an autogenerated mock type for the Broadcaster type
type MockBroadcaster struct {
mock.Mock
}
type MockBroadcaster_Expecter struct {
mock *mock.Mock
}
func (_m *MockBroadcaster) EXPECT() *MockBroadcaster_Expecter {
return &MockBroadcaster_Expecter{mock: &_m.Mock}
}
// Ack provides a mock function with given fields: ctx, msg
func (_m *MockBroadcaster) Ack(ctx context.Context, msg message.ImmutableMessage) error {
ret := _m.Called(ctx, msg)
if len(ret) == 0 {
panic("no return value specified for Ack")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, message.ImmutableMessage) error); ok {
r0 = rf(ctx, msg)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockBroadcaster_Ack_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Ack'
type MockBroadcaster_Ack_Call struct {
*mock.Call
}
// Ack is a helper method to define mock.On call
// - ctx context.Context
// - msg message.ImmutableMessage
func (_e *MockBroadcaster_Expecter) Ack(ctx interface{}, msg interface{}) *MockBroadcaster_Ack_Call {
return &MockBroadcaster_Ack_Call{Call: _e.mock.On("Ack", ctx, msg)}
}
func (_c *MockBroadcaster_Ack_Call) Run(run func(ctx context.Context, msg message.ImmutableMessage)) *MockBroadcaster_Ack_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(message.ImmutableMessage))
})
return _c
}
func (_c *MockBroadcaster_Ack_Call) Return(_a0 error) *MockBroadcaster_Ack_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockBroadcaster_Ack_Call) RunAndReturn(run func(context.Context, message.ImmutableMessage) error) *MockBroadcaster_Ack_Call {
_c.Call.Return(run)
return _c
}
// Broadcast provides a mock function with given fields: ctx, msg
func (_m *MockBroadcaster) Broadcast(ctx context.Context, msg message.BroadcastMutableMessage) (*types.BroadcastAppendResult, error) {
ret := _m.Called(ctx, msg)
if len(ret) == 0 {
panic("no return value specified for Broadcast")
}
var r0 *types.BroadcastAppendResult
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, message.BroadcastMutableMessage) (*types.BroadcastAppendResult, error)); ok {
return rf(ctx, msg)
}
if rf, ok := ret.Get(0).(func(context.Context, message.BroadcastMutableMessage) *types.BroadcastAppendResult); ok {
r0 = rf(ctx, msg)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*types.BroadcastAppendResult)
}
}
if rf, ok := ret.Get(1).(func(context.Context, message.BroadcastMutableMessage) error); ok {
r1 = rf(ctx, msg)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockBroadcaster_Broadcast_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Broadcast'
type MockBroadcaster_Broadcast_Call struct {
*mock.Call
}
// Broadcast is a helper method to define mock.On call
// - ctx context.Context
// - msg message.BroadcastMutableMessage
func (_e *MockBroadcaster_Expecter) Broadcast(ctx interface{}, msg interface{}) *MockBroadcaster_Broadcast_Call {
return &MockBroadcaster_Broadcast_Call{Call: _e.mock.On("Broadcast", ctx, msg)}
}
func (_c *MockBroadcaster_Broadcast_Call) Run(run func(ctx context.Context, msg message.BroadcastMutableMessage)) *MockBroadcaster_Broadcast_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(message.BroadcastMutableMessage))
})
return _c
}
func (_c *MockBroadcaster_Broadcast_Call) Return(_a0 *types.BroadcastAppendResult, _a1 error) *MockBroadcaster_Broadcast_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockBroadcaster_Broadcast_Call) RunAndReturn(run func(context.Context, message.BroadcastMutableMessage) (*types.BroadcastAppendResult, error)) *MockBroadcaster_Broadcast_Call {
_c.Call.Return(run)
return _c
}
// Close provides a mock function with no fields
func (_m *MockBroadcaster) Close() {
_m.Called()
}
// MockBroadcaster_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
type MockBroadcaster_Close_Call struct {
*mock.Call
}
// Close is a helper method to define mock.On call
func (_e *MockBroadcaster_Expecter) Close() *MockBroadcaster_Close_Call {
return &MockBroadcaster_Close_Call{Call: _e.mock.On("Close")}
}
func (_c *MockBroadcaster_Close_Call) Run(run func()) *MockBroadcaster_Close_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockBroadcaster_Close_Call) Return() *MockBroadcaster_Close_Call {
_c.Call.Return()
return _c
}
func (_c *MockBroadcaster_Close_Call) RunAndReturn(run func()) *MockBroadcaster_Close_Call {
_c.Run(run)
return _c
}
// LegacyAck provides a mock function with given fields: ctx, broadcastID, vchannel
func (_m *MockBroadcaster) LegacyAck(ctx context.Context, broadcastID uint64, vchannel string) error {
ret := _m.Called(ctx, broadcastID, vchannel)
if len(ret) == 0 {
panic("no return value specified for LegacyAck")
}
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, uint64, string) error); ok {
r0 = rf(ctx, broadcastID, vchannel)
} else {
r0 = ret.Error(0)
}
return r0
}
// MockBroadcaster_LegacyAck_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LegacyAck'
type MockBroadcaster_LegacyAck_Call struct {
*mock.Call
}
// LegacyAck is a helper method to define mock.On call
// - ctx context.Context
// - broadcastID uint64
// - vchannel string
func (_e *MockBroadcaster_Expecter) LegacyAck(ctx interface{}, broadcastID interface{}, vchannel interface{}) *MockBroadcaster_LegacyAck_Call {
return &MockBroadcaster_LegacyAck_Call{Call: _e.mock.On("LegacyAck", ctx, broadcastID, vchannel)}
}
func (_c *MockBroadcaster_LegacyAck_Call) Run(run func(ctx context.Context, broadcastID uint64, vchannel string)) *MockBroadcaster_LegacyAck_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(uint64), args[2].(string))
})
return _c
}
func (_c *MockBroadcaster_LegacyAck_Call) Return(_a0 error) *MockBroadcaster_LegacyAck_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockBroadcaster_LegacyAck_Call) RunAndReturn(run func(context.Context, uint64, string) error) *MockBroadcaster_LegacyAck_Call {
_c.Call.Return(run)
return _c
}
// NewMockBroadcaster creates a new instance of MockBroadcaster. 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 NewMockBroadcaster(t interface {
mock.TestingT
Cleanup(func())
}) *MockBroadcaster {
mock := &MockBroadcaster{}
mock.Mock.Test(t)
t.Cleanup(func() { mock.AssertExpectations(t) })
return mock
}

View File

@ -9,6 +9,8 @@ import (
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"
)
// MockHandlerClient is an autogenerated mock type for the HandlerClient type
@ -231,6 +233,65 @@ func (_c *MockHandlerClient_GetLatestMVCCTimestampIfLocal_Call) RunAndReturn(run
return _c
}
// GetReplicateCheckpoint provides a mock function with given fields: ctx, channelName
func (_m *MockHandlerClient) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
ret := _m.Called(ctx, channelName)
if len(ret) == 0 {
panic("no return value specified for GetReplicateCheckpoint")
}
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_GetReplicateCheckpoint_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetReplicateCheckpoint'
type MockHandlerClient_GetReplicateCheckpoint_Call struct {
*mock.Call
}
// GetReplicateCheckpoint is a helper method to define mock.On call
// - ctx context.Context
// - channelName string
func (_e *MockHandlerClient_Expecter) GetReplicateCheckpoint(ctx interface{}, channelName interface{}) *MockHandlerClient_GetReplicateCheckpoint_Call {
return &MockHandlerClient_GetReplicateCheckpoint_Call{Call: _e.mock.On("GetReplicateCheckpoint", ctx, channelName)}
}
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) Run(run func(ctx context.Context, channelName string)) *MockHandlerClient_GetReplicateCheckpoint_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) Return(_a0 *wal.ReplicateCheckpoint, _a1 error) *MockHandlerClient_GetReplicateCheckpoint_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockHandlerClient_GetReplicateCheckpoint_Call) RunAndReturn(run func(context.Context, string) (*wal.ReplicateCheckpoint, error)) *MockHandlerClient_GetReplicateCheckpoint_Call {
_c.Call.Return(run)
return _c
}
// GetWALMetricsIfLocal provides a mock function with given fields: ctx
func (_m *MockHandlerClient) GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error) {
ret := _m.Called(ctx)

View File

@ -3,8 +3,10 @@
package mock_wal
import (
wal "github.com/milvus-io/milvus/internal/streamingnode/server/wal"
message "github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
mock "github.com/stretchr/testify/mock"
wal "github.com/milvus-io/milvus/internal/streamingnode/server/wal"
)
// MockOpenerBuilder is an autogenerated mock type for the OpenerBuilder type
@ -78,18 +80,18 @@ func (_c *MockOpenerBuilder_Build_Call) RunAndReturn(run func() (wal.Opener, err
}
// Name provides a mock function with no fields
func (_m *MockOpenerBuilder) Name() string {
func (_m *MockOpenerBuilder) Name() message.WALName {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for Name")
}
var r0 string
if rf, ok := ret.Get(0).(func() string); ok {
var r0 message.WALName
if rf, ok := ret.Get(0).(func() message.WALName); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(string)
r0 = ret.Get(0).(message.WALName)
}
return r0
@ -112,12 +114,12 @@ func (_c *MockOpenerBuilder_Name_Call) Run(run func()) *MockOpenerBuilder_Name_C
return _c
}
func (_c *MockOpenerBuilder_Name_Call) Return(_a0 string) *MockOpenerBuilder_Name_Call {
func (_c *MockOpenerBuilder_Name_Call) Return(_a0 message.WALName) *MockOpenerBuilder_Name_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockOpenerBuilder_Name_Call) RunAndReturn(run func() string) *MockOpenerBuilder_Name_Call {
func (_c *MockOpenerBuilder_Name_Call) RunAndReturn(run func() message.WALName) *MockOpenerBuilder_Name_Call {
_c.Call.Return(run)
return _c
}

View File

@ -453,18 +453,18 @@ func (_c *MockWAL_Read_Call) RunAndReturn(run func(context.Context, wal.ReadOpti
}
// WALName provides a mock function with no fields
func (_m *MockWAL) WALName() string {
func (_m *MockWAL) WALName() message.WALName {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for WALName")
}
var r0 string
if rf, ok := ret.Get(0).(func() string); ok {
var r0 message.WALName
if rf, ok := ret.Get(0).(func() message.WALName); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(string)
r0 = ret.Get(0).(message.WALName)
}
return r0
@ -487,12 +487,12 @@ func (_c *MockWAL_WALName_Call) Run(run func()) *MockWAL_WALName_Call {
return _c
}
func (_c *MockWAL_WALName_Call) Return(_a0 string) *MockWAL_WALName_Call {
func (_c *MockWAL_WALName_Call) Return(_a0 message.WALName) *MockWAL_WALName_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *MockWAL_WALName_Call) RunAndReturn(run func() string) *MockWAL_WALName_Call {
func (_c *MockWAL_WALName_Call) RunAndReturn(run func() message.WALName) *MockWAL_WALName_Call {
_c.Call.Return(run)
return _c
}

View File

@ -40,8 +40,10 @@ import (
"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-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/http"
"github.com/milvus-io/milvus/internal/proxy/connection"
"github.com/milvus-io/milvus/internal/proxy/replicate"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/analyzer"
"github.com/milvus-io/milvus/internal/util/hookutil"
@ -62,6 +64,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/util/metricsinfo"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/ratelimitutil"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/requestutil"
"github.com/milvus-io/milvus/pkg/v2/util/retry"
"github.com/milvus-io/milvus/pkg/v2/util/timerecord"
@ -1737,7 +1740,7 @@ func (node *Proxy) GetLoadingProgress(ctx context.Context, request *milvuspb.Get
}, nil
}
func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (resp *milvuspb.GetLoadStateResponse, err error) {
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return &milvuspb.GetLoadStateResponse{Status: merr.Status(err)}, nil
}
@ -1771,7 +1774,10 @@ func (node *Proxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadSt
defer func() {
log.Debug(
rpcDone(method),
zap.Any("request", request))
zap.Any("request", request),
zap.Any("response", resp),
zap.Error(err),
)
metrics.ProxyReqLatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), method).Observe(float64(tr.ElapseSpan().Milliseconds()))
}()
@ -6445,3 +6451,86 @@ func (node *Proxy) ListFileResources(ctx context.Context, req *milvuspb.ListFile
log.Info("ListFileResources success", zap.Int("count", len(resp.GetResources())))
return resp, nil
}
// UpdateReplicateConfiguration applies a full replacement of the current replication configuration across Milvus clusters.
func (node *Proxy) UpdateReplicateConfiguration(ctx context.Context, req *milvuspb.UpdateReplicateConfigurationRequest) (*commonpb.Status, error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-UpdateReplicateConfiguration")
defer sp.End()
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return merr.Status(err), nil
}
log.Ctx(ctx).Info("UpdateReplicateConfiguration received", replicateutil.ConfigLogFields(req.GetReplicateConfiguration())...)
err := streaming.WAL().Replicate().UpdateReplicateConfiguration(ctx, req.GetReplicateConfiguration())
if err != nil {
log.Ctx(ctx).Warn("UpdateReplicateConfiguration fail", zap.Error(err))
return merr.Status(err), nil
}
log.Ctx(ctx).Info("UpdateReplicateConfiguration success", replicateutil.ConfigLogFields(req.GetReplicateConfiguration())...)
return merr.Status(nil), nil
}
// GetReplicateInfo retrieves replication-related metadata from a target Milvus cluster.
// TODO: sheep, only get target checkpoint
func (node *Proxy) GetReplicateInfo(ctx context.Context, req *milvuspb.GetReplicateInfoRequest) (resp *milvuspb.GetReplicateInfoResponse, err error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-GetReplicateInfo")
defer sp.End()
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return nil, err
}
logger := log.Ctx(ctx).With(zap.String("sourceClusterID", req.GetSourceClusterId()))
logger.Info("GetReplicateInfo received")
defer func() {
if err != nil {
logger.Warn("GetReplicateInfo fail", zap.Error(err))
} else {
logger.Info("GetReplicateInfo success", zap.Any("checkpoints", resp.GetCheckpoints()))
}
}()
configHelper, err := streaming.WAL().Replicate().GetReplicateConfiguration(ctx)
if err != nil {
return nil, err
}
currentCluster := configHelper.GetCurrentCluster()
checkpoints := make([]*commonpb.ReplicateCheckpoint, 0, len(currentCluster.GetPchannels()))
for _, pchannel := range currentCluster.GetPchannels() {
checkpoint, err := streaming.WAL().Replicate().GetReplicateCheckpoint(ctx, pchannel)
if err != nil {
return nil, err
}
checkpoints = append(checkpoints, checkpoint.IntoProto())
}
return &milvuspb.GetReplicateInfoResponse{
Checkpoints: checkpoints,
}, nil
}
// CreateReplicateStream establishes a replication stream on the target Milvus cluster.
func (node *Proxy) CreateReplicateStream(stream milvuspb.MilvusService_CreateReplicateStreamServer) (err error) {
ctx := stream.Context()
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-CreateReplicateStream")
defer sp.End()
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return err
}
log.Ctx(ctx).Info("replicate stream created")
defer func() {
if err != nil {
log.Ctx(ctx).Warn("replicate stream closed with error", zap.Error(err))
} else {
log.Ctx(ctx).Info("replicate stream closed")
}
}()
s, err := replicate.CreateReplicateServer(stream)
if err != nil {
return err
}
return s.Execute()
}

View File

@ -41,7 +41,6 @@ import (
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/metrics"
"github.com/milvus-io/milvus/pkg/v2/proto/internalpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/util/expr"
"github.com/milvus-io/milvus/pkg/v2/util/logutil"
"github.com/milvus-io/milvus/pkg/v2/util/metricsinfo"
@ -263,10 +262,6 @@ func (node *Proxy) Init() error {
uuid.EnableRandPool()
log.Debug("enable rand pool for UUIDv4 generation")
if hookutil.IsClusterEncyptionEnabled() {
message.RegisterCipher(hookutil.GetCipher())
}
log.Info("init proxy done", zap.Int64("nodeID", paramtable.GetNodeID()), zap.String("Address", node.address))
return nil
}

View File

@ -0,0 +1,159 @@
package replicate
import (
"io"
"sync"
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
)
const replicateRespChanLength = 128
func CreateReplicateServer(streamServer milvuspb.MilvusService_CreateReplicateStreamServer) (*ReplicateStreamServer, error) {
clusterID, err := contextutil.GetClusterID(streamServer.Context())
if err != nil {
return nil, err
}
return &ReplicateStreamServer{
clusterID: clusterID,
streamServer: streamServer,
replicateRespCh: make(chan *milvuspb.ReplicateResponse, replicateRespChanLength),
wg: sync.WaitGroup{},
}, nil
}
// ReplicateStreamServer is a ReplicateStreamServer of replicate messages.
type ReplicateStreamServer struct {
clusterID string
streamServer milvuspb.MilvusService_CreateReplicateStreamServer
replicateRespCh chan *milvuspb.ReplicateResponse // All processing messages result should sent from theses channel.
wg sync.WaitGroup
}
// Execute starts the replicate server.
func (p *ReplicateStreamServer) Execute() error {
// Start a recv arm to handle the control message from client.
go func() {
// recv loop will be blocked until the stream is closed.
_ = p.recvLoop()
}()
// Start a send loop on current main goroutine.
// the loop will be blocked until the stream is closed.
err := p.sendLoop()
return err
}
// sendLoop sends the message to client.
func (p *ReplicateStreamServer) sendLoop() (err error) {
defer func() {
if err != nil {
log.Warn("send arm of stream closed by unexpected error", zap.Error(err))
return
}
log.Info("send arm of stream closed")
}()
for {
select {
case resp, ok := <-p.replicateRespCh:
if !ok {
return nil
}
if err := p.streamServer.Send(resp); err != nil {
return err
}
case <-p.streamServer.Context().Done():
return errors.Wrap(p.streamServer.Context().Err(), "cancel send loop by stream server")
}
}
}
// recvLoop receives the message from client.
func (p *ReplicateStreamServer) recvLoop() (err error) {
defer func() {
p.wg.Wait()
close(p.replicateRespCh)
if err != nil {
log.Warn("recv arm of stream closed by unexpected error", zap.Error(err))
return
}
log.Info("recv arm of stream closed")
}()
for {
req, err := p.streamServer.Recv()
if err == io.EOF {
return nil
}
if err != nil {
return err
}
switch req := req.Request.(type) {
case *milvuspb.ReplicateRequest_ReplicateMessage:
err := p.handleReplicateMessage(req)
if err != nil {
return err
}
default:
log.Warn("unknown request type", zap.Any("request", req))
}
}
}
// handleReplicateMessage handles the replicate message request.
func (p *ReplicateStreamServer) handleReplicateMessage(req *milvuspb.ReplicateRequest_ReplicateMessage) error {
// TODO: sheep, update metrics.
p.wg.Add(1)
defer p.wg.Done()
reqMsg := req.ReplicateMessage.GetMessage()
msg := message.NewReplicateMessage(req.ReplicateMessage.SourceClusterId, reqMsg)
sourceTs := msg.ReplicateHeader().TimeTick
log.Debug("recv replicate message from client",
zap.String("messageID", reqMsg.GetId().GetId()),
zap.Uint64("sourceTimeTick", sourceTs),
log.FieldMessage(msg),
)
// Append message to wal.
_, err := streaming.WAL().Replicate().Append(p.streamServer.Context(), msg)
if err == nil {
p.sendReplicateResult(sourceTs)
return nil
}
if status.AsStreamingError(err).IsIgnoredOperation() {
log.Info("append replicate message to wal ignored", log.FieldMessage(msg), zap.Error(err))
return nil
}
// unexpected error, will close the stream and wait for client to reconnect.
log.Warn("append replicate message to wal failed", log.FieldMessage(msg), zap.Error(err))
return err
}
// sendReplicateResult sends the replicate result to client.
func (p *ReplicateStreamServer) sendReplicateResult(sourceTimeTick uint64) {
resp := &milvuspb.ReplicateResponse{
Response: &milvuspb.ReplicateResponse_ReplicateConfirmedMessageInfo{
ReplicateConfirmedMessageInfo: &milvuspb.ReplicateConfirmedMessageInfo{
ConfirmedTimeTick: sourceTimeTick,
},
},
}
// If server context is canceled, it means the stream has been closed.
// all pending response message should be dropped, client side will handle it.
select {
case p.replicateRespCh <- resp:
log.Debug("send replicate message response to client", zap.Uint64("confirmedTimeTick", sourceTimeTick))
case <-p.streamServer.Context().Done():
log.Warn("stream closed before replicate message response sent", zap.Uint64("confirmedTimeTick", sourceTimeTick))
return
}
}

View File

@ -0,0 +1,237 @@
package replicate
import (
"context"
"io"
"sync"
"testing"
"time"
"github.com/apache/pulsar-client-go/pulsar"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"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-proto/go-api/v2/msgpb"
"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/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
pulsar2 "github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/pulsar"
)
// createContextWithClusterID creates a context with cluster ID in metadata (simulating incoming context)
func createContextWithClusterID(clusterID string) context.Context {
if clusterID == "" {
return context.Background()
}
md := metadata.New(map[string]string{
"cluster-id": clusterID,
})
return metadata.NewIncomingContext(context.Background(), md)
}
func TestReplicateStreamServer_Execute(t *testing.T) {
ctx := createContextWithClusterID("test-cluster")
mockStreamServer := newMockReplicateStreamServer(ctx)
const msgCount = replicateRespChanLength * 10
// Setup WAL mock
replicateService := mock_streaming.NewMockReplicateService(t)
tt := uint64(1)
replicateService.EXPECT().Append(mock.Anything, mock.Anything).
RunAndReturn(func(ctx context.Context, msg message.ReplicateMutableMessage) (*types.AppendResult, error) {
defer func() { tt++ }()
return &types.AppendResult{
TimeTick: tt,
}, nil
})
mockWAL := mock_streaming.NewMockWALAccesser(t)
mockWAL.EXPECT().Replicate().Return(replicateService)
streaming.SetWALForTest(mockWAL)
server, err := CreateReplicateServer(mockStreamServer)
assert.NoError(t, err)
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
err := server.Execute()
assert.NoError(t, err)
}()
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < msgCount; i++ {
tt := uint64(i + 1)
messageID := pulsar2.NewPulsarID(pulsar.EarliestMessageID())
msg := message.NewInsertMessageBuilderV1().
WithVChannel("test-vchannel").
WithHeader(&messagespb.InsertMessageHeader{}).
WithBody(&msgpb.InsertRequest{}).
MustBuildMutable().WithTimeTick(tt).
WithLastConfirmed(messageID)
milvusMsg := message.ImmutableMessageToMilvusMessage(commonpb.WALName_Pulsar.String(), msg.IntoImmutableMessage(messageID))
mockStreamServer.SendRequest(&milvuspb.ReplicateRequest{
Request: &milvuspb.ReplicateRequest_ReplicateMessage{
ReplicateMessage: &milvuspb.ReplicateMessage{
Message: milvusMsg,
},
},
})
}
}()
for i := 0; i < msgCount; i++ {
tt := uint64(i + 1)
sentResp := mockStreamServer.GetSentResponse()
assert.NotNil(t, sentResp)
assert.Equal(t, tt, sentResp.GetReplicateConfirmedMessageInfo().GetConfirmedTimeTick())
}
// Close the stream to stop execution
mockStreamServer.CloseSend()
wg.Wait()
}
func TestReplicateStreamServer_ContextCanceled(t *testing.T) {
ctx, cancel := context.WithCancel(createContextWithClusterID("test-cluster"))
mockStreamServer := newMockReplicateStreamServer(ctx)
server, err := CreateReplicateServer(mockStreamServer)
assert.NoError(t, err)
// Test send loop with canceled context
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
err := server.Execute()
assert.Error(t, err)
assert.True(t, errors.Is(err, context.Canceled))
}()
// Cancel context
cancel()
wg.Wait()
}
func TestReplicateStreamServer_recvLoop_RecvError(t *testing.T) {
ctx := createContextWithClusterID("test-cluster")
mockStreamServer := newMockReplicateStreamServer(ctx)
mockStreamServer.recvError = errors.New("recv error")
server, err := CreateReplicateServer(mockStreamServer)
assert.NoError(t, err)
// Test recv loop with recv error
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
err := server.recvLoop()
assert.Error(t, err)
assert.Equal(t, "recv error", err.Error())
}()
wg.Wait()
}
// mockReplicateStreamServer implements the milvuspb.MilvusService_CreateReplicateStreamServer interface
type mockReplicateStreamServer struct {
ctx context.Context
sendError error
recvError error
recvRequests chan *milvuspb.ReplicateRequest
sentResponses chan *milvuspb.ReplicateResponse
closeCh chan struct{}
timeout time.Duration
}
func newMockReplicateStreamServer(ctx context.Context) *mockReplicateStreamServer {
return &mockReplicateStreamServer{
ctx: ctx,
recvRequests: make(chan *milvuspb.ReplicateRequest, 128),
sentResponses: make(chan *milvuspb.ReplicateResponse, 128),
closeCh: make(chan struct{}, 1),
timeout: 10 * time.Second,
}
}
func (m *mockReplicateStreamServer) Send(resp *milvuspb.ReplicateResponse) error {
if m.sendError != nil {
return m.sendError
}
m.sentResponses <- resp
return nil
}
func (m *mockReplicateStreamServer) Recv() (*milvuspb.ReplicateRequest, error) {
if m.recvError != nil {
return nil, m.recvError
}
select {
case <-m.closeCh:
return nil, io.EOF
case req := <-m.recvRequests:
return req, nil
case <-time.After(m.timeout):
return nil, errors.New("recv timeout")
}
}
func (m *mockReplicateStreamServer) RecvMsg(msg interface{}) error {
return nil
}
func (m *mockReplicateStreamServer) SendMsg(msg interface{}) error {
return nil
}
func (m *mockReplicateStreamServer) Header() (metadata.MD, error) {
return nil, nil
}
func (m *mockReplicateStreamServer) Trailer() metadata.MD {
return nil
}
func (m *mockReplicateStreamServer) SendHeader(md metadata.MD) error {
return nil
}
func (m *mockReplicateStreamServer) SetHeader(md metadata.MD) error {
return nil
}
func (m *mockReplicateStreamServer) SetTrailer(md metadata.MD) {
}
func (m *mockReplicateStreamServer) Context() context.Context {
return m.ctx
}
func (m *mockReplicateStreamServer) CloseSend() error {
close(m.closeCh)
return nil
}
func (m *mockReplicateStreamServer) SendRequest(req *milvuspb.ReplicateRequest) {
m.recvRequests <- req
}
func (m *mockReplicateStreamServer) GetSentResponse() *milvuspb.ReplicateResponse {
select {
case resp := <-m.sentResponses:
return resp
case <-time.After(m.timeout):
return nil
}
}

View File

@ -72,7 +72,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/metrics"
"github.com/milvus-io/milvus/pkg/v2/mq/msgdispatcher"
"github.com/milvus-io/milvus/pkg/v2/proto/datapb"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/rmq"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/util/expr"
"github.com/milvus-io/milvus/pkg/v2/util/hardware"
"github.com/milvus-io/milvus/pkg/v2/util/lifetime"
@ -539,7 +539,7 @@ func (node *QueryNode) Stop() error {
err := node.session.GoingStop()
if err != nil {
log.Warn("session fail to go stopping state", zap.Error(err))
} else if util.MustSelectWALName() != rmq.WALName { // rocksmq cannot support querynode graceful stop because of using local storage.
} else if util.MustSelectWALName() != message.WALNameRocksmq { // rocksmq cannot support querynode graceful stop because of using local storage.
metrics.StoppingBalanceNodeNum.WithLabelValues().Set(1)
// TODO: Redundant timeout control, graceful stop timeout is controlled by outside by `component`.
// Integration test is still using it, Remove it in future.

View File

@ -46,6 +46,7 @@ import (
"github.com/milvus-io/milvus/internal/querynodev2/segments"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/internal/util/streamrpc"
"github.com/milvus-io/milvus/pkg/v2/common"
"github.com/milvus-io/milvus/pkg/v2/log"
@ -53,8 +54,8 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v2/proto/internalpb"
"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/types"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/rmq"
"github.com/milvus-io/milvus/pkg/v2/util/conc"
"github.com/milvus-io/milvus/pkg/v2/util/etcd"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
@ -2420,12 +2421,14 @@ func TestQueryNodeService(t *testing.T) {
local.EXPECT().GetLatestMVCCTimestampIfLocal(mock.Anything, mock.Anything).Return(0, nil).Maybe()
local.EXPECT().GetMetricsIfLocal(mock.Anything).Return(&types.StreamingNodeMetrics{}, nil).Maybe()
wal.EXPECT().Local().Return(local).Maybe()
wal.EXPECT().WALName().Return(rmq.WALName).Maybe()
scanner := mock_streaming.NewMockScanner(t)
scanner.EXPECT().Done().Return(make(chan struct{})).Maybe()
scanner.EXPECT().Error().Return(nil).Maybe()
scanner.EXPECT().Close().Return().Maybe()
wal.EXPECT().Read(mock.Anything, mock.Anything).Return(scanner).Maybe()
paramtable.SetRole(typeutil.StandaloneRole)
paramtable.Get().MQCfg.Type.SwapTempValue(message.WALNameRocksmq.String())
util.InitAndSelectWALName()
streaming.SetWALForTest(wal)
defer streaming.RecoverWALForTest()

View File

@ -8,11 +8,13 @@ import (
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -96,6 +98,32 @@ func (c *AssignmentServiceImpl) ReportAssignmentError(ctx context.Context, pchan
return nil
}
// UpdateReplicateConfiguration updates the replicate configuration to the milvus cluster.
func (c *AssignmentServiceImpl) UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error {
if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
return status.NewOnShutdownError("assignment service client is closing")
}
defer c.lifetime.Done()
service, err := c.service.GetService(c.ctx)
if err != nil {
return err
}
_, err = service.UpdateReplicateConfiguration(ctx, &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: config,
})
return err
}
func (c *AssignmentServiceImpl) GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
if !c.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, status.NewOnShutdownError("assignment service client is closing")
}
defer c.lifetime.Done()
return c.watcher.GetLatestReplicateConfiguration(ctx)
}
// Close closes the assignment service.
func (c *AssignmentServiceImpl) Close() {
c.lifetime.SetState(typeutil.LifetimeStateStopped)

View File

@ -15,10 +15,13 @@ import (
"github.com/milvus-io/milvus/pkg/v2/mocks/proto/mock_streamingpb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
func TestAssignmentService(t *testing.T) {
paramtable.Init()
s := mock_lazygrpc.NewMockService[streamingpb.StreamingCoordAssignmentServiceClient](t)
c := mock_streamingpb.NewMockStreamingCoordAssignmentServiceClient(t)
s.EXPECT().GetService(mock.Anything).Return(c, nil)

View File

@ -8,6 +8,8 @@ import (
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -22,6 +24,7 @@ func newAssignmentDiscoverClient(w *watcher, streamClient streamingpb.StreamingC
exitCh: make(chan struct{}),
wg: sync.WaitGroup{},
lastErrorReportedTerm: make(map[string]int64),
clusterID: paramtable.Get().CommonCfg.ClusterPrefix.GetValue(),
}
c.executeBackgroundTask()
return c
@ -37,6 +40,7 @@ type assignmentDiscoverClient struct {
wg sync.WaitGroup
streamClient streamingpb.StreamingCoordAssignmentService_AssignmentDiscoverClient
lastErrorReportedTerm map[string]int64
clusterID string
}
// ReportAssignmentError reports the assignment error to server.
@ -162,6 +166,9 @@ func (c *assignmentDiscoverClient) recvLoop() (err error) {
Version: newIncomingVersion,
Assignments: newIncomingAssignments,
CChannel: resp.FullAssignment.Cchannel,
ReplicateConfigHelper: replicateutil.MustNewConfigHelper(
c.clusterID,
resp.FullAssignment.ReplicateConfiguration),
})
case *streamingpb.AssignmentDiscoverResponse_Close:
// nothing to do now, just wait io.EOF.

View File

@ -7,6 +7,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -18,8 +19,9 @@ func newWatcher() *watcher {
return &watcher{
cond: syncutil.NewContextCond(&sync.Mutex{}),
lastVersionedAssignment: types.VersionedStreamingNodeAssignments{
Version: typeutil.VersionInt64Pair{Global: -1, Local: -1},
Assignments: make(map[int64]types.StreamingNodeAssignment),
Version: typeutil.VersionInt64Pair{Global: -1, Local: -1},
Assignments: make(map[int64]types.StreamingNodeAssignment),
ReplicateConfigHelper: nil,
},
}
}
@ -42,6 +44,19 @@ func (w *watcher) GetLatestDiscover(ctx context.Context) (*types.VersionedStream
return &last, nil
}
func (w *watcher) GetLatestReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error) {
w.cond.L.Lock()
for (w.lastVersionedAssignment.Version.Global == -1 && w.lastVersionedAssignment.Version.Local == -1) ||
w.lastVersionedAssignment.ReplicateConfigHelper == nil {
if err := w.cond.Wait(ctx); err != nil {
return nil, err
}
}
last := w.lastVersionedAssignment
w.cond.L.Unlock()
return last.ReplicateConfigHelper, nil
}
// AssignmentDiscover watches the assignment discovery.
func (w *watcher) AssignmentDiscover(ctx context.Context, cb func(*types.VersionedStreamingNodeAssignments) error) error {
w.cond.L.Lock()

View File

@ -10,9 +10,8 @@ import (
)
// NewGRPCBroadcastService creates a new broadcast service with grpc.
func NewGRPCBroadcastService(walName string, service lazygrpc.Service[streamingpb.StreamingCoordBroadcastServiceClient]) *GRPCBroadcastServiceImpl {
func NewGRPCBroadcastService(service lazygrpc.Service[streamingpb.StreamingCoordBroadcastServiceClient]) *GRPCBroadcastServiceImpl {
return &GRPCBroadcastServiceImpl{
walName: walName,
service: service,
}
}
@ -20,7 +19,6 @@ func NewGRPCBroadcastService(walName string, service lazygrpc.Service[streamingp
// GRPCBroadcastServiceImpl is the implementation of BroadcastService based on grpc service.
// If the streaming coord is not deployed at current node, these implementation will be used.
type GRPCBroadcastServiceImpl struct {
walName string
service lazygrpc.Service[streamingpb.StreamingCoordBroadcastServiceClient]
}
@ -37,7 +35,7 @@ func (c *GRPCBroadcastServiceImpl) Broadcast(ctx context.Context, msg message.Br
}
results := make(map[string]*types.AppendResult, len(resp.Results))
for channel, result := range resp.Results {
msgID, err := message.UnmarshalMessageID(c.walName, result.Id.Id)
msgID, err := message.UnmarshalMessageID(result.Id)
if err != nil {
return nil, err
}

View File

@ -12,7 +12,6 @@ import (
"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_lazygrpc"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
"github.com/milvus-io/milvus/pkg/v2/mocks/proto/mock_streamingpb"
"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/rmq"
@ -21,7 +20,7 @@ import (
func TestBroadcast(t *testing.T) {
s := newMockServer(t, 0)
bs := NewGRPCBroadcastService(walimplstest.WALName, s)
bs := NewGRPCBroadcastService(s)
msg := message.NewDropCollectionMessageBuilderV1().
WithHeader(&message.DropCollectionMessageHeader{}).
WithBody(&msgpb.DropCollectionRequest{}).
@ -42,9 +41,7 @@ func newMockServer(t *testing.T, sendDelay time.Duration) lazygrpc.Service[strea
c.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&streamingpb.BroadcastResponse{
Results: map[string]*streamingpb.ProduceMessageResponseResult{
"v1": {
Id: &messagespb.MessageID{
Id: walimplstest.NewTestMessageID(1).Marshal(),
},
Id: walimplstest.NewTestMessageID(1).IntoProto(),
},
},
BroadcastId: 1,

View File

@ -8,6 +8,7 @@ import (
"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/json"
"github.com/milvus-io/milvus/internal/streamingcoord/client/assignment"
"github.com/milvus-io/milvus/internal/streamingcoord/client/broadcast"
@ -16,13 +17,13 @@ import (
streamingserviceinterceptor "github.com/milvus-io/milvus/internal/util/streamingutil/service/interceptor"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"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/tracer"
"github.com/milvus-io/milvus/pkg/v2/util/interceptor"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -33,6 +34,12 @@ type AssignmentService interface {
// AssignmentDiscover is used to watches the assignment discovery.
types.AssignmentDiscoverWatcher
// UpdateReplicateConfiguration updates the replicate configuration to the milvus cluster.
UpdateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) error
// GetReplicateConfiguration returns the replicate configuration of the milvus cluster.
GetReplicateConfiguration(ctx context.Context) (*replicateutil.ConfigHelper, error)
// GetLatestAssignments returns the latest assignment discovery result.
GetLatestAssignments(ctx context.Context) (*types.VersionedStreamingNodeAssignments, error)
@ -89,7 +96,7 @@ func NewClient(etcdCli *clientv3.Client) Client {
conn: conn,
rb: rb,
assignmentService: assignmentServiceImpl,
broadcastService: broadcast.NewGRPCBroadcastService(util.MustSelectWALName(), broadcastService),
broadcastService: broadcast.NewGRPCBroadcastService(broadcastService),
}
}

View File

@ -7,6 +7,7 @@ import (
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"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/util/syncutil"
)
@ -26,6 +27,9 @@ type (
// Balancer is a local component, it should promise all channel can be assigned, and reach the final consistency.
// Balancer should be thread safe.
type Balancer interface {
// GetLatestChannelAssignment returns the latest channel assignment.
GetLatestChannelAssignment() (*WatchChannelAssignmentsCallbackParam, error)
// GetAllStreamingNodes fetches all streaming node info.
GetAllStreamingNodes(ctx context.Context) (map[int64]*types.StreamingNodeInfo, error)
@ -51,6 +55,9 @@ type Balancer interface {
// MarkAsAvailable marks the pchannels as available, and trigger a rebalance.
MarkAsUnavailable(ctx context.Context, pChannels []types.PChannelInfo) error
// UpdateReplicateConfiguration updates the replicate configuration.
UpdateReplicateConfiguration(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2) error
// Trigger is a hint to trigger a balance.
Trigger(ctx context.Context) error

View File

@ -15,6 +15,7 @@ import (
"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
"github.com/milvus-io/milvus/internal/util/streamingutil/status"
"github.com/milvus-io/milvus/pkg/v2/log"
"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/util/contextutil"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
@ -83,6 +84,15 @@ func (b *balancerImpl) RegisterStreamingEnabledNotifier(notifier *syncutil.Async
b.channelMetaManager.RegisterStreamingEnabledNotifier(notifier)
}
func (b *balancerImpl) GetLatestChannelAssignment() (*WatchChannelAssignmentsCallbackParam, error) {
if !b.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, status.NewOnShutdownError("balancer is closing")
}
defer b.lifetime.Done()
return b.channelMetaManager.GetLatestChannelAssignment()
}
// GetAllStreamingNodes fetches all streaming node info.
func (b *balancerImpl) GetAllStreamingNodes(ctx context.Context) (map[int64]*types.StreamingNodeInfo, error) {
return resource.Resource().StreamingNodeManagerClient().GetAllStreamingNodes(ctx)
@ -105,6 +115,22 @@ func (b *balancerImpl) WatchChannelAssignments(ctx context.Context, cb WatchChan
return b.channelMetaManager.WatchAssignmentResult(ctx, cb)
}
// UpdateReplicateConfiguration updates the replicate configuration.
func (b *balancerImpl) UpdateReplicateConfiguration(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2) error {
if !b.lifetime.Add(typeutil.LifetimeStateWorking) {
return status.NewOnShutdownError("balancer is closing")
}
defer b.lifetime.Done()
ctx, cancel := contextutil.MergeContext(ctx, b.ctx)
defer cancel()
if err := b.channelMetaManager.UpdateReplicateConfiguration(ctx, msgs...); err != nil {
return err
}
return nil
}
// UpdateBalancePolicy update the balance policy.
func (b *balancerImpl) UpdateBalancePolicy(ctx context.Context, req *types.UpdateWALBalancePolicyRequest) (*types.UpdateWALBalancePolicyResponse, error) {
if !b.lifetime.Add(typeutil.LifetimeStateWorking) {

View File

@ -106,6 +106,7 @@ func TestBalancer(t *testing.T) {
}, nil
})
catalog.EXPECT().SavePChannels(mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, nil)
// Test for lower datanode and proxy version protection.
metaRoot := paramtable.Get().EtcdCfg.MetaRootPath.GetValue()
@ -331,6 +332,7 @@ func TestBalancer_WithRecoveryLag(t *testing.T) {
}, nil
})
catalog.EXPECT().SavePChannels(mock.Anything, mock.Anything).Return(nil).Maybe()
catalog.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, nil)
ctx := context.Background()
b, err := balancer.RecoverBalancer(ctx, "test-channel-1")

View File

@ -7,10 +7,14 @@ 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/streamingcoord/server/resource"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"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/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/retry"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
@ -20,9 +24,11 @@ var ErrChannelNotExist = errors.New("channel not exist")
type (
WatchChannelAssignmentsCallbackParam struct {
Version typeutil.VersionInt64Pair
CChannelAssignment *streamingpb.CChannelAssignment
Relations []types.PChannelInfoAssigned
Version typeutil.VersionInt64Pair
CChannelAssignment *streamingpb.CChannelAssignment
PChannelView *PChannelView
Relations []types.PChannelInfoAssigned
ReplicateConfiguration *commonpb.ReplicateConfiguration
}
WatchChannelAssignmentsCallback func(param WatchChannelAssignmentsCallbackParam) error
)
@ -39,7 +45,10 @@ func RecoverChannelManager(ctx context.Context, incomingChannel ...string) (*Cha
if err != nil {
return nil, err
}
replicateConfig, err := recoverReplicateConfiguration(ctx)
if err != nil {
return nil, err
}
channels, metrics, err := recoverFromConfigurationAndMeta(ctx, streamingVersion, incomingChannel...)
if err != nil {
return nil, err
@ -56,6 +65,7 @@ func RecoverChannelManager(ctx context.Context, incomingChannel ...string) (*Cha
metrics: metrics,
cchannelMeta: cchannelMeta,
streamingVersion: streamingVersion,
replicateConfig: replicateConfig,
}, nil
}
@ -116,6 +126,14 @@ func recoverFromConfigurationAndMeta(ctx context.Context, streamingVersion *stre
return channels, metrics, nil
}
func recoverReplicateConfiguration(ctx context.Context) (*replicateConfigHelper, error) {
config, err := resource.Resource().StreamingCatalog().GetReplicateConfiguration(ctx)
if err != nil {
return nil, err
}
return newReplicateConfigHelper(config), nil
}
// ChannelManager manages the channels.
// ChannelManager is the `wal` of channel assignment and unassignment.
// Every operation applied to the streaming node should be recorded in ChannelManager first.
@ -129,6 +147,7 @@ type ChannelManager struct {
// null if no streaming service has been run.
// 1 if streaming service has been run once.
streamingEnableNotifiers []*syncutil.AsyncTaskNotifier[struct{}]
replicateConfig *replicateConfigHelper
}
// RegisterStreamingEnabledNotifier registers a notifier into the balancer.
@ -320,6 +339,18 @@ func (cm *ChannelManager) GetLatestWALLocated(ctx context.Context, pchannel stri
return 0, false
}
// GetLatestChannelAssignment returns the latest channel assignment.
func (cm *ChannelManager) GetLatestChannelAssignment() (*WatchChannelAssignmentsCallbackParam, error) {
var result WatchChannelAssignmentsCallbackParam
if _, err := cm.applyAssignments(func(param WatchChannelAssignmentsCallbackParam) error {
result = param
return nil
}); err != nil {
return nil, err
}
return &result, nil
}
func (cm *ChannelManager) WatchAssignmentResult(ctx context.Context, cb WatchChannelAssignmentsCallback) error {
// push the first balance result to watcher callback function if balance result is ready.
version, err := cm.applyAssignments(cb)
@ -337,6 +368,52 @@ func (cm *ChannelManager) WatchAssignmentResult(ctx context.Context, cb WatchCha
}
}
// UpdateReplicateConfiguration updates the in-memory replicate configuration.
func (cm *ChannelManager) UpdateReplicateConfiguration(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2) error {
config := replicateutil.MustNewConfigHelper(paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), msgs[0].Header().ReplicateConfiguration)
pchannels := make([]types.AckedCheckpoint, 0, len(msgs))
for _, msg := range msgs {
pchannels = append(pchannels, types.AckedCheckpoint{
Channel: funcutil.ToPhysicalChannel(msg.VChannel()),
MessageID: msg.LastConfirmedMessageID(),
LastConfirmedMessageID: msg.LastConfirmedMessageID(),
TimeTick: msg.TimeTick(),
})
}
cm.cond.L.Lock()
defer cm.cond.L.Unlock()
if cm.replicateConfig == nil {
cm.replicateConfig = newReplicateConfigHelperFromMessage(msgs[0])
} else {
// StartUpdating starts the updating process.
if !cm.replicateConfig.StartUpdating(config.GetReplicateConfiguration(), msgs[0].BroadcastHeader().VChannels) {
return nil
}
}
cm.replicateConfig.Apply(config.GetReplicateConfiguration(), pchannels)
dirtyConfig, dirtyCDCTasks, dirty := cm.replicateConfig.ConsumeIfDirty(config.GetReplicateConfiguration())
if !dirty {
// the meta is not dirty, so nothing updated, return it directly.
return nil
}
if err := resource.Resource().StreamingCatalog().SaveReplicateConfiguration(ctx, dirtyConfig, dirtyCDCTasks); err != nil {
return err
}
// If the acked result is nil, it means the all the channels are acked,
// so we can update the version and push the new replicate configuration into client.
if dirtyConfig.AckedResult == nil {
// update metrics.
cm.cond.UnsafeBroadcast()
cm.version.Local++
cm.metrics.UpdateAssignmentVersion(cm.version.Local)
}
return nil
}
// applyAssignments applies the assignments.
func (cm *ChannelManager) applyAssignments(cb WatchChannelAssignmentsCallback) (typeutil.VersionInt64Pair, error) {
cm.cond.L.Lock()
@ -348,13 +425,21 @@ func (cm *ChannelManager) applyAssignments(cb WatchChannelAssignmentsCallback) (
}
version := cm.version
cchannelAssignment := proto.Clone(cm.cchannelMeta).(*streamingpb.CChannelMeta)
pchannelViews := newPChannelView(cm.channels)
cm.cond.L.Unlock()
var replicateConfig *commonpb.ReplicateConfiguration
if cm.replicateConfig != nil {
replicateConfig = cm.replicateConfig.GetReplicateConfiguration()
}
return version, cb(WatchChannelAssignmentsCallbackParam{
Version: version,
CChannelAssignment: &streamingpb.CChannelAssignment{
Meta: cchannelAssignment,
},
Relations: assignments,
PChannelView: pchannelViews,
Relations: assignments,
ReplicateConfiguration: replicateConfig,
})
}

View File

@ -32,6 +32,7 @@ func TestChannelManager(t *testing.T) {
Version: 1,
}, nil)
catalog.EXPECT().ListPChannel(mock.Anything).Return(nil, errors.New("recover failure"))
catalog.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, nil)
m, err := RecoverChannelManager(ctx)
assert.Nil(t, m)
assert.Error(t, err)
@ -131,6 +132,7 @@ func TestStreamingEnableChecker(t *testing.T) {
catalog.EXPECT().GetVersion(mock.Anything).Return(nil, nil)
catalog.EXPECT().SaveVersion(mock.Anything, mock.Anything).Return(nil)
catalog.EXPECT().ListPChannel(mock.Anything).Return(nil, nil)
catalog.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, nil)
m, err := RecoverChannelManager(ctx, "test-channel")
assert.NoError(t, err)
@ -183,6 +185,7 @@ func TestChannelManagerWatch(t *testing.T) {
}, nil
})
catalog.EXPECT().SavePChannels(mock.Anything, mock.Anything).Return(nil)
catalog.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, nil)
manager, err := RecoverChannelManager(context.Background())
assert.NoError(t, err)

View File

@ -22,7 +22,7 @@ func newPChannelView(metas map[ChannelID]*PChannelMeta) *PChannelView {
panic(fmt.Sprintf("duplicate rw channel: %s", id.String()))
}
view.Channels[id] = meta
stat := StaticPChannelStatsManager.MustGet().GetPChannelStats(id).View()
stat := StaticPChannelStatsManager.Get().GetPChannelStats(id).View()
stat.LastAssignTimestamp = meta.LastAssignTimestamp()
view.Stats[id] = stat
}

View File

@ -0,0 +1,117 @@
package channel
import (
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"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/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
)
// replicateConfigHelper is a helper to manage the replicate configuration.
type replicateConfigHelper struct {
*replicateutil.ConfigHelper
ackedPendings *types.AckedResult
dirty bool
}
// newReplicateConfigHelperFromMessage creates a new replicate config helper from message.
func newReplicateConfigHelperFromMessage(replicateConfig message.ImmutablePutReplicateConfigMessageV2) *replicateConfigHelper {
return newReplicateConfigHelper(&streamingpb.ReplicateConfigurationMeta{
ReplicateConfiguration: nil,
AckedResult: types.NewAckedPendings(replicateConfig.BroadcastHeader().VChannels).AckedResult,
})
}
// newReplicateConfigHelper creates a new replicate config helper from proto.
func newReplicateConfigHelper(replicateConfig *streamingpb.ReplicateConfigurationMeta) *replicateConfigHelper {
if replicateConfig == nil {
return nil
}
return &replicateConfigHelper{
ConfigHelper: replicateutil.MustNewConfigHelper(paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), replicateConfig.GetReplicateConfiguration()),
ackedPendings: types.NewAckedPendingsFromProto(replicateConfig.GetAckedResult()),
dirty: false,
}
}
// StartUpdating starts the updating process.
// return true if the replicate configuration is changed, false otherwise.
func (rc *replicateConfigHelper) StartUpdating(config *commonpb.ReplicateConfiguration, pchannels []string) bool {
if rc.ConfigHelper != nil && proto.Equal(config, rc.GetReplicateConfiguration()) {
return false
}
if rc.ackedPendings == nil {
rc.ackedPendings = types.NewAckedPendings(pchannels)
}
return true
}
// Apply applies the replicate configuration to the wal.
func (rc *replicateConfigHelper) Apply(config *commonpb.ReplicateConfiguration, cp []types.AckedCheckpoint) {
if rc.ackedPendings == nil {
panic("ackedPendings is nil when applying replicate configuration")
}
for _, cp := range cp {
if rc.ackedPendings.Ack(cp) {
rc.dirty = true
}
}
}
// ConsumeIfDirty consumes the dirty part of the replicate configuration.
func (rc *replicateConfigHelper) ConsumeIfDirty(incoming *commonpb.ReplicateConfiguration) (config *streamingpb.ReplicateConfigurationMeta, replicatingTasks []*streamingpb.ReplicatePChannelMeta, dirty bool) {
if !rc.dirty {
return nil, nil, false
}
rc.dirty = false
if !rc.ackedPendings.IsAllAcked() {
// not all the channels are acked, return the current replicate configuration and acked result.
var cfg *commonpb.ReplicateConfiguration
if rc.ConfigHelper != nil {
cfg = rc.ConfigHelper.GetReplicateConfiguration()
}
return &streamingpb.ReplicateConfigurationMeta{
ReplicateConfiguration: cfg,
AckedResult: rc.ackedPendings.AckedResult,
}, nil, true
}
// all the channels are acked, return the new replicate configuration and acked result.
newConfig := replicateutil.MustNewConfigHelper(paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), incoming)
newIncomingCDCTasks := rc.getNewIncomingTask(newConfig)
rc.ConfigHelper = newConfig
rc.ackedPendings = nil
return &streamingpb.ReplicateConfigurationMeta{
ReplicateConfiguration: incoming,
AckedResult: nil,
}, newIncomingCDCTasks, true
}
// getNewIncomingTask gets the new incoming task from replicatingTasks.
func (cm *replicateConfigHelper) getNewIncomingTask(newConfig *replicateutil.ConfigHelper) []*streamingpb.ReplicatePChannelMeta {
incoming := newConfig.GetCurrentCluster()
var current *replicateutil.MilvusCluster
if cm.ConfigHelper != nil {
current = cm.ConfigHelper.GetCurrentCluster()
}
incomingReplicatingTasks := make([]*streamingpb.ReplicatePChannelMeta, 0, len(incoming.TargetClusters()))
for _, targetCluster := range incoming.TargetClusters() {
if current != nil && current.TargetCluster(targetCluster.GetClusterId()) != nil {
// target already exists, skip it.
continue
}
for _, pchannel := range targetCluster.GetPchannels() {
incomingReplicatingTasks = append(incomingReplicatingTasks, &streamingpb.ReplicatePChannelMeta{
SourceChannelName: targetCluster.MustGetSourceChannel(pchannel),
TargetChannelName: pchannel,
TargetCluster: targetCluster.MilvusCluster,
})
}
}
return incomingReplicatingTasks
}

View File

@ -0,0 +1,129 @@
package channel
import (
"testing"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
)
type ReplicateConfigHelperSuite struct {
suite.Suite
helper *replicateConfigHelper
}
func TestReplicateConfigHelperSuite(t *testing.T) {
suite.Run(t, new(ReplicateConfigHelperSuite))
}
func (s *ReplicateConfigHelperSuite) SetupTest() {
s.helper = nil
}
func (s *ReplicateConfigHelperSuite) TestNewReplicateConfigHelper() {
// Test nil input
helper := newReplicateConfigHelper(nil)
s.Nil(helper)
// Test valid input
meta := &streamingpb.ReplicateConfigurationMeta{
ReplicateConfiguration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
AckedResult: types.NewAckedPendings([]string{"p1", "p2"}).AckedResult,
}
helper = newReplicateConfigHelper(meta)
s.NotNil(helper)
s.NotNil(helper.ConfigHelper)
s.NotNil(helper.ackedPendings)
s.False(helper.dirty)
}
func (s *ReplicateConfigHelperSuite) TestStartUpdating() {
s.helper = &replicateConfigHelper{
ConfigHelper: nil,
ackedPendings: nil,
dirty: false,
}
config := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
}
pchannels := []string{"p1", "p2"}
// First update should return true
changed := s.helper.StartUpdating(config, pchannels)
s.True(changed)
s.NotNil(s.helper.ackedPendings)
s.helper.Apply(config, []types.AckedCheckpoint{
{Channel: "p1", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1},
{Channel: "p2", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1},
})
s.helper.ConsumeIfDirty(config)
// Same config should return false
changed = s.helper.StartUpdating(config, pchannels)
s.False(changed)
}
func (s *ReplicateConfigHelperSuite) TestApply() {
s.helper = &replicateConfigHelper{
ConfigHelper: nil,
ackedPendings: types.NewAckedPendings([]string{"p1", "p2"}),
dirty: false,
}
config := &commonpb.ReplicateConfiguration{}
checkpoints := []types.AckedCheckpoint{
{Channel: "p1", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1},
{Channel: "p2", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1},
}
s.helper.Apply(config, checkpoints)
s.True(s.helper.dirty)
s.True(s.helper.ackedPendings.IsAllAcked())
}
func (s *ReplicateConfigHelperSuite) TestConsumeIfDirty() {
s.helper = &replicateConfigHelper{
ConfigHelper: nil,
ackedPendings: types.NewAckedPendings([]string{"p1", "p2"}),
dirty: true,
}
// Not all acked case
config, tasks, dirty := s.helper.ConsumeIfDirty(&commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
})
s.NotNil(config)
s.Nil(tasks)
s.True(dirty)
s.False(s.helper.dirty)
// All acked case
s.helper.dirty = true
s.helper.ackedPendings.Ack(types.AckedCheckpoint{Channel: "p1", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1})
s.helper.ackedPendings.Ack(types.AckedCheckpoint{Channel: "p2", MessageID: walimplstest.NewTestMessageID(1), LastConfirmedMessageID: walimplstest.NewTestMessageID(1), TimeTick: 1})
config, tasks, dirty = s.helper.ConsumeIfDirty(&commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
})
s.NotNil(config)
s.NotNil(tasks)
s.True(dirty)
s.False(s.helper.dirty)
s.Nil(s.helper.ackedPendings)
}

View File

@ -4,12 +4,20 @@ import (
"context"
"github.com/prometheus/client_golang/prometheus"
"github.com/samber/lo"
"go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
"github.com/milvus-io/milvus/internal/streamingcoord/server/service/discover"
"github.com/milvus-io/milvus/pkg/v2/log"
"github.com/milvus-io/milvus/pkg/v2/metrics"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/rmq"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/replicateutil"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
)
@ -19,10 +27,13 @@ var _ streamingpb.StreamingCoordAssignmentServiceServer = (*assignmentServiceImp
func NewAssignmentService(
balancer *syncutil.Future[balancer.Balancer],
) streamingpb.StreamingCoordAssignmentServiceServer {
return &assignmentServiceImpl{
assignmentService := &assignmentServiceImpl{
balancer: balancer,
listenerTotal: metrics.StreamingCoordAssignmentListenerTotal.WithLabelValues(paramtable.GetStringNodeID()),
}
// TODO: after recovering from wal, add it to here.
// registry.RegisterPutReplicateConfigV2AckCallback(assignmentService.putReplicateConfiguration)
return assignmentService
}
type AssignmentService interface {
@ -31,6 +42,8 @@ type AssignmentService interface {
// assignmentServiceImpl is the implementation of the assignment service.
type assignmentServiceImpl struct {
streamingpb.UnimplementedStreamingCoordAssignmentServiceServer
balancer *syncutil.Future[balancer.Balancer]
listenerTotal prometheus.Gauge
}
@ -47,6 +60,86 @@ func (s *assignmentServiceImpl) AssignmentDiscover(server streamingpb.StreamingC
return discover.NewAssignmentDiscoverServer(balancer, server).Execute()
}
// UpdateReplicateConfiguration updates the replicate configuration to the milvus cluster.
func (s *assignmentServiceImpl) UpdateReplicateConfiguration(ctx context.Context, req *streamingpb.UpdateReplicateConfigurationRequest) (*streamingpb.UpdateReplicateConfigurationResponse, error) {
config := req.GetConfiguration()
log.Ctx(ctx).Info("UpdateReplicateConfiguration received", replicateutil.ConfigLogFields(config)...)
// TODO: after recovering from wal, do a broadcast operation here.
msg, err := s.validateReplicateConfiguration(ctx, config)
if err != nil {
return nil, err
}
// TODO: After recovering from wal, we can get the immutable message from wal system.
// Now, we just mock the immutable message here.
mutableMsg := msg.SplitIntoMutableMessage()
mockMessages := make([]message.ImmutablePutReplicateConfigMessageV2, 0)
for _, msg := range mutableMsg {
mockMessages = append(mockMessages,
message.MustAsImmutablePutReplicateConfigMessageV2(msg.WithTimeTick(0).WithLastConfirmedUseMessageID().IntoImmutableMessage(rmq.NewRmqID(1))),
)
}
// TODO: After recovering from wal, remove the operation here.
if err := s.putReplicateConfiguration(ctx, mockMessages...); err != nil {
return nil, err
}
return &streamingpb.UpdateReplicateConfigurationResponse{}, nil
}
// validateReplicateConfiguration validates the replicate configuration.
func (s *assignmentServiceImpl) validateReplicateConfiguration(ctx context.Context, config *commonpb.ReplicateConfiguration) (message.BroadcastMutableMessage, error) {
balancer, err := s.balancer.GetWithContext(ctx)
if err != nil {
return nil, err
}
// get all pchannels
latestAssignment, err := balancer.GetLatestChannelAssignment()
if err != nil {
return nil, err
}
pchannels := lo.MapToSlice(latestAssignment.PChannelView.Channels, func(_ channel.ChannelID, channel *channel.PChannelMeta) string {
return channel.Name()
})
// validate the configuration itself
currentClusterID := paramtable.Get().CommonCfg.ClusterPrefix.GetValue()
validator := replicateutil.NewReplicateConfigValidator(config, currentClusterID, pchannels)
if err := validator.Validate(); err != nil {
log.Ctx(ctx).Warn("UpdateReplicateConfiguration fail", zap.Error(err))
return nil, err
}
// TODO: validate the incoming configuration is compatible with the current config.
if _, err := replicateutil.NewConfigHelper(paramtable.Get().CommonCfg.ClusterPrefix.GetValue(), config); err != nil {
return nil, err
}
b := message.NewPutReplicateConfigMessageBuilderV2().
WithHeader(&message.PutReplicateConfigMessageHeader{
ReplicateConfiguration: config,
}).
WithBody(&message.PutReplicateConfigMessageBody{}).
WithBroadcast(pchannels).
MustBuildBroadcast()
// TODO: After recovering from wal, remove the operation here.
b.WithBroadcastID(1)
return b, nil
}
// putReplicateConfiguration puts the replicate configuration into the balancer.
// It's a callback function of the broadcast service.
func (s *assignmentServiceImpl) putReplicateConfiguration(ctx context.Context, msgs ...message.ImmutablePutReplicateConfigMessageV2) error {
balancer, err := s.balancer.GetWithContext(ctx)
if err != nil {
return err
}
return balancer.UpdateReplicateConfiguration(ctx, msgs...)
}
// UpdateWALBalancePolicy is used to update the WAL balance policy.
func (s *assignmentServiceImpl) UpdateWALBalancePolicy(ctx context.Context, req *streamingpb.UpdateWALBalancePolicyRequest) (*streamingpb.UpdateWALBalancePolicyResponse, error) {
balancer, err := s.balancer.GetWithContext(ctx)

View File

@ -4,7 +4,6 @@ import (
"context"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster"
"github.com/milvus-io/milvus/internal/util/streamingutil/util"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
@ -19,14 +18,12 @@ type BroadcastService interface {
func NewBroadcastService(bc *syncutil.Future[broadcaster.Broadcaster]) BroadcastService {
return &broadcastServceImpl{
broadcaster: bc,
walName: util.MustSelectWALName(),
}
}
// broadcastServiceeeeImpl is the implementation of the broadcast service.
type broadcastServceImpl struct {
broadcaster *syncutil.Future[broadcaster.Broadcaster]
walName string
}
// Broadcast broadcasts the message to all channels.
@ -63,7 +60,11 @@ func (s *broadcastServceImpl) Ack(ctx context.Context, req *streamingpb.Broadcas
}
return &streamingpb.BroadcastAckResponse{}, nil
}
if err := broadcaster.Ack(ctx, message.NewImmutableMessageFromProto(s.walName, req.Message)); err != nil {
if err := broadcaster.Ack(ctx, message.NewImmutableMesasge(
message.MustUnmarshalMessageID(req.Message.Id),
req.Message.Payload,
req.Message.Properties,
)); err != nil {
return nil, err
}
return &streamingpb.BroadcastAckResponse{}, nil

View File

@ -0,0 +1,46 @@
package service
import (
"context"
"testing"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster"
"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v2/streaming/walimpls/impls/walimplstest"
"github.com/milvus-io/milvus/pkg/v2/util/syncutil"
)
func TestBroadcastService(t *testing.T) {
fb := syncutil.NewFuture[broadcaster.Broadcaster]()
mb := mock_broadcaster.NewMockBroadcaster(t)
fb.Set(mb)
mb.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&types.BroadcastAppendResult{}, nil)
mb.EXPECT().Ack(mock.Anything, mock.Anything).Return(nil)
mb.EXPECT().LegacyAck(mock.Anything, mock.Anything, mock.Anything).Return(nil)
service := NewBroadcastService(fb)
service.Broadcast(context.Background(), &streamingpb.BroadcastRequest{
Message: &messagespb.Message{
Payload: []byte("payload"),
Properties: map[string]string{"_bh": "1"},
},
})
service.Ack(context.Background(), &streamingpb.BroadcastAckRequest{
BroadcastId: 1,
Vchannel: "v1",
})
service.Ack(context.Background(), &streamingpb.BroadcastAckRequest{
BroadcastId: 1,
Vchannel: "v1",
Message: &commonpb.ImmutableMessage{
Id: walimplstest.NewTestMessageID(1).IntoProto(),
Payload: []byte("payload"),
Properties: map[string]string{"key": "value"},
},
})
}

View File

@ -36,8 +36,9 @@ func (h *discoverGrpcServerHelper) SendFullAssignment(param balancer.WatchChanne
Global: param.Version.Global,
Local: param.Version.Local,
},
Assignments: assignments,
Cchannel: param.CChannelAssignment,
Assignments: assignments,
Cchannel: param.CChannelAssignment,
ReplicateConfiguration: param.ReplicateConfiguration,
},
},
})

View File

@ -8,6 +8,7 @@ import (
"github.com/stretchr/testify/assert"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v2/util/etcd"
@ -20,6 +21,8 @@ func TestServer(t *testing.T) {
params := paramtable.Get()
channel.RecoverPChannelStatsManager([]string{})
endpoints := params.EtcdCfg.Endpoints.GetValue()
etcdEndpoints := strings.Split(endpoints, ",")
c, err := etcd.GetRemoteEtcdClient(etcdEndpoints)

View File

@ -172,7 +172,7 @@ func (c *consumerImpl) recvLoop() (err error) {
}
switch resp := resp.Response.(type) {
case *streamingpb.ConsumeResponse_Consume:
msgID, err := message.UnmarshalMessageID(c.walName, resp.Consume.GetMessage().GetId().GetId())
msgID, err := message.UnmarshalMessageID(resp.Consume.Message.Id)
if err != nil {
return err
}

View File

@ -153,9 +153,7 @@ func newMockedConsumerImpl(t *testing.T, ctx context.Context, h message.Handler)
recvCh <- &streamingpb.ConsumeResponse{
Response: &streamingpb.ConsumeResponse_Create{
Create: &streamingpb.CreateConsumerResponse{
WalName: walimplstest.WALName,
},
Create: &streamingpb.CreateConsumerResponse{},
},
}
recvCh <- &streamingpb.ConsumeResponse{

View File

@ -12,6 +12,7 @@ import (
"github.com/milvus-io/milvus/internal/streamingnode/client/handler/assignment"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
"github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker"
streamingserviceinterceptor "github.com/milvus-io/milvus/internal/util/streamingutil/service/interceptor"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
@ -70,6 +71,9 @@ type HandlerClient interface {
// If the wal is located at remote, it will return 0, error.
GetLatestMVCCTimestampIfLocal(ctx context.Context, vchannel string) (uint64, error)
// GetReplicateCheckpoint returns the WAL checkpoint that will be used to create scanner.
GetReplicateCheckpoint(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)

View File

@ -65,6 +65,18 @@ func (hc *handlerClientImpl) GetLatestMVCCTimestampIfLocal(ctx context.Context,
return w.GetLatestMVCCTimestamp(ctx, vchannel)
}
// GetReplicateCheckpoint returns the WAL checkpoint that will be used to create scanner.
func (hc *handlerClientImpl) GetReplicateCheckpoint(ctx context.Context, channelName string) (*wal.ReplicateCheckpoint, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {
return nil, ErrClientClosed
}
defer hc.lifetime.Done()
return nil, nil
// TODO: sheep, implement it.
}
// GetWALMetricsIfLocal gets the metrics of the local wal.
func (hc *handlerClientImpl) GetWALMetricsIfLocal(ctx context.Context) (*types.StreamingNodeMetrics, error) {
if !hc.lifetime.Add(typeutil.LifetimeStateWorking) {

View File

@ -298,10 +298,7 @@ func (p *producerImpl) recvLoop() (err error) {
var result produceResponse
switch produceResp := resp.Produce.Response.(type) {
case *streamingpb.ProduceMessageResponse_Result:
msgID, err := message.UnmarshalMessageID(
p.walName,
produceResp.Result.GetId().GetId(),
)
msgID, err := message.UnmarshalMessageID(produceResp.Result.GetId())
if err != nil {
return err
}

View File

@ -10,7 +10,6 @@ import (
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/pkg/v2/mocks/proto/mock_streamingpb"
"github.com/milvus-io/milvus/pkg/v2/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
@ -50,9 +49,7 @@ func TestProducer(t *testing.T) {
recvCh <- &streamingpb.ProduceResponse{
Response: &streamingpb.ProduceResponse_Create{
Create: &streamingpb.CreateProducerResponse{
WalName: walimplstest.WALName,
},
Create: &streamingpb.CreateProducerResponse{},
},
}
producer, err := CreateProducer(ctx, opts, c)
@ -89,7 +86,7 @@ func TestProducer(t *testing.T) {
RequestId: 2,
Response: &streamingpb.ProduceMessageResponse_Result{
Result: &streamingpb.ProduceMessageResponseResult{
Id: &messagespb.MessageID{Id: walimplstest.NewTestMessageID(1).Marshal()},
Id: walimplstest.NewTestMessageID(1).IntoProto(),
},
},
},

View File

@ -44,7 +44,7 @@ func (impl *WALFlusherImpl) getRecoveryInfos(ctx context.Context, vchannel []str
var checkpoint message.MessageID
for _, info := range recoveryInfos {
messageID := adaptor.MustGetMessageIDFromMQWrapperIDBytes(impl.wal.Get().WALName(), info.GetInfo().GetSeekPosition().GetMsgID())
messageID := adaptor.MustGetMessageIDFromMQWrapperIDBytes(info.GetInfo().GetSeekPosition().GetMsgID())
if checkpoint == nil || messageID.LT(checkpoint) {
checkpoint = messageID
}

Some files were not shown because too many files have changed in this diff Show More