admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: streaming node client implementation (#34653)

Browse Source 
issue: #33285

- add streaming node grpc client wrapper
- add unittest for streaming node grpc client side
- fix binary unsafe bug for message

---------

Signed-off-by: chyezh <chyezh@outlook.com>
This commit is contained in:
chyezh 2024-07-19 17:37:40 +08:00 committed by GitHub
parent ed057e6fce
commit 86eff6e589
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
42 changed files with 4501 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
configs/milvus.yaml
View File

@ -580,6 +580,17 @@ dataNode:
memoryBufferRatio: 0.1 # The ratio of memory buffer of clustering compaction. Data larger than threshold will be spilled to storage.
workPoolSize: 8
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'.

17
internal/.mockery.yaml
View File

@ -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:

178
internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go Normal file
View File

@ -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
}

398
internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go Normal file
View File

@ -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
}

398
internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go Normal file
View File

@ -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
}

250
internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go Normal file
View File

@ -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
}

158
internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go Normal file
View File

@ -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
}

148
internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go Normal file
View File

@ -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
}

251
internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go Normal file
View File

@ -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
}

20
internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go
View File

@ -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))
})

176
internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go Normal file
View File

@ -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
}

209
internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go Normal file
View File

@ -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
}

22
internal/streamingnode/client/handler/assignment/watcher.go Normal file
View File

@ -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()
}

108
internal/streamingnode/client/handler/assignment/watcher_impl.go Normal file
View File

@ -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()
}

91
internal/streamingnode/client/handler/assignment/watcher_test.go Normal file
View File

@ -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)
}

18
internal/streamingnode/client/handler/consumer/consumer.go Normal file
View File

@ -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()
}

181
internal/streamingnode/client/handler/consumer/consumer_impl.go Normal file
View File

@ -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))
}
}
}

85
internal/streamingnode/client/handler/consumer/consumer_test.go Normal file
View File

@ -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())
}

151
internal/streamingnode/client/handler/handler_client.go Normal file
View File

@ -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 := &paramtable.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
}

216
internal/streamingnode/client/handler/handler_client_impl.go Normal file
View File

@ -0,0 +1,216 @@
package handler
import (
"context"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/cockroachdb/errors"
"go.uber.org/zap"
"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.createHandlerAfterStreamingNodeReady(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.createHandlerAfterStreamingNodeReady(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
}
// createHandlerAfterStreamingNodeReady 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) createHandlerAfterStreamingNodeReady(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")
}
start := time.Now()
nextBackoff := backoff.NextBackOff()
logger.Info("wait for next backoff", zap.Duration("nextBackoff", nextBackoff))
isAssignemtChange, err := hc.waitForNextBackoff(ctx, pchannel, assign, nextBackoff)
cost := time.Since(start)
if err != nil {
logger.Warn("wait for next backoff failed", zap.Error(err), zap.Duration("cost", cost))
return nil, err
}
logger.Info("wait for next backoff done", zap.Bool("isAssignmentChange", isAssignemtChange), zap.Duration("cost", cost))
}
}
// 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()
}

145
internal/streamingnode/client/handler/handler_client_test.go Normal file
View File

@ -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()
}

35
internal/streamingnode/client/handler/producer/produce_grpc_client.go Normal file
View File

@ -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{},
},
})
}

30
internal/streamingnode/client/handler/producer/producer.go Normal file
View File

@ -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()
}

312
internal/streamingnode/client/handler/producer/producer_impl.go Normal file
View File

@ -0,0 +1,312 @@
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
}
}
// 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
}
}

112
internal/streamingnode/client/handler/producer/producer_test.go Normal file
View File

@ -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())
}

51
internal/streamingnode/client/handler/shared_producer.go Normal file
View File

@ -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()
}

24
internal/streamingnode/client/manager/manager.go
View File

@ -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()
}

112
internal/streamingnode/client/manager/manager_client.go Normal file
View File

@ -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 := &paramtable.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
}

191
internal/streamingnode/client/manager/manager_client_impl.go Normal file
View File

@ -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()
}

174
internal/streamingnode/client/manager/manager_test.go Normal file
View File

@ -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()
}

14
internal/streamingnode/server/resource/idalloc/test_mock_root_coord_client.go
View File

@ -12,9 +12,9 @@ import (
"go.uber.org/atomic"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/util/merr"
)
func NewMockRootCoordClient(t *testing.T) *mocks.MockRootCoordClient {
@ -27,9 +27,7 @@ func NewMockRootCoordClient(t *testing.T) *mocks.MockRootCoordClient {
}
c := counter.Add(uint64(atr.Count))
return &rootcoordpb.AllocTimestampResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
Status: merr.Success(),
Timestamp: c - uint64(atr.Count),
Count: atr.Count,
}, nil
@ -42,11 +40,9 @@ func NewMockRootCoordClient(t *testing.T) *mocks.MockRootCoordClient {
}
c := counter.Add(uint64(atr.Count))
return &rootcoordpb.AllocIDResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
},
ID: int64(c - uint64(atr.Count)),
Count: atr.Count,
Status: merr.Success(),
ID: int64(c - uint64(atr.Count)),
Count: atr.Count,
}, nil
},
).Maybe()

4
internal/util/streamingutil/service/attributes/attributes_test.go
View File

@ -3,9 +3,10 @@ package attributes
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/streaming/util/types"
"github.com/stretchr/testify/assert"
)
func TestAttributes(t *testing.T) {
@ -38,6 +39,5 @@ func TestAttributes(t *testing.T) {
attr = new(Attributes)
attr = WithServerID(attr, 1)
serverID = GetServerID(attr)
assert.Equal(t, int64(1), *GetServerID(attr))
}

2
internal/util/streamingutil/service/balancer/balancer.go
View File

@ -23,9 +23,9 @@
package balancer
import (
"errors"
"fmt"
"github.com/cockroachdb/errors"
"google.golang.org/grpc/balancer"
"google.golang.org/grpc/balancer/base"
"google.golang.org/grpc/connectivity"

4
internal/util/streamingutil/service/interceptor/server.go
View File

@ -15,7 +15,7 @@ func NewStreamingServiceUnaryServerInterceptor() grpc.UnaryServerInterceptor {
return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
resp, err := handler(ctx, req)
if err == nil {
return resp, err
return resp, nil
}
// Streaming Service Method should be overwrite the response error code.
if strings.HasPrefix(info.FullMethod, streamingpb.ServiceMethodPrefix) {
@ -35,7 +35,7 @@ func NewStreamingServiceStreamServerInterceptor() grpc.StreamServerInterceptor {
return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
err := handler(srv, ss)
if err == nil {
return err
return nil
}
// Streaming Service Method should be overwrite the response error code.

2
internal/util/streamingutil/service/resolver/builder.go
View File

@ -1,9 +1,9 @@
package resolver
import (
"errors"
"time"
"github.com/cockroachdb/errors"
clientv3 "go.etcd.io/etcd/client/v3"
"go.uber.org/zap"
"google.golang.org/grpc/resolver"

7
internal/util/streamingutil/service/resolver/builder_test.go
View File

@ -4,13 +4,14 @@ import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/grpc/resolver"
"github.com/milvus-io/milvus/internal/mocks/google.golang.org/grpc/mock_resolver"
"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_discoverer"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
"github.com/milvus-io/milvus/pkg/util/typeutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/grpc/resolver"
)
func TestNewBuilder(t *testing.T) {

2
internal/util/streamingutil/service/resolver/resolver.go
View File

@ -2,8 +2,8 @@ package resolver
import (
"context"
"errors"
"github.com/cockroachdb/errors"
"google.golang.org/grpc/resolver"
"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"

6
internal/util/streamingutil/service/resolver/resolver_with_discoverer.go
View File

@ -160,9 +160,11 @@ func (r *resolverWithDiscoverer) doDiscover() {
// Update all grpc resolver.
r.logger.Info("service discover update, update resolver", zap.Any("state", state), zap.Int("resolver_count", len(grpcResolvers)))
for watcher := range grpcResolvers {
// update operation do not block.
// Update operation do not block.
// Only return error if the resolver is closed, so just print a info log and delete the resolver.
if err := watcher.Update(state); err != nil {
r.logger.Info("resolver is closed, unregister the resolver", zap.Error(err))
// updateError is always context.Canceled.
r.logger.Info("resolver is closed, unregister the resolver", zap.NamedError("updateError", err))
delete(grpcResolvers, watcher)
}
}

18
pkg/util/paramtable/component_param.go
View File

@ -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)
}

52
pkg/util/paramtable/grpc_param.go
View File

@ -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(),
}
}

99
pkg/util/typeutil/shared_reference.go Normal file
View File

@ -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()
}
}

81
pkg/util/typeutil/shared_reference_test.go Normal file
View File

@ -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()
})
}