enhance: streaming node client implementation (#34653)

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>

configs/milvus.yaml

@@ -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'.

internal/.mockery.yaml

@@ -11,6 +11,15 @@ packages:
   github.com/milvus-io/milvus/internal/streamingnode/client/manager:
     interfaces:
       ManagerClient:
+  github.com/milvus-io/milvus/internal/streamingnode/client/handler/assignment:
+    interfaces:
+      Watcher:
+  github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer:
+    interfaces:
+      Producer:
+  github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer:
+    interfaces:
+      Consumer:
   github.com/milvus-io/milvus/internal/streamingnode/server/wal:
     interfaces:
       OpenerBuilder:
@@ -30,6 +39,10 @@ packages:
       StreamingNodeHandlerService_ConsumeServer:
       StreamingNodeHandlerService_ProduceServer:
       StreamingCoordAssignmentService_AssignmentDiscoverServer:
+      StreamingNodeManagerServiceClient:
+      StreamingNodeHandlerServiceClient:
+      StreamingNodeHandlerService_ConsumeClient:
+      StreamingNodeHandlerService_ProduceClient:
   github.com/milvus-io/milvus/internal/streamingnode/server/walmanager:
     interfaces:
       Manager:
@@ -40,9 +53,13 @@ packages:
     interfaces:
       Discoverer:
       AssignmentDiscoverWatcher:
+  github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc:
+    interfaces:
+      Service:
   github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver:
     interfaces:
       Resolver:
+      Builder:
   google.golang.org/grpc/resolver:
     interfaces:
       ClientConn:

internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerServiceClient.go

@@ -0,0 +1,178 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_streamingpb
+
+import (
+	context "context"
+
+	grpc "google.golang.org/grpc"
+
+	mock "github.com/stretchr/testify/mock"
+
+	streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb"
+)
+
+// MockStreamingNodeHandlerServiceClient is an autogenerated mock type for the StreamingNodeHandlerServiceClient type
+type MockStreamingNodeHandlerServiceClient struct {
+	mock.Mock
+}
+
+type MockStreamingNodeHandlerServiceClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockStreamingNodeHandlerServiceClient) EXPECT() *MockStreamingNodeHandlerServiceClient_Expecter {
+	return &MockStreamingNodeHandlerServiceClient_Expecter{mock: &_m.Mock}
+}
+
+// Consume provides a mock function with given fields: ctx, opts
+func (_m *MockStreamingNodeHandlerServiceClient) Consume(ctx context.Context, opts ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 streamingpb.StreamingNodeHandlerService_ConsumeClient
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, error)); ok {
+		return rf(ctx, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) streamingpb.StreamingNodeHandlerService_ConsumeClient); ok {
+		r0 = rf(ctx, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(streamingpb.StreamingNodeHandlerService_ConsumeClient)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerServiceClient_Consume_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Consume'
+type MockStreamingNodeHandlerServiceClient_Consume_Call struct {
+	*mock.Call
+}
+
+// Consume is a helper method to define mock.On call
+//   - ctx context.Context
+//   - opts ...grpc.CallOption
+func (_e *MockStreamingNodeHandlerServiceClient_Expecter) Consume(ctx interface{}, opts ...interface{}) *MockStreamingNodeHandlerServiceClient_Consume_Call {
+	return &MockStreamingNodeHandlerServiceClient_Consume_Call{Call: _e.mock.On("Consume",
+		append([]interface{}{ctx}, opts...)...)}
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Consume_Call) Run(run func(ctx context.Context, opts ...grpc.CallOption)) *MockStreamingNodeHandlerServiceClient_Consume_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-1)
+		for i, a := range args[1:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Consume_Call) Return(_a0 streamingpb.StreamingNodeHandlerService_ConsumeClient, _a1 error) *MockStreamingNodeHandlerServiceClient_Consume_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Consume_Call) RunAndReturn(run func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ConsumeClient, error)) *MockStreamingNodeHandlerServiceClient_Consume_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Produce provides a mock function with given fields: ctx, opts
+func (_m *MockStreamingNodeHandlerServiceClient) Produce(ctx context.Context, opts ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 streamingpb.StreamingNodeHandlerService_ProduceClient
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, error)); ok {
+		return rf(ctx, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, ...grpc.CallOption) streamingpb.StreamingNodeHandlerService_ProduceClient); ok {
+		r0 = rf(ctx, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(streamingpb.StreamingNodeHandlerService_ProduceClient)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerServiceClient_Produce_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Produce'
+type MockStreamingNodeHandlerServiceClient_Produce_Call struct {
+	*mock.Call
+}
+
+// Produce is a helper method to define mock.On call
+//   - ctx context.Context
+//   - opts ...grpc.CallOption
+func (_e *MockStreamingNodeHandlerServiceClient_Expecter) Produce(ctx interface{}, opts ...interface{}) *MockStreamingNodeHandlerServiceClient_Produce_Call {
+	return &MockStreamingNodeHandlerServiceClient_Produce_Call{Call: _e.mock.On("Produce",
+		append([]interface{}{ctx}, opts...)...)}
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Produce_Call) Run(run func(ctx context.Context, opts ...grpc.CallOption)) *MockStreamingNodeHandlerServiceClient_Produce_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-1)
+		for i, a := range args[1:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Produce_Call) Return(_a0 streamingpb.StreamingNodeHandlerService_ProduceClient, _a1 error) *MockStreamingNodeHandlerServiceClient_Produce_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerServiceClient_Produce_Call) RunAndReturn(run func(context.Context, ...grpc.CallOption) (streamingpb.StreamingNodeHandlerService_ProduceClient, error)) *MockStreamingNodeHandlerServiceClient_Produce_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockStreamingNodeHandlerServiceClient creates a new instance of MockStreamingNodeHandlerServiceClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockStreamingNodeHandlerServiceClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockStreamingNodeHandlerServiceClient {
+	mock := &MockStreamingNodeHandlerServiceClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ConsumeClient.go

@@ -0,0 +1,398 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_streamingpb
+
+import (
+	context "context"
+
+	mock "github.com/stretchr/testify/mock"
+	metadata "google.golang.org/grpc/metadata"
+
+	streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb"
+)
+
+// MockStreamingNodeHandlerService_ConsumeClient is an autogenerated mock type for the StreamingNodeHandlerService_ConsumeClient type
+type MockStreamingNodeHandlerService_ConsumeClient struct {
+	mock.Mock
+}
+
+type MockStreamingNodeHandlerService_ConsumeClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) EXPECT() *MockStreamingNodeHandlerService_ConsumeClient_Expecter {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseSend provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) CloseSend() error {
+	ret := _m.Called()
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) CloseSend() *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call) RunAndReturn(run func() error) *MockStreamingNodeHandlerService_ConsumeClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) Context() context.Context {
+	ret := _m.Called()
+
+	var r0 context.Context
+	if rf, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type MockStreamingNodeHandlerService_ConsumeClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Context() *MockStreamingNodeHandlerService_ConsumeClient_Context_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) Return(_a0 context.Context) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Context_Call) RunAndReturn(run func() context.Context) *MockStreamingNodeHandlerService_ConsumeClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) Header() (metadata.MD, error) {
+	ret := _m.Called()
+
+	var r0 metadata.MD
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type MockStreamingNodeHandlerService_ConsumeClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Header() *MockStreamingNodeHandlerService_ConsumeClient_Header_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) Return(_a0 metadata.MD, _a1 error) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *MockStreamingNodeHandlerService_ConsumeClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Recv provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) Recv() (*streamingpb.ConsumeResponse, error) {
+	ret := _m.Called()
+
+	var r0 *streamingpb.ConsumeResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (*streamingpb.ConsumeResponse, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() *streamingpb.ConsumeResponse); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*streamingpb.ConsumeResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
+type MockStreamingNodeHandlerService_ConsumeClient_Recv_Call struct {
+	*mock.Call
+}
+
+// Recv is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Recv() *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Recv_Call{Call: _e.mock.On("Recv")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) Return(_a0 *streamingpb.ConsumeResponse, _a1 error) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call) RunAndReturn(run func() (*streamingpb.ConsumeResponse, error)) *MockStreamingNodeHandlerService_ConsumeClient_Recv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function with given fields: m
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) RecvMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m interface{}
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) RecvMsg(m interface{}) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(interface{}))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ConsumeClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function with given fields: _a0
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) Send(_a0 *streamingpb.ConsumeRequest) error {
+	ret := _m.Called(_a0)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(*streamingpb.ConsumeRequest) error); ok {
+		r0 = rf(_a0)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type MockStreamingNodeHandlerService_ConsumeClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - _a0 *streamingpb.ConsumeRequest
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Send(_a0 interface{}) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Send_Call{Call: _e.mock.On("Send", _a0)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) Run(run func(_a0 *streamingpb.ConsumeRequest)) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(*streamingpb.ConsumeRequest))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Send_Call) RunAndReturn(run func(*streamingpb.ConsumeRequest) error) *MockStreamingNodeHandlerService_ConsumeClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function with given fields: m
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) SendMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m interface{}
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) SendMsg(m interface{}) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(interface{}))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ConsumeClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ConsumeClient) Trailer() metadata.MD {
+	ret := _m.Called()
+
+	var r0 metadata.MD
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ConsumeClient_Expecter) Trailer() *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call {
+	return &MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) Run(run func()) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) Return(_a0 metadata.MD) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *MockStreamingNodeHandlerService_ConsumeClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockStreamingNodeHandlerService_ConsumeClient creates a new instance of MockStreamingNodeHandlerService_ConsumeClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockStreamingNodeHandlerService_ConsumeClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockStreamingNodeHandlerService_ConsumeClient {
+	mock := &MockStreamingNodeHandlerService_ConsumeClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/proto/mock_streamingpb/mock_StreamingNodeHandlerService_ProduceClient.go

@@ -0,0 +1,398 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_streamingpb
+
+import (
+	context "context"
+
+	mock "github.com/stretchr/testify/mock"
+	metadata "google.golang.org/grpc/metadata"
+
+	streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb"
+)
+
+// MockStreamingNodeHandlerService_ProduceClient is an autogenerated mock type for the StreamingNodeHandlerService_ProduceClient type
+type MockStreamingNodeHandlerService_ProduceClient struct {
+	mock.Mock
+}
+
+type MockStreamingNodeHandlerService_ProduceClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockStreamingNodeHandlerService_ProduceClient) EXPECT() *MockStreamingNodeHandlerService_ProduceClient_Expecter {
+	return &MockStreamingNodeHandlerService_ProduceClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseSend provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ProduceClient) CloseSend() error {
+	ret := _m.Called()
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) CloseSend() *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call) RunAndReturn(run func() error) *MockStreamingNodeHandlerService_ProduceClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ProduceClient) Context() context.Context {
+	ret := _m.Called()
+
+	var r0 context.Context
+	if rf, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type MockStreamingNodeHandlerService_ProduceClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Context() *MockStreamingNodeHandlerService_ProduceClient_Context_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) Return(_a0 context.Context) *MockStreamingNodeHandlerService_ProduceClient_Context_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Context_Call) RunAndReturn(run func() context.Context) *MockStreamingNodeHandlerService_ProduceClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ProduceClient) Header() (metadata.MD, error) {
+	ret := _m.Called()
+
+	var r0 metadata.MD
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type MockStreamingNodeHandlerService_ProduceClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Header() *MockStreamingNodeHandlerService_ProduceClient_Header_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) Return(_a0 metadata.MD, _a1 error) *MockStreamingNodeHandlerService_ProduceClient_Header_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *MockStreamingNodeHandlerService_ProduceClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Recv provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ProduceClient) Recv() (*streamingpb.ProduceResponse, error) {
+	ret := _m.Called()
+
+	var r0 *streamingpb.ProduceResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (*streamingpb.ProduceResponse, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() *streamingpb.ProduceResponse); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*streamingpb.ProduceResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
+type MockStreamingNodeHandlerService_ProduceClient_Recv_Call struct {
+	*mock.Call
+}
+
+// Recv is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Recv() *MockStreamingNodeHandlerService_ProduceClient_Recv_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_Recv_Call{Call: _e.mock.On("Recv")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) Return(_a0 *streamingpb.ProduceResponse, _a1 error) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Recv_Call) RunAndReturn(run func() (*streamingpb.ProduceResponse, error)) *MockStreamingNodeHandlerService_ProduceClient_Recv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function with given fields: m
+func (_m *MockStreamingNodeHandlerService_ProduceClient) RecvMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m interface{}
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) RecvMsg(m interface{}) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(interface{}))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ProduceClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function with given fields: _a0
+func (_m *MockStreamingNodeHandlerService_ProduceClient) Send(_a0 *streamingpb.ProduceRequest) error {
+	ret := _m.Called(_a0)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(*streamingpb.ProduceRequest) error); ok {
+		r0 = rf(_a0)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type MockStreamingNodeHandlerService_ProduceClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - _a0 *streamingpb.ProduceRequest
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Send(_a0 interface{}) *MockStreamingNodeHandlerService_ProduceClient_Send_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_Send_Call{Call: _e.mock.On("Send", _a0)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) Run(run func(_a0 *streamingpb.ProduceRequest)) *MockStreamingNodeHandlerService_ProduceClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(*streamingpb.ProduceRequest))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_Send_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Send_Call) RunAndReturn(run func(*streamingpb.ProduceRequest) error) *MockStreamingNodeHandlerService_ProduceClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function with given fields: m
+func (_m *MockStreamingNodeHandlerService_ProduceClient) SendMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m interface{}
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) SendMsg(m interface{}) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) Run(run func(m interface{})) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(interface{}))
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) Return(_a0 error) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call) RunAndReturn(run func(interface{}) error) *MockStreamingNodeHandlerService_ProduceClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function with given fields:
+func (_m *MockStreamingNodeHandlerService_ProduceClient) Trailer() metadata.MD {
+	ret := _m.Called()
+
+	var r0 metadata.MD
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	return r0
+}
+
+// MockStreamingNodeHandlerService_ProduceClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type MockStreamingNodeHandlerService_ProduceClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *MockStreamingNodeHandlerService_ProduceClient_Expecter) Trailer() *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call {
+	return &MockStreamingNodeHandlerService_ProduceClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) Run(run func()) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) Return(_a0 metadata.MD) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *MockStreamingNodeHandlerService_ProduceClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockStreamingNodeHandlerService_ProduceClient creates a new instance of MockStreamingNodeHandlerService_ProduceClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockStreamingNodeHandlerService_ProduceClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockStreamingNodeHandlerService_ProduceClient {
+	mock := &MockStreamingNodeHandlerService_ProduceClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/proto/mock_streamingpb/mock_StreamingNodeManagerServiceClient.go

@@ -0,0 +1,250 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_streamingpb
+
+import (
+	context "context"
+
+	grpc "google.golang.org/grpc"
+
+	mock "github.com/stretchr/testify/mock"
+
+	streamingpb "github.com/milvus-io/milvus/internal/proto/streamingpb"
+)
+
+// MockStreamingNodeManagerServiceClient is an autogenerated mock type for the StreamingNodeManagerServiceClient type
+type MockStreamingNodeManagerServiceClient struct {
+	mock.Mock
+}
+
+type MockStreamingNodeManagerServiceClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockStreamingNodeManagerServiceClient) EXPECT() *MockStreamingNodeManagerServiceClient_Expecter {
+	return &MockStreamingNodeManagerServiceClient_Expecter{mock: &_m.Mock}
+}
+
+// Assign provides a mock function with given fields: ctx, in, opts
+func (_m *MockStreamingNodeManagerServiceClient) Assign(ctx context.Context, in *streamingpb.StreamingNodeManagerAssignRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *streamingpb.StreamingNodeManagerAssignResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerAssignResponse); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerAssignResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeManagerServiceClient_Assign_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Assign'
+type MockStreamingNodeManagerServiceClient_Assign_Call struct {
+	*mock.Call
+}
+
+// Assign is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *streamingpb.StreamingNodeManagerAssignRequest
+//   - opts ...grpc.CallOption
+func (_e *MockStreamingNodeManagerServiceClient_Expecter) Assign(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_Assign_Call {
+	return &MockStreamingNodeManagerServiceClient_Assign_Call{Call: _e.mock.On("Assign",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerAssignRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_Assign_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*streamingpb.StreamingNodeManagerAssignRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) Return(_a0 *streamingpb.StreamingNodeManagerAssignResponse, _a1 error) *MockStreamingNodeManagerServiceClient_Assign_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Assign_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerAssignRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error)) *MockStreamingNodeManagerServiceClient_Assign_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CollectStatus provides a mock function with given fields: ctx, in, opts
+func (_m *MockStreamingNodeManagerServiceClient) CollectStatus(ctx context.Context, in *streamingpb.StreamingNodeManagerCollectStatusRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *streamingpb.StreamingNodeManagerCollectStatusResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerCollectStatusResponse); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerCollectStatusResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeManagerServiceClient_CollectStatus_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CollectStatus'
+type MockStreamingNodeManagerServiceClient_CollectStatus_Call struct {
+	*mock.Call
+}
+
+// CollectStatus is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *streamingpb.StreamingNodeManagerCollectStatusRequest
+//   - opts ...grpc.CallOption
+func (_e *MockStreamingNodeManagerServiceClient_Expecter) CollectStatus(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_CollectStatus_Call {
+	return &MockStreamingNodeManagerServiceClient_CollectStatus_Call{Call: _e.mock.On("CollectStatus",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerCollectStatusRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_CollectStatus_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*streamingpb.StreamingNodeManagerCollectStatusRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) Return(_a0 *streamingpb.StreamingNodeManagerCollectStatusResponse, _a1 error) *MockStreamingNodeManagerServiceClient_CollectStatus_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_CollectStatus_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerCollectStatusRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error)) *MockStreamingNodeManagerServiceClient_CollectStatus_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Remove provides a mock function with given fields: ctx, in, opts
+func (_m *MockStreamingNodeManagerServiceClient) Remove(ctx context.Context, in *streamingpb.StreamingNodeManagerRemoveRequest, opts ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	var r0 *streamingpb.StreamingNodeManagerRemoveResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) *streamingpb.StreamingNodeManagerRemoveResponse); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*streamingpb.StreamingNodeManagerRemoveResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockStreamingNodeManagerServiceClient_Remove_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Remove'
+type MockStreamingNodeManagerServiceClient_Remove_Call struct {
+	*mock.Call
+}
+
+// Remove is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *streamingpb.StreamingNodeManagerRemoveRequest
+//   - opts ...grpc.CallOption
+func (_e *MockStreamingNodeManagerServiceClient_Expecter) Remove(ctx interface{}, in interface{}, opts ...interface{}) *MockStreamingNodeManagerServiceClient_Remove_Call {
+	return &MockStreamingNodeManagerServiceClient_Remove_Call{Call: _e.mock.On("Remove",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) Run(run func(ctx context.Context, in *streamingpb.StreamingNodeManagerRemoveRequest, opts ...grpc.CallOption)) *MockStreamingNodeManagerServiceClient_Remove_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*streamingpb.StreamingNodeManagerRemoveRequest), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) Return(_a0 *streamingpb.StreamingNodeManagerRemoveResponse, _a1 error) *MockStreamingNodeManagerServiceClient_Remove_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockStreamingNodeManagerServiceClient_Remove_Call) RunAndReturn(run func(context.Context, *streamingpb.StreamingNodeManagerRemoveRequest, ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error)) *MockStreamingNodeManagerServiceClient_Remove_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockStreamingNodeManagerServiceClient creates a new instance of MockStreamingNodeManagerServiceClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockStreamingNodeManagerServiceClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockStreamingNodeManagerServiceClient {
+	mock := &MockStreamingNodeManagerServiceClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/streamingnode/client/handler/mock_assignment/mock_Watcher.go

@@ -0,0 +1,158 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_assignment
+
+import (
+	context "context"
+
+	types "github.com/milvus-io/milvus/pkg/streaming/util/types"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// MockWatcher is an autogenerated mock type for the Watcher type
+type MockWatcher struct {
+	mock.Mock
+}
+
+type MockWatcher_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockWatcher) EXPECT() *MockWatcher_Expecter {
+	return &MockWatcher_Expecter{mock: &_m.Mock}
+}
+
+// Close provides a mock function with given fields:
+func (_m *MockWatcher) Close() {
+	_m.Called()
+}
+
+// MockWatcher_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
+type MockWatcher_Close_Call struct {
+	*mock.Call
+}
+
+// Close is a helper method to define mock.On call
+func (_e *MockWatcher_Expecter) Close() *MockWatcher_Close_Call {
+	return &MockWatcher_Close_Call{Call: _e.mock.On("Close")}
+}
+
+func (_c *MockWatcher_Close_Call) Run(run func()) *MockWatcher_Close_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockWatcher_Close_Call) Return() *MockWatcher_Close_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockWatcher_Close_Call) RunAndReturn(run func()) *MockWatcher_Close_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Get provides a mock function with given fields: ctx, channel
+func (_m *MockWatcher) Get(ctx context.Context, channel string) *types.PChannelInfoAssigned {
+	ret := _m.Called(ctx, channel)
+
+	var r0 *types.PChannelInfoAssigned
+	if rf, ok := ret.Get(0).(func(context.Context, string) *types.PChannelInfoAssigned); ok {
+		r0 = rf(ctx, channel)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*types.PChannelInfoAssigned)
+		}
+	}
+
+	return r0
+}
+
+// MockWatcher_Get_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Get'
+type MockWatcher_Get_Call struct {
+	*mock.Call
+}
+
+// Get is a helper method to define mock.On call
+//   - ctx context.Context
+//   - channel string
+func (_e *MockWatcher_Expecter) Get(ctx interface{}, channel interface{}) *MockWatcher_Get_Call {
+	return &MockWatcher_Get_Call{Call: _e.mock.On("Get", ctx, channel)}
+}
+
+func (_c *MockWatcher_Get_Call) Run(run func(ctx context.Context, channel string)) *MockWatcher_Get_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(string))
+	})
+	return _c
+}
+
+func (_c *MockWatcher_Get_Call) Return(_a0 *types.PChannelInfoAssigned) *MockWatcher_Get_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockWatcher_Get_Call) RunAndReturn(run func(context.Context, string) *types.PChannelInfoAssigned) *MockWatcher_Get_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Watch provides a mock function with given fields: ctx, channel, previous
+func (_m *MockWatcher) Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error {
+	ret := _m.Called(ctx, channel, previous)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, string, *types.PChannelInfoAssigned) error); ok {
+		r0 = rf(ctx, channel, previous)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockWatcher_Watch_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Watch'
+type MockWatcher_Watch_Call struct {
+	*mock.Call
+}
+
+// Watch is a helper method to define mock.On call
+//   - ctx context.Context
+//   - channel string
+//   - previous *types.PChannelInfoAssigned
+func (_e *MockWatcher_Expecter) Watch(ctx interface{}, channel interface{}, previous interface{}) *MockWatcher_Watch_Call {
+	return &MockWatcher_Watch_Call{Call: _e.mock.On("Watch", ctx, channel, previous)}
+}
+
+func (_c *MockWatcher_Watch_Call) Run(run func(ctx context.Context, channel string, previous *types.PChannelInfoAssigned)) *MockWatcher_Watch_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(string), args[2].(*types.PChannelInfoAssigned))
+	})
+	return _c
+}
+
+func (_c *MockWatcher_Watch_Call) Return(_a0 error) *MockWatcher_Watch_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockWatcher_Watch_Call) RunAndReturn(run func(context.Context, string, *types.PChannelInfoAssigned) error) *MockWatcher_Watch_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockWatcher creates a new instance of MockWatcher. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockWatcher(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockWatcher {
+	mock := &MockWatcher{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/streamingnode/client/handler/mock_consumer/mock_Consumer.go

@@ -0,0 +1,148 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_consumer
+
+import mock "github.com/stretchr/testify/mock"
+
+// MockConsumer is an autogenerated mock type for the Consumer type
+type MockConsumer struct {
+	mock.Mock
+}
+
+type MockConsumer_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockConsumer) EXPECT() *MockConsumer_Expecter {
+	return &MockConsumer_Expecter{mock: &_m.Mock}
+}
+
+// Close provides a mock function with given fields:
+func (_m *MockConsumer) Close() {
+	_m.Called()
+}
+
+// MockConsumer_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
+type MockConsumer_Close_Call struct {
+	*mock.Call
+}
+
+// Close is a helper method to define mock.On call
+func (_e *MockConsumer_Expecter) Close() *MockConsumer_Close_Call {
+	return &MockConsumer_Close_Call{Call: _e.mock.On("Close")}
+}
+
+func (_c *MockConsumer_Close_Call) Run(run func()) *MockConsumer_Close_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockConsumer_Close_Call) Return() *MockConsumer_Close_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockConsumer_Close_Call) RunAndReturn(run func()) *MockConsumer_Close_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Done provides a mock function with given fields:
+func (_m *MockConsumer) Done() <-chan struct{} {
+	ret := _m.Called()
+
+	var r0 <-chan struct{}
+	if rf, ok := ret.Get(0).(func() <-chan struct{}); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(<-chan struct{})
+		}
+	}
+
+	return r0
+}
+
+// MockConsumer_Done_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Done'
+type MockConsumer_Done_Call struct {
+	*mock.Call
+}
+
+// Done is a helper method to define mock.On call
+func (_e *MockConsumer_Expecter) Done() *MockConsumer_Done_Call {
+	return &MockConsumer_Done_Call{Call: _e.mock.On("Done")}
+}
+
+func (_c *MockConsumer_Done_Call) Run(run func()) *MockConsumer_Done_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockConsumer_Done_Call) Return(_a0 <-chan struct{}) *MockConsumer_Done_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockConsumer_Done_Call) RunAndReturn(run func() <-chan struct{}) *MockConsumer_Done_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Error provides a mock function with given fields:
+func (_m *MockConsumer) Error() error {
+	ret := _m.Called()
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// MockConsumer_Error_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Error'
+type MockConsumer_Error_Call struct {
+	*mock.Call
+}
+
+// Error is a helper method to define mock.On call
+func (_e *MockConsumer_Expecter) Error() *MockConsumer_Error_Call {
+	return &MockConsumer_Error_Call{Call: _e.mock.On("Error")}
+}
+
+func (_c *MockConsumer_Error_Call) Run(run func()) *MockConsumer_Error_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockConsumer_Error_Call) Return(_a0 error) *MockConsumer_Error_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockConsumer_Error_Call) RunAndReturn(run func() error) *MockConsumer_Error_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockConsumer creates a new instance of MockConsumer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockConsumer(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockConsumer {
+	mock := &MockConsumer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/streamingnode/client/handler/mock_producer/mock_Producer.go

@@ -0,0 +1,251 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_producer
+
+import (
+	context "context"
+
+	message "github.com/milvus-io/milvus/pkg/streaming/util/message"
+	mock "github.com/stretchr/testify/mock"
+
+	types "github.com/milvus-io/milvus/pkg/streaming/util/types"
+)
+
+// MockProducer is an autogenerated mock type for the Producer type
+type MockProducer struct {
+	mock.Mock
+}
+
+type MockProducer_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockProducer) EXPECT() *MockProducer_Expecter {
+	return &MockProducer_Expecter{mock: &_m.Mock}
+}
+
+// Assignment provides a mock function with given fields:
+func (_m *MockProducer) Assignment() types.PChannelInfoAssigned {
+	ret := _m.Called()
+
+	var r0 types.PChannelInfoAssigned
+	if rf, ok := ret.Get(0).(func() types.PChannelInfoAssigned); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(types.PChannelInfoAssigned)
+	}
+
+	return r0
+}
+
+// MockProducer_Assignment_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Assignment'
+type MockProducer_Assignment_Call struct {
+	*mock.Call
+}
+
+// Assignment is a helper method to define mock.On call
+func (_e *MockProducer_Expecter) Assignment() *MockProducer_Assignment_Call {
+	return &MockProducer_Assignment_Call{Call: _e.mock.On("Assignment")}
+}
+
+func (_c *MockProducer_Assignment_Call) Run(run func()) *MockProducer_Assignment_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockProducer_Assignment_Call) Return(_a0 types.PChannelInfoAssigned) *MockProducer_Assignment_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockProducer_Assignment_Call) RunAndReturn(run func() types.PChannelInfoAssigned) *MockProducer_Assignment_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Available provides a mock function with given fields:
+func (_m *MockProducer) Available() <-chan struct{} {
+	ret := _m.Called()
+
+	var r0 <-chan struct{}
+	if rf, ok := ret.Get(0).(func() <-chan struct{}); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(<-chan struct{})
+		}
+	}
+
+	return r0
+}
+
+// MockProducer_Available_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Available'
+type MockProducer_Available_Call struct {
+	*mock.Call
+}
+
+// Available is a helper method to define mock.On call
+func (_e *MockProducer_Expecter) Available() *MockProducer_Available_Call {
+	return &MockProducer_Available_Call{Call: _e.mock.On("Available")}
+}
+
+func (_c *MockProducer_Available_Call) Run(run func()) *MockProducer_Available_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockProducer_Available_Call) Return(_a0 <-chan struct{}) *MockProducer_Available_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockProducer_Available_Call) RunAndReturn(run func() <-chan struct{}) *MockProducer_Available_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Close provides a mock function with given fields:
+func (_m *MockProducer) Close() {
+	_m.Called()
+}
+
+// MockProducer_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
+type MockProducer_Close_Call struct {
+	*mock.Call
+}
+
+// Close is a helper method to define mock.On call
+func (_e *MockProducer_Expecter) Close() *MockProducer_Close_Call {
+	return &MockProducer_Close_Call{Call: _e.mock.On("Close")}
+}
+
+func (_c *MockProducer_Close_Call) Run(run func()) *MockProducer_Close_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockProducer_Close_Call) Return() *MockProducer_Close_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockProducer_Close_Call) RunAndReturn(run func()) *MockProducer_Close_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// IsAvailable provides a mock function with given fields:
+func (_m *MockProducer) IsAvailable() bool {
+	ret := _m.Called()
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// MockProducer_IsAvailable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsAvailable'
+type MockProducer_IsAvailable_Call struct {
+	*mock.Call
+}
+
+// IsAvailable is a helper method to define mock.On call
+func (_e *MockProducer_Expecter) IsAvailable() *MockProducer_IsAvailable_Call {
+	return &MockProducer_IsAvailable_Call{Call: _e.mock.On("IsAvailable")}
+}
+
+func (_c *MockProducer_IsAvailable_Call) Run(run func()) *MockProducer_IsAvailable_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockProducer_IsAvailable_Call) Return(_a0 bool) *MockProducer_IsAvailable_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockProducer_IsAvailable_Call) RunAndReturn(run func() bool) *MockProducer_IsAvailable_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Produce provides a mock function with given fields: ctx, msg
+func (_m *MockProducer) Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) {
+	ret := _m.Called(ctx, msg)
+
+	var r0 message.MessageID
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, message.MutableMessage) (message.MessageID, error)); ok {
+		return rf(ctx, msg)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, message.MutableMessage) message.MessageID); ok {
+		r0 = rf(ctx, msg)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(message.MessageID)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, message.MutableMessage) error); ok {
+		r1 = rf(ctx, msg)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockProducer_Produce_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Produce'
+type MockProducer_Produce_Call struct {
+	*mock.Call
+}
+
+// Produce is a helper method to define mock.On call
+//   - ctx context.Context
+//   - msg message.MutableMessage
+func (_e *MockProducer_Expecter) Produce(ctx interface{}, msg interface{}) *MockProducer_Produce_Call {
+	return &MockProducer_Produce_Call{Call: _e.mock.On("Produce", ctx, msg)}
+}
+
+func (_c *MockProducer_Produce_Call) Run(run func(ctx context.Context, msg message.MutableMessage)) *MockProducer_Produce_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(message.MutableMessage))
+	})
+	return _c
+}
+
+func (_c *MockProducer_Produce_Call) Return(_a0 message.MessageID, _a1 error) *MockProducer_Produce_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockProducer_Produce_Call) RunAndReturn(run func(context.Context, message.MutableMessage) (message.MessageID, error)) *MockProducer_Produce_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockProducer creates a new instance of MockProducer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockProducer(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockProducer {
+	mock := &MockProducer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/streamingnode/server/mock_walmanager/mock_Manager.go

@@ -109,17 +109,17 @@ func (_c *MockManager_GetAllAvailableChannels_Call) RunAndReturn(run func() ([]t
 	return _c
 }
 
-// GetAvailableWAL provides a mock function with given fields: channel
+// GetAvailableWAL provides a mock function with given fields: _a0
-func (_m *MockManager) GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, error) {
+func (_m *MockManager) GetAvailableWAL(_a0 types.PChannelInfo) (wal.WAL, error) {
-	ret := _m.Called(channel)
+	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
+//   - _a0 types.PChannelInfo
-func (_e *MockManager_Expecter) GetAvailableWAL(channel interface{}) *MockManager_GetAvailableWAL_Call {
+func (_e *MockManager_Expecter) GetAvailableWAL(_a0 interface{}) *MockManager_GetAvailableWAL_Call {
-	return &MockManager_GetAvailableWAL_Call{Call: _e.mock.On("GetAvailableWAL", channel)}
+	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))
 	})

internal/mocks/util/streamingutil/service/mock_lazygrpc/mock_Service.go

@@ -0,0 +1,176 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_lazygrpc
+
+import (
+	context "context"
+
+	grpc "google.golang.org/grpc"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// MockService is an autogenerated mock type for the Service type
+type MockService[T interface{}] struct {
+	mock.Mock
+}
+
+type MockService_Expecter[T interface{}] struct {
+	mock *mock.Mock
+}
+
+func (_m *MockService[T]) EXPECT() *MockService_Expecter[T] {
+	return &MockService_Expecter[T]{mock: &_m.Mock}
+}
+
+// Close provides a mock function with given fields:
+func (_m *MockService[T]) Close() {
+	_m.Called()
+}
+
+// MockService_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
+type MockService_Close_Call[T interface{}] struct {
+	*mock.Call
+}
+
+// Close is a helper method to define mock.On call
+func (_e *MockService_Expecter[T]) Close() *MockService_Close_Call[T] {
+	return &MockService_Close_Call[T]{Call: _e.mock.On("Close")}
+}
+
+func (_c *MockService_Close_Call[T]) Run(run func()) *MockService_Close_Call[T] {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockService_Close_Call[T]) Return() *MockService_Close_Call[T] {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockService_Close_Call[T]) RunAndReturn(run func()) *MockService_Close_Call[T] {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetConn provides a mock function with given fields: ctx
+func (_m *MockService[T]) GetConn(ctx context.Context) (*grpc.ClientConn, error) {
+	ret := _m.Called(ctx)
+
+	var r0 *grpc.ClientConn
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) (*grpc.ClientConn, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) *grpc.ClientConn); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*grpc.ClientConn)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockService_GetConn_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetConn'
+type MockService_GetConn_Call[T interface{}] struct {
+	*mock.Call
+}
+
+// GetConn is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *MockService_Expecter[T]) GetConn(ctx interface{}) *MockService_GetConn_Call[T] {
+	return &MockService_GetConn_Call[T]{Call: _e.mock.On("GetConn", ctx)}
+}
+
+func (_c *MockService_GetConn_Call[T]) Run(run func(ctx context.Context)) *MockService_GetConn_Call[T] {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *MockService_GetConn_Call[T]) Return(_a0 *grpc.ClientConn, _a1 error) *MockService_GetConn_Call[T] {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockService_GetConn_Call[T]) RunAndReturn(run func(context.Context) (*grpc.ClientConn, error)) *MockService_GetConn_Call[T] {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetService provides a mock function with given fields: ctx
+func (_m *MockService[T]) GetService(ctx context.Context) (T, error) {
+	ret := _m.Called(ctx)
+
+	var r0 T
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) (T, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) T); ok {
+		r0 = rf(ctx)
+	} else {
+		r0 = ret.Get(0).(T)
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockService_GetService_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetService'
+type MockService_GetService_Call[T interface{}] struct {
+	*mock.Call
+}
+
+// GetService is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *MockService_Expecter[T]) GetService(ctx interface{}) *MockService_GetService_Call[T] {
+	return &MockService_GetService_Call[T]{Call: _e.mock.On("GetService", ctx)}
+}
+
+func (_c *MockService_GetService_Call[T]) Run(run func(ctx context.Context)) *MockService_GetService_Call[T] {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *MockService_GetService_Call[T]) Return(_a0 T, _a1 error) *MockService_GetService_Call[T] {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockService_GetService_Call[T]) RunAndReturn(run func(context.Context) (T, error)) *MockService_GetService_Call[T] {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockService creates a new instance of MockService. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockService[T interface{}](t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockService[T] {
+	mock := &MockService[T]{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/mocks/util/streamingutil/service/mock_resolver/mock_Builder.go

@@ -0,0 +1,209 @@
+// Code generated by mockery v2.32.4. DO NOT EDIT.
+
+package mock_resolver
+
+import (
+	mock "github.com/stretchr/testify/mock"
+	resolver "google.golang.org/grpc/resolver"
+
+	serviceresolver "github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
+)
+
+// MockBuilder is an autogenerated mock type for the Builder type
+type MockBuilder struct {
+	mock.Mock
+}
+
+type MockBuilder_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockBuilder) EXPECT() *MockBuilder_Expecter {
+	return &MockBuilder_Expecter{mock: &_m.Mock}
+}
+
+// Build provides a mock function with given fields: target, cc, opts
+func (_m *MockBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
+	ret := _m.Called(target, cc, opts)
+
+	var r0 resolver.Resolver
+	var r1 error
+	if rf, ok := ret.Get(0).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) (resolver.Resolver, error)); ok {
+		return rf(target, cc, opts)
+	}
+	if rf, ok := ret.Get(0).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) resolver.Resolver); ok {
+		r0 = rf(target, cc, opts)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(resolver.Resolver)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) error); ok {
+		r1 = rf(target, cc, opts)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MockBuilder_Build_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Build'
+type MockBuilder_Build_Call struct {
+	*mock.Call
+}
+
+// Build is a helper method to define mock.On call
+//   - target resolver.Target
+//   - cc resolver.ClientConn
+//   - opts resolver.BuildOptions
+func (_e *MockBuilder_Expecter) Build(target interface{}, cc interface{}, opts interface{}) *MockBuilder_Build_Call {
+	return &MockBuilder_Build_Call{Call: _e.mock.On("Build", target, cc, opts)}
+}
+
+func (_c *MockBuilder_Build_Call) Run(run func(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions)) *MockBuilder_Build_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(resolver.Target), args[1].(resolver.ClientConn), args[2].(resolver.BuildOptions))
+	})
+	return _c
+}
+
+func (_c *MockBuilder_Build_Call) Return(_a0 resolver.Resolver, _a1 error) *MockBuilder_Build_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MockBuilder_Build_Call) RunAndReturn(run func(resolver.Target, resolver.ClientConn, resolver.BuildOptions) (resolver.Resolver, error)) *MockBuilder_Build_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Close provides a mock function with given fields:
+func (_m *MockBuilder) Close() {
+	_m.Called()
+}
+
+// MockBuilder_Close_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Close'
+type MockBuilder_Close_Call struct {
+	*mock.Call
+}
+
+// Close is a helper method to define mock.On call
+func (_e *MockBuilder_Expecter) Close() *MockBuilder_Close_Call {
+	return &MockBuilder_Close_Call{Call: _e.mock.On("Close")}
+}
+
+func (_c *MockBuilder_Close_Call) Run(run func()) *MockBuilder_Close_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockBuilder_Close_Call) Return() *MockBuilder_Close_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockBuilder_Close_Call) RunAndReturn(run func()) *MockBuilder_Close_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Resolver provides a mock function with given fields:
+func (_m *MockBuilder) Resolver() serviceresolver.Resolver {
+	ret := _m.Called()
+
+	var r0 serviceresolver.Resolver
+	if rf, ok := ret.Get(0).(func() serviceresolver.Resolver); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(serviceresolver.Resolver)
+		}
+	}
+
+	return r0
+}
+
+// MockBuilder_Resolver_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Resolver'
+type MockBuilder_Resolver_Call struct {
+	*mock.Call
+}
+
+// Resolver is a helper method to define mock.On call
+func (_e *MockBuilder_Expecter) Resolver() *MockBuilder_Resolver_Call {
+	return &MockBuilder_Resolver_Call{Call: _e.mock.On("Resolver")}
+}
+
+func (_c *MockBuilder_Resolver_Call) Run(run func()) *MockBuilder_Resolver_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockBuilder_Resolver_Call) Return(_a0 serviceresolver.Resolver) *MockBuilder_Resolver_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockBuilder_Resolver_Call) RunAndReturn(run func() serviceresolver.Resolver) *MockBuilder_Resolver_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Scheme provides a mock function with given fields:
+func (_m *MockBuilder) Scheme() string {
+	ret := _m.Called()
+
+	var r0 string
+	if rf, ok := ret.Get(0).(func() string); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+
+	return r0
+}
+
+// MockBuilder_Scheme_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Scheme'
+type MockBuilder_Scheme_Call struct {
+	*mock.Call
+}
+
+// Scheme is a helper method to define mock.On call
+func (_e *MockBuilder_Expecter) Scheme() *MockBuilder_Scheme_Call {
+	return &MockBuilder_Scheme_Call{Call: _e.mock.On("Scheme")}
+}
+
+func (_c *MockBuilder_Scheme_Call) Run(run func()) *MockBuilder_Scheme_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockBuilder_Scheme_Call) Return(_a0 string) *MockBuilder_Scheme_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockBuilder_Scheme_Call) RunAndReturn(run func() string) *MockBuilder_Scheme_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockBuilder creates a new instance of MockBuilder. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockBuilder(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockBuilder {
+	mock := &MockBuilder{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/streamingnode/client/handler/assignment/watcher.go

@@ -0,0 +1,22 @@
+package assignment
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+)
+
+var _ Watcher = (*watcherImpl)(nil)
+
+// Watcher is the interface for the channel assignment.
+type Watcher interface {
+	// Get gets the channel assignment.
+	Get(ctx context.Context, channel string) *types.PChannelInfoAssigned
+
+	// Watch watches the channel assignment.
+	// Block until new term is coming.
+	Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error
+
+	// Close stop the watcher.
+	Close()
+}

internal/streamingnode/client/handler/assignment/watcher_impl.go

@@ -0,0 +1,108 @@
+package assignment
+
+import (
+	"context"
+	"sync"
+
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
+	"github.com/milvus-io/milvus/pkg/log"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/syncutil"
+)
+
+func NewWatcher(r resolver.Resolver) Watcher {
+	ctx, cancel := context.WithCancel(context.Background())
+	w := &watcherImpl{
+		ctx:         ctx,
+		cancel:      cancel,
+		r:           r,
+		cond:        *syncutil.NewContextCond(&sync.Mutex{}),
+		assignments: make(map[string]types.PChannelInfoAssigned),
+	}
+	go w.execute()
+	return w
+}
+
+// watcherImpl is the implementation of the assignment watcher.
+type watcherImpl struct {
+	ctx         context.Context
+	cancel      context.CancelFunc
+	r           resolver.Resolver
+	cond        syncutil.ContextCond
+	assignments map[string]types.PChannelInfoAssigned // map pchannel to node.
+}
+
+// execute starts the watcher.
+func (w *watcherImpl) execute() {
+	log.Info("assignment watcher start")
+	var err error
+	defer func() {
+		// error can be ignored here, so use info level log here.
+		log.Info("assignment watcher close", zap.Error(err))
+	}()
+
+	// error can be ignored here, error is always cancel by watcher's close as expected.
+	// otherwise, the resolver's close is unexpected.
+	err = w.r.Watch(w.ctx, func(state discoverer.VersionedState) error {
+		w.updateAssignment(state)
+		return nil
+	})
+}
+
+// updateAssignment updates the assignment.
+func (w *watcherImpl) updateAssignment(state discoverer.VersionedState) {
+	newAssignments := make(map[string]types.PChannelInfoAssigned)
+	for _, assignments := range state.ChannelAssignmentInfo() {
+		for _, pChannelInfo := range assignments.Channels {
+			newAssignments[pChannelInfo.Name] = types.PChannelInfoAssigned{
+				Channel: pChannelInfo,
+				Node:    assignments.NodeInfo,
+			}
+		}
+	}
+	w.cond.LockAndBroadcast()
+	w.assignments = newAssignments
+	w.cond.L.Unlock()
+}
+
+// Get gets the current pchannel assignment.
+func (w *watcherImpl) Get(ctx context.Context, channel string) *types.PChannelInfoAssigned {
+	w.cond.L.Lock()
+	defer w.cond.L.Unlock()
+
+	if info, ok := w.assignments[channel]; ok {
+		return &info
+	}
+	return nil
+}
+
+// Watch watches the channel assignment.
+func (w *watcherImpl) Watch(ctx context.Context, channel string, previous *types.PChannelInfoAssigned) error {
+	w.cond.L.Lock()
+
+	term := types.InitialTerm
+	if previous != nil {
+		term = previous.Channel.Term
+	}
+
+	for {
+		if info, ok := w.assignments[channel]; ok {
+			if info.Channel.Term > term {
+				break
+			}
+		}
+		if err := w.cond.Wait(ctx); err != nil {
+			return err
+		}
+	}
+	w.cond.L.Unlock()
+	return nil
+}
+
+// Close closes the watcher.
+func (w *watcherImpl) Close() {
+	w.cancel()
+}

internal/streamingnode/client/handler/assignment/watcher_test.go

@@ -0,0 +1,91 @@
+package assignment
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"google.golang.org/grpc/resolver"
+
+	"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/attributes"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+func TestWatcher(t *testing.T) {
+	r := mock_resolver.NewMockResolver(t)
+
+	ch := make(chan discoverer.VersionedState)
+	r.EXPECT().Watch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(s discoverer.VersionedState) error) error {
+		for {
+			select {
+			case v, ok := <-ch:
+				if !ok {
+					return nil
+				}
+				f(v)
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+		}
+	})
+	w := NewWatcher(r)
+	defer w.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	a := w.Get(ctx, "test_pchannel")
+	assert.Nil(t, a)
+	err := w.Watch(ctx, "test_pchannel", nil)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+
+	ch <- discoverer.VersionedState{
+		Version: typeutil.VersionInt64(1),
+		State: resolver.State{
+			Addresses: []resolver.Address{
+				{
+					Addr: "test_addr",
+					BalancerAttributes: attributes.WithChannelAssignmentInfo(
+						new(attributes.Attributes),
+						&types.StreamingNodeAssignment{
+							NodeInfo: types.StreamingNodeInfo{
+								ServerID: 1,
+								Address:  "test_addr",
+							},
+							Channels: map[string]types.PChannelInfo{
+								"test_pchannel": {
+									Name: "test_pchannel",
+									Term: 1,
+								},
+								"test_pchannel_2": {
+									Name: "test_pchannel_2",
+									Term: 2,
+								},
+							},
+						},
+					),
+				},
+			},
+		},
+	}
+	err = w.Watch(context.Background(), "test_pchannel", nil)
+	assert.NoError(t, err)
+	a = w.Get(ctx, "test_pchannel")
+	assert.NotNil(t, a)
+	assert.Equal(t, int64(1), a.Channel.Term)
+
+	err = w.Watch(context.Background(), "test_pchannel_2", nil)
+	assert.NoError(t, err)
+	a = w.Get(ctx, "test_pchannel_2")
+	assert.NotNil(t, a)
+	assert.Equal(t, int64(2), a.Channel.Term)
+
+	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+	err = w.Watch(ctx, "test_pchannel", a)
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}

internal/streamingnode/client/handler/consumer/consumer.go

@@ -0,0 +1,18 @@
+package consumer
+
+var _ Consumer = (*consumerImpl)(nil)
+
+// Consumer is the interface that wraps the basic consume method on grpc stream.
+// Consumer is work on a single stream on grpc,
+// so Consumer cannot recover from failure because of the stream is broken.
+type Consumer interface {
+	// Error returns the error of scanner failed.
+	// Will block until scanner is closed or Chan is dry out.
+	Error() error
+
+	// Done returns a channel which will be closed when scanner is finished or closed.
+	Done() <-chan struct{}
+
+	// Close the consumer, release the underlying resources.
+	Close()
+}

internal/streamingnode/client/handler/consumer/consumer_impl.go

@@ -0,0 +1,181 @@
+package consumer
+
+import (
+	"context"
+	"io"
+
+	"github.com/cockroachdb/errors"
+	"go.uber.org/zap"
+	"google.golang.org/grpc"
+
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/typeconverter"
+	"github.com/milvus-io/milvus/pkg/log"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/syncutil"
+)
+
+// ConsumerOptions is the options for creating a consumer.
+type ConsumerOptions struct {
+	// The cosume target
+	Assignment *types.PChannelInfoAssigned
+
+	// DeliverPolicy is the deliver policy of the consumer.
+	DeliverPolicy options.DeliverPolicy
+
+	// DeliverFilters is the deliver filters of the consumer.
+	DeliverFilters []options.DeliverFilter
+
+	// Handler is the message handler used to handle message after recv from consumer.
+	MessageHandler message.Handler
+}
+
+// CreateConsumer creates a new consumer client.
+func CreateConsumer(
+	ctx context.Context,
+	opts *ConsumerOptions,
+	handlerClient streamingpb.StreamingNodeHandlerServiceClient,
+) (Consumer, error) {
+	ctx, err := createConsumeRequest(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO: configurable or auto adjust grpc.MaxCallRecvMsgSize
+	streamClient, err := handlerClient.Consume(ctx, grpc.MaxCallRecvMsgSize(8388608))
+	if err != nil {
+		return nil, err
+	}
+
+	// Recv the first response from server.
+	// It must be a create response.
+	resp, err := streamClient.Recv()
+	if err != nil {
+		return nil, err
+	}
+	createResp := resp.GetCreate()
+	if createResp == nil {
+		return nil, status.NewInvalidRequestSeq("first message arrive must be create response")
+	}
+	cli := &consumerImpl{
+		walName:          createResp.GetWalName(),
+		assignment:       *opts.Assignment,
+		grpcStreamClient: streamClient,
+		handlerClient:    handlerClient,
+		logger: log.With(
+			zap.String("walName", createResp.GetWalName()),
+			zap.String("pchannel", opts.Assignment.Channel.Name),
+			zap.Int64("term", opts.Assignment.Channel.Term),
+			zap.Int64("streamingNodeID", opts.Assignment.Node.ServerID)),
+		msgHandler: opts.MessageHandler,
+		finishErr:  syncutil.NewFuture[error](),
+	}
+	go cli.execute()
+	return cli, nil
+}
+
+// createConsumeRequest creates the consume request.
+func createConsumeRequest(ctx context.Context, opts *ConsumerOptions) (context.Context, error) {
+	// select server to consume.
+	ctx = contextutil.WithPickServerID(ctx, opts.Assignment.Node.ServerID)
+	// create the consumer request.
+	deliverPolicy, err := typeconverter.NewProtoFromDeliverPolicy(opts.DeliverPolicy)
+	if err != nil {
+		return nil, errors.Wrap(err, "at convert deliver policy")
+	}
+	deliverFilters, err := typeconverter.NewProtosFromDeliverFilters(opts.DeliverFilters)
+	if err != nil {
+		return nil, errors.Wrap(err, "at convert deliver filters")
+	}
+	return contextutil.WithCreateConsumer(ctx, &streamingpb.CreateConsumerRequest{
+		Pchannel:       typeconverter.NewProtoFromPChannelInfo(opts.Assignment.Channel),
+		DeliverPolicy:  deliverPolicy,
+		DeliverFilters: deliverFilters,
+	}), nil
+}
+
+type consumerImpl struct {
+	walName          string
+	assignment       types.PChannelInfoAssigned
+	grpcStreamClient streamingpb.StreamingNodeHandlerService_ConsumeClient
+	handlerClient    streamingpb.StreamingNodeHandlerServiceClient
+	logger           *log.MLogger
+	msgHandler       message.Handler
+	finishErr        *syncutil.Future[error]
+}
+
+// Close close the consumer client.
+func (c *consumerImpl) Close() {
+	// Send the close request to server.
+	if err := c.grpcStreamClient.Send(&streamingpb.ConsumeRequest{
+		Request: &streamingpb.ConsumeRequest_Close{},
+	}); err != nil {
+		c.logger.Warn("send close request failed", zap.Error(err))
+	}
+	// close the grpc client stream.
+	if err := c.grpcStreamClient.CloseSend(); err != nil {
+		c.logger.Warn("close grpc stream failed", zap.Error(err))
+	}
+	<-c.finishErr.Done()
+}
+
+// Error returns the error of the consumer client.
+func (c *consumerImpl) Error() error {
+	return c.finishErr.Get()
+}
+
+// Done returns a channel that closes when the consumer client is closed.
+func (c *consumerImpl) Done() <-chan struct{} {
+	return c.finishErr.Done()
+}
+
+// execute starts the recv loop.
+func (c *consumerImpl) execute() {
+	c.recvLoop()
+}
+
+// recvLoop is the recv arm of the grpc stream.
+// Throughput of the grpc framework should be ok to use single stream to receive message.
+// Once throughput is not enough, look at https://grpc.io/docs/guides/performance/ to find the solution.
+func (c *consumerImpl) recvLoop() (err error) {
+	defer func() {
+		if err != nil {
+			c.logger.Warn("recv arm of stream closed with unexpected error", zap.Error(err))
+		} else {
+			c.logger.Info("recv arm of stream closed")
+		}
+		c.finishErr.Set(err)
+		c.msgHandler.Close()
+	}()
+
+	for {
+		resp, err := c.grpcStreamClient.Recv()
+		if errors.Is(err, io.EOF) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		switch resp := resp.Response.(type) {
+		case *streamingpb.ConsumeResponse_Consume:
+			msgID, err := message.UnmarshalMessageID(c.walName, resp.Consume.GetId().GetId())
+			if err != nil {
+				return err
+			}
+			c.msgHandler.Handle(message.NewImmutableMesasge(
+				msgID,
+				resp.Consume.GetMessage().GetPayload(),
+				resp.Consume.GetMessage().GetProperties(),
+			))
+		case *streamingpb.ConsumeResponse_Close:
+			// Should receive io.EOF after that.
+			// Do nothing at current implementation.
+		default:
+			c.logger.Warn("unknown response type", zap.Any("response", resp))
+		}
+	}
+}

internal/streamingnode/client/handler/consumer/consumer_test.go

@@ -0,0 +1,85 @@
+package consumer
+
+import (
+	"context"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/streaming/walimpls/impls/walimplstest"
+)
+
+func TestConsumer(t *testing.T) {
+	c := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t)
+	cc := mock_streamingpb.NewMockStreamingNodeHandlerService_ConsumeClient(t)
+	recvCh := make(chan *streamingpb.ConsumeResponse, 10)
+	cc.EXPECT().Recv().RunAndReturn(func() (*streamingpb.ConsumeResponse, error) {
+		msg, ok := <-recvCh
+		if !ok {
+			return nil, io.EOF
+		}
+		return msg, nil
+	})
+	sendCh := make(chan *streamingpb.ConsumeRequest, 10)
+	cc.EXPECT().Send(mock.Anything).RunAndReturn(func(cr *streamingpb.ConsumeRequest) error {
+		sendCh <- cr
+		return nil
+	})
+	c.EXPECT().Consume(mock.Anything, mock.Anything).Return(cc, nil)
+	cc.EXPECT().CloseSend().RunAndReturn(func() error {
+		recvCh <- &streamingpb.ConsumeResponse{Response: &streamingpb.ConsumeResponse_Close{}}
+		close(recvCh)
+		return nil
+	})
+
+	ctx := context.Background()
+	resultCh := make(message.ChanMessageHandler, 1)
+	opts := &ConsumerOptions{
+		Assignment: &types.PChannelInfoAssigned{
+			Channel: types.PChannelInfo{Name: "test", Term: 1},
+			Node:    types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
+		},
+		DeliverPolicy: options.DeliverPolicyAll(),
+		DeliverFilters: []options.DeliverFilter{
+			options.DeliverFilterVChannel("test-1"),
+			options.DeliverFilterTimeTickGT(100),
+		},
+		MessageHandler: resultCh,
+	}
+
+	recvCh <- &streamingpb.ConsumeResponse{
+		Response: &streamingpb.ConsumeResponse_Create{
+			Create: &streamingpb.CreateConsumerResponse{
+				WalName: "test",
+			},
+		},
+	}
+	recvCh <- &streamingpb.ConsumeResponse{
+		Response: &streamingpb.ConsumeResponse_Consume{
+			Consume: &streamingpb.ConsumeMessageReponse{
+				Id: &streamingpb.MessageID{
+					Id: walimplstest.NewTestMessageID(1).Marshal(),
+				},
+				Message: &streamingpb.Message{
+					Payload:    []byte{},
+					Properties: make(map[string]string),
+				},
+			},
+		},
+	}
+	consumer, err := CreateConsumer(ctx, opts, c)
+	assert.NoError(t, err)
+	assert.NotNil(t, consumer)
+	consumer.Close()
+	msg := <-resultCh
+	assert.True(t, msg.MessageID().EQ(walimplstest.NewTestMessageID(1)))
+	<-consumer.Done()
+	assert.NoError(t, consumer.Error())
+}

internal/streamingnode/client/handler/handler_client.go

@@ -0,0 +1,151 @@
+package handler
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/assignment"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker"
+	streamingserviceinterceptor "github.com/milvus-io/milvus/internal/util/streamingutil/service/interceptor"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/tracer"
+	"github.com/milvus-io/milvus/pkg/util/interceptor"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+	"github.com/milvus-io/milvus/pkg/util/lock"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+var _ HandlerClient = (*handlerClientImpl)(nil)
+
+type (
+	Producer = producer.Producer
+	Consumer = consumer.Consumer
+)
+
+// ProducerOptions is the options for creating a producer.
+type ProducerOptions struct {
+	// PChannel is the pchannel of the producer.
+	PChannel string
+}
+
+// ConsumerOptions is the options for creating a consumer.
+type ConsumerOptions struct {
+	// PChannel is the pchannel of the consumer.
+	PChannel string
+
+	// DeliverPolicy is the deliver policy of the consumer.
+	DeliverPolicy options.DeliverPolicy
+
+	// DeliverFilters is the deliver filters of the consumer.
+	DeliverFilters []options.DeliverFilter
+
+	// Handler is the message handler used to handle message after recv from consumer.
+	MessageHandler message.Handler
+}
+
+// HandlerClient is the interface that wraps streamingpb.StreamingNodeHandlerServiceClient.
+// HandlerClient wraps the PChannel Assignment Service Discovery.
+// Provides the ability to create pchannel-level producer and consumer.
+type HandlerClient interface {
+	// CreateProducer creates a producer.
+	// Producer is a stream client without keep alive promise.
+	// It will be available until context canceled, active close, streaming error or remote server wal closing.
+	// Because of there's no more ProducerOptions except PChannel, so a producer of same PChannel is shared by reference count.
+	CreateProducer(ctx context.Context, opts *ProducerOptions) (Producer, error)
+
+	// CreateConsumer creates a consumer.
+	// Consumer is a stream client without keep alive promise.
+	// It will be available until context canceled, active close, streaming error or remote server wal closing.
+	// A consumer will not share stream connection with other consumers.
+	CreateConsumer(ctx context.Context, opts *ConsumerOptions) (Consumer, error)
+
+	// Close closes the handler client.
+	// It will only stop the underlying service discovery, but don't stop the producer and consumer created by it.
+	// So please close Producer and Consumer created by it before close the handler client.
+	Close()
+}
+
+// NewHandlerClient creates a new handler client.
+func NewHandlerClient(w types.AssignmentDiscoverWatcher) HandlerClient {
+	rb := resolver.NewChannelAssignmentBuilder(w)
+	dialTimeout := paramtable.Get().StreamingNodeGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond)
+	dialOptions := getDialOptions(rb)
+	conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) {
+		ctx, cancel := context.WithTimeout(ctx, dialTimeout)
+		defer cancel()
+		return grpc.DialContext(
+			ctx,
+			resolver.ChannelAssignmentResolverScheme+":///"+typeutil.StreamingNodeRole,
+			dialOptions..., // TODO: we should use dynamic service config in future by add it to resolver.
+		)
+	})
+	watcher := assignment.NewWatcher(rb.Resolver())
+	return &handlerClientImpl{
+		lifetime:              lifetime.NewLifetime(lifetime.Working),
+		service:               lazygrpc.WithServiceCreator(conn, streamingpb.NewStreamingNodeHandlerServiceClient),
+		rb:                    rb,
+		watcher:               watcher,
+		rebalanceTrigger:      w,
+		sharedProducers:       make(map[string]*typeutil.WeakReference[Producer]),
+		sharedProducerKeyLock: lock.NewKeyLock[string](),
+		newProducer:           producer.CreateProducer,
+		newConsumer:           consumer.CreateConsumer,
+	}
+}
+
+// getDialOptions returns grpc dial options.
+func getDialOptions(rb resolver.Builder) []grpc.DialOption {
+	cfg := &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
+}

internal/streamingnode/client/handler/handler_client_impl.go

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

internal/streamingnode/client/handler/handler_client_test.go

@@ -0,0 +1,145 @@
+package handler
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb"
+	"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_assignment"
+	"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_consumer"
+	"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_producer"
+	"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_lazygrpc"
+	"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
+	"github.com/milvus-io/milvus/pkg/mocks/streaming/util/mock_types"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+	"github.com/milvus-io/milvus/pkg/util/lock"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+func TestHandlerClient(t *testing.T) {
+	assignment := &types.PChannelInfoAssigned{
+		Channel: types.PChannelInfo{Name: "pchannel", Term: 1},
+		Node:    types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
+	}
+
+	service := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeHandlerServiceClient](t)
+	handlerServiceClient := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t)
+	service.EXPECT().GetService(mock.Anything).Return(handlerServiceClient, nil)
+	rb := mock_resolver.NewMockBuilder(t)
+	rb.EXPECT().Close().Run(func() {})
+	w := mock_assignment.NewMockWatcher(t)
+	w.EXPECT().Close().Run(func() {})
+
+	p := mock_producer.NewMockProducer(t)
+	p.EXPECT().IsAvailable().Return(true)
+	p.EXPECT().Assignment().Return(*assignment)
+	p.EXPECT().Close().Run(func() {})
+	c := mock_consumer.NewMockConsumer(t)
+	c.EXPECT().Close().Run(func() {})
+
+	rebalanceTrigger := mock_types.NewMockAssignmentRebalanceTrigger(t)
+	rebalanceTrigger.EXPECT().ReportAssignmentError(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+	pK := 0
+	handler := &handlerClientImpl{
+		lifetime:              lifetime.NewLifetime(lifetime.Working),
+		service:               service,
+		rb:                    rb,
+		watcher:               w,
+		rebalanceTrigger:      rebalanceTrigger,
+		sharedProducers:       make(map[string]*typeutil.WeakReference[producer.Producer]),
+		sharedProducerKeyLock: lock.NewKeyLock[string](),
+		newProducer: func(ctx context.Context, opts *producer.ProducerOptions, handler streamingpb.StreamingNodeHandlerServiceClient) (Producer, error) {
+			if pK == 0 {
+				pK++
+				return nil, status.NewUnmatchedChannelTerm("pchannel", 1, 2)
+			}
+			return p, nil
+		},
+		newConsumer: func(ctx context.Context, opts *consumer.ConsumerOptions, handlerClient streamingpb.StreamingNodeHandlerServiceClient) (Consumer, error) {
+			return c, nil
+		},
+	}
+	ctx := context.Background()
+
+	k := 0
+	w.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, s string) *types.PChannelInfoAssigned {
+		if k == 0 {
+			k++
+			return nil
+		}
+		return assignment
+	})
+	w.EXPECT().Watch(mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+	producer, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"})
+	assert.NoError(t, err)
+	assert.NotNil(t, producer)
+	producer2, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"})
+	assert.NoError(t, err)
+	assert.NotNil(t, producer)
+	p.EXPECT().IsAvailable().Unset()
+	p.EXPECT().IsAvailable().Return(false)
+	producer3, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"})
+	assert.NoError(t, err)
+	assert.NotNil(t, producer3)
+	producer.Close()
+	producer2.Close()
+	producer3.Close()
+
+	producer4, err := handler.CreateProducer(ctx, &ProducerOptions{PChannel: "pchannel"})
+	assert.NoError(t, err)
+	assert.NotNil(t, producer4)
+	producer4.Close()
+
+	consumer, err := handler.CreateConsumer(ctx, &ConsumerOptions{
+		PChannel:      "pchannel",
+		DeliverPolicy: options.DeliverPolicyAll(),
+		DeliverFilters: []options.DeliverFilter{
+			options.DeliverFilterTimeTickGT(10),
+			options.DeliverFilterTimeTickGTE(10),
+			options.DeliverFilterVChannel("vchannel"),
+		},
+		MessageHandler: make(message.ChanMessageHandler),
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, consumer)
+	consumer.Close()
+
+	service.EXPECT().Close().Return()
+	handler.Close()
+	producer, err = handler.CreateProducer(ctx, nil)
+	assert.Error(t, err)
+	assert.Nil(t, producer)
+
+	consumer, err = handler.CreateConsumer(ctx, nil)
+	assert.Error(t, err)
+	assert.Nil(t, consumer)
+}
+
+func TestDial(t *testing.T) {
+	paramtable.Init()
+
+	w := mock_types.NewMockAssignmentDiscoverWatcher(t)
+	w.EXPECT().AssignmentDiscover(mock.Anything, mock.Anything).RunAndReturn(
+		func(ctx context.Context, f func(*types.VersionedStreamingNodeAssignments) error) error {
+			return context.Canceled
+		},
+	)
+	handler := NewHandlerClient(w)
+	assert.NotNil(t, handler)
+	time.Sleep(100 * time.Millisecond)
+	handler.Close()
+}

internal/streamingnode/client/handler/producer/produce_grpc_client.go

@@ -0,0 +1,35 @@
+package producer
+
+import (
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+)
+
+// produceGrpcClient is a wrapped producer server of log messages.
+type produceGrpcClient struct {
+	streamingpb.StreamingNodeHandlerService_ProduceClient
+}
+
+// SendProduceMessage sends the produce message to server.
+func (p *produceGrpcClient) SendProduceMessage(requestID int64, msg message.MutableMessage) error {
+	return p.Send(&streamingpb.ProduceRequest{
+		Request: &streamingpb.ProduceRequest_Produce{
+			Produce: &streamingpb.ProduceMessageRequest{
+				RequestId: requestID,
+				Message: &streamingpb.Message{
+					Payload:    msg.Payload(),
+					Properties: msg.Properties().ToRawMap(),
+				},
+			},
+		},
+	})
+}
+
+// SendClose sends the close request to server.
+func (p *produceGrpcClient) SendClose() error {
+	return p.Send(&streamingpb.ProduceRequest{
+		Request: &streamingpb.ProduceRequest_Close{
+			Close: &streamingpb.CloseProducerRequest{},
+		},
+	})
+}

internal/streamingnode/client/handler/producer/producer.go

@@ -0,0 +1,30 @@
+package producer
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+)
+
+var _ Producer = (*producerImpl)(nil)
+
+// Producer is the interface that wraps the basic produce method on grpc stream.
+// Producer is work on a single stream on grpc,
+// so Producer cannot recover from failure because of the stream is broken.
+type Producer interface {
+	// Assignment returns the assignment of the producer.
+	Assignment() types.PChannelInfoAssigned
+
+	// Produce sends the produce message to server.
+	Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error)
+
+	// Check if a producer is available.
+	IsAvailable() bool
+
+	// Available returns a channel that will be closed when the producer is unavailable.
+	Available() <-chan struct{}
+
+	// Close close the producer client.
+	Close()
+}

internal/streamingnode/client/handler/producer/producer_impl.go

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

internal/streamingnode/client/handler/producer/producer_test.go

@@ -0,0 +1,112 @@
+package producer
+
+import (
+	"context"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/streaming/walimpls/impls/walimplstest"
+)
+
+func TestProducer(t *testing.T) {
+	c := mock_streamingpb.NewMockStreamingNodeHandlerServiceClient(t)
+	cc := mock_streamingpb.NewMockStreamingNodeHandlerService_ProduceClient(t)
+	recvCh := make(chan *streamingpb.ProduceResponse, 10)
+	cc.EXPECT().Recv().RunAndReturn(func() (*streamingpb.ProduceResponse, error) {
+		msg, ok := <-recvCh
+		if !ok {
+			return nil, io.EOF
+		}
+		return msg, nil
+	})
+	sendCh := make(chan struct{}, 5)
+	cc.EXPECT().Send(mock.Anything).RunAndReturn(func(pr *streamingpb.ProduceRequest) error {
+		sendCh <- struct{}{}
+		return nil
+	})
+	c.EXPECT().Produce(mock.Anything, mock.Anything).Return(cc, nil)
+	cc.EXPECT().CloseSend().RunAndReturn(func() error {
+		recvCh <- &streamingpb.ProduceResponse{Response: &streamingpb.ProduceResponse_Close{}}
+		close(recvCh)
+		return nil
+	})
+
+	ctx := context.Background()
+	opts := &ProducerOptions{
+		Assignment: &types.PChannelInfoAssigned{
+			Channel: types.PChannelInfo{Name: "test", Term: 1},
+			Node:    types.StreamingNodeInfo{ServerID: 1, Address: "localhost"},
+		},
+	}
+
+	recvCh <- &streamingpb.ProduceResponse{
+		Response: &streamingpb.ProduceResponse_Create{
+			Create: &streamingpb.CreateProducerResponse{
+				WalName: "test",
+			},
+		},
+	}
+	producer, err := CreateProducer(ctx, opts, c)
+	assert.NoError(t, err)
+	assert.NotNil(t, producer)
+	ch := make(chan struct{})
+	go func() {
+		msgID, err := producer.Produce(ctx, message.NewMutableMessageBuilder().
+			WithMessageType(message.MessageTypeUnknown).
+			WithPayload([]byte{}).
+			BuildMutable())
+		assert.Error(t, err)
+		assert.Nil(t, msgID)
+		msgID, err = producer.Produce(ctx, message.NewMutableMessageBuilder().
+			WithMessageType(message.MessageTypeUnknown).
+			WithPayload([]byte{}).
+			BuildMutable())
+		assert.NoError(t, err)
+		assert.NotNil(t, msgID)
+		close(ch)
+	}()
+	<-sendCh
+	recvCh <- &streamingpb.ProduceResponse{
+		Response: &streamingpb.ProduceResponse_Produce{
+			Produce: &streamingpb.ProduceMessageResponse{
+				RequestId: 1,
+				Response: &streamingpb.ProduceMessageResponse_Error{
+					Error: &streamingpb.StreamingError{Code: 1},
+				},
+			},
+		},
+	}
+	<-sendCh
+	recvCh <- &streamingpb.ProduceResponse{
+		Response: &streamingpb.ProduceResponse_Produce{
+			Produce: &streamingpb.ProduceMessageResponse{
+				RequestId: 2,
+				Response: &streamingpb.ProduceMessageResponse_Result{
+					Result: &streamingpb.ProduceMessageResponseResult{
+						Id: &streamingpb.MessageID{Id: walimplstest.NewTestMessageID(1).Marshal()},
+					},
+				},
+			},
+		},
+	}
+	<-ch
+
+	ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+	defer cancel()
+	_, err = producer.Produce(ctx, message.NewMutableMessageBuilder().
+		WithMessageType(message.MessageTypeUnknown).
+		WithPayload([]byte{}).
+		BuildMutable())
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	assert.True(t, producer.IsAvailable())
+	producer.Close()
+	assert.False(t, producer.IsAvailable())
+}

internal/streamingnode/client/handler/shared_producer.go

@@ -0,0 +1,51 @@
+package handler
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+var _ producer.Producer = (*sharedProducer)(nil)
+
+// newSharedProducer creates a shared producer.
+func newSharedProducer(ref *typeutil.SharedReference[Producer]) sharedProducer {
+	return sharedProducer{
+		SharedReference: ref,
+	}
+}
+
+// sharedProducer is a shared producer.
+type sharedProducer struct {
+	*typeutil.SharedReference[Producer]
+}
+
+// Clone clones the shared producer.
+func (sp sharedProducer) Clone() *sharedProducer {
+	return &sharedProducer{
+		SharedReference: sp.SharedReference.Clone(),
+	}
+}
+
+// Assignment returns the assignment of the producer.
+func (sp sharedProducer) Assignment() types.PChannelInfoAssigned {
+	return sp.Deref().Assignment()
+}
+
+// Produce sends the produce message to server.
+func (sp sharedProducer) Produce(ctx context.Context, msg message.MutableMessage) (message.MessageID, error) {
+	return sp.Deref().Produce(ctx, msg)
+}
+
+// Check if a producer is available.
+func (sp sharedProducer) IsAvailable() bool {
+	return sp.Deref().IsAvailable()
+}
+
+// Available returns a channel that will be closed when the producer is unavailable.
+func (sp sharedProducer) Available() <-chan struct{} {
+	return sp.Deref().Available()
+}

internal/streamingnode/client/manager/manager.go

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

internal/streamingnode/client/manager/manager_client.go

@@ -0,0 +1,112 @@
+package manager
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker"
+	streamingserviceinterceptor "github.com/milvus-io/milvus/internal/util/streamingutil/service/interceptor"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/tracer"
+	"github.com/milvus-io/milvus/pkg/util/interceptor"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+// ManagerClient is the client to manage wal instances in all streamingnode.
+// ManagerClient wraps the Session Service Discovery.
+// Provides the ability to assign and remove wal instances for the channel on streaming node.
+type ManagerClient interface {
+	// WatchNodeChanged returns a channel that receive the signal that a streaming node change.
+	WatchNodeChanged(ctx context.Context) (<-chan struct{}, error)
+
+	// CollectAllStatus collects status of all streamingnode, such as load balance attributes.
+	CollectAllStatus(ctx context.Context) (map[int64]*types.StreamingNodeStatus, error)
+
+	// Assign a wal instance for the channel on streaming node of given server id.
+	Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error
+
+	// Remove the wal instance for the channel on streaming node of given server id.
+	Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error
+
+	// Close closes the manager client.
+	// It close the underlying connection, stop the node watcher and release all resources.
+	Close()
+}
+
+// NewManagerClient creates a new manager client.
+func NewManagerClient(etcdCli *clientv3.Client) ManagerClient {
+	role := sessionutil.GetSessionPrefixByRole(typeutil.StreamingNodeRole)
+	rb := resolver.NewSessionBuilder(etcdCli, role)
+	dialTimeout := paramtable.Get().StreamingNodeGrpcClientCfg.DialTimeout.GetAsDuration(time.Millisecond)
+	dialOptions := getDialOptions(rb)
+	conn := lazygrpc.NewConn(func(ctx context.Context) (*grpc.ClientConn, error) {
+		ctx, cancel := context.WithTimeout(ctx, dialTimeout)
+		defer cancel()
+		return grpc.DialContext(
+			ctx,
+			resolver.SessionResolverScheme+":///"+typeutil.StreamingNodeRole,
+			dialOptions...,
+		)
+	})
+	return &managerClientImpl{
+		lifetime: lifetime.NewLifetime(lifetime.Working),
+		rb:       rb,
+		service:  lazygrpc.WithServiceCreator(conn, streamingpb.NewStreamingNodeManagerServiceClient),
+	}
+}
+
+// getDialOptions returns grpc dial options.
+func getDialOptions(rb resolver.Builder) []grpc.DialOption {
+	cfg := &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
+}

internal/streamingnode/client/manager/manager_client_impl.go

@@ -0,0 +1,191 @@
+package manager
+
+import (
+	"context"
+	"sync"
+
+	"go.uber.org/zap"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/balancer/picker"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/lazygrpc"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/resolver"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/typeconverter"
+	"github.com/milvus-io/milvus/pkg/log"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+)
+
+var _ ManagerClient = (*managerClientImpl)(nil)
+
+// managerClientImpl implements ManagerClient.
+type managerClientImpl struct {
+	lifetime lifetime.Lifetime[lifetime.State]
+
+	rb      resolver.Builder
+	service lazygrpc.Service[streamingpb.StreamingNodeManagerServiceClient]
+}
+
+func (c *managerClientImpl) WatchNodeChanged(ctx context.Context) (<-chan struct{}, error) {
+	if c.lifetime.Add(lifetime.IsWorking) != nil {
+		return nil, status.NewOnShutdownError("manager client is closing")
+	}
+	defer c.lifetime.Done()
+
+	resultCh := make(chan struct{}, 1)
+	go func() {
+		defer close(resultCh)
+		c.rb.Resolver().Watch(ctx, func(state resolver.VersionedState) error {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-c.lifetime.CloseCh():
+				return status.NewOnShutdownError("manager client is closing")
+			case resultCh <- struct{}{}:
+			}
+			return nil
+		})
+	}()
+	return resultCh, nil
+}
+
+// CollectAllStatus collects status in all underlying streamingnode.
+func (c *managerClientImpl) CollectAllStatus(ctx context.Context) (map[int64]*types.StreamingNodeStatus, error) {
+	if c.lifetime.Add(lifetime.IsWorking) != nil {
+		return nil, status.NewOnShutdownError("manager client is closing")
+	}
+	defer c.lifetime.Done()
+
+	// Get all discovered streamingnode.
+	state := c.rb.Resolver().GetLatestState()
+	if len(state.State.Addresses) == 0 {
+		return make(map[int64]*types.StreamingNodeStatus), nil
+	}
+
+	// Collect status of all streamingnode.
+	result, err := c.getAllStreamingNodeStatus(ctx, state)
+	if err != nil {
+		return nil, err
+	}
+
+	// Collect status may cost some time, so we need to check the lifetime again.
+	newState := c.rb.Resolver().GetLatestState()
+	if newState.Version.GT(state.Version) {
+		newSession := newState.Sessions()
+		for serverID := range result {
+			if session, ok := newSession[serverID]; !ok {
+				result[serverID].Err = types.ErrNotAlive
+			} else if session.Stopping {
+				result[serverID].Err = types.ErrStopping
+			}
+		}
+	}
+	return result, nil
+}
+
+func (c *managerClientImpl) getAllStreamingNodeStatus(ctx context.Context, state discoverer.VersionedState) (map[int64]*types.StreamingNodeStatus, error) {
+	// wait for manager service ready.
+	manager, err := c.service.GetService(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	g, _ := errgroup.WithContext(ctx)
+	g.SetLimit(10)
+	var mu sync.Mutex
+	result := make(map[int64]*types.StreamingNodeStatus, len(state.State.Addresses))
+	for serverID, session := range state.Sessions() {
+		serverID := serverID
+		address := session.Address
+		g.Go(func() error {
+			ctx := contextutil.WithPickServerID(ctx, serverID)
+			resp, err := manager.CollectStatus(ctx, &streamingpb.StreamingNodeManagerCollectStatusRequest{})
+			mu.Lock()
+			defer mu.Unlock()
+			result[serverID] = &types.StreamingNodeStatus{
+				StreamingNodeInfo: types.StreamingNodeInfo{
+					ServerID: serverID,
+					Address:  address,
+				},
+				Err: err,
+			}
+
+			if err != nil {
+				log.Warn("collect status failed, skip", zap.Int64("serverID", serverID), zap.Error(err))
+				return err
+			}
+			log.Debug("collect status success", zap.Int64("serverID", serverID), zap.Any("status", resp))
+			return nil
+		})
+	}
+	g.Wait()
+
+	return result, nil
+}
+
+// Assign a wal instance for the channel on log node of given server id.
+func (c *managerClientImpl) Assign(ctx context.Context, pchannel types.PChannelInfoAssigned) error {
+	if c.lifetime.Add(lifetime.IsWorking) != nil {
+		return status.NewOnShutdownError("manager client is closing")
+	}
+	defer c.lifetime.Done()
+
+	// wait for manager service ready.
+	manager, err := c.service.GetService(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Select a log node to assign the wal instance.
+	ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID)
+	_, err = manager.Assign(ctx, &streamingpb.StreamingNodeManagerAssignRequest{
+		Pchannel: typeconverter.NewProtoFromPChannelInfo(pchannel.Channel),
+	})
+	return err
+}
+
+// Remove the wal instance for the channel on log node of given server id.
+func (c *managerClientImpl) Remove(ctx context.Context, pchannel types.PChannelInfoAssigned) error {
+	if c.lifetime.Add(lifetime.IsWorking) != nil {
+		return status.NewOnShutdownError("manager client is closing")
+	}
+	defer c.lifetime.Done()
+
+	// wait for manager service ready.
+	manager, err := c.service.GetService(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Select a streaming node to remove the wal instance.
+	ctx = contextutil.WithPickServerID(ctx, pchannel.Node.ServerID)
+	_, err = manager.Remove(ctx, &streamingpb.StreamingNodeManagerRemoveRequest{
+		Pchannel: typeconverter.NewProtoFromPChannelInfo(pchannel.Channel),
+	})
+	// The following error can be treated as success.
+	// 1. err is nil, a real remove operation at streaming node has been happened.
+	// 2. err is ErrSubConnNoExist, the streaming node is not alive at view of session, so the wal on it is already removed.
+	// 3. err is SkippedOperation, the streaming node is not the owner of the wal, so the wal on it is already removed.
+	if err == nil || picker.IsErrSubConnNoExist(err) {
+		return nil
+	}
+	statusErr := status.AsStreamingError(err)
+	if statusErr == nil || statusErr.IsSkippedOperation() {
+		return nil
+	}
+	return err
+}
+
+// Close closes the manager client.
+func (c *managerClientImpl) Close() {
+	c.lifetime.SetState(lifetime.Stopped)
+	c.lifetime.Wait()
+	c.lifetime.Close()
+
+	c.service.Close()
+	c.rb.Close()
+}

internal/streamingnode/client/manager/manager_test.go

@@ -0,0 +1,174 @@
+package manager
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/resolver"
+
+	"github.com/milvus-io/milvus/internal/mocks/proto/mock_streamingpb"
+	"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_lazygrpc"
+	"github.com/milvus-io/milvus/internal/mocks/util/streamingutil/service/mock_resolver"
+	"github.com/milvus-io/milvus/internal/proto/streamingpb"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/attributes"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/contextutil"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/service/discoverer"
+	"github.com/milvus-io/milvus/pkg/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/util/etcd"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/typeutil"
+)
+
+func TestManager(t *testing.T) {
+	rb := mock_resolver.NewMockBuilder(t)
+	managerService := mock_lazygrpc.NewMockService[streamingpb.StreamingNodeManagerServiceClient](t)
+	m := &managerClientImpl{
+		lifetime: lifetime.NewLifetime(lifetime.Working),
+		rb:       rb,
+		service:  managerService,
+	}
+	r := mock_resolver.NewMockResolver(t)
+	rb.EXPECT().Resolver().Return(r)
+	r.EXPECT().Watch(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, f func(discoverer.VersionedState) error) error {
+		f(discoverer.VersionedState{})
+		return context.Canceled
+	})
+
+	// Test WatchNodeChanged.
+	resultCh, err := m.WatchNodeChanged(context.Background())
+	assert.NoError(t, err)
+	_, ok := <-resultCh
+	assert.True(t, ok)
+	_, ok = <-resultCh
+	assert.False(t, ok)
+
+	// Test CollectAllStatus
+	r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState {
+		return discoverer.VersionedState{
+			Version: typeutil.VersionInt64(1),
+			State:   resolver.State{},
+		}
+	})
+	// Not address here.
+	nodes, err := m.CollectAllStatus(context.Background())
+	assert.NoError(t, err)
+	assert.Len(t, nodes, 0)
+
+	// Has address.
+	managerServiceClient := mock_streamingpb.NewMockStreamingNodeManagerServiceClient(t)
+	managerService.EXPECT().GetService(mock.Anything).RunAndReturn(func(ctx context.Context) (streamingpb.StreamingNodeManagerServiceClient, error) {
+		return managerServiceClient, nil
+	})
+	managerServiceClient.EXPECT().CollectStatus(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, snmcsr *streamingpb.StreamingNodeManagerCollectStatusRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerCollectStatusResponse, error) {
+		return &streamingpb.StreamingNodeManagerCollectStatusResponse{}, nil
+	})
+
+	i := 0
+	states := []map[uint64]bool{
+		{1: false, 2: false, 3: true},
+		{1: true, 2: false},
+	}
+	r.EXPECT().GetLatestState().Unset()
+	r.EXPECT().GetLatestState().RunAndReturn(func() discoverer.VersionedState {
+		s := newVersionedState(int64(i), states[i])
+		i++
+		return s
+	})
+
+	nodes, err = m.CollectAllStatus(context.Background())
+	assert.NoError(t, err)
+	assert.Len(t, nodes, 3)
+	assert.ErrorIs(t, nodes[3].Err, types.ErrNotAlive)
+	assert.ErrorIs(t, nodes[1].Err, types.ErrStopping)
+
+	// Test Assign
+	serverID := int64(2)
+	managerServiceClient.EXPECT().Assign(mock.Anything, mock.Anything).RunAndReturn(
+		func(ctx context.Context, snmcsr *streamingpb.StreamingNodeManagerAssignRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerAssignResponse, error) {
+			pickedServerID, ok := contextutil.GetPickServerID(ctx)
+			assert.True(t, ok)
+			assert.Equal(t, serverID, pickedServerID)
+			return nil, nil
+		})
+	err = m.Assign(context.Background(), types.PChannelInfoAssigned{
+		Channel: types.PChannelInfo{Name: "p", Term: 1},
+		Node:    types.StreamingNodeInfo{ServerID: serverID},
+	})
+	assert.NoError(t, err)
+
+	// Test Remove
+	serverID = int64(2)
+	managerServiceClient.EXPECT().Remove(mock.Anything, mock.Anything).RunAndReturn(
+		func(ctx context.Context, snmrr *streamingpb.StreamingNodeManagerRemoveRequest, co ...grpc.CallOption) (*streamingpb.StreamingNodeManagerRemoveResponse, error) {
+			pickedServerID, ok := contextutil.GetPickServerID(ctx)
+			assert.True(t, ok)
+			assert.Equal(t, serverID, pickedServerID)
+			return nil, nil
+		})
+	err = m.Remove(context.Background(), types.PChannelInfoAssigned{
+		Channel: types.PChannelInfo{Name: "p", Term: 1},
+		Node:    types.StreamingNodeInfo{ServerID: serverID},
+	})
+	assert.NoError(t, err)
+
+	// Test Close
+	managerService.EXPECT().Close().Return()
+	rb.EXPECT().Close().Return()
+	m.Close()
+
+	nodes, err = m.CollectAllStatus(context.Background())
+	assert.Nil(t, nodes)
+	assert.Error(t, err)
+	err = m.Assign(context.Background(), types.PChannelInfoAssigned{})
+	assert.Error(t, err)
+	err = m.Remove(context.Background(), types.PChannelInfoAssigned{})
+	assert.Error(t, err)
+	resultCh, err = m.WatchNodeChanged(context.Background())
+	assert.Nil(t, resultCh)
+	assert.Error(t, err)
+}
+
+func newVersionedState(version int64, serverIDs map[uint64]bool) discoverer.VersionedState {
+	state := discoverer.VersionedState{
+		Version: typeutil.VersionInt64(version),
+		State: resolver.State{
+			Addresses: make([]resolver.Address, 0, len(serverIDs)),
+		},
+	}
+
+	for serverID, stopping := range serverIDs {
+		state.State.Addresses = append(state.State.Addresses, resolver.Address{
+			Addr: fmt.Sprintf("localhost:%d", serverID),
+			BalancerAttributes: attributes.WithSession(
+				new(attributes.Attributes), &sessionutil.SessionRaw{
+					ServerID: int64(serverID),
+					Stopping: stopping,
+				},
+			),
+		})
+	}
+	return state
+}
+
+func TestDial(t *testing.T) {
+	paramtable.Init()
+
+	err := etcd.InitEtcdServer(true, "", t.TempDir(), "stdout", "info")
+	assert.NoError(t, err)
+	defer etcd.StopEtcdServer()
+	c, err := etcd.GetEmbedEtcdClient()
+	assert.NoError(t, err)
+	assert.NotNil(t, c)
+
+	client := NewManagerClient(c)
+	assert.NotNil(t, client)
+	time.Sleep(100 * time.Millisecond)
+	client.Close()
+}

internal/streamingnode/server/resource/idalloc/test_mock_root_coord_client.go

@@ -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{
+				Status:    merr.Success(),
-					ErrorCode: commonpb.ErrorCode_Success,
-				},
 				Timestamp: c - uint64(atr.Count),
 				Count:     atr.Count,
 			}, nil
@@ -42,9 +40,7 @@ func NewMockRootCoordClient(t *testing.T) *mocks.MockRootCoordClient {
 			}
 			c := counter.Add(uint64(atr.Count))
 			return &rootcoordpb.AllocIDResponse{
-				Status: &commonpb.Status{
+				Status: merr.Success(),
-					ErrorCode: commonpb.ErrorCode_Success,
-				},
 				ID:     int64(c - uint64(atr.Count)),
 				Count:  atr.Count,
 			}, nil

internal/util/streamingutil/service/attributes/attributes_test.go

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

internal/util/streamingutil/service/balancer/balancer.go

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

internal/util/streamingutil/service/interceptor/server.go

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

internal/util/streamingutil/service/resolver/builder.go

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

internal/util/streamingutil/service/resolver/builder_test.go

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

internal/util/streamingutil/service/resolver/resolver.go

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

internal/util/streamingutil/service/resolver/resolver_with_discoverer.go

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

pkg/util/paramtable/component_param.go

@@ -97,6 +97,8 @@ type ComponentParam struct {
 	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)
 }

pkg/util/paramtable/grpc_param.go

@@ -19,8 +19,12 @@ package paramtable
 import (
 	"fmt"
 	"strconv"
+	"time"
 
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/keepalive"
 
 	"github.com/milvus-io/milvus/pkg/log"
 	"github.com/milvus-io/milvus/pkg/util/funcutil"
@@ -219,6 +223,7 @@ type GrpcClientConfig struct {
 	MaxAttempts             ParamItem `refreshable:"false"`
 	InitialBackoff          ParamItem `refreshable:"false"`
 	MaxBackoff              ParamItem `refreshable:"false"`
+	BackoffMultiplier       ParamItem `refreshable:"false"`
 	MinResetInterval        ParamItem `refreshable:"false"`
 	MaxCancelError          ParamItem `refreshable:"false"`
 	MinSessionCheckInterval ParamItem `refreshable:"false"`
@@ -397,6 +402,14 @@ func (p *GrpcClientConfig) Init(domain string, base *BaseTable) {
 	}
 	p.MaxBackoff.Init(base.mgr)
 
+	p.BackoffMultiplier = ParamItem{
+		Key:          "grpc.client.backoffMultiplier",
+		Version:      "2.5.0",
+		DefaultValue: "2.0",
+		Export:       true,
+	}
+	p.BackoffMultiplier.Init(base.mgr)
+
 	compressionEnabled := fmt.Sprintf("%t", DefaultCompressionEnabled)
 	p.CompressionEnabled = ParamItem{
 		Key:     "grpc.client.compressionEnabled",
@@ -478,3 +491,42 @@ func (p *GrpcClientConfig) Init(domain string, base *BaseTable) {
 	}
 	p.MaxCancelError.Init(base.mgr)
 }
+
+// GetDialOptionsFromConfig returns grpc dial options from config.
+func (p *GrpcClientConfig) GetDialOptionsFromConfig() []grpc.DialOption {
+	compress := ""
+	if p.CompressionEnabled.GetAsBool() {
+		compress = "zstd"
+	}
+	return []grpc.DialOption{
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(p.ClientMaxRecvSize.GetAsInt()),
+			grpc.MaxCallSendMsgSize(p.ClientMaxSendSize.GetAsInt()),
+			grpc.UseCompressor(compress),
+		),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:                p.KeepAliveTime.GetAsDuration(time.Millisecond),
+			Timeout:             p.KeepAliveTimeout.GetAsDuration(time.Millisecond),
+			PermitWithoutStream: true,
+		}),
+		grpc.WithConnectParams(grpc.ConnectParams{
+			Backoff: backoff.Config{
+				BaseDelay:  100 * time.Millisecond,
+				Multiplier: 1.6,
+				Jitter:     0.2,
+				MaxDelay:   3 * time.Second,
+			},
+			MinConnectTimeout: p.DialTimeout.GetAsDuration(time.Millisecond),
+		}),
+	}
+}
+
+// GetDefaultRetryPolicy returns default grpc retry policy.
+func (p *GrpcClientConfig) GetDefaultRetryPolicy() map[string]interface{} {
+	return map[string]interface{}{
+		"maxAttempts":       p.MaxAttempts.GetAsInt(),
+		"initialBackoff":    fmt.Sprintf("%fs", p.InitialBackoff.GetAsFloat()),
+		"maxBackoff":        fmt.Sprintf("%fs", p.MaxBackoff.GetAsFloat()),
+		"backoffMultiplier": p.BackoffMultiplier.GetAsFloat(),
+	}
+}

pkg/util/typeutil/shared_reference.go

@@ -0,0 +1,99 @@
+package typeutil
+
+import (
+	"go.uber.org/atomic"
+)
+
+// Closable is an interface that can be closed or release the resource.
+type Closable interface {
+	Close()
+}
+
+// SharedReference is a reference type that can be shared with only one close operation.
+// Without supported of determined destructor, A safe SharedReference is rely on user's behavior.
+type SharedReference[T Closable] struct {
+	inner     T
+	strongCnt *atomic.Int32
+}
+
+// NewSharedReference creates a new SharedReference with the inner object.
+func NewSharedReference[T Closable](inner T) *SharedReference[T] {
+	return &SharedReference[T]{
+		inner:     inner,
+		strongCnt: atomic.NewInt32(1),
+	}
+}
+
+// Deref returns the inner object.
+// Deref should only be called if the reference is not closed.
+// Otherwise UB happens.
+func (sr *SharedReference[T]) Deref() T {
+	return sr.inner
+}
+
+// Clone returns a new SharedReference with the same inner object.
+// Clone should only be called if the reference is not closed.
+// Otherwise UB happens.
+func (sr *SharedReference[T]) Clone() *SharedReference[T] {
+	n := sr.strongCnt.Load()
+
+	for {
+		if n == 0 {
+			panic("SharedReference: Clone on a closed reference")
+		}
+		if sr.strongCnt.CompareAndSwap(n, n+1) {
+			break
+		}
+		n = sr.strongCnt.Load()
+	}
+	return sr
+}
+
+// Downgrade returns a new WeakReference with the same inner object.
+// After downgrade, the SharedReference can not be used.
+func (sr *SharedReference[T]) Downgrade() *WeakReference[T] {
+	w := &WeakReference[T]{
+		inner:     sr.inner,
+		strongCnt: sr.strongCnt,
+	}
+	sr.Close()
+	return w
+}
+
+// Close closes the reference, should only be called once at a sharedReference.
+// After called this method, all other method of the SharedReference should not be called any more.
+// if the reference count is 0, the inner object will be closed.
+func (sr *SharedReference[T]) Close() {
+	sr.strongCnt.Load()
+	refcnt := sr.strongCnt.Dec()
+	if refcnt < 0 {
+		panic("SharedReference: Close on a closed reference")
+	}
+	if refcnt == 0 {
+		sr.inner.Close()
+	}
+}
+
+// WeakReference is a weak reference type that can be shared with only one close operation.
+type WeakReference[T Closable] struct {
+	inner     T
+	strongCnt *atomic.Int32
+	// TODO: add weak ref count in future.
+}
+
+func (wr *WeakReference[T]) Upgrade() *SharedReference[T] {
+	n := wr.strongCnt.Load()
+	for {
+		// There's no strong reference of dLogImpl, so it's impossible to create a DLogRef.
+		if n == 0 {
+			return nil
+		}
+		if wr.strongCnt.CompareAndSwap(n, n+1) {
+			return &SharedReference[T]{
+				inner:     wr.inner,
+				strongCnt: wr.strongCnt,
+			}
+		}
+		n = wr.strongCnt.Load()
+	}
+}

pkg/util/typeutil/shared_reference_test.go

@@ -0,0 +1,81 @@
+package typeutil
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/atomic"
+)
+
+type closer struct {
+	t      *testing.T
+	closed atomic.Bool
+}
+
+func (c *closer) Close() {
+	assert.True(c.t, c.closed.CompareAndSwap(false, true))
+}
+
+func TestSharedReference(t *testing.T) {
+	// test downgrade and upgrade
+	closer1 := &closer{t: t}
+	r := NewSharedReference(closer1)
+	w := r.Downgrade()
+	assert.Nil(t, w.Upgrade())
+	assert.True(t, closer1.closed.Load())
+
+	// test concurrent
+	closer1 = &closer{t: t}
+
+	r = NewSharedReference(closer1)
+
+	wg := &sync.WaitGroup{}
+	wg.Add(4)
+	r1 := r.Clone()
+	go func() {
+		defer wg.Done()
+		defer r1.Close()
+		for i := 0; i < 1000; i++ {
+			r3 := r1.Clone()
+			r3.Close()
+		}
+	}()
+	r2 := r.Clone()
+	go func() {
+		defer wg.Done()
+		defer r2.Close()
+		for i := 0; i < 1000; i++ {
+			r := r2.Clone()
+			r.Close()
+		}
+	}()
+	r3 := r.Clone().Downgrade()
+	go func() {
+		defer wg.Done()
+		r := r3.Upgrade()
+		r.Close()
+	}()
+	r4 := r.Clone().Downgrade()
+	go func() {
+		defer wg.Done()
+		r := r4.Upgrade()
+		r.Close()
+	}()
+	wg.Wait()
+
+	assert.False(t, r.Deref().closed.Load())
+	r.Close()
+	assert.True(t, closer1.closed.Load())
+
+	assert.Panics(t, func() {
+		r.Clone()
+	})
+
+	closer1 = &closer{t: t}
+	r = NewSharedReference(closer1)
+	r.Close()
+	assert.Panics(t, func() {
+		r.Close()
+	})
+}