From aa50e6281f6a7a6d9c9c1ed6e51fb435f25b43f5 Mon Sep 17 00:00:00 2001 From: chyezh Date: Thu, 4 Jul 2024 13:05:47 +0800 Subject: [PATCH] enhance: streaming node client implementation - add streaming node grpc client wrapper - add unittest for streaming node grpc client side - fix binary unsafe bug for message Signed-off-by: chyezh --- configs/milvus.yaml | 11 + internal/.mockery.yaml | 17 + .../mock_StreamingNodeHandlerServiceClient.go | 178 ++++++++ ...reamingNodeHandlerService_ConsumeClient.go | 398 ++++++++++++++++++ ...reamingNodeHandlerService_ProduceClient.go | 398 ++++++++++++++++++ .../mock_StreamingNodeManagerServiceClient.go | 250 +++++++++++ .../handler/mock_assignment/mock_Watcher.go | 158 +++++++ .../handler/mock_consumer/mock_Consumer.go | 148 +++++++ .../handler/mock_producer/mock_Producer.go | 251 +++++++++++ .../server/mock_walmanager/mock_Manager.go | 20 +- .../service/mock_lazygrpc/mock_Service.go | 176 ++++++++ .../service/mock_resolver/mock_Builder.go | 209 +++++++++ .../client/handler/assignment/watcher.go | 22 + .../client/handler/assignment/watcher_impl.go | 108 +++++ .../client/handler/assignment/watcher_test.go | 91 ++++ .../client/handler/consumer/consumer.go | 18 + .../client/handler/consumer/consumer_impl.go | 181 ++++++++ .../client/handler/consumer/consumer_test.go | 85 ++++ .../client/handler/handler_client.go | 151 +++++++ .../client/handler/handler_client_impl.go | 214 ++++++++++ .../client/handler/handler_client_test.go | 145 +++++++ .../handler/producer/produce_grpc_client.go | 35 ++ .../client/handler/producer/producer.go | 30 ++ .../client/handler/producer/producer_impl.go | 313 ++++++++++++++ .../client/handler/producer/producer_test.go | 112 +++++ .../client/handler/shared_producer.go | 51 +++ .../streamingnode/client/manager/manager.go | 24 -- .../client/manager/manager_client.go | 112 +++++ .../client/manager/manager_client_impl.go | 191 +++++++++ .../client/manager/manager_test.go | 174 ++++++++ .../server/walmanager/manager.go | 2 +- pkg/util/paramtable/component_param.go | 18 +- pkg/util/paramtable/grpc_param.go | 52 +++ pkg/util/typeutil/shared_reference.go | 99 +++++ pkg/util/typeutil/shared_reference_test.go | 81 ++++ 35 files changed, 4481 insertions(+), 42 deletions(-) create mode 100644 internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go create mode 100644 internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go create mode 100644 internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go create mode 100644 internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go create mode 100644 internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go create mode 100644 internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go create mode 100644 internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go create mode 100644 internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go create mode 100644 internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go create mode 100644 internal/streamingnode/client/handler/assignment/watcher.go create mode 100644 internal/streamingnode/client/handler/assignment/watcher_impl.go create mode 100644 internal/streamingnode/client/handler/assignment/watcher_test.go create mode 100644 internal/streamingnode/client/handler/consumer/consumer.go create mode 100644 internal/streamingnode/client/handler/consumer/consumer_impl.go create mode 100644 internal/streamingnode/client/handler/consumer/consumer_test.go create mode 100644 internal/streamingnode/client/handler/handler_client.go create mode 100644 internal/streamingnode/client/handler/handler_client_impl.go create mode 100644 internal/streamingnode/client/handler/handler_client_test.go create mode 100644 internal/streamingnode/client/handler/producer/produce_grpc_client.go create mode 100644 internal/streamingnode/client/handler/producer/producer.go create mode 100644 internal/streamingnode/client/handler/producer/producer_impl.go create mode 100644 internal/streamingnode/client/handler/producer/producer_test.go create mode 100644 internal/streamingnode/client/handler/shared_producer.go delete mode 100644 internal/streamingnode/client/manager/manager.go create mode 100644 internal/streamingnode/client/manager/manager_client.go create mode 100644 internal/streamingnode/client/manager/manager_client_impl.go create mode 100644 internal/streamingnode/client/manager/manager_test.go create mode 100644 pkg/util/typeutil/shared_reference.go create mode 100644 pkg/util/typeutil/shared_reference_test.go diff --git a/configs/milvus.yaml b/configs/milvus.yaml index e4cd6d1193dd6..3632172580826 100644 --- a/configs/milvus.yaml +++ b/configs/milvus.yaml @@ -572,6 +572,17 @@ dataNode: clusteringCompaction: memoryBufferRatio: 0.1 # The ratio of memory buffer of clustering compaction. Data larger than threshold will be spilled to storage. +streamingNode: + # can specify ip for example + # ip: 127.0.0.1 + ip: # if not specify address, will use the first unicastable address as local ip + port: 19532 + grpc: + serverMaxSendSize: 536870912 + serverMaxRecvSize: 536870912 + clientMaxSendSize: 268435456 + clientMaxRecvSize: 268435456 + # Configures the system log output. log: level: info # Only supports debug, info, warn, error, panic, or fatal. Default 'info'. diff --git a/internal/.mockery.yaml b/internal/.mockery.yaml index 3840c369a9e60..0d1de8c737e1e 100644 --- a/internal/.mockery.yaml +++ b/internal/.mockery.yaml @@ -11,6 +11,15 @@ packages: github.com/milvus-io/milvus/internal/streamingnode/client/manager: interfaces: ManagerClient: + github.com/milvus-io/milvus/internal/streamingnode/client/handler/assignment: + interfaces: + Watcher: + github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer: + interfaces: + Producer: + github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer: + interfaces: + Consumer: github.com/milvus-io/milvus/internal/streamingnode/server/wal: interfaces: OpenerBuilder: @@ -30,6 +39,10 @@ packages: StreamingNodeHandlerService_ConsumeServer: StreamingNodeHandlerService_ProduceServer: StreamingCoordAssignmentService_AssignmentDiscoverServer: + StreamingNodeManagerServiceClient: + StreamingNodeHandlerServiceClient: + StreamingNodeHandlerService_ConsumeClient: + StreamingNodeHandlerService_ProduceClient: github.com/milvus-io/milvus/internal/streamingnode/server/walmanager: interfaces: Manager: @@ -40,9 +53,13 @@ packages: interfaces: Discoverer: AssignmentDiscoverWatcher: + github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc: + interfaces: + Service: github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver: interfaces: Resolver: + Builder: google.golang.org/grpc/resolver: interfaces: ClientConn: diff --git a/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go new file mode 100644 index 0000000000000..671519bec6161 --- /dev/null +++ b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go @@ -0,0 +1,178 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_streamingpb + +import ( + context "context" + + grpc "google.golang.org/grpc" + + mock "github.com/stretchr/testify/mock" + + streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb" +) + +// MockStreamingNodeHandlerServiceClient is an autogenerated mock type for the StreamingNodeHandlerServiceClient type +type MockStreamingNodeHandlerServiceClient struct { + mock.Mock +} + +type MockStreamingNodeHandlerServiceClient_Expecter struct { + mock *mock.Mock +} + +func (_m *MockStreamingNodeHandlerServiceClient) EXPECT() *MockStreamingNodeHandlerServiceClient_Expecter { + return &MockStreamingNodeHandlerServiceClient_Expecter{mock: &_m.Mock} +} + +// Consume provides a mock function with given fields: ctx, opts +func (_m *MockStreamingNodeHandlerServiceClient) Consume(ctx context.Context, opts ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, 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...) + + var r0 streamingpb.StreamingNodeHandlerService_ConsumeClient + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, error)); ok { + return rf(ctx, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) streamingpb.StreamingNodeHandlerService_ConsumeClient); ok { + r0 = rf(ctx, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(streamingpb.StreamingNodeHandlerService_ConsumeClient) + } + } + + 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 +} + +// MockStreamingNodeHandlerServiceClient_Consume_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Consume' +type MockStreamingNodeHandlerServiceClient_Consume_Call struct { + *mock.Call +} + +// Consume is a helper method to define mock.On call +// - ctx context.Context +// - opts ...grpc.CallOption +func (_e *MockStreamingNodeHandlerServiceClient_Expecter) Consume(ctx interface{}, opts ...interface{}) *MockStreamingNodeHandlerServiceClient_Consume_Call { + return &MockStreamingNodeHandlerServiceClient_Consume_Call{Call: _e.mock.On("Consume", + append([]interface{}{ctx}, opts...)...)} +} + +func (_c *MockStreamingNodeHandlerServiceClient_Consume_Call) Run(run func(ctx context.Context, opts ...grpc.CallOption)) *MockStreamingNodeHandlerServiceClient_Consume_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 *MockStreamingNodeHandlerServiceClient_Consume_Call) Return(_a0 streamingpb.StreamingNodeHandlerService_ConsumeClient, _a1 error) *MockStreamingNodeHandlerServiceClient_Consume_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerServiceClient_Consume_Call) RunAndReturn(run func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, error)) *MockStreamingNodeHandlerServiceClient_Consume_Call { + _c.Call.Return(run) + return _c +} + +// Produce provides a mock function with given fields: ctx, opts +func (_m *MockStreamingNodeHandlerServiceClient) Produce(ctx context.Context, opts ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, 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...) + + var r0 streamingpb.StreamingNodeHandlerService_ProduceClient + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, error)); ok { + return rf(ctx, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) streamingpb.StreamingNodeHandlerService_ProduceClient); ok { + r0 = rf(ctx, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(streamingpb.StreamingNodeHandlerService_ProduceClient) + } + } + + 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 +} + +// MockStreamingNodeHandlerServiceClient_Produce_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Produce' +type MockStreamingNodeHandlerServiceClient_Produce_Call struct { + *mock.Call +} + +// Produce is a helper method to define mock.On call +// - ctx context.Context +// - opts ...grpc.CallOption +func (_e *MockStreamingNodeHandlerServiceClient_Expecter) Produce(ctx interface{}, opts ...interface{}) *MockStreamingNodeHandlerServiceClient_Produce_Call { + return &MockStreamingNodeHandlerServiceClient_Produce_Call{Call: _e.mock.On("Produce", + append([]interface{}{ctx}, opts...)...)} +} + +func (_c *MockStreamingNodeHandlerServiceClient_Produce_Call) Run(run func(ctx context.Context, opts ...grpc.CallOption)) *MockStreamingNodeHandlerServiceClient_Produce_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 *MockStreamingNodeHandlerServiceClient_Produce_Call) Return(_a0 streamingpb.StreamingNodeHandlerService_ProduceClient, _a1 error) *MockStreamingNodeHandlerServiceClient_Produce_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerServiceClient_Produce_Call) RunAndReturn(run func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, error)) *MockStreamingNodeHandlerServiceClient_Produce_Call { + _c.Call.Return(run) + return _c +} + +// NewMockStreamingNodeHandlerServiceClient creates a new instance of MockStreamingNodeHandlerServiceClient. 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 NewMockStreamingNodeHandlerServiceClient(t interface { + mock.TestingT + Cleanup(func()) +}) *MockStreamingNodeHandlerServiceClient { + mock := &MockStreamingNodeHandlerServiceClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go new file mode 100644 index 0000000000000..3889b9457840e --- /dev/null +++ b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go @@ -0,0 +1,398 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_streamingpb + +import ( + context "context" + + mock "github.com/stretchr/testify/mock" + metadata "google.golang.org/grpc/metadata" + + streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb" +) + +// MockStreamingNodeHandlerService_ConsumeClient is an autogenerated mock type for the StreamingNodeHandlerService_ConsumeClient type +type MockStreamingNodeHandlerService_ConsumeClient struct { + mock.Mock +} + +type MockStreamingNodeHandlerService_ConsumeClient_Expecter struct { + mock *mock.Mock +} + +func (_m *MockStreamingNodeHandlerService_ConsumeClient) EXPECT() *MockStreamingNodeHandlerService_ConsumeClient_Expecter { + return &MockStreamingNodeHandlerService_ConsumeClient_Expecter{mock: &_m.Mock} +} + +// CloseSend provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ConsumeClient) CloseSend() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend' +type MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call struct { + *mock.Call +} + +// CloseSend is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) CloseSend() *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call{Call: _e.mock.On("CloseSend")} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) RunAndReturn(run func() error) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call { + _c.Call.Return(run) + return _c +} + +// Context provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ConsumeClient) Context() context.Context { + ret := _m.Called() + + var r0 context.Context + if rf, ok := ret.Get(0).(func() context.Context); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(context.Context) + } + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context' +type MockStreamingNodeHandlerService_ConsumeClient_Context_Call struct { + *mock.Call +} + +// Context is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Context() *MockStreamingNodeHandlerService_ConsumeClient_Context_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_Context_Call{Call: _e.mock.On("Context")} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) Return(_a0 context.Context) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) RunAndReturn(run func() context.Context) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call { + _c.Call.Return(run) + return _c +} + +// Header provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ConsumeClient) Header() (metadata.MD, error) { + ret := _m.Called() + + var r0 metadata.MD + var r1 error + if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeHandlerService_ConsumeClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header' +type MockStreamingNodeHandlerService_ConsumeClient_Header_Call struct { + *mock.Call +} + +// Header is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Header() *MockStreamingNodeHandlerService_ConsumeClient_Header_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_Header_Call{Call: _e.mock.On("Header")} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) Return(_a0 metadata.MD, _a1 error) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call { + _c.Call.Return(run) + return _c +} + +// Recv provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ConsumeClient) Recv() (*streamingpb.ConsumeResponse, error) { + ret := _m.Called() + + var r0 *streamingpb.ConsumeResponse + var r1 error + if rf, ok := ret.Get(0).(func() (*streamingpb.ConsumeResponse, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() *streamingpb.ConsumeResponse); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*streamingpb.ConsumeResponse) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeHandlerService_ConsumeClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv' +type MockStreamingNodeHandlerService_ConsumeClient_Recv_Call struct { + *mock.Call +} + +// Recv is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Recv() *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_Recv_Call{Call: _e.mock.On("Recv")} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) Return(_a0 *streamingpb.ConsumeResponse, _a1 error) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) RunAndReturn(run func() (*streamingpb.ConsumeResponse, error)) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call { + _c.Call.Return(run) + return _c +} + +// RecvMsg provides a mock function with given fields: m +func (_m *MockStreamingNodeHandlerService_ConsumeClient) RecvMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg' +type MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call struct { + *mock.Call +} + +// RecvMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) RecvMsg(m interface{}) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call { + _c.Call.Return(run) + return _c +} + +// Send provides a mock function with given fields: _a0 +func (_m *MockStreamingNodeHandlerService_ConsumeClient) Send(_a0 *streamingpb.ConsumeRequest) error { + ret := _m.Called(_a0) + + var r0 error + if rf, ok := ret.Get(0).(func(*streamingpb.ConsumeRequest) error); ok { + r0 = rf(_a0) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send' +type MockStreamingNodeHandlerService_ConsumeClient_Send_Call struct { + *mock.Call +} + +// Send is a helper method to define mock.On call +// - _a0 *streamingpb.ConsumeRequest +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Send(_a0 interface{}) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_Send_Call{Call: _e.mock.On("Send", _a0)} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) Run(run func(_a0 *streamingpb.ConsumeRequest)) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(*streamingpb.ConsumeRequest)) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) RunAndReturn(run func(*streamingpb.ConsumeRequest) error) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call { + _c.Call.Return(run) + return _c +} + +// SendMsg provides a mock function with given fields: m +func (_m *MockStreamingNodeHandlerService_ConsumeClient) SendMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg' +type MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call struct { + *mock.Call +} + +// SendMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) SendMsg(m interface{}) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call { + _c.Call.Return(run) + return _c +} + +// Trailer provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ConsumeClient) Trailer() metadata.MD { + ret := _m.Called() + + var r0 metadata.MD + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + return r0 +} + +// MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer' +type MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call struct { + *mock.Call +} + +// Trailer is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Trailer() *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call { + return &MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call{Call: _e.mock.On("Trailer")} +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) Return(_a0 metadata.MD) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call { + _c.Call.Return(run) + return _c +} + +// NewMockStreamingNodeHandlerService_ConsumeClient creates a new instance of MockStreamingNodeHandlerService_ConsumeClient. 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 NewMockStreamingNodeHandlerService_ConsumeClient(t interface { + mock.TestingT + Cleanup(func()) +}) *MockStreamingNodeHandlerService_ConsumeClient { + mock := &MockStreamingNodeHandlerService_ConsumeClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go new file mode 100644 index 0000000000000..ad20971185761 --- /dev/null +++ b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go @@ -0,0 +1,398 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_streamingpb + +import ( + context "context" + + mock "github.com/stretchr/testify/mock" + metadata "google.golang.org/grpc/metadata" + + streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb" +) + +// MockStreamingNodeHandlerService_ProduceClient is an autogenerated mock type for the StreamingNodeHandlerService_ProduceClient type +type MockStreamingNodeHandlerService_ProduceClient struct { + mock.Mock +} + +type MockStreamingNodeHandlerService_ProduceClient_Expecter struct { + mock *mock.Mock +} + +func (_m *MockStreamingNodeHandlerService_ProduceClient) EXPECT() *MockStreamingNodeHandlerService_ProduceClient_Expecter { + return &MockStreamingNodeHandlerService_ProduceClient_Expecter{mock: &_m.Mock} +} + +// CloseSend provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ProduceClient) CloseSend() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend' +type MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call struct { + *mock.Call +} + +// CloseSend is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) CloseSend() *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call { + return &MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call{Call: _e.mock.On("CloseSend")} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) RunAndReturn(run func() error) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call { + _c.Call.Return(run) + return _c +} + +// Context provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ProduceClient) Context() context.Context { + ret := _m.Called() + + var r0 context.Context + if rf, ok := ret.Get(0).(func() context.Context); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(context.Context) + } + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context' +type MockStreamingNodeHandlerService_ProduceClient_Context_Call struct { + *mock.Call +} + +// Context is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Context() *MockStreamingNodeHandlerService_ProduceClient_Context_Call { + return &MockStreamingNodeHandlerService_ProduceClient_Context_Call{Call: _e.mock.On("Context")} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Context_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) Return(_a0 context.Context) *MockStreamingNodeHandlerService_ProduceClient_Context_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) RunAndReturn(run func() context.Context) *MockStreamingNodeHandlerService_ProduceClient_Context_Call { + _c.Call.Return(run) + return _c +} + +// Header provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ProduceClient) Header() (metadata.MD, error) { + ret := _m.Called() + + var r0 metadata.MD + var r1 error + if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeHandlerService_ProduceClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header' +type MockStreamingNodeHandlerService_ProduceClient_Header_Call struct { + *mock.Call +} + +// Header is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Header() *MockStreamingNodeHandlerService_ProduceClient_Header_Call { + return &MockStreamingNodeHandlerService_ProduceClient_Header_Call{Call: _e.mock.On("Header")} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Header_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) Return(_a0 metadata.MD, _a1 error) *MockStreamingNodeHandlerService_ProduceClient_Header_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *MockStreamingNodeHandlerService_ProduceClient_Header_Call { + _c.Call.Return(run) + return _c +} + +// Recv provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ProduceClient) Recv() (*streamingpb.ProduceResponse, error) { + ret := _m.Called() + + var r0 *streamingpb.ProduceResponse + var r1 error + if rf, ok := ret.Get(0).(func() (*streamingpb.ProduceResponse, error)); ok { + return rf() + } + if rf, ok := ret.Get(0).(func() *streamingpb.ProduceResponse); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*streamingpb.ProduceResponse) + } + } + + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeHandlerService_ProduceClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv' +type MockStreamingNodeHandlerService_ProduceClient_Recv_Call struct { + *mock.Call +} + +// Recv is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Recv() *MockStreamingNodeHandlerService_ProduceClient_Recv_Call { + return &MockStreamingNodeHandlerService_ProduceClient_Recv_Call{Call: _e.mock.On("Recv")} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) Return(_a0 *streamingpb.ProduceResponse, _a1 error) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) RunAndReturn(run func() (*streamingpb.ProduceResponse, error)) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call { + _c.Call.Return(run) + return _c +} + +// RecvMsg provides a mock function with given fields: m +func (_m *MockStreamingNodeHandlerService_ProduceClient) RecvMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg' +type MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call struct { + *mock.Call +} + +// RecvMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) RecvMsg(m interface{}) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call { + return &MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call { + _c.Call.Return(run) + return _c +} + +// Send provides a mock function with given fields: _a0 +func (_m *MockStreamingNodeHandlerService_ProduceClient) Send(_a0 *streamingpb.ProduceRequest) error { + ret := _m.Called(_a0) + + var r0 error + if rf, ok := ret.Get(0).(func(*streamingpb.ProduceRequest) error); ok { + r0 = rf(_a0) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send' +type MockStreamingNodeHandlerService_ProduceClient_Send_Call struct { + *mock.Call +} + +// Send is a helper method to define mock.On call +// - _a0 *streamingpb.ProduceRequest +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Send(_a0 interface{}) *MockStreamingNodeHandlerService_ProduceClient_Send_Call { + return &MockStreamingNodeHandlerService_ProduceClient_Send_Call{Call: _e.mock.On("Send", _a0)} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) Run(run func(_a0 *streamingpb.ProduceRequest)) *MockStreamingNodeHandlerService_ProduceClient_Send_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(*streamingpb.ProduceRequest)) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_Send_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) RunAndReturn(run func(*streamingpb.ProduceRequest) error) *MockStreamingNodeHandlerService_ProduceClient_Send_Call { + _c.Call.Return(run) + return _c +} + +// SendMsg provides a mock function with given fields: m +func (_m *MockStreamingNodeHandlerService_ProduceClient) SendMsg(m interface{}) error { + ret := _m.Called(m) + + var r0 error + if rf, ok := ret.Get(0).(func(interface{}) error); ok { + r0 = rf(m) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg' +type MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call struct { + *mock.Call +} + +// SendMsg is a helper method to define mock.On call +// - m interface{} +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) SendMsg(m interface{}) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call { + return &MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(interface{})) + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call { + _c.Call.Return(run) + return _c +} + +// Trailer provides a mock function with given fields: +func (_m *MockStreamingNodeHandlerService_ProduceClient) Trailer() metadata.MD { + ret := _m.Called() + + var r0 metadata.MD + if rf, ok := ret.Get(0).(func() metadata.MD); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(metadata.MD) + } + } + + return r0 +} + +// MockStreamingNodeHandlerService_ProduceClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer' +type MockStreamingNodeHandlerService_ProduceClient_Trailer_Call struct { + *mock.Call +} + +// Trailer is a helper method to define mock.On call +func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Trailer() *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call { + return &MockStreamingNodeHandlerService_ProduceClient_Trailer_Call{Call: _e.mock.On("Trailer")} +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) Return(_a0 metadata.MD) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call { + _c.Call.Return(run) + return _c +} + +// NewMockStreamingNodeHandlerService_ProduceClient creates a new instance of MockStreamingNodeHandlerService_ProduceClient. 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 NewMockStreamingNodeHandlerService_ProduceClient(t interface { + mock.TestingT + Cleanup(func()) +}) *MockStreamingNodeHandlerService_ProduceClient { + mock := &MockStreamingNodeHandlerService_ProduceClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go new file mode 100644 index 0000000000000..5354f7b9a1ae2 --- /dev/null +++ b/internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go @@ -0,0 +1,250 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_streamingpb + +import ( + context "context" + + grpc "google.golang.org/grpc" + + mock "github.com/stretchr/testify/mock" + + streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb" +) + +// MockStreamingNodeManagerServiceClient is an autogenerated mock type for the StreamingNodeManagerServiceClient type +type MockStreamingNodeManagerServiceClient struct { + mock.Mock +} + +type MockStreamingNodeManagerServiceClient_Expecter struct { + mock *mock.Mock +} + +func (_m *MockStreamingNodeManagerServiceClient) EXPECT() *MockStreamingNodeManagerServiceClient_Expecter { + return &MockStreamingNodeManagerServiceClient_Expecter{mock: &_m.Mock} +} + +// Assign provides a mock function with given fields: ctx, in, opts +func (_m *MockStreamingNodeManagerServiceClient) Assign(ctx context.Context, in *streamingpb.StreamingNodeManagerAssignRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + var r0 *streamingpb.StreamingNodeManagerAssignResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerAssignResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerAssignResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeManagerServiceClient_Assign_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Assign' +type MockStreamingNodeManagerServiceClient_Assign_Call struct { + *mock.Call +} + +// Assign is a helper method to define mock.On call +// - ctx context.Context +// - in *streamingpb.StreamingNodeManagerAssignRequest +// - opts ...grpc.CallOption +func (_e *MockStreamingNodeManagerServiceClient_Expecter) Assign(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_Assign_Call { + return &MockStreamingNodeManagerServiceClient_Assign_Call{Call: _e.mock.On("Assign", + append([]interface{}{ctx, in}, opts...)...)} +} + +func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerAssignRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_Assign_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].(*streamingpb.StreamingNodeManagerAssignRequest), variadicArgs...) + }) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) Return(_a0 *streamingpb.StreamingNodeManagerAssignResponse, _a1 error) *MockStreamingNodeManagerServiceClient_Assign_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error)) *MockStreamingNodeManagerServiceClient_Assign_Call { + _c.Call.Return(run) + return _c +} + +// CollectStatus provides a mock function with given fields: ctx, in, opts +func (_m *MockStreamingNodeManagerServiceClient) CollectStatus(ctx context.Context, in *streamingpb.StreamingNodeManagerCollectStatusRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + var r0 *streamingpb.StreamingNodeManagerCollectStatusResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerCollectStatusResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerCollectStatusResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeManagerServiceClient_CollectStatus_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CollectStatus' +type MockStreamingNodeManagerServiceClient_CollectStatus_Call struct { + *mock.Call +} + +// CollectStatus is a helper method to define mock.On call +// - ctx context.Context +// - in *streamingpb.StreamingNodeManagerCollectStatusRequest +// - opts ...grpc.CallOption +func (_e *MockStreamingNodeManagerServiceClient_Expecter) CollectStatus(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_CollectStatus_Call { + return &MockStreamingNodeManagerServiceClient_CollectStatus_Call{Call: _e.mock.On("CollectStatus", + append([]interface{}{ctx, in}, opts...)...)} +} + +func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerCollectStatusRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_CollectStatus_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].(*streamingpb.StreamingNodeManagerCollectStatusRequest), variadicArgs...) + }) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) Return(_a0 *streamingpb.StreamingNodeManagerCollectStatusResponse, _a1 error) *MockStreamingNodeManagerServiceClient_CollectStatus_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error)) *MockStreamingNodeManagerServiceClient_CollectStatus_Call { + _c.Call.Return(run) + return _c +} + +// Remove provides a mock function with given fields: ctx, in, opts +func (_m *MockStreamingNodeManagerServiceClient) Remove(ctx context.Context, in *streamingpb.StreamingNodeManagerRemoveRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error) { + _va := make([]interface{}, len(opts)) + for _i := range opts { + _va[_i] = opts[_i] + } + var _ca []interface{} + _ca = append(_ca, ctx, in) + _ca = append(_ca, _va...) + ret := _m.Called(_ca...) + + var r0 *streamingpb.StreamingNodeManagerRemoveResponse + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error)); ok { + return rf(ctx, in, opts...) + } + if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerRemoveResponse); ok { + r0 = rf(ctx, in, opts...) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerRemoveResponse) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) error); ok { + r1 = rf(ctx, in, opts...) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockStreamingNodeManagerServiceClient_Remove_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Remove' +type MockStreamingNodeManagerServiceClient_Remove_Call struct { + *mock.Call +} + +// Remove is a helper method to define mock.On call +// - ctx context.Context +// - in *streamingpb.StreamingNodeManagerRemoveRequest +// - opts ...grpc.CallOption +func (_e *MockStreamingNodeManagerServiceClient_Expecter) Remove(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_Remove_Call { + return &MockStreamingNodeManagerServiceClient_Remove_Call{Call: _e.mock.On("Remove", + append([]interface{}{ctx, in}, opts...)...)} +} + +func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerRemoveRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_Remove_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].(*streamingpb.StreamingNodeManagerRemoveRequest), variadicArgs...) + }) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) Return(_a0 *streamingpb.StreamingNodeManagerRemoveResponse, _a1 error) *MockStreamingNodeManagerServiceClient_Remove_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error)) *MockStreamingNodeManagerServiceClient_Remove_Call { + _c.Call.Return(run) + return _c +} + +// NewMockStreamingNodeManagerServiceClient creates a new instance of MockStreamingNodeManagerServiceClient. 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 NewMockStreamingNodeManagerServiceClient(t interface { + mock.TestingT + Cleanup(func()) +}) *MockStreamingNodeManagerServiceClient { + mock := &MockStreamingNodeManagerServiceClient{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go b/internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go new file mode 100644 index 0000000000000..f51798aeea86c --- /dev/null +++ b/internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go @@ -0,0 +1,158 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_assignment + +import ( + context "context" + + types "github.com/milvus-io/milvus/pkg/streaming/util/types" + mock "github.com/stretchr/testify/mock" +) + +// MockWatcher is an autogenerated mock type for the Watcher type +type MockWatcher struct { + mock.Mock +} + +type MockWatcher_Expecter struct { + mock *mock.Mock +} + +func (_m *MockWatcher) EXPECT() *MockWatcher_Expecter { + return &MockWatcher_Expecter{mock: &_m.Mock} +} + +// Close provides a mock function with given fields: +func (_m *MockWatcher) Close() { + _m.Called() +} + +// MockWatcher_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockWatcher_Close_Call struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockWatcher_Expecter) Close() *MockWatcher_Close_Call { + return &MockWatcher_Close_Call{Call: _e.mock.On("Close")} +} + +func (_c *MockWatcher_Close_Call) Run(run func()) *MockWatcher_Close_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockWatcher_Close_Call) Return() *MockWatcher_Close_Call { + _c.Call.Return() + return _c +} + +func (_c *MockWatcher_Close_Call) RunAndReturn(run func()) *MockWatcher_Close_Call { + _c.Call.Return(run) + return _c +} + +// Get provides a mock function with given fields: ctx, channel +func (_m *MockWatcher) Get(ctx context.Context, channel string) *types.PChannelInfoAssigned { + ret := _m.Called(ctx, channel) + + var r0 *types.PChannelInfoAssigned + if rf, ok := ret.Get(0).(func(context.Context, string) *types.PChannelInfoAssigned); ok { + r0 = rf(ctx, channel) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*types.PChannelInfoAssigned) + } + } + + return r0 +} + +// MockWatcher_Get_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Get' +type MockWatcher_Get_Call struct { + *mock.Call +} + +// Get is a helper method to define mock.On call +// - ctx context.Context +// - channel string +func (_e *MockWatcher_Expecter) Get(ctx interface{}, channel interface{}) *MockWatcher_Get_Call { + return &MockWatcher_Get_Call{Call: _e.mock.On("Get", ctx, channel)} +} + +func (_c *MockWatcher_Get_Call) Run(run func(ctx context.Context, channel string)) *MockWatcher_Get_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string)) + }) + return _c +} + +func (_c *MockWatcher_Get_Call) Return(_a0 *types.PChannelInfoAssigned) *MockWatcher_Get_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockWatcher_Get_Call) RunAndReturn(run func(context.Context, string) *types.PChannelInfoAssigned) *MockWatcher_Get_Call { + _c.Call.Return(run) + return _c +} + +// Watch provides a mock function with given fields: ctx, channel, previous +func (_m *MockWatcher) Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error { + ret := _m.Called(ctx, channel, previous) + + var r0 error + if rf, ok := ret.Get(0).(func(context.Context, string, *types.PChannelInfoAssigned) error); ok { + r0 = rf(ctx, channel, previous) + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockWatcher_Watch_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Watch' +type MockWatcher_Watch_Call struct { + *mock.Call +} + +// Watch is a helper method to define mock.On call +// - ctx context.Context +// - channel string +// - previous *types.PChannelInfoAssigned +func (_e *MockWatcher_Expecter) Watch(ctx interface{}, channel interface{}, previous interface{}) *MockWatcher_Watch_Call { + return &MockWatcher_Watch_Call{Call: _e.mock.On("Watch", ctx, channel, previous)} +} + +func (_c *MockWatcher_Watch_Call) Run(run func(ctx context.Context, channel string, previous *types.PChannelInfoAssigned)) *MockWatcher_Watch_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(string), args[2].(*types.PChannelInfoAssigned)) + }) + return _c +} + +func (_c *MockWatcher_Watch_Call) Return(_a0 error) *MockWatcher_Watch_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockWatcher_Watch_Call) RunAndReturn(run func(context.Context, string, *types.PChannelInfoAssigned) error) *MockWatcher_Watch_Call { + _c.Call.Return(run) + return _c +} + +// NewMockWatcher creates a new instance of MockWatcher. 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 NewMockWatcher(t interface { + mock.TestingT + Cleanup(func()) +}) *MockWatcher { + mock := &MockWatcher{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go b/internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go new file mode 100644 index 0000000000000..f5c8ca5925223 --- /dev/null +++ b/internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go @@ -0,0 +1,148 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_consumer + +import mock "github.com/stretchr/testify/mock" + +// MockConsumer is an autogenerated mock type for the Consumer type +type MockConsumer struct { + mock.Mock +} + +type MockConsumer_Expecter struct { + mock *mock.Mock +} + +func (_m *MockConsumer) EXPECT() *MockConsumer_Expecter { + return &MockConsumer_Expecter{mock: &_m.Mock} +} + +// Close provides a mock function with given fields: +func (_m *MockConsumer) Close() { + _m.Called() +} + +// MockConsumer_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockConsumer_Close_Call struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockConsumer_Expecter) Close() *MockConsumer_Close_Call { + return &MockConsumer_Close_Call{Call: _e.mock.On("Close")} +} + +func (_c *MockConsumer_Close_Call) Run(run func()) *MockConsumer_Close_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockConsumer_Close_Call) Return() *MockConsumer_Close_Call { + _c.Call.Return() + return _c +} + +func (_c *MockConsumer_Close_Call) RunAndReturn(run func()) *MockConsumer_Close_Call { + _c.Call.Return(run) + return _c +} + +// Done provides a mock function with given fields: +func (_m *MockConsumer) Done() <-chan struct{} { + ret := _m.Called() + + var r0 <-chan struct{} + if rf, ok := ret.Get(0).(func() <-chan struct{}); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(<-chan struct{}) + } + } + + return r0 +} + +// MockConsumer_Done_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Done' +type MockConsumer_Done_Call struct { + *mock.Call +} + +// Done is a helper method to define mock.On call +func (_e *MockConsumer_Expecter) Done() *MockConsumer_Done_Call { + return &MockConsumer_Done_Call{Call: _e.mock.On("Done")} +} + +func (_c *MockConsumer_Done_Call) Run(run func()) *MockConsumer_Done_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockConsumer_Done_Call) Return(_a0 <-chan struct{}) *MockConsumer_Done_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockConsumer_Done_Call) RunAndReturn(run func() <-chan struct{}) *MockConsumer_Done_Call { + _c.Call.Return(run) + return _c +} + +// Error provides a mock function with given fields: +func (_m *MockConsumer) Error() error { + ret := _m.Called() + + var r0 error + if rf, ok := ret.Get(0).(func() error); ok { + r0 = rf() + } else { + r0 = ret.Error(0) + } + + return r0 +} + +// MockConsumer_Error_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Error' +type MockConsumer_Error_Call struct { + *mock.Call +} + +// Error is a helper method to define mock.On call +func (_e *MockConsumer_Expecter) Error() *MockConsumer_Error_Call { + return &MockConsumer_Error_Call{Call: _e.mock.On("Error")} +} + +func (_c *MockConsumer_Error_Call) Run(run func()) *MockConsumer_Error_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockConsumer_Error_Call) Return(_a0 error) *MockConsumer_Error_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockConsumer_Error_Call) RunAndReturn(run func() error) *MockConsumer_Error_Call { + _c.Call.Return(run) + return _c +} + +// NewMockConsumer creates a new instance of MockConsumer. 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 NewMockConsumer(t interface { + mock.TestingT + Cleanup(func()) +}) *MockConsumer { + mock := &MockConsumer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go b/internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go new file mode 100644 index 0000000000000..61dfda479cfdb --- /dev/null +++ b/internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go @@ -0,0 +1,251 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_producer + +import ( + context "context" + + message "github.com/milvus-io/milvus/pkg/streaming/util/message" + mock "github.com/stretchr/testify/mock" + + types "github.com/milvus-io/milvus/pkg/streaming/util/types" +) + +// MockProducer is an autogenerated mock type for the Producer type +type MockProducer struct { + mock.Mock +} + +type MockProducer_Expecter struct { + mock *mock.Mock +} + +func (_m *MockProducer) EXPECT() *MockProducer_Expecter { + return &MockProducer_Expecter{mock: &_m.Mock} +} + +// Assignment provides a mock function with given fields: +func (_m *MockProducer) Assignment() types.PChannelInfoAssigned { + ret := _m.Called() + + var r0 types.PChannelInfoAssigned + if rf, ok := ret.Get(0).(func() types.PChannelInfoAssigned); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(types.PChannelInfoAssigned) + } + + return r0 +} + +// MockProducer_Assignment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Assignment' +type MockProducer_Assignment_Call struct { + *mock.Call +} + +// Assignment is a helper method to define mock.On call +func (_e *MockProducer_Expecter) Assignment() *MockProducer_Assignment_Call { + return &MockProducer_Assignment_Call{Call: _e.mock.On("Assignment")} +} + +func (_c *MockProducer_Assignment_Call) Run(run func()) *MockProducer_Assignment_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockProducer_Assignment_Call) Return(_a0 types.PChannelInfoAssigned) *MockProducer_Assignment_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockProducer_Assignment_Call) RunAndReturn(run func() types.PChannelInfoAssigned) *MockProducer_Assignment_Call { + _c.Call.Return(run) + return _c +} + +// Available provides a mock function with given fields: +func (_m *MockProducer) Available() <-chan struct{} { + ret := _m.Called() + + var r0 <-chan struct{} + if rf, ok := ret.Get(0).(func() <-chan struct{}); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(<-chan struct{}) + } + } + + return r0 +} + +// MockProducer_Available_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Available' +type MockProducer_Available_Call struct { + *mock.Call +} + +// Available is a helper method to define mock.On call +func (_e *MockProducer_Expecter) Available() *MockProducer_Available_Call { + return &MockProducer_Available_Call{Call: _e.mock.On("Available")} +} + +func (_c *MockProducer_Available_Call) Run(run func()) *MockProducer_Available_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockProducer_Available_Call) Return(_a0 <-chan struct{}) *MockProducer_Available_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockProducer_Available_Call) RunAndReturn(run func() <-chan struct{}) *MockProducer_Available_Call { + _c.Call.Return(run) + return _c +} + +// Close provides a mock function with given fields: +func (_m *MockProducer) Close() { + _m.Called() +} + +// MockProducer_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockProducer_Close_Call struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockProducer_Expecter) Close() *MockProducer_Close_Call { + return &MockProducer_Close_Call{Call: _e.mock.On("Close")} +} + +func (_c *MockProducer_Close_Call) Run(run func()) *MockProducer_Close_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockProducer_Close_Call) Return() *MockProducer_Close_Call { + _c.Call.Return() + return _c +} + +func (_c *MockProducer_Close_Call) RunAndReturn(run func()) *MockProducer_Close_Call { + _c.Call.Return(run) + return _c +} + +// IsAvailable provides a mock function with given fields: +func (_m *MockProducer) IsAvailable() bool { + ret := _m.Called() + + var r0 bool + if rf, ok := ret.Get(0).(func() bool); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(bool) + } + + return r0 +} + +// MockProducer_IsAvailable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsAvailable' +type MockProducer_IsAvailable_Call struct { + *mock.Call +} + +// IsAvailable is a helper method to define mock.On call +func (_e *MockProducer_Expecter) IsAvailable() *MockProducer_IsAvailable_Call { + return &MockProducer_IsAvailable_Call{Call: _e.mock.On("IsAvailable")} +} + +func (_c *MockProducer_IsAvailable_Call) Run(run func()) *MockProducer_IsAvailable_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockProducer_IsAvailable_Call) Return(_a0 bool) *MockProducer_IsAvailable_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockProducer_IsAvailable_Call) RunAndReturn(run func() bool) *MockProducer_IsAvailable_Call { + _c.Call.Return(run) + return _c +} + +// Produce provides a mock function with given fields: ctx, msg +func (_m *MockProducer) Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) { + ret := _m.Called(ctx, msg) + + var r0 message.MessageID + var r1 error + if rf, ok := ret.Get(0).(func(context.Context, message.MutableMessage) (message.MessageID, error)); ok { + return rf(ctx, msg) + } + if rf, ok := ret.Get(0).(func(context.Context, message.MutableMessage) message.MessageID); ok { + r0 = rf(ctx, msg) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(message.MessageID) + } + } + + if rf, ok := ret.Get(1).(func(context.Context, message.MutableMessage) error); ok { + r1 = rf(ctx, msg) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockProducer_Produce_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Produce' +type MockProducer_Produce_Call struct { + *mock.Call +} + +// Produce is a helper method to define mock.On call +// - ctx context.Context +// - msg message.MutableMessage +func (_e *MockProducer_Expecter) Produce(ctx interface{}, msg interface{}) *MockProducer_Produce_Call { + return &MockProducer_Produce_Call{Call: _e.mock.On("Produce", ctx, msg)} +} + +func (_c *MockProducer_Produce_Call) Run(run func(ctx context.Context, msg message.MutableMessage)) *MockProducer_Produce_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context), args[1].(message.MutableMessage)) + }) + return _c +} + +func (_c *MockProducer_Produce_Call) Return(_a0 message.MessageID, _a1 error) *MockProducer_Produce_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockProducer_Produce_Call) RunAndReturn(run func(context.Context, message.MutableMessage) (message.MessageID, error)) *MockProducer_Produce_Call { + _c.Call.Return(run) + return _c +} + +// NewMockProducer creates a new instance of MockProducer. 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 NewMockProducer(t interface { + mock.TestingT + Cleanup(func()) +}) *MockProducer { + mock := &MockProducer{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go b/internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go index 4c12954e6c705..16b3bbe6f1dca 100644 --- a/internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go +++ b/internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go @@ -109,17 +109,17 @@ func (_c *MockManager_GetAllAvailableChannels_Call) RunAndReturn(run func() ([]t return _c } -// GetAvailableWAL provides a mock function with given fields: channel -func (_m *MockManager) GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, error) { - ret := _m.Called(channel) +// GetAvailableWAL provides a mock function with given fields: _a0 +func (_m *MockManager) GetAvailableWAL(_a0 types.PChannelInfo) (wal.WAL, error) { + ret := _m.Called(_a0) var r0 wal.WAL var r1 error if rf, ok := ret.Get(0).(func(types.PChannelInfo) (wal.WAL, error)); ok { - return rf(channel) + return rf(_a0) } if rf, ok := ret.Get(0).(func(types.PChannelInfo) wal.WAL); ok { - r0 = rf(channel) + r0 = rf(_a0) } else { if ret.Get(0) != nil { r0 = ret.Get(0).(wal.WAL) @@ -127,7 +127,7 @@ func (_m *MockManager) GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, err } if rf, ok := ret.Get(1).(func(types.PChannelInfo) error); ok { - r1 = rf(channel) + r1 = rf(_a0) } else { r1 = ret.Error(1) } @@ -141,12 +141,12 @@ type MockManager_GetAvailableWAL_Call struct { } // GetAvailableWAL is a helper method to define mock.On call -// - channel types.PChannelInfo -func (_e *MockManager_Expecter) GetAvailableWAL(channel interface{}) *MockManager_GetAvailableWAL_Call { - return &MockManager_GetAvailableWAL_Call{Call: _e.mock.On("GetAvailableWAL", channel)} +// - _a0 types.PChannelInfo +func (_e *MockManager_Expecter) GetAvailableWAL(_a0 interface{}) *MockManager_GetAvailableWAL_Call { + return &MockManager_GetAvailableWAL_Call{Call: _e.mock.On("GetAvailableWAL", _a0)} } -func (_c *MockManager_GetAvailableWAL_Call) Run(run func(channel types.PChannelInfo)) *MockManager_GetAvailableWAL_Call { +func (_c *MockManager_GetAvailableWAL_Call) Run(run func(_a0 types.PChannelInfo)) *MockManager_GetAvailableWAL_Call { _c.Call.Run(func(args mock.Arguments) { run(args[0].(types.PChannelInfo)) }) diff --git a/internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go b/internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go new file mode 100644 index 0000000000000..2c8d838751979 --- /dev/null +++ b/internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go @@ -0,0 +1,176 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_lazygrpc + +import ( + context "context" + + grpc "google.golang.org/grpc" + + mock "github.com/stretchr/testify/mock" +) + +// MockService is an autogenerated mock type for the Service type +type MockService[T interface{}] struct { + mock.Mock +} + +type MockService_Expecter[T interface{}] struct { + mock *mock.Mock +} + +func (_m *MockService[T]) EXPECT() *MockService_Expecter[T] { + return &MockService_Expecter[T]{mock: &_m.Mock} +} + +// Close provides a mock function with given fields: +func (_m *MockService[T]) Close() { + _m.Called() +} + +// MockService_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockService_Close_Call[T interface{}] struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockService_Expecter[T]) Close() *MockService_Close_Call[T] { + return &MockService_Close_Call[T]{Call: _e.mock.On("Close")} +} + +func (_c *MockService_Close_Call[T]) Run(run func()) *MockService_Close_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockService_Close_Call[T]) Return() *MockService_Close_Call[T] { + _c.Call.Return() + return _c +} + +func (_c *MockService_Close_Call[T]) RunAndReturn(run func()) *MockService_Close_Call[T] { + _c.Call.Return(run) + return _c +} + +// GetConn provides a mock function with given fields: ctx +func (_m *MockService[T]) GetConn(ctx context.Context) (*grpc.ClientConn, error) { + ret := _m.Called(ctx) + + var r0 *grpc.ClientConn + var r1 error + if rf, ok := ret.Get(0).(func(context.Context) (*grpc.ClientConn, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) *grpc.ClientConn); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*grpc.ClientConn) + } + } + + if rf, ok := ret.Get(1).(func(context.Context) error); ok { + r1 = rf(ctx) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockService_GetConn_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetConn' +type MockService_GetConn_Call[T interface{}] struct { + *mock.Call +} + +// GetConn is a helper method to define mock.On call +// - ctx context.Context +func (_e *MockService_Expecter[T]) GetConn(ctx interface{}) *MockService_GetConn_Call[T] { + return &MockService_GetConn_Call[T]{Call: _e.mock.On("GetConn", ctx)} +} + +func (_c *MockService_GetConn_Call[T]) Run(run func(ctx context.Context)) *MockService_GetConn_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context)) + }) + return _c +} + +func (_c *MockService_GetConn_Call[T]) Return(_a0 *grpc.ClientConn, _a1 error) *MockService_GetConn_Call[T] { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockService_GetConn_Call[T]) RunAndReturn(run func(context.Context) (*grpc.ClientConn, error)) *MockService_GetConn_Call[T] { + _c.Call.Return(run) + return _c +} + +// GetService provides a mock function with given fields: ctx +func (_m *MockService[T]) GetService(ctx context.Context) (T, error) { + ret := _m.Called(ctx) + + var r0 T + var r1 error + if rf, ok := ret.Get(0).(func(context.Context) (T, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) T); ok { + r0 = rf(ctx) + } else { + r0 = ret.Get(0).(T) + } + + if rf, ok := ret.Get(1).(func(context.Context) error); ok { + r1 = rf(ctx) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockService_GetService_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetService' +type MockService_GetService_Call[T interface{}] struct { + *mock.Call +} + +// GetService is a helper method to define mock.On call +// - ctx context.Context +func (_e *MockService_Expecter[T]) GetService(ctx interface{}) *MockService_GetService_Call[T] { + return &MockService_GetService_Call[T]{Call: _e.mock.On("GetService", ctx)} +} + +func (_c *MockService_GetService_Call[T]) Run(run func(ctx context.Context)) *MockService_GetService_Call[T] { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(context.Context)) + }) + return _c +} + +func (_c *MockService_GetService_Call[T]) Return(_a0 T, _a1 error) *MockService_GetService_Call[T] { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockService_GetService_Call[T]) RunAndReturn(run func(context.Context) (T, error)) *MockService_GetService_Call[T] { + _c.Call.Return(run) + return _c +} + +// NewMockService creates a new instance of MockService. 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 NewMockService[T interface{}](t interface { + mock.TestingT + Cleanup(func()) +}) *MockService[T] { + mock := &MockService[T]{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go b/internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go new file mode 100644 index 0000000000000..774a548cde626 --- /dev/null +++ b/internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go @@ -0,0 +1,209 @@ +// Code generated by mockery v2.32.4. DO NOT EDIT. + +package mock_resolver + +import ( + mock "github.com/stretchr/testify/mock" + resolver "google.golang.org/grpc/resolver" + + serviceresolver "github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver" +) + +// MockBuilder is an autogenerated mock type for the Builder type +type MockBuilder struct { + mock.Mock +} + +type MockBuilder_Expecter struct { + mock *mock.Mock +} + +func (_m *MockBuilder) EXPECT() *MockBuilder_Expecter { + return &MockBuilder_Expecter{mock: &_m.Mock} +} + +// Build provides a mock function with given fields: target, cc, opts +func (_m *MockBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) { + ret := _m.Called(target, cc, opts) + + var r0 resolver.Resolver + var r1 error + if rf, ok := ret.Get(0).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) (resolver.Resolver, error)); ok { + return rf(target, cc, opts) + } + if rf, ok := ret.Get(0).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) resolver.Resolver); ok { + r0 = rf(target, cc, opts) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(resolver.Resolver) + } + } + + if rf, ok := ret.Get(1).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) error); ok { + r1 = rf(target, cc, opts) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// MockBuilder_Build_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Build' +type MockBuilder_Build_Call struct { + *mock.Call +} + +// Build is a helper method to define mock.On call +// - target resolver.Target +// - cc resolver.ClientConn +// - opts resolver.BuildOptions +func (_e *MockBuilder_Expecter) Build(target interface{}, cc interface{}, opts interface{}) *MockBuilder_Build_Call { + return &MockBuilder_Build_Call{Call: _e.mock.On("Build", target, cc, opts)} +} + +func (_c *MockBuilder_Build_Call) Run(run func(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions)) *MockBuilder_Build_Call { + _c.Call.Run(func(args mock.Arguments) { + run(args[0].(resolver.Target), args[1].(resolver.ClientConn), args[2].(resolver.BuildOptions)) + }) + return _c +} + +func (_c *MockBuilder_Build_Call) Return(_a0 resolver.Resolver, _a1 error) *MockBuilder_Build_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +func (_c *MockBuilder_Build_Call) RunAndReturn(run func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) (resolver.Resolver, error)) *MockBuilder_Build_Call { + _c.Call.Return(run) + return _c +} + +// Close provides a mock function with given fields: +func (_m *MockBuilder) Close() { + _m.Called() +} + +// MockBuilder_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close' +type MockBuilder_Close_Call struct { + *mock.Call +} + +// Close is a helper method to define mock.On call +func (_e *MockBuilder_Expecter) Close() *MockBuilder_Close_Call { + return &MockBuilder_Close_Call{Call: _e.mock.On("Close")} +} + +func (_c *MockBuilder_Close_Call) Run(run func()) *MockBuilder_Close_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockBuilder_Close_Call) Return() *MockBuilder_Close_Call { + _c.Call.Return() + return _c +} + +func (_c *MockBuilder_Close_Call) RunAndReturn(run func()) *MockBuilder_Close_Call { + _c.Call.Return(run) + return _c +} + +// Resolver provides a mock function with given fields: +func (_m *MockBuilder) Resolver() serviceresolver.Resolver { + ret := _m.Called() + + var r0 serviceresolver.Resolver + if rf, ok := ret.Get(0).(func() serviceresolver.Resolver); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(serviceresolver.Resolver) + } + } + + return r0 +} + +// MockBuilder_Resolver_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Resolver' +type MockBuilder_Resolver_Call struct { + *mock.Call +} + +// Resolver is a helper method to define mock.On call +func (_e *MockBuilder_Expecter) Resolver() *MockBuilder_Resolver_Call { + return &MockBuilder_Resolver_Call{Call: _e.mock.On("Resolver")} +} + +func (_c *MockBuilder_Resolver_Call) Run(run func()) *MockBuilder_Resolver_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockBuilder_Resolver_Call) Return(_a0 serviceresolver.Resolver) *MockBuilder_Resolver_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockBuilder_Resolver_Call) RunAndReturn(run func() serviceresolver.Resolver) *MockBuilder_Resolver_Call { + _c.Call.Return(run) + return _c +} + +// Scheme provides a mock function with given fields: +func (_m *MockBuilder) Scheme() string { + ret := _m.Called() + + var r0 string + if rf, ok := ret.Get(0).(func() string); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(string) + } + + return r0 +} + +// MockBuilder_Scheme_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Scheme' +type MockBuilder_Scheme_Call struct { + *mock.Call +} + +// Scheme is a helper method to define mock.On call +func (_e *MockBuilder_Expecter) Scheme() *MockBuilder_Scheme_Call { + return &MockBuilder_Scheme_Call{Call: _e.mock.On("Scheme")} +} + +func (_c *MockBuilder_Scheme_Call) Run(run func()) *MockBuilder_Scheme_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockBuilder_Scheme_Call) Return(_a0 string) *MockBuilder_Scheme_Call { + _c.Call.Return(_a0) + return _c +} + +func (_c *MockBuilder_Scheme_Call) RunAndReturn(run func() string) *MockBuilder_Scheme_Call { + _c.Call.Return(run) + return _c +} + +// NewMockBuilder creates a new instance of MockBuilder. 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 NewMockBuilder(t interface { + mock.TestingT + Cleanup(func()) +}) *MockBuilder { + mock := &MockBuilder{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/internal/streamingnode/client/handler/assignment/watcher.go b/internal/streamingnode/client/handler/assignment/watcher.go new file mode 100644 index 0000000000000..ca80e8ac207a1 --- /dev/null +++ b/internal/streamingnode/client/handler/assignment/watcher.go @@ -0,0 +1,22 @@ +package assignment + +import ( + "context" + + "github.com/milvus-io/milvus/pkg/streaming/util/types" +) + +var _ Watcher = (*watcherImpl)(nil) + +// Watcher is the interface for the channel assignment. +type Watcher interface { + // Get gets the channel assignment. + Get(ctx context.Context, channel string) *types.PChannelInfoAssigned + + // Watch watches the channel assignment. + // Block until new term is coming. + Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error + + // Close stop the watcher. + Close() +} diff --git a/internal/streamingnode/client/handler/assignment/watcher_impl.go b/internal/streamingnode/client/handler/assignment/watcher_impl.go new file mode 100644 index 0000000000000..aa276a5b4672b --- /dev/null +++ b/internal/streamingnode/client/handler/assignment/watcher_impl.go @@ -0,0 +1,108 @@ +package assignment + +import ( + "context" + "sync" + + "go.uber.org/zap" + + "github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver" + "github.com/milvus-io/milvus/pkg/log" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/syncutil" +) + +func NewWatcher(r resolver.Resolver) Watcher { + ctx, cancel := context.WithCancel(context.Background()) + w := &watcherImpl{ + ctx: ctx, + cancel: cancel, + r: r, + cond: *syncutil.NewContextCond(&sync.Mutex{}), + assignments: make(map[string]types.PChannelInfoAssigned), + } + go w.execute() + return w +} + +// watcherImpl is the implementation of the assignment watcher. +type watcherImpl struct { + ctx context.Context + cancel context.CancelFunc + r resolver.Resolver + cond syncutil.ContextCond + assignments map[string]types.PChannelInfoAssigned // map pchannel to node. +} + +// execute starts the watcher. +func (w *watcherImpl) execute() { + log.Info("assignment watcher start") + var err error + defer func() { + // error can be ignored here, so use info level log here. + log.Info("assignment watcher close", zap.Error(err)) + }() + + // error can be ignored here, error is always cancel by watcher's close as expected. + // otherwise, the resolver's close is unexpected. + err = w.r.Watch(w.ctx, func(state discoverer.VersionedState) error { + w.updateAssignment(state) + return nil + }) +} + +// updateAssignment updates the assignment. +func (w *watcherImpl) updateAssignment(state discoverer.VersionedState) { + newAssignments := make(map[string]types.PChannelInfoAssigned) + for _, assignments := range state.ChannelAssignmentInfo() { + for _, pChannelInfo := range assignments.Channels { + newAssignments[pChannelInfo.Name] = types.PChannelInfoAssigned{ + Channel: pChannelInfo, + Node: assignments.NodeInfo, + } + } + } + w.cond.LockAndBroadcast() + w.assignments = newAssignments + w.cond.L.Unlock() +} + +// Get gets the current pchannel assignment. +func (w *watcherImpl) Get(ctx context.Context, channel string) *types.PChannelInfoAssigned { + w.cond.L.Lock() + defer w.cond.L.Unlock() + + if info, ok := w.assignments[channel]; ok { + return &info + } + return nil +} + +// Watch watches the channel assignment. +func (w *watcherImpl) Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error { + w.cond.L.Lock() + + term := types.InitialTerm + if previous != nil { + term = previous.Channel.Term + } + + for { + if info, ok := w.assignments[channel]; ok { + if info.Channel.Term > term { + break + } + } + if err := w.cond.Wait(ctx); err != nil { + return err + } + } + w.cond.L.Unlock() + return nil +} + +// Close closes the watcher. +func (w *watcherImpl) Close() { + w.cancel() +} diff --git a/internal/streamingnode/client/handler/assignment/watcher_test.go b/internal/streamingnode/client/handler/assignment/watcher_test.go new file mode 100644 index 0000000000000..01d3e069830c1 --- /dev/null +++ b/internal/streamingnode/client/handler/assignment/watcher_test.go @@ -0,0 +1,91 @@ +package assignment + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "google.golang.org/grpc/resolver" + + "github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/attributes" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +func TestWatcher(t *testing.T) { + r := mock_resolver.NewMockResolver(t) + + ch := make(chan discoverer.VersionedState) + r.EXPECT().Watch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(s discoverer.VersionedState) error) error { + for { + select { + case v, ok := <-ch: + if !ok { + return nil + } + f(v) + case <-ctx.Done(): + return ctx.Err() + } + } + }) + w := NewWatcher(r) + defer w.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + a := w.Get(ctx, "test_pchannel") + assert.Nil(t, a) + err := w.Watch(ctx, "test_pchannel", nil) + assert.ErrorIs(t, err, context.DeadlineExceeded) + + ch <- discoverer.VersionedState{ + Version: typeutil.VersionInt64(1), + State: resolver.State{ + Addresses: []resolver.Address{ + { + Addr: "test_addr", + BalancerAttributes: attributes.WithChannelAssignmentInfo( + new(attributes.Attributes), + &types.StreamingNodeAssignment{ + NodeInfo: types.StreamingNodeInfo{ + ServerID: 1, + Address: "test_addr", + }, + Channels: map[string]types.PChannelInfo{ + "test_pchannel": { + Name: "test_pchannel", + Term: 1, + }, + "test_pchannel_2": { + Name: "test_pchannel_2", + Term: 2, + }, + }, + }, + ), + }, + }, + }, + } + err = w.Watch(context.Background(), "test_pchannel", nil) + assert.NoError(t, err) + a = w.Get(ctx, "test_pchannel") + assert.NotNil(t, a) + assert.Equal(t, int64(1), a.Channel.Term) + + err = w.Watch(context.Background(), "test_pchannel_2", nil) + assert.NoError(t, err) + a = w.Get(ctx, "test_pchannel_2") + assert.NotNil(t, a) + assert.Equal(t, int64(2), a.Channel.Term) + + ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + err = w.Watch(ctx, "test_pchannel", a) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} diff --git a/internal/streamingnode/client/handler/consumer/consumer.go b/internal/streamingnode/client/handler/consumer/consumer.go new file mode 100644 index 0000000000000..d9fd1deb1abd1 --- /dev/null +++ b/internal/streamingnode/client/handler/consumer/consumer.go @@ -0,0 +1,18 @@ +package consumer + +var _ Consumer = (*consumerImpl)(nil) + +// Consumer is the interface that wraps the basic consume method on grpc stream. +// Consumer is work on a single stream on grpc, +// so Consumer cannot recover from failure because of the stream is broken. +type Consumer interface { + // Error returns the error of scanner failed. + // Will block until scanner is closed or Chan is dry out. + Error() error + + // Done returns a channel which will be closed when scanner is finished or closed. + Done() <-chan struct{} + + // Close the consumer, release the underlying resources. + Close() +} diff --git a/internal/streamingnode/client/handler/consumer/consumer_impl.go b/internal/streamingnode/client/handler/consumer/consumer_impl.go new file mode 100644 index 0000000000000..8c8b9485a2d43 --- /dev/null +++ b/internal/streamingnode/client/handler/consumer/consumer_impl.go @@ -0,0 +1,181 @@ +package consumer + +import ( + "context" + "io" + + "github.com/cockroachdb/errors" + "go.uber.org/zap" + "google.golang.org/grpc" + + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil" + "github.com/milvus-io/milvus/internal/util/streamingutil/status" + "github.com/milvus-io/milvus/internal/util/streamingutil/typeconverter" + "github.com/milvus-io/milvus/pkg/log" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/options" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/syncutil" +) + +// ConsumerOptions is the options for creating a consumer. +type ConsumerOptions struct { + // The cosume target + Assignment *types.PChannelInfoAssigned + + // DeliverPolicy is the deliver policy of the consumer. + DeliverPolicy options.DeliverPolicy + + // DeliverFilters is the deliver filters of the consumer. + DeliverFilters []options.DeliverFilter + + // Handler is the message handler used to handle message after recv from consumer. + MessageHandler message.Handler +} + +// CreateConsumer creates a new consumer client. +func CreateConsumer( + ctx context.Context, + opts *ConsumerOptions, + handlerClient streamingpb.StreamingNodeHandlerServiceClient, +) (Consumer, error) { + ctx, err := createConsumeRequest(ctx, opts) + if err != nil { + return nil, err + } + + // TODO: configurable or auto adjust grpc.MaxCallRecvMsgSize + streamClient, err := handlerClient.Consume(ctx, grpc.MaxCallRecvMsgSize(8388608)) + if err != nil { + return nil, err + } + + // Recv the first response from server. + // It must be a create response. + resp, err := streamClient.Recv() + if err != nil { + return nil, err + } + createResp := resp.GetCreate() + if createResp == nil { + return nil, status.NewInvalidRequestSeq("first message arrive must be create response") + } + cli := &consumerImpl{ + walName: createResp.GetWalName(), + assignment: *opts.Assignment, + grpcStreamClient: streamClient, + handlerClient: handlerClient, + logger: log.With( + zap.String("walName", createResp.GetWalName()), + zap.String("pchannel", opts.Assignment.Channel.Name), + zap.Int64("term", opts.Assignment.Channel.Term), + zap.Int64("streamingNodeID", opts.Assignment.Node.ServerID)), + msgHandler: opts.MessageHandler, + finishErr: syncutil.NewFuture[error](), + } + go cli.execute() + return cli, nil +} + +// createConsumeRequest creates the consume request. +func createConsumeRequest(ctx context.Context, opts *ConsumerOptions) (context.Context, error) { + // select server to consume. + ctx = contextutil.WithPickServerID(ctx, opts.Assignment.Node.ServerID) + // create the consumer request. + deliverPolicy, err := typeconverter.NewProtoFromDeliverPolicy(opts.DeliverPolicy) + if err != nil { + return nil, errors.Wrap(err, "at convert deliver policy") + } + deliverFilters, err := typeconverter.NewProtosFromDeliverFilters(opts.DeliverFilters) + if err != nil { + return nil, errors.Wrap(err, "at convert deliver filters") + } + return contextutil.WithCreateConsumer(ctx, &streamingpb.CreateConsumerRequest{ + Pchannel: typeconverter.NewProtoFromPChannelInfo(opts.Assignment.Channel), + DeliverPolicy: deliverPolicy, + DeliverFilters: deliverFilters, + }), nil +} + +type consumerImpl struct { + walName string + assignment types.PChannelInfoAssigned + grpcStreamClient streamingpb.StreamingNodeHandlerService_ConsumeClient + handlerClient streamingpb.StreamingNodeHandlerServiceClient + logger *log.MLogger + msgHandler message.Handler + finishErr *syncutil.Future[error] +} + +// Close close the consumer client. +func (c *consumerImpl) Close() { + // Send the close request to server. + if err := c.grpcStreamClient.Send(&streamingpb.ConsumeRequest{ + Request: &streamingpb.ConsumeRequest_Close{}, + }); err != nil { + c.logger.Warn("send close request failed", zap.Error(err)) + } + // close the grpc client stream. + if err := c.grpcStreamClient.CloseSend(); err != nil { + c.logger.Warn("close grpc stream failed", zap.Error(err)) + } + <-c.finishErr.Done() +} + +// Error returns the error of the consumer client. +func (c *consumerImpl) Error() error { + return c.finishErr.Get() +} + +// Done returns a channel that closes when the consumer client is closed. +func (c *consumerImpl) Done() <-chan struct{} { + return c.finishErr.Done() +} + +// execute starts the recv loop. +func (c *consumerImpl) execute() { + c.recvLoop() +} + +// recvLoop is the recv arm of the grpc stream. +// Throughput of the grpc framework should be ok to use single stream to receive message. +// Once throughput is not enough, look at https://grpc.io/docs/guides/performance/ to find the solution. +func (c *consumerImpl) recvLoop() (err error) { + defer func() { + if err != nil { + c.logger.Warn("recv arm of stream closed with unexpected error", zap.Error(err)) + } else { + c.logger.Info("recv arm of stream closed") + } + c.finishErr.Set(err) + c.msgHandler.Close() + }() + + for { + resp, err := c.grpcStreamClient.Recv() + if errors.Is(err, io.EOF) { + return nil + } + if err != nil { + return err + } + switch resp := resp.Response.(type) { + case *streamingpb.ConsumeResponse_Consume: + msgID, err := message.UnmarshalMessageID(c.walName, resp.Consume.GetId().GetId()) + if err != nil { + return err + } + c.msgHandler.Handle(message.NewImmutableMesasge( + msgID, + resp.Consume.GetMessage().GetPayload(), + resp.Consume.GetMessage().GetProperties(), + )) + case *streamingpb.ConsumeResponse_Close: + // Should receive io.EOF after that. + // Do nothing at current implementation. + default: + c.logger.Warn("unknown response type", zap.Any("response", resp)) + } + } +} diff --git a/internal/streamingnode/client/handler/consumer/consumer_test.go b/internal/streamingnode/client/handler/consumer/consumer_test.go new file mode 100644 index 0000000000000..c9440eca1d893 --- /dev/null +++ b/internal/streamingnode/client/handler/consumer/consumer_test.go @@ -0,0 +1,85 @@ +package consumer + +import ( + "context" + "io" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + + "github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb" + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/options" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/streaming/walimpls/impls/walimplstest" +) + +func TestConsumer(t *testing.T) { + c := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t) + cc := mock_streamingpb.NewMockStreamingNodeHandlerService_ConsumeClient(t) + recvCh := make(chan *streamingpb.ConsumeResponse, 10) + cc.EXPECT().Recv().RunAndReturn(func() (*streamingpb.ConsumeResponse, error) { + msg, ok := <-recvCh + if !ok { + return nil, io.EOF + } + return msg, nil + }) + sendCh := make(chan *streamingpb.ConsumeRequest, 10) + cc.EXPECT().Send(mock.Anything).RunAndReturn(func(cr *streamingpb.ConsumeRequest) error { + sendCh <- cr + return nil + }) + c.EXPECT().Consume(mock.Anything, mock.Anything).Return(cc, nil) + cc.EXPECT().CloseSend().RunAndReturn(func() error { + recvCh <- &streamingpb.ConsumeResponse{Response: &streamingpb.ConsumeResponse_Close{}} + close(recvCh) + return nil + }) + + ctx := context.Background() + resultCh := make(message.ChanMessageHandler, 1) + opts := &ConsumerOptions{ + Assignment: &types.PChannelInfoAssigned{ + Channel: types.PChannelInfo{Name: "test", Term: 1}, + Node: types.StreamingNodeInfo{ServerID: 1, Address: "localhost"}, + }, + DeliverPolicy: options.DeliverPolicyAll(), + DeliverFilters: []options.DeliverFilter{ + options.DeliverFilterVChannel("test-1"), + options.DeliverFilterTimeTickGT(100), + }, + MessageHandler: resultCh, + } + + recvCh <- &streamingpb.ConsumeResponse{ + Response: &streamingpb.ConsumeResponse_Create{ + Create: &streamingpb.CreateConsumerResponse{ + WalName: "test", + }, + }, + } + recvCh <- &streamingpb.ConsumeResponse{ + Response: &streamingpb.ConsumeResponse_Consume{ + Consume: &streamingpb.ConsumeMessageReponse{ + Id: &streamingpb.MessageID{ + Id: walimplstest.NewTestMessageID(1).Marshal(), + }, + Message: &streamingpb.Message{ + Payload: []byte{}, + Properties: make(map[string]string), + }, + }, + }, + } + consumer, err := CreateConsumer(ctx, opts, c) + assert.NoError(t, err) + assert.NotNil(t, consumer) + consumer.Close() + msg := <-resultCh + assert.True(t, msg.MessageID().EQ(walimplstest.NewTestMessageID(1))) + <-consumer.Done() + assert.NoError(t, consumer.Error()) +} diff --git a/internal/streamingnode/client/handler/handler_client.go b/internal/streamingnode/client/handler/handler_client.go new file mode 100644 index 0000000000000..1d9d5ae7bbd77 --- /dev/null +++ b/internal/streamingnode/client/handler/handler_client.go @@ -0,0 +1,151 @@ +package handler + +import ( + "context" + "encoding/json" + "time" + + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "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/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" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/options" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/tracer" + "github.com/milvus-io/milvus/pkg/util/interceptor" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/lock" + "github.com/milvus-io/milvus/pkg/util/paramtable" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +var _ HandlerClient = (*handlerClientImpl)(nil) + +type ( + Producer = producer.Producer + Consumer = consumer.Consumer +) + +// ProducerOptions is the options for creating a producer. +type ProducerOptions struct { + // PChannel is the pchannel of the producer. + PChannel string +} + +// ConsumerOptions is the options for creating a consumer. +type ConsumerOptions struct { + // PChannel is the pchannel of the consumer. + PChannel string + + // DeliverPolicy is the deliver policy of the consumer. + DeliverPolicy options.DeliverPolicy + + // DeliverFilters is the deliver filters of the consumer. + DeliverFilters []options.DeliverFilter + + // Handler is the message handler used to handle message after recv from consumer. + MessageHandler message.Handler +} + +// HandlerClient is the interface that wraps streamingpb.StreamingNodeHandlerServiceClient. +// HandlerClient wraps the PChannel Assignment Service Discovery. +// Provides the ability to create pchannel-level producer and consumer. +type HandlerClient interface { + // CreateProducer creates a producer. + // Producer is a stream client without keep alive promise. + // It will be available until context canceled, active close, streaming error or remote server wal closing. + // Because of there's no more ProducerOptions except PChannel, so a producer of same PChannel is shared by reference count. + CreateProducer(ctx context.Context, opts *ProducerOptions) (Producer, error) + + // CreateConsumer creates a consumer. + // Consumer is a stream client without keep alive promise. + // It will be available until context canceled, active close, streaming error or remote server wal closing. + // A consumer will not share stream connection with other consumers. + CreateConsumer(ctx context.Context, opts *ConsumerOptions) (Consumer, error) + + // Close closes the handler client. + // It will only stop the underlying service discovery, but don't stop the producer and consumer created by it. + // So please close Producer and Consumer created by it before close the handler client. + Close() +} + +// NewHandlerClient creates a new handler client. +func NewHandlerClient(w types.AssignmentDiscoverWatcher) HandlerClient { + rb := resolver.NewChannelAssignmentBuilder(w) + dialTimeout := paramtable.Get().StreamingNodeGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond) + dialOptions := getDialOptions(rb) + conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) { + ctx, cancel := context.WithTimeout(ctx, dialTimeout) + defer cancel() + return grpc.DialContext( + ctx, + resolver.ChannelAssignmentResolverScheme+":///"+typeutil.StreamingNodeRole, + dialOptions..., // TODO: we should use dynamic service config in future by add it to resolver. + ) + }) + watcher := assignment.NewWatcher(rb.Resolver()) + return &handlerClientImpl{ + lifetime: lifetime.NewLifetime(lifetime.Working), + service: lazygrpc.WithServiceCreator(conn, streamingpb.NewStreamingNodeHandlerServiceClient), + rb: rb, + watcher: watcher, + rebalanceTrigger: w, + sharedProducers: make(map[string]*typeutil.WeakReference[Producer]), + sharedProducerKeyLock: lock.NewKeyLock[string](), + newProducer: producer.CreateProducer, + newConsumer: consumer.CreateConsumer, + } +} + +// getDialOptions returns grpc dial options. +func getDialOptions(rb resolver.Builder) []grpc.DialOption { + cfg := ¶mtable.Get().StreamingNodeGrpcClientCfg + retryPolicy := cfg.GetDefaultRetryPolicy() + retryPolicy["retryableStatusCodes"] = []string{"UNAVAILABLE"} + defaultServiceConfig := map[string]interface{}{ + "loadBalancingConfig": []map[string]interface{}{ + {picker.ServerIDPickerBalancerName: map[string]interface{}{}}, + }, + "methodConfig": []map[string]interface{}{ + { + "name": []map[string]string{ + {"service": "milvus.proto.streaming.StreamingNodeHandlerService"}, + }, + "waitForReady": true, + "retryPolicy": retryPolicy, + }, + }, + } + defaultServiceConfigJSON, err := json.Marshal(defaultServiceConfig) + if err != nil { + panic(err) + } + dialOptions := cfg.GetDialOptionsFromConfig() + dialOptions = append(dialOptions, + grpc.WithBlock(), + grpc.WithResolvers(rb), + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithChainUnaryInterceptor( + otelgrpc.UnaryClientInterceptor(tracer.GetInterceptorOpts()...), + interceptor.ClusterInjectionUnaryClientInterceptor(), + streamingserviceinterceptor.NewStreamingServiceUnaryClientInterceptor(), + ), + grpc.WithChainStreamInterceptor( + otelgrpc.StreamClientInterceptor(tracer.GetInterceptorOpts()...), + interceptor.ClusterInjectionStreamClientInterceptor(), + streamingserviceinterceptor.NewStreamingServiceStreamClientInterceptor(), + ), + grpc.WithReturnConnectionError(), + grpc.WithDefaultServiceConfig(string(defaultServiceConfigJSON)), + ) + return dialOptions +} diff --git a/internal/streamingnode/client/handler/handler_client_impl.go b/internal/streamingnode/client/handler/handler_client_impl.go new file mode 100644 index 0000000000000..23f1f0f047568 --- /dev/null +++ b/internal/streamingnode/client/handler/handler_client_impl.go @@ -0,0 +1,214 @@ +package handler + +import ( + "context" + "time" + + "go.uber.org/zap" + + "github.com/cenkalti/backoff/v4" + "github.com/cockroachdb/errors" + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "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/util/streamingutil/service/balancer/picker" + "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/status" + "github.com/milvus-io/milvus/pkg/log" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/lock" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +var errWaitNextBackoff = errors.New("wait for next backoff") + +type handlerClientImpl struct { + lifetime lifetime.Lifetime[lifetime.State] + service lazygrpc.Service[streamingpb.StreamingNodeHandlerServiceClient] + rb resolver.Builder + watcher assignment.Watcher + rebalanceTrigger types.AssignmentRebalanceTrigger + sharedProducers map[string]*typeutil.WeakReference[Producer] // map the pchannel to shared producer. + sharedProducerKeyLock *lock.KeyLock[string] + newProducer func(ctx context.Context, opts *producer.ProducerOptions, handler streamingpb.StreamingNodeHandlerServiceClient) (Producer, error) + newConsumer func(ctx context.Context, opts *consumer.ConsumerOptions, handlerClient streamingpb.StreamingNodeHandlerServiceClient) (Consumer, error) +} + +// CreateProducer creates a producer. +func (hc *handlerClientImpl) CreateProducer(ctx context.Context, opts *ProducerOptions) (Producer, error) { + if hc.lifetime.Add(lifetime.IsWorking) != nil { + return nil, status.NewOnShutdownError("handler client is closed") + } + defer hc.lifetime.Done() + + p, err := hc.createHandlerUntilStreamingNodeReady(ctx, opts.PChannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) { + // Wait for handler service is ready. + handlerService, err := hc.service.GetService(ctx) + if err != nil { + return nil, err + } + return hc.createOrGetSharedProducer(ctx, &producer.ProducerOptions{Assignment: assign}, handlerService) + }) + if err != nil { + return nil, err + } + return p.(Producer), nil +} + +// CreateConsumer creates a consumer. +func (hc *handlerClientImpl) CreateConsumer(ctx context.Context, opts *ConsumerOptions) (Consumer, error) { + if hc.lifetime.Add(lifetime.IsWorking) != nil { + return nil, status.NewOnShutdownError("handler client is closed") + } + defer hc.lifetime.Done() + + c, err := hc.createHandlerUntilStreamingNodeReady(ctx, opts.PChannel, func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error) { + // Wait for handler service is ready. + handlerService, err := hc.service.GetService(ctx) + if err != nil { + return nil, err + } + return hc.newConsumer(ctx, &consumer.ConsumerOptions{ + Assignment: assign, + DeliverPolicy: opts.DeliverPolicy, + DeliverFilters: opts.DeliverFilters, + MessageHandler: opts.MessageHandler, + }, handlerService) + }) + if err != nil { + return nil, err + } + return c.(Consumer), nil +} + +// createHandlerUntilStreamingNodeReady creates a handler until streaming node ready. +// If streaming node is not ready, it will block until new assignment term is coming or context timeout. +func (hc *handlerClientImpl) createHandlerUntilStreamingNodeReady(ctx context.Context, pchannel string, create func(ctx context.Context, assign *types.PChannelInfoAssigned) (any, error)) (any, error) { + logger := log.With(zap.String("pchannel", pchannel)) + // TODO: backoff should be configurable. + backoff := backoff.NewExponentialBackOff() + for { + assign := hc.watcher.Get(ctx, pchannel) + if assign != nil { + // Find assignment, try to create producer on this assignment. + c, err := create(ctx, assign) + if err == nil { + return c, nil + } + logger.Warn("create handler failed", zap.Any("assignment", assign), zap.Error(err)) + + // Check if the error is permanent failure until new assignment. + if isPermanentFailureUntilNewAssignment(err) { + reportErr := hc.rebalanceTrigger.ReportAssignmentError(ctx, assign.Channel, err) + logger.Info("report assignment error", zap.NamedError("assignmentError", err), zap.Error(reportErr)) + } + } else { + log.Warn("assignment not found") + } + + nextBackoff := backoff.NextBackOff() + logger.Info("wait for next backoff", zap.Duration("nextBackoff", nextBackoff)) + isAssignemtChange, err := hc.waitForNextBackoff(ctx, pchannel, assign, nextBackoff) + if err != nil { + logger.Warn("wait for next backoff failed", zap.Error(err)) + return nil, err + } + logger.Info("wait for next backoff done", zap.Bool("isAssignmentChange", isAssignemtChange)) + } +} + +// waitForNextBackoff waits for next backoff. +func (hc *handlerClientImpl) waitForNextBackoff(ctx context.Context, pchannel string, assign *types.PChannelInfoAssigned, nextBackoff time.Duration) (bool, error) { + ctx, cancel := context.WithTimeoutCause(ctx, nextBackoff, errWaitNextBackoff) + defer cancel() + // Block until new assignment term is coming. + err := hc.watcher.Watch(ctx, pchannel, assign) + if err == nil || errors.Is(context.Cause(ctx), errWaitNextBackoff) { + return err == nil, nil + } + return false, err +} + +// getFromSharedProducers gets a shared producer from shared producers.A +func (hc *handlerClientImpl) getFromSharedProducers(channelInfo types.PChannelInfo) Producer { + weakProducerRef, ok := hc.sharedProducers[channelInfo.Name] + if !ok { + return nil + } + + strongProducerRef := weakProducerRef.Upgrade() + if strongProducerRef == nil { + // upgrade failure means the outer producer is all closed. + // remove the weak ref and create again. + delete(hc.sharedProducers, channelInfo.Name) + return nil + } + + p := newSharedProducer(strongProducerRef) + if !p.IsAvailable() || p.Assignment().Channel.Term < channelInfo.Term { + // if the producer is not available or the term is less than expected. + // close it and return to create new one. + p.Close() + delete(hc.sharedProducers, channelInfo.Name) + return nil + } + return p +} + +// createOrGetSharedProducer creates or get a shared producer. +// because vchannel in same pchannel can share the same producer. +func (hc *handlerClientImpl) createOrGetSharedProducer( + ctx context.Context, + opts *producer.ProducerOptions, + handlerService streamingpb.StreamingNodeHandlerServiceClient, +) (Producer, error) { + hc.sharedProducerKeyLock.Lock(opts.Assignment.Channel.Name) + defer hc.sharedProducerKeyLock.Unlock(opts.Assignment.Channel.Name) + + // check if shared producer is created within key lock. + if p := hc.getFromSharedProducers(opts.Assignment.Channel); p != nil { + return p, nil + } + + // create a new producer and insert it into shared producers. + newProducer, err := hc.newProducer(ctx, opts, handlerService) + if err != nil { + return nil, err + } + newStrongProducerRef := typeutil.NewSharedReference(newProducer) + // store a weak ref and return a strong ref. + returned := newStrongProducerRef.Clone() + stored := newStrongProducerRef.Downgrade() + hc.sharedProducers[opts.Assignment.Channel.Name] = stored + return newSharedProducer(returned), nil +} + +// Close closes the handler client. +func (hc *handlerClientImpl) Close() { + hc.lifetime.SetState(lifetime.Stopped) + hc.lifetime.Wait() + hc.lifetime.Close() + + hc.watcher.Close() + hc.service.Close() + hc.rb.Close() +} + +// isPermanentFailureUntilNewAssignment checks if the error is permanent failure until new assignment. +// If the encounter this error, client should notify the assignment service to rebalance the assignment and update discovery result. +// block until new assignment term is coming or context timeout. +func isPermanentFailureUntilNewAssignment(err error) bool { + if err == nil { + return false + } + // The error is reported by grpc balancer at client that the sub connection is not exist (remote server is down at view of session). + if picker.IsErrSubConnNoExist(err) { + return true + } + // The error is reported by remote server that the wal is not exist at remote server. + streamingServiceErr := status.AsStreamingError(err) + return streamingServiceErr.IsWrongStreamingNode() +} diff --git a/internal/streamingnode/client/handler/handler_client_test.go b/internal/streamingnode/client/handler/handler_client_test.go new file mode 100644 index 0000000000000..08045c75477b7 --- /dev/null +++ b/internal/streamingnode/client/handler/handler_client_test.go @@ -0,0 +1,145 @@ +package handler + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + + "github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb" + "github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_assignment" + "github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_consumer" + "github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_producer" + "github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_lazygrpc" + "github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver" + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "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/util/streamingutil/status" + "github.com/milvus-io/milvus/pkg/mocks/streaming/util/mock_types" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/options" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/lock" + "github.com/milvus-io/milvus/pkg/util/paramtable" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +func TestHandlerClient(t *testing.T) { + assignment := &types.PChannelInfoAssigned{ + Channel: types.PChannelInfo{Name: "pchannel", Term: 1}, + Node: types.StreamingNodeInfo{ServerID: 1, Address: "localhost"}, + } + + service := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeHandlerServiceClient](t) + handlerServiceClient := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t) + service.EXPECT().GetService(mock.Anything).Return(handlerServiceClient, nil) + rb := mock_resolver.NewMockBuilder(t) + rb.EXPECT().Close().Run(func() {}) + w := mock_assignment.NewMockWatcher(t) + w.EXPECT().Close().Run(func() {}) + + p := mock_producer.NewMockProducer(t) + p.EXPECT().IsAvailable().Return(true) + p.EXPECT().Assignment().Return(*assignment) + p.EXPECT().Close().Run(func() {}) + c := mock_consumer.NewMockConsumer(t) + c.EXPECT().Close().Run(func() {}) + + rebalanceTrigger := mock_types.NewMockAssignmentRebalanceTrigger(t) + rebalanceTrigger.EXPECT().ReportAssignmentError(mock.Anything, mock.Anything, mock.Anything).Return(nil) + + pK := 0 + handler := &handlerClientImpl{ + lifetime: lifetime.NewLifetime(lifetime.Working), + service: service, + rb: rb, + watcher: w, + rebalanceTrigger: rebalanceTrigger, + sharedProducers: make(map[string]*typeutil.WeakReference[producer.Producer]), + sharedProducerKeyLock: lock.NewKeyLock[string](), + newProducer: func(ctx context.Context, opts *producer.ProducerOptions, handler streamingpb.StreamingNodeHandlerServiceClient) (Producer, error) { + if pK == 0 { + pK++ + return nil, status.NewUnmatchedChannelTerm("pchannel", 1, 2) + } + return p, nil + }, + newConsumer: func(ctx context.Context, opts *consumer.ConsumerOptions, handlerClient streamingpb.StreamingNodeHandlerServiceClient) (Consumer, error) { + return c, nil + }, + } + ctx := context.Background() + + k := 0 + w.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) *types.PChannelInfoAssigned { + if k == 0 { + k++ + return nil + } + return assignment + }) + w.EXPECT().Watch(mock.Anything, mock.Anything, mock.Anything).Return(nil) + + producer, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"}) + assert.NoError(t, err) + assert.NotNil(t, producer) + producer2, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"}) + assert.NoError(t, err) + assert.NotNil(t, producer) + p.EXPECT().IsAvailable().Unset() + p.EXPECT().IsAvailable().Return(false) + producer3, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"}) + assert.NoError(t, err) + assert.NotNil(t, producer3) + producer.Close() + producer2.Close() + producer3.Close() + + producer4, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"}) + assert.NoError(t, err) + assert.NotNil(t, producer4) + producer4.Close() + + consumer, err := handler.CreateConsumer(ctx, &ConsumerOptions{ + PChannel: "pchannel", + DeliverPolicy: options.DeliverPolicyAll(), + DeliverFilters: []options.DeliverFilter{ + options.DeliverFilterTimeTickGT(10), + options.DeliverFilterTimeTickGTE(10), + options.DeliverFilterVChannel("vchannel"), + }, + MessageHandler: make(message.ChanMessageHandler), + }) + assert.NoError(t, err) + assert.NotNil(t, consumer) + consumer.Close() + + service.EXPECT().Close().Return() + handler.Close() + producer, err = handler.CreateProducer(ctx, nil) + assert.Error(t, err) + assert.Nil(t, producer) + + consumer, err = handler.CreateConsumer(ctx, nil) + assert.Error(t, err) + assert.Nil(t, consumer) +} + +func TestDial(t *testing.T) { + paramtable.Init() + + w := mock_types.NewMockAssignmentDiscoverWatcher(t) + w.EXPECT().AssignmentDiscover(mock.Anything, mock.Anything).RunAndReturn( + func(ctx context.Context, f func(*types.VersionedStreamingNodeAssignments) error) error { + return context.Canceled + }, + ) + handler := NewHandlerClient(w) + assert.NotNil(t, handler) + time.Sleep(100 * time.Millisecond) + handler.Close() +} diff --git a/internal/streamingnode/client/handler/producer/produce_grpc_client.go b/internal/streamingnode/client/handler/producer/produce_grpc_client.go new file mode 100644 index 0000000000000..2bee5cc376bc7 --- /dev/null +++ b/internal/streamingnode/client/handler/producer/produce_grpc_client.go @@ -0,0 +1,35 @@ +package producer + +import ( + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/streaming/util/message" +) + +// produceGrpcClient is a wrapped producer server of log messages. +type produceGrpcClient struct { + streamingpb.StreamingNodeHandlerService_ProduceClient +} + +// SendProduceMessage sends the produce message to server. +func (p *produceGrpcClient) SendProduceMessage(requestID int64, msg message.MutableMessage) error { + return p.Send(&streamingpb.ProduceRequest{ + Request: &streamingpb.ProduceRequest_Produce{ + Produce: &streamingpb.ProduceMessageRequest{ + RequestId: requestID, + Message: &streamingpb.Message{ + Payload: msg.Payload(), + Properties: msg.Properties().ToRawMap(), + }, + }, + }, + }) +} + +// SendClose sends the close request to server. +func (p *produceGrpcClient) SendClose() error { + return p.Send(&streamingpb.ProduceRequest{ + Request: &streamingpb.ProduceRequest_Close{ + Close: &streamingpb.CloseProducerRequest{}, + }, + }) +} diff --git a/internal/streamingnode/client/handler/producer/producer.go b/internal/streamingnode/client/handler/producer/producer.go new file mode 100644 index 0000000000000..b611931c86ac0 --- /dev/null +++ b/internal/streamingnode/client/handler/producer/producer.go @@ -0,0 +1,30 @@ +package producer + +import ( + "context" + + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/types" +) + +var _ Producer = (*producerImpl)(nil) + +// Producer is the interface that wraps the basic produce method on grpc stream. +// Producer is work on a single stream on grpc, +// so Producer cannot recover from failure because of the stream is broken. +type Producer interface { + // Assignment returns the assignment of the producer. + Assignment() types.PChannelInfoAssigned + + // Produce sends the produce message to server. + Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) + + // Check if a producer is available. + IsAvailable() bool + + // Available returns a channel that will be closed when the producer is unavailable. + Available() <-chan struct{} + + // Close close the producer client. + Close() +} diff --git a/internal/streamingnode/client/handler/producer/producer_impl.go b/internal/streamingnode/client/handler/producer/producer_impl.go new file mode 100644 index 0000000000000..15783bc406adc --- /dev/null +++ b/internal/streamingnode/client/handler/producer/producer_impl.go @@ -0,0 +1,313 @@ +package producer + +import ( + "context" + "fmt" + "io" + "sync" + + "github.com/cockroachdb/errors" + "go.uber.org/zap" + + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil" + "github.com/milvus-io/milvus/internal/util/streamingutil/status" + "github.com/milvus-io/milvus/internal/util/streamingutil/typeconverter" + "github.com/milvus-io/milvus/pkg/log" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +// ProducerOptions is the options for creating a producer. +type ProducerOptions struct { + // The produce target + Assignment *types.PChannelInfoAssigned +} + +// CreateProducer create a new producer client. +func CreateProducer( + ctx context.Context, + opts *ProducerOptions, + handler streamingpb.StreamingNodeHandlerServiceClient, +) (Producer, error) { + ctx = createProduceRequest(ctx, opts) + streamClient, err := handler.Produce(ctx) + if err != nil { + return nil, err + } + + // Initialize the producer client. + produceClient := &produceGrpcClient{ + streamClient, + } + + // Recv the first response from server. + // It must be a create response. + resp, err := produceClient.Recv() + if err != nil { + return nil, err + } + createResp := resp.GetCreate() + if createResp == nil { + return nil, status.NewInvalidRequestSeq("first message arrive must be create response") + } + + // Initialize the producer client finished. + cli := &producerImpl{ + assignment: *opts.Assignment, + walName: createResp.GetWalName(), + logger: log.With( + zap.String("walName", createResp.GetWalName()), + zap.String("pchannel", opts.Assignment.Channel.Name), + zap.Int64("term", opts.Assignment.Channel.Term), + zap.Int64("streamingNodeID", opts.Assignment.Node.ServerID)), + lifetime: lifetime.NewLifetime[lifetime.State](lifetime.Working), + idAllocator: typeutil.NewIDAllocator(), + grpcStreamClient: produceClient, + pendingRequests: sync.Map{}, + requestCh: make(chan *produceRequest), + sendExitCh: make(chan struct{}), + finishedCh: make(chan struct{}), + } + + // Start the producer client. + go cli.execute() + return cli, nil +} + +// createProduceRequest creates the produce request. +func createProduceRequest(ctx context.Context, opts *ProducerOptions) context.Context { + // select server to consume. + ctx = contextutil.WithPickServerID(ctx, opts.Assignment.Node.ServerID) + // select channel to consume. + return contextutil.WithCreateProducer(ctx, &streamingpb.CreateProducerRequest{ + Pchannel: typeconverter.NewProtoFromPChannelInfo(opts.Assignment.Channel), + }) +} + +// Expected message sequence: +// CreateProducer +// ProduceRequest 1 -> ProduceResponse Or Error 1 +// ProduceRequest 2 -> ProduceResponse Or Error 2 +// ProduceRequest 3 -> ProduceResponse Or Error 3 +// CloseProducer +type producerImpl struct { + assignment types.PChannelInfoAssigned + walName string + logger *log.MLogger + lifetime lifetime.Lifetime[lifetime.State] + idAllocator *typeutil.IDAllocator + grpcStreamClient *produceGrpcClient + + pendingRequests sync.Map + requestCh chan *produceRequest + sendExitCh chan struct{} + finishedCh chan struct{} +} + +type produceRequest struct { + ctx context.Context + msg message.MutableMessage + respCh chan produceResponse +} + +type produceResponse struct { + id message.MessageID + err error +} + +// Assignment returns the assignment of the producer. +func (p *producerImpl) Assignment() types.PChannelInfoAssigned { + return p.assignment +} + +// Produce sends the produce message to server. +func (p *producerImpl) Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) { + if p.lifetime.Add(lifetime.IsWorking) != nil { + return nil, status.NewOnShutdownError("producer client is shutting down") + } + defer p.lifetime.Done() + + respCh := make(chan produceResponse, 1) + req := &produceRequest{ + ctx: ctx, + msg: msg, + respCh: respCh, + } + + // Send the produce message to server. + select { + case <-ctx.Done(): + return nil, ctx.Err() + case p.requestCh <- req: + case <-p.sendExitCh: + return nil, status.NewInner("producer stream client is closed") + } + + // Wait for the response from server or context timeout. + select { + case <-ctx.Done(): + return nil, ctx.Err() + case resp := <-respCh: + return resp.id, resp.err + } +} + +// execute executes the producer client. +func (p *producerImpl) execute() { + defer close(p.finishedCh) + + errSendCh := p.startSend() + errRecvCh := p.startRecv() + + // Wait for send and recv arm to exit. + err := <-errRecvCh + <-errSendCh + + // Clear all pending request. + p.pendingRequests.Range(func(key, value interface{}) bool { + value.(*produceRequest).respCh <- produceResponse{ + err: status.NewUnknownError(fmt.Sprintf("request sent but no response returned, %v", err)), + } + return true + }) +} + +// IsAvailable returns whether the producer is available. +func (p *producerImpl) IsAvailable() bool { + select { + case <-p.Available(): + return false + default: + return true + } +} + +// Available returns a channel that will be closed when the producer is unavailable. +func (p *producerImpl) Available() <-chan struct{} { + return p.sendExitCh +} + +// Close close the producer client. +func (p *producerImpl) Close() { + // Wait for all message has been sent. + p.lifetime.SetState(lifetime.Stopped) + p.lifetime.Wait() + close(p.requestCh) + + // Wait for send and recv arm to exit. + <-p.finishedCh +} + +// startSend starts the send loop. +func (p *producerImpl) startSend() <-chan struct{} { + ch := make(chan struct{}) + go func() { + _ = p.sendLoop() + close(ch) + }() + return ch +} + +// startRecv starts the recv loop. +func (p *producerImpl) startRecv() <-chan error { + errCh := make(chan error, 1) + go func() { + errCh <- p.recvLoop() + }() + return errCh +} + +// sendLoop sends the produce message to server. +func (p *producerImpl) sendLoop() (err error) { + defer func() { + if err != nil { + p.logger.Warn("send arm of stream closed by unexpected error", zap.Error(err)) + } else { + p.logger.Info("send arm of stream closed") + } + close(p.sendExitCh) + if err := p.grpcStreamClient.CloseSend(); err != nil { + p.logger.Warn("failed to close send", zap.Error(err)) + } + }() + + for req := range p.requestCh { + requestID := p.idAllocator.Allocate() + // Store the request to pending request map. + p.pendingRequests.Store(requestID, req) + // Send the produce message to server. + if err := p.grpcStreamClient.SendProduceMessage(requestID, req.msg); err != nil { + // If send failed, remove the request from pending request map and return error to client. + p.notifyRequest(requestID, produceResponse{ + err: err, + }) + return err + } + p.logger.Debug("send produce message to server", zap.Int64("requestID", requestID)) + } + // all message has been sent, sent close response. + return p.grpcStreamClient.SendClose() +} + +// recvLoop receives the produce response from server. +func (p *producerImpl) recvLoop() (err error) { + defer func() { + if err != nil { + p.logger.Warn("recv arm of stream closed by unexpected error", zap.Error(err)) + return + } + p.logger.Info("recv arm of stream closed") + }() + + for { + resp, err := p.grpcStreamClient.Recv() + if errors.Is(err, io.EOF) { + p.logger.Debug("stream closed successful") + return nil + } + if err != nil { + return err + } + switch resp := resp.Response.(type) { + case *streamingpb.ProduceResponse_Produce: + var result produceResponse + switch produceResp := resp.Produce.Response.(type) { + case *streamingpb.ProduceMessageResponse_Result: + msgID, err := message.UnmarshalMessageID( + p.walName, + produceResp.Result.GetId().GetId(), + ) + if err != nil { + return err + } + result = produceResponse{ + id: msgID, + } + case *streamingpb.ProduceMessageResponse_Error: + result = produceResponse{ + err: status.New(produceResp.Error.Code, produceResp.Error.Cause), + } + default: + panic("unreachable") + } + p.notifyRequest(resp.Produce.RequestId, result) + case *streamingpb.ProduceResponse_Close: + // recv io.EOF after this message. + default: + // skip message here. + p.logger.Error("unknown response type", zap.Any("response", resp)) + } + } +} + +// notifyRequest notify the request has been returned from server. +func (p *producerImpl) notifyRequest(requestID int64, resp produceResponse) { + pendingRequest, loaded := p.pendingRequests.LoadAndDelete(requestID) + if loaded { + p.logger.Debug("recv send produce message from server", zap.Int64("requestID", requestID)) + pendingRequest.(*produceRequest).respCh <- resp + } +} diff --git a/internal/streamingnode/client/handler/producer/producer_test.go b/internal/streamingnode/client/handler/producer/producer_test.go new file mode 100644 index 0000000000000..db0092c5be3a7 --- /dev/null +++ b/internal/streamingnode/client/handler/producer/producer_test.go @@ -0,0 +1,112 @@ +package producer + +import ( + "context" + "io" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + + "github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb" + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/streaming/walimpls/impls/walimplstest" +) + +func TestProducer(t *testing.T) { + c := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t) + cc := mock_streamingpb.NewMockStreamingNodeHandlerService_ProduceClient(t) + recvCh := make(chan *streamingpb.ProduceResponse, 10) + cc.EXPECT().Recv().RunAndReturn(func() (*streamingpb.ProduceResponse, error) { + msg, ok := <-recvCh + if !ok { + return nil, io.EOF + } + return msg, nil + }) + sendCh := make(chan struct{}, 5) + cc.EXPECT().Send(mock.Anything).RunAndReturn(func(pr *streamingpb.ProduceRequest) error { + sendCh <- struct{}{} + return nil + }) + c.EXPECT().Produce(mock.Anything, mock.Anything).Return(cc, nil) + cc.EXPECT().CloseSend().RunAndReturn(func() error { + recvCh <- &streamingpb.ProduceResponse{Response: &streamingpb.ProduceResponse_Close{}} + close(recvCh) + return nil + }) + + ctx := context.Background() + opts := &ProducerOptions{ + Assignment: &types.PChannelInfoAssigned{ + Channel: types.PChannelInfo{Name: "test", Term: 1}, + Node: types.StreamingNodeInfo{ServerID: 1, Address: "localhost"}, + }, + } + + recvCh <- &streamingpb.ProduceResponse{ + Response: &streamingpb.ProduceResponse_Create{ + Create: &streamingpb.CreateProducerResponse{ + WalName: "test", + }, + }, + } + producer, err := CreateProducer(ctx, opts, c) + assert.NoError(t, err) + assert.NotNil(t, producer) + ch := make(chan struct{}) + go func() { + msgID, err := producer.Produce(ctx, message.NewMutableMessageBuilder(). + WithMessageType(message.MessageTypeUnknown). + WithPayload([]byte{}). + BuildMutable()) + assert.Error(t, err) + assert.Nil(t, msgID) + msgID, err = producer.Produce(ctx, message.NewMutableMessageBuilder(). + WithMessageType(message.MessageTypeUnknown). + WithPayload([]byte{}). + BuildMutable()) + assert.NoError(t, err) + assert.NotNil(t, msgID) + close(ch) + }() + <-sendCh + recvCh <- &streamingpb.ProduceResponse{ + Response: &streamingpb.ProduceResponse_Produce{ + Produce: &streamingpb.ProduceMessageResponse{ + RequestId: 1, + Response: &streamingpb.ProduceMessageResponse_Error{ + Error: &streamingpb.StreamingError{Code: 1}, + }, + }, + }, + } + <-sendCh + recvCh <- &streamingpb.ProduceResponse{ + Response: &streamingpb.ProduceResponse_Produce{ + Produce: &streamingpb.ProduceMessageResponse{ + RequestId: 2, + Response: &streamingpb.ProduceMessageResponse_Result{ + Result: &streamingpb.ProduceMessageResponseResult{ + Id: &streamingpb.MessageID{Id: walimplstest.NewTestMessageID(1).Marshal()}, + }, + }, + }, + }, + } + <-ch + + ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + _, err = producer.Produce(ctx, message.NewMutableMessageBuilder(). + WithMessageType(message.MessageTypeUnknown). + WithPayload([]byte{}). + BuildMutable()) + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.True(t, producer.IsAvailable()) + producer.Close() + assert.False(t, producer.IsAvailable()) +} diff --git a/internal/streamingnode/client/handler/shared_producer.go b/internal/streamingnode/client/handler/shared_producer.go new file mode 100644 index 0000000000000..7e6d8da2e1465 --- /dev/null +++ b/internal/streamingnode/client/handler/shared_producer.go @@ -0,0 +1,51 @@ +package handler + +import ( + "context" + + "github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer" + "github.com/milvus-io/milvus/pkg/streaming/util/message" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +var _ producer.Producer = (*sharedProducer)(nil) + +// newSharedProducer creates a shared producer. +func newSharedProducer(ref *typeutil.SharedReference[Producer]) sharedProducer { + return sharedProducer{ + SharedReference: ref, + } +} + +// sharedProducer is a shared producer. +type sharedProducer struct { + *typeutil.SharedReference[Producer] +} + +// Clone clones the shared producer. +func (sp sharedProducer) Clone() *sharedProducer { + return &sharedProducer{ + SharedReference: sp.SharedReference.Clone(), + } +} + +// Assignment returns the assignment of the producer. +func (sp sharedProducer) Assignment() types.PChannelInfoAssigned { + return sp.Deref().Assignment() +} + +// Produce sends the produce message to server. +func (sp sharedProducer) Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) { + return sp.Deref().Produce(ctx, msg) +} + +// Check if a producer is available. +func (sp sharedProducer) IsAvailable() bool { + return sp.Deref().IsAvailable() +} + +// Available returns a channel that will be closed when the producer is unavailable. +func (sp sharedProducer) Available() <-chan struct{} { + return sp.Deref().Available() +} diff --git a/internal/streamingnode/client/manager/manager.go b/internal/streamingnode/client/manager/manager.go deleted file mode 100644 index 8582c18f364c2..0000000000000 --- a/internal/streamingnode/client/manager/manager.go +++ /dev/null @@ -1,24 +0,0 @@ -package manager - -import ( - "context" - - "github.com/milvus-io/milvus/pkg/streaming/util/types" -) - -type ManagerClient interface { - // WatchNodeChanged returns a channel that receive a node change. - WatchNodeChanged(ctx context.Context) (<-chan struct{}, error) - - // CollectStatus collects status of all wal instances in all streamingnode. - CollectAllStatus(ctx context.Context) (map[int64]*types.StreamingNodeStatus, error) - - // Assign a wal instance for the channel on log node of given server id. - Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error - - // Remove the wal instance for the channel on log node of given server id. - Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error - - // Close closes the manager client. - Close() -} diff --git a/internal/streamingnode/client/manager/manager_client.go b/internal/streamingnode/client/manager/manager_client.go new file mode 100644 index 0000000000000..2012ec70a6dcc --- /dev/null +++ b/internal/streamingnode/client/manager/manager_client.go @@ -0,0 +1,112 @@ +package manager + +import ( + "context" + "encoding/json" + "time" + + clientv3 "go.etcd.io/etcd/client/v3" + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/internal/util/sessionutil" + "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" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/tracer" + "github.com/milvus-io/milvus/pkg/util/interceptor" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/paramtable" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +// ManagerClient is the client to manage wal instances in all streamingnode. +// ManagerClient wraps the Session Service Discovery. +// Provides the ability to assign and remove wal instances for the channel on streaming node. +type ManagerClient interface { + // WatchNodeChanged returns a channel that receive the signal that a streaming node change. + WatchNodeChanged(ctx context.Context) (<-chan struct{}, error) + + // CollectAllStatus collects status of all streamingnode, such as load balance attributes. + CollectAllStatus(ctx context.Context) (map[int64]*types.StreamingNodeStatus, error) + + // Assign a wal instance for the channel on streaming node of given server id. + Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error + + // Remove the wal instance for the channel on streaming node of given server id. + Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error + + // Close closes the manager client. + // It close the underlying connection, stop the node watcher and release all resources. + Close() +} + +// NewManagerClient creates a new manager client. +func NewManagerClient(etcdCli *clientv3.Client) ManagerClient { + role := sessionutil.GetSessionPrefixByRole(typeutil.StreamingNodeRole) + rb := resolver.NewSessionBuilder(etcdCli, role) + dialTimeout := paramtable.Get().StreamingNodeGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond) + dialOptions := getDialOptions(rb) + conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) { + ctx, cancel := context.WithTimeout(ctx, dialTimeout) + defer cancel() + return grpc.DialContext( + ctx, + resolver.SessionResolverScheme+":///"+typeutil.StreamingNodeRole, + dialOptions..., + ) + }) + return &managerClientImpl{ + lifetime: lifetime.NewLifetime(lifetime.Working), + rb: rb, + service: lazygrpc.WithServiceCreator(conn, streamingpb.NewStreamingNodeManagerServiceClient), + } +} + +// getDialOptions returns grpc dial options. +func getDialOptions(rb resolver.Builder) []grpc.DialOption { + cfg := ¶mtable.Get().StreamingNodeGrpcClientCfg + retryPolicy := cfg.GetDefaultRetryPolicy() + retryPolicy["retryableStatusCodes"] = []string{"UNAVAILABLE"} + defaultServiceConfig := map[string]interface{}{ + "loadBalancingConfig": []map[string]interface{}{ + {picker.ServerIDPickerBalancerName: map[string]interface{}{}}, + }, + "methodConfig": []map[string]interface{}{ + { + "name": []map[string]string{ + {"service": "milvus.proto.streaming.StreamingNodeManagerService"}, + }, + "waitForReady": true, + "retryPolicy": retryPolicy, + }, + }, + } + defaultServiceConfigJSON, err := json.Marshal(defaultServiceConfig) + if err != nil { + panic(err) + } + dialOptions := cfg.GetDialOptionsFromConfig() + dialOptions = append(dialOptions, + grpc.WithBlock(), + grpc.WithResolvers(rb), + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithChainUnaryInterceptor( + otelgrpc.UnaryClientInterceptor(tracer.GetInterceptorOpts()...), + interceptor.ClusterInjectionUnaryClientInterceptor(), + streamingserviceinterceptor.NewStreamingServiceUnaryClientInterceptor(), + ), + grpc.WithChainStreamInterceptor( + otelgrpc.StreamClientInterceptor(tracer.GetInterceptorOpts()...), + interceptor.ClusterInjectionStreamClientInterceptor(), + streamingserviceinterceptor.NewStreamingServiceStreamClientInterceptor(), + ), + grpc.WithReturnConnectionError(), + grpc.WithDefaultServiceConfig(string(defaultServiceConfigJSON)), + ) + return dialOptions +} diff --git a/internal/streamingnode/client/manager/manager_client_impl.go b/internal/streamingnode/client/manager/manager_client_impl.go new file mode 100644 index 0000000000000..490056cb2473d --- /dev/null +++ b/internal/streamingnode/client/manager/manager_client_impl.go @@ -0,0 +1,191 @@ +package manager + +import ( + "context" + "sync" + + "go.uber.org/zap" + "golang.org/x/sync/errgroup" + + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer" + "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/status" + "github.com/milvus-io/milvus/internal/util/streamingutil/typeconverter" + "github.com/milvus-io/milvus/pkg/log" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/lifetime" +) + +var _ ManagerClient = (*managerClientImpl)(nil) + +// managerClientImpl implements ManagerClient. +type managerClientImpl struct { + lifetime lifetime.Lifetime[lifetime.State] + + rb resolver.Builder + service lazygrpc.Service[streamingpb.StreamingNodeManagerServiceClient] +} + +func (c *managerClientImpl) WatchNodeChanged(ctx context.Context) (<-chan struct{}, error) { + if c.lifetime.Add(lifetime.IsWorking) != nil { + return nil, status.NewOnShutdownError("manager client is closing") + } + defer c.lifetime.Done() + + resultCh := make(chan struct{}, 1) + go func() { + defer close(resultCh) + c.rb.Resolver().Watch(ctx, func(state resolver.VersionedState) error { + select { + case <-ctx.Done(): + return ctx.Err() + case <-c.lifetime.CloseCh(): + return status.NewOnShutdownError("manager client is closing") + case resultCh <- struct{}{}: + } + return nil + }) + }() + return resultCh, nil +} + +// CollectAllStatus collects status in all underlying streamingnode. +func (c *managerClientImpl) CollectAllStatus(ctx context.Context) (map[int64]*types.StreamingNodeStatus, error) { + if c.lifetime.Add(lifetime.IsWorking) != nil { + return nil, status.NewOnShutdownError("manager client is closing") + } + defer c.lifetime.Done() + + // Get all discovered streamingnode. + state := c.rb.Resolver().GetLatestState() + if len(state.State.Addresses) == 0 { + return make(map[int64]*types.StreamingNodeStatus), nil + } + + // Collect status of all streamingnode. + result, err := c.getAllStreamingNodeStatus(ctx, state) + if err != nil { + return nil, err + } + + // Collect status may cost some time, so we need to check the lifetime again. + newState := c.rb.Resolver().GetLatestState() + if newState.Version.GT(state.Version) { + newSession := newState.Sessions() + for serverID := range result { + if session, ok := newSession[serverID]; !ok { + result[serverID].Err = types.ErrNotAlive + } else if session.Stopping { + result[serverID].Err = types.ErrStopping + } + } + } + return result, nil +} + +func (c *managerClientImpl) getAllStreamingNodeStatus(ctx context.Context, state discoverer.VersionedState) (map[int64]*types.StreamingNodeStatus, error) { + // wait for manager service ready. + manager, err := c.service.GetService(ctx) + if err != nil { + return nil, err + } + + g, _ := errgroup.WithContext(ctx) + g.SetLimit(10) + var mu sync.Mutex + result := make(map[int64]*types.StreamingNodeStatus, len(state.State.Addresses)) + for serverID, session := range state.Sessions() { + serverID := serverID + address := session.Address + g.Go(func() error { + ctx := contextutil.WithPickServerID(ctx, serverID) + resp, err := manager.CollectStatus(ctx, &streamingpb.StreamingNodeManagerCollectStatusRequest{}) + mu.Lock() + defer mu.Unlock() + result[serverID] = &types.StreamingNodeStatus{ + StreamingNodeInfo: types.StreamingNodeInfo{ + ServerID: serverID, + Address: address, + }, + Err: err, + } + + if err != nil { + log.Warn("collect status failed, skip", zap.Int64("serverID", serverID), zap.Error(err)) + return err + } + log.Debug("collect status success", zap.Int64("serverID", serverID), zap.Any("status", resp)) + return nil + }) + } + g.Wait() + + return result, nil +} + +// Assign a wal instance for the channel on log node of given server id. +func (c *managerClientImpl) Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error { + if c.lifetime.Add(lifetime.IsWorking) != nil { + return status.NewOnShutdownError("manager client is closing") + } + defer c.lifetime.Done() + + // wait for manager service ready. + manager, err := c.service.GetService(ctx) + if err != nil { + return err + } + + // Select a log node to assign the wal instance. + ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID) + _, err = manager.Assign(ctx, &streamingpb.StreamingNodeManagerAssignRequest{ + Pchannel: typeconverter.NewProtoFromPChannelInfo(pchannel.Channel), + }) + return err +} + +// Remove the wal instance for the channel on log node of given server id. +func (c *managerClientImpl) Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error { + if c.lifetime.Add(lifetime.IsWorking) != nil { + return status.NewOnShutdownError("manager client is closing") + } + defer c.lifetime.Done() + + // wait for manager service ready. + manager, err := c.service.GetService(ctx) + if err != nil { + return err + } + + // Select a streaming node to remove the wal instance. + ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID) + _, err = manager.Remove(ctx, &streamingpb.StreamingNodeManagerRemoveRequest{ + Pchannel: typeconverter.NewProtoFromPChannelInfo(pchannel.Channel), + }) + // The following error can be treated as success. + // 1. err is nil, a real remove operation at streaming node has been happened. + // 2. err is ErrSubConnNoExist, the streaming node is not alive at view of session, so the wal on it is already removed. + // 3. err is SkippedOperation, the streaming node is not the owner of the wal, so the wal on it is already removed. + if err == nil || picker.IsErrSubConnNoExist(err) { + return nil + } + statusErr := status.AsStreamingError(err) + if statusErr == nil || statusErr.IsSkippedOperation() { + return nil + } + return err +} + +// Close closes the manager client. +func (c *managerClientImpl) Close() { + c.lifetime.SetState(lifetime.Stopped) + c.lifetime.Wait() + c.lifetime.Close() + + c.service.Close() + c.rb.Close() +} diff --git a/internal/streamingnode/client/manager/manager_test.go b/internal/streamingnode/client/manager/manager_test.go new file mode 100644 index 0000000000000..d2d05e6e0e2eb --- /dev/null +++ b/internal/streamingnode/client/manager/manager_test.go @@ -0,0 +1,174 @@ +package manager + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "google.golang.org/grpc" + "google.golang.org/grpc/resolver" + + "github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb" + "github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_lazygrpc" + "github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver" + "github.com/milvus-io/milvus/internal/proto/streamingpb" + "github.com/milvus-io/milvus/internal/util/sessionutil" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/attributes" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil" + "github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer" + "github.com/milvus-io/milvus/pkg/streaming/util/types" + "github.com/milvus-io/milvus/pkg/util/etcd" + "github.com/milvus-io/milvus/pkg/util/lifetime" + "github.com/milvus-io/milvus/pkg/util/paramtable" + "github.com/milvus-io/milvus/pkg/util/typeutil" +) + +func TestManager(t *testing.T) { + rb := mock_resolver.NewMockBuilder(t) + managerService := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeManagerServiceClient](t) + m := &managerClientImpl{ + lifetime: lifetime.NewLifetime(lifetime.Working), + rb: rb, + service: managerService, + } + r := mock_resolver.NewMockResolver(t) + rb.EXPECT().Resolver().Return(r) + r.EXPECT().Watch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(discoverer.VersionedState) error) error { + f(discoverer.VersionedState{}) + return context.Canceled + }) + + // Test WatchNodeChanged. + resultCh, err := m.WatchNodeChanged(context.Background()) + assert.NoError(t, err) + _, ok := <-resultCh + assert.True(t, ok) + _, ok = <-resultCh + assert.False(t, ok) + + // Test CollectAllStatus + r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState { + return discoverer.VersionedState{ + Version: typeutil.VersionInt64(1), + State: resolver.State{}, + } + }) + // Not address here. + nodes, err := m.CollectAllStatus(context.Background()) + assert.NoError(t, err) + assert.Len(t, nodes, 0) + + // Has address. + managerServiceClient := mock_streamingpb.NewMockStreamingNodeManagerServiceClient(t) + managerService.EXPECT().GetService(mock.Anything).RunAndReturn(func(ctx context.Context) (streamingpb.StreamingNodeManagerServiceClient, error) { + return managerServiceClient, nil + }) + managerServiceClient.EXPECT().CollectStatus(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, snmcsr *streamingpb.StreamingNodeManagerCollectStatusRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error) { + return &streamingpb.StreamingNodeManagerCollectStatusResponse{}, nil + }) + + i := 0 + states := []map[uint64]bool{ + {1: false, 2: false, 3: true}, + {1: true, 2: false}, + } + r.EXPECT().GetLatestState().Unset() + r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState { + s := newVersionedState(int64(i), states[i]) + i++ + return s + }) + + nodes, err = m.CollectAllStatus(context.Background()) + assert.NoError(t, err) + assert.Len(t, nodes, 3) + assert.ErrorIs(t, nodes[3].Err, types.ErrNotAlive) + assert.ErrorIs(t, nodes[1].Err, types.ErrStopping) + + // Test Assign + serverID := int64(2) + managerServiceClient.EXPECT().Assign(mock.Anything, mock.Anything).RunAndReturn( + func(ctx context.Context, snmcsr *streamingpb.StreamingNodeManagerAssignRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error) { + pickedServerID, ok := contextutil.GetPickServerID(ctx) + assert.True(t, ok) + assert.Equal(t, serverID, pickedServerID) + return nil, nil + }) + err = m.Assign(context.Background(), types.PChannelInfoAssigned{ + Channel: types.PChannelInfo{Name: "p", Term: 1}, + Node: types.StreamingNodeInfo{ServerID: serverID}, + }) + assert.NoError(t, err) + + // Test Remove + serverID = int64(2) + managerServiceClient.EXPECT().Remove(mock.Anything, mock.Anything).RunAndReturn( + func(ctx context.Context, snmrr *streamingpb.StreamingNodeManagerRemoveRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error) { + pickedServerID, ok := contextutil.GetPickServerID(ctx) + assert.True(t, ok) + assert.Equal(t, serverID, pickedServerID) + return nil, nil + }) + err = m.Remove(context.Background(), types.PChannelInfoAssigned{ + Channel: types.PChannelInfo{Name: "p", Term: 1}, + Node: types.StreamingNodeInfo{ServerID: serverID}, + }) + assert.NoError(t, err) + + // Test Close + managerService.EXPECT().Close().Return() + rb.EXPECT().Close().Return() + m.Close() + + nodes, err = m.CollectAllStatus(context.Background()) + assert.Nil(t, nodes) + assert.Error(t, err) + err = m.Assign(context.Background(), types.PChannelInfoAssigned{}) + assert.Error(t, err) + err = m.Remove(context.Background(), types.PChannelInfoAssigned{}) + assert.Error(t, err) + resultCh, err = m.WatchNodeChanged(context.Background()) + assert.Nil(t, resultCh) + assert.Error(t, err) +} + +func newVersionedState(version int64, serverIDs map[uint64]bool) discoverer.VersionedState { + state := discoverer.VersionedState{ + Version: typeutil.VersionInt64(version), + State: resolver.State{ + Addresses: make([]resolver.Address, 0, len(serverIDs)), + }, + } + + for serverID, stopping := range serverIDs { + state.State.Addresses = append(state.State.Addresses, resolver.Address{ + Addr: fmt.Sprintf("localhost:%d", serverID), + BalancerAttributes: attributes.WithSession( + new(attributes.Attributes), &sessionutil.SessionRaw{ + ServerID: int64(serverID), + Stopping: stopping, + }, + ), + }) + } + return state +} + +func TestDial(t *testing.T) { + paramtable.Init() + + err := etcd.InitEtcdServer(true, "", t.TempDir(), "stdout", "info") + assert.NoError(t, err) + defer etcd.StopEtcdServer() + c, err := etcd.GetEmbedEtcdClient() + assert.NoError(t, err) + assert.NotNil(t, c) + + client := NewManagerClient(c) + assert.NotNil(t, client) + time.Sleep(100 * time.Millisecond) + client.Close() +} diff --git a/internal/streamingnode/server/walmanager/manager.go b/internal/streamingnode/server/walmanager/manager.go index 142461ef617c0..445f1c68bdabf 100644 --- a/internal/streamingnode/server/walmanager/manager.go +++ b/internal/streamingnode/server/walmanager/manager.go @@ -18,7 +18,7 @@ type Manager interface { // GetAvailableWAL returns a available wal instance for the channel. // Return nil if the wal instance is not found. - GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, error) + GetAvailableWAL(types.PChannelInfo) (wal.WAL, error) // GetAllAvailableWALInfo returns all available channel info. GetAllAvailableChannels() ([]types.PChannelInfo, error) diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 80703b478928c..88f6a9294faee 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -90,13 +90,15 @@ type ComponentParam struct { DataNodeGrpcServerCfg GrpcServerConfig IndexNodeGrpcServerCfg GrpcServerConfig - RootCoordGrpcClientCfg GrpcClientConfig - ProxyGrpcClientCfg GrpcClientConfig - QueryCoordGrpcClientCfg GrpcClientConfig - QueryNodeGrpcClientCfg GrpcClientConfig - DataCoordGrpcClientCfg GrpcClientConfig - DataNodeGrpcClientCfg GrpcClientConfig - IndexNodeGrpcClientCfg GrpcClientConfig + RootCoordGrpcClientCfg GrpcClientConfig + ProxyGrpcClientCfg GrpcClientConfig + QueryCoordGrpcClientCfg GrpcClientConfig + QueryNodeGrpcClientCfg GrpcClientConfig + DataCoordGrpcClientCfg GrpcClientConfig + DataNodeGrpcClientCfg GrpcClientConfig + IndexNodeGrpcClientCfg GrpcClientConfig + StreamingCoordGrpcClientCfg GrpcClientConfig + StreamingNodeGrpcClientCfg GrpcClientConfig IntegrationTestCfg integrationTestConfig @@ -151,6 +153,8 @@ func (p *ComponentParam) init(bt *BaseTable) { p.DataCoordGrpcClientCfg.Init("dataCoord", bt) p.DataNodeGrpcClientCfg.Init("dataNode", bt) p.IndexNodeGrpcClientCfg.Init("indexNode", bt) + p.StreamingCoordGrpcClientCfg.Init("streamingCoord", bt) + p.StreamingNodeGrpcClientCfg.Init("streamingNode", bt) p.IntegrationTestCfg.init(bt) } diff --git a/pkg/util/paramtable/grpc_param.go b/pkg/util/paramtable/grpc_param.go index 31340b6b01987..d75f6f61ffadb 100644 --- a/pkg/util/paramtable/grpc_param.go +++ b/pkg/util/paramtable/grpc_param.go @@ -19,8 +19,12 @@ package paramtable import ( "fmt" "strconv" + "time" "go.uber.org/zap" + "google.golang.org/grpc" + "google.golang.org/grpc/backoff" + "google.golang.org/grpc/keepalive" "github.com/milvus-io/milvus/pkg/log" "github.com/milvus-io/milvus/pkg/util/funcutil" @@ -219,6 +223,7 @@ type GrpcClientConfig struct { MaxAttempts ParamItem `refreshable:"false"` InitialBackoff ParamItem `refreshable:"false"` MaxBackoff ParamItem `refreshable:"false"` + BackoffMultiplier ParamItem `refreshable:"false"` MinResetInterval ParamItem `refreshable:"false"` MaxCancelError ParamItem `refreshable:"false"` MinSessionCheckInterval ParamItem `refreshable:"false"` @@ -397,6 +402,14 @@ func (p *GrpcClientConfig) Init(domain string, base *BaseTable) { } p.MaxBackoff.Init(base.mgr) + p.BackoffMultiplier = ParamItem{ + Key: "grpc.client.backoffMultiplier", + Version: "2.5.0", + DefaultValue: "2.0", + Export: true, + } + p.BackoffMultiplier.Init(base.mgr) + compressionEnabled := fmt.Sprintf("%t", DefaultCompressionEnabled) p.CompressionEnabled = ParamItem{ Key: "grpc.client.compressionEnabled", @@ -478,3 +491,42 @@ func (p *GrpcClientConfig) Init(domain string, base *BaseTable) { } p.MaxCancelError.Init(base.mgr) } + +// GetDialOptionsFromConfig returns grpc dial options from config. +func (p *GrpcClientConfig) GetDialOptionsFromConfig() []grpc.DialOption { + compress := "" + if p.CompressionEnabled.GetAsBool() { + compress = "zstd" + } + return []grpc.DialOption{ + grpc.WithDefaultCallOptions( + grpc.MaxCallRecvMsgSize(p.ClientMaxRecvSize.GetAsInt()), + grpc.MaxCallSendMsgSize(p.ClientMaxSendSize.GetAsInt()), + grpc.UseCompressor(compress), + ), + grpc.WithKeepaliveParams(keepalive.ClientParameters{ + Time: p.KeepAliveTime.GetAsDuration(time.Millisecond), + Timeout: p.KeepAliveTimeout.GetAsDuration(time.Millisecond), + PermitWithoutStream: true, + }), + grpc.WithConnectParams(grpc.ConnectParams{ + Backoff: backoff.Config{ + BaseDelay: 100 * time.Millisecond, + Multiplier: 1.6, + Jitter: 0.2, + MaxDelay: 3 * time.Second, + }, + MinConnectTimeout: p.DialTimeout.GetAsDuration(time.Millisecond), + }), + } +} + +// GetDefaultRetryPolicy returns default grpc retry policy. +func (p *GrpcClientConfig) GetDefaultRetryPolicy() map[string]interface{} { + return map[string]interface{}{ + "maxAttempts": p.MaxAttempts.GetAsInt(), + "initialBackoff": fmt.Sprintf("%fs", p.InitialBackoff.GetAsFloat()), + "maxBackoff": fmt.Sprintf("%fs", p.MaxBackoff.GetAsFloat()), + "backoffMultiplier": p.BackoffMultiplier.GetAsFloat(), + } +} diff --git a/pkg/util/typeutil/shared_reference.go b/pkg/util/typeutil/shared_reference.go new file mode 100644 index 0000000000000..97610859b37e2 --- /dev/null +++ b/pkg/util/typeutil/shared_reference.go @@ -0,0 +1,99 @@ +package typeutil + +import ( + "go.uber.org/atomic" +) + +// Closable is an interface that can be closed or release the resource. +type Closable interface { + Close() +} + +// SharedReference is a reference type that can be shared with only one close operation. +// Without supported of determined destructor, A safe SharedReference is rely on user's behavior. +type SharedReference[T Closable] struct { + inner T + strongCnt *atomic.Int32 +} + +// NewSharedReference creates a new SharedReference with the inner object. +func NewSharedReference[T Closable](inner T) *SharedReference[T] { + return &SharedReference[T]{ + inner: inner, + strongCnt: atomic.NewInt32(1), + } +} + +// Deref returns the inner object. +// Deref should only be called if the reference is not closed. +// Otherwise UB happens. +func (sr *SharedReference[T]) Deref() T { + return sr.inner +} + +// Clone returns a new SharedReference with the same inner object. +// Clone should only be called if the reference is not closed. +// Otherwise UB happens. +func (sr *SharedReference[T]) Clone() *SharedReference[T] { + n := sr.strongCnt.Load() + + for { + if n == 0 { + panic("SharedReference: Clone on a closed reference") + } + if sr.strongCnt.CompareAndSwap(n, n+1) { + break + } + n = sr.strongCnt.Load() + } + return sr +} + +// Downgrade returns a new WeakReference with the same inner object. +// After downgrade, the SharedReference can not be used. +func (sr *SharedReference[T]) Downgrade() *WeakReference[T] { + w := &WeakReference[T]{ + inner: sr.inner, + strongCnt: sr.strongCnt, + } + sr.Close() + return w +} + +// Close closes the reference, should only be called once at a sharedReference. +// After called this method, all other method of the SharedReference should not be called any more. +// if the reference count is 0, the inner object will be closed. +func (sr *SharedReference[T]) Close() { + sr.strongCnt.Load() + refcnt := sr.strongCnt.Dec() + if refcnt < 0 { + panic("SharedReference: Close on a closed reference") + } + if refcnt == 0 { + sr.inner.Close() + } +} + +// WeakReference is a weak reference type that can be shared with only one close operation. +type WeakReference[T Closable] struct { + inner T + strongCnt *atomic.Int32 + // TODO: add weak ref count in future. +} + +func (wr *WeakReference[T]) Upgrade() *SharedReference[T] { + n := wr.strongCnt.Load() + for { + // There's no strong reference of dLogImpl, so it's impossible to create a DLogRef. + if n == 0 { + return nil + } + if wr.strongCnt.CompareAndSwap(n, n+1) { + return &SharedReference[T]{ + inner: wr.inner, + strongCnt: wr.strongCnt, + } + } + n = wr.strongCnt.Load() + } +} diff --git a/pkg/util/typeutil/shared_reference_test.go b/pkg/util/typeutil/shared_reference_test.go new file mode 100644 index 0000000000000..2770cc4fca617 --- /dev/null +++ b/pkg/util/typeutil/shared_reference_test.go @@ -0,0 +1,81 @@ +package typeutil + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/atomic" +) + +type closer struct { + t *testing.T + closed atomic.Bool +} + +func (c *closer) Close() { + assert.True(c.t, c.closed.CompareAndSwap(false, true)) +} + +func TestSharedReference(t *testing.T) { + // test downgrade and upgrade + closer1 := &closer{t: t} + r := NewSharedReference(closer1) + w := r.Downgrade() + assert.Nil(t, w.Upgrade()) + assert.True(t, closer1.closed.Load()) + + // test concurrent + closer1 = &closer{t: t} + + r = NewSharedReference(closer1) + + wg := &sync.WaitGroup{} + wg.Add(4) + r1 := r.Clone() + go func() { + defer wg.Done() + defer r1.Close() + for i := 0; i < 1000; i++ { + r3 := r1.Clone() + r3.Close() + } + }() + r2 := r.Clone() + go func() { + defer wg.Done() + defer r2.Close() + for i := 0; i < 1000; i++ { + r := r2.Clone() + r.Close() + } + }() + r3 := r.Clone().Downgrade() + go func() { + defer wg.Done() + r := r3.Upgrade() + r.Close() + }() + r4 := r.Clone().Downgrade() + go func() { + defer wg.Done() + r := r4.Upgrade() + r.Close() + }() + wg.Wait() + + assert.False(t, r.Deref().closed.Load()) + r.Close() + assert.True(t, closer1.closed.Load()) + + assert.Panics(t, func() { + r.Clone() + }) + + closer1 = &closer{t: t} + r = NewSharedReference(closer1) + r.Close() + assert.Panics(t, func() { + r.Close() + }) +}