mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-01-07 19:31:51 +08:00
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:
parent
98d23de36c
commit
51f69f32d0
7
Makefile
7
Makefile
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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=
|
||||
|
||||
@ -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)
|
||||
|
||||
57
client/milvusclient/replicate.go
Normal file
57
client/milvusclient/replicate.go
Normal 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
|
||||
}
|
||||
113
client/milvusclient/replicate_builder.go
Normal file
113
client/milvusclient/replicate_builder.go
Normal 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
|
||||
}
|
||||
116
client/milvusclient/replicate_example_test.go
Normal file
116
client/milvusclient/replicate_example_test.go
Normal 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
44
cmd/components/cdc.go
Normal 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
|
||||
}
|
||||
@ -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)
|
||||
|
||||
@ -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
7
go.mod
@ -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
4
go.sum
@ -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=
|
||||
|
||||
@ -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:
|
||||
|
||||
24
internal/cdc/.mockery.yaml
Normal file
24
internal/cdc/.mockery.yaml
Normal 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:
|
||||
51
internal/cdc/cluster/cluster_client.go
Normal file
51
internal/cdc/cluster/cluster_client.go
Normal 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
|
||||
}
|
||||
34
internal/cdc/cluster/milvus_client.go
Normal file
34
internal/cdc/cluster/milvus_client.go
Normal 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
|
||||
}
|
||||
97
internal/cdc/cluster/mock_cluster_client.go
Normal file
97
internal/cdc/cluster/mock_cluster_client.go
Normal 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
|
||||
}
|
||||
233
internal/cdc/cluster/mock_milvus_client.go
Normal file
233
internal/cdc/cluster/mock_milvus_client.go
Normal 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
|
||||
}
|
||||
24
internal/cdc/controller/controller.go
Normal file
24
internal/cdc/controller/controller.go
Normal 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()
|
||||
}
|
||||
83
internal/cdc/controller/controllerimpl/controller_impl.go
Normal file
83
internal/cdc/controller/controllerimpl/controller_impl.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@ -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()
|
||||
}
|
||||
96
internal/cdc/controller/mock_controller.go
Normal file
96
internal/cdc/controller/mock_controller.go
Normal 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
|
||||
}
|
||||
68
internal/cdc/replication/mock_replicate_manager_client.go
Normal file
68
internal/cdc/replication/mock_replicate_manager_client.go
Normal 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
|
||||
}
|
||||
25
internal/cdc/replication/replicate_manager_client.go
Normal file
25
internal/cdc/replication/replicate_manager_client.go
Normal 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)
|
||||
}
|
||||
237
internal/cdc/replication/replicatemanager/channel_replicator.go
Normal file
237
internal/cdc/replication/replicatemanager/channel_replicator.go
Normal 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()
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
@ -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
|
||||
}
|
||||
@ -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")
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
121
internal/cdc/replication/replicatestream/metrics.go
Normal file
121
internal/cdc/replication/replicatestream/metrics.go
Normal 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()
|
||||
}
|
||||
@ -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
|
||||
}
|
||||
211
internal/cdc/replication/replicatestream/msg_queue.go
Normal file
211
internal/cdc/replication/replicatestream/msg_queue.go
Normal 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)
|
||||
}
|
||||
287
internal/cdc/replication/replicatestream/msg_queue_test.go
Normal file
287
internal/cdc/replication/replicatestream/msg_queue_test.go
Normal 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)
|
||||
}
|
||||
@ -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
|
||||
@ -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()
|
||||
}
|
||||
@ -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()
|
||||
}
|
||||
143
internal/cdc/resource/resource.go
Normal file
143
internal/cdc/resource/resource.go
Normal 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")
|
||||
}
|
||||
}
|
||||
}
|
||||
37
internal/cdc/resource/test_utility.go
Normal file
37
internal/cdc/resource/test_utility.go
Normal 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
55
internal/cdc/server.go
Normal 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
|
||||
}
|
||||
@ -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)
|
||||
|
||||
|
||||
194
internal/distributed/cdc/service.go
Normal file
194
internal/distributed/cdc/service.go
Normal 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(¶mtable.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
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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")
|
||||
)
|
||||
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -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()),
|
||||
|
||||
149
internal/distributed/streaming/replicate_service.go
Normal file
149
internal/distributed/streaming/replicate_service.go
Normal 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
|
||||
}
|
||||
124
internal/distributed/streaming/replicate_service_test.go
Normal file
124
internal/distributed/streaming/replicate_service_test.go
Normal 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
|
||||
}
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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{} {
|
||||
|
||||
@ -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 {
|
||||
|
||||
@ -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
|
||||
|
||||
@ -6,4 +6,8 @@ const (
|
||||
BroadcastTaskPrefix = MetaPrefix + "broadcast-task/"
|
||||
VersionPrefix = MetaPrefix + "version/"
|
||||
CChannelMetaPrefix = MetaPrefix + "cchannel/"
|
||||
|
||||
// Replicate
|
||||
ReplicatePChannelMetaPrefix = MetaPrefix + "replicating-pchannel/"
|
||||
ReplicateConfigurationKey = MetaPrefix + "replicate-configuration"
|
||||
)
|
||||
|
||||
@ -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)
|
||||
}
|
||||
|
||||
@ -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")
|
||||
}
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
@ -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 {
|
||||
|
||||
144
internal/mocks/mock_metastore/mock_ReplicationCatalog.go
Normal file
144
internal/mocks/mock_metastore/mock_ReplicationCatalog.go
Normal 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
|
||||
}
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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()
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
159
internal/proxy/replicate/replicate_stream_server.go
Normal file
159
internal/proxy/replicate/replicate_stream_server.go
Normal 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
|
||||
}
|
||||
}
|
||||
237
internal/proxy/replicate/replicate_stream_server_test.go
Normal file
237
internal/proxy/replicate/replicate_stream_server_test.go
Normal 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
|
||||
}
|
||||
}
|
||||
@ -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.
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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")
|
||||
|
||||
@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
46
internal/streamingcoord/server/service/broadcast_test.go
Normal file
46
internal/streamingcoord/server/service/broadcast_test.go
Normal 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"},
|
||||
},
|
||||
})
|
||||
}
|
||||
@ -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,
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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{
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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(),
|
||||
},
|
||||
},
|
||||
},
|
||||
|
||||
@ -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
Loading…
x
Reference in New Issue
Block a user