enhance: streaming service client (#34656)

issue: #33285

- implement streaming service client.
- implement producing and consuming service client by streaming coord
client and streaming node client.

Signed-off-by: chyezh <chyezh@outlook.com>

internal/distributed/streaming/OWNERS

@@ -0,0 +1,5 @@
+reviewers:
+  - chyezh
+
+approvers:
+  - maintainers

internal/distributed/streaming/append.go

@@ -0,0 +1,171 @@
+package streaming
+
+import (
+	"context"
+	"sync"
+
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/producer"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/util/funcutil"
+)
+
+// newAppendResponseN creates a new append response.
+func newAppendResponseN(n int) AppendResponses {
+	return AppendResponses{
+		Responses: make([]AppendResponse, n),
+	}
+}
+
+// AppendResponse is the response of one append operation.
+type AppendResponse struct {
+	MessageID message.MessageID
+	Error     error
+}
+
+// AppendResponses is the response of append operation.
+type AppendResponses struct {
+	Responses []AppendResponse
+}
+
+// IsAnyError returns the first error in the responses.
+func (a AppendResponses) IsAnyError() error {
+	for _, r := range a.Responses {
+		if r.Error != nil {
+			return r.Error
+		}
+	}
+	return nil
+}
+
+// fillAllError fills all the responses with the same error.
+func (a *AppendResponses) fillAllError(err error) {
+	for i := range a.Responses {
+		a.Responses[i].Error = err
+	}
+}
+
+// fillResponseAtIdx fill the response at idx
+func (a *AppendResponses) fillResponseAtIdx(resp AppendResponse, idx int) {
+	a.Responses[idx] = resp
+}
+
+// dispatchByPChannel dispatches the message into different pchannel.
+func (w *walAccesserImpl) dispatchByPChannel(ctx context.Context, msgs ...message.MutableMessage) AppendResponses {
+	if len(msgs) == 0 {
+		return newAppendResponseN(0)
+	}
+
+	// dispatch the messages into different pchannel.
+	dispatchedMessages, indexes := w.dispatchMessages(msgs...)
+
+	// only one pchannel, append it directly, no more goroutine needed.
+	if len(dispatchedMessages) == 1 {
+		for pchannel, msgs := range dispatchedMessages {
+			return w.appendToPChannel(ctx, pchannel, msgs...)
+		}
+	}
+
+	// otherwise, start multiple goroutine to append to different pchannel.
+	resp := newAppendResponseN(len(msgs))
+	wg := sync.WaitGroup{}
+	wg.Add(len(dispatchedMessages))
+
+	mu := sync.Mutex{}
+	for pchannel, msgs := range dispatchedMessages {
+		pchannel := pchannel
+		msgs := msgs
+		idxes := indexes[pchannel]
+		go func() {
+			defer wg.Done()
+			singleResp := w.appendToPChannel(ctx, pchannel, msgs...)
+			mu.Lock()
+			for i, idx := range idxes {
+				resp.fillResponseAtIdx(singleResp.Responses[i], idx)
+			}
+			mu.Unlock()
+		}()
+	}
+	wg.Wait()
+	return resp
+}
+
+// dispatchMessages dispatches the messages into different pchannel.
+func (w *walAccesserImpl) dispatchMessages(msgs ...message.MutableMessage) (map[string][]message.MutableMessage, map[string][]int) {
+	dispatchedMessages := make(map[string][]message.MutableMessage, 0)
+	// record the index of the message in the msgs, used to fill back response.
+	indexes := make(map[string][]int, 0)
+	for idx, msg := range msgs {
+		pchannel := funcutil.ToPhysicalChannel(msg.VChannel())
+		if _, ok := dispatchedMessages[pchannel]; !ok {
+			dispatchedMessages[pchannel] = make([]message.MutableMessage, 0)
+			indexes[pchannel] = make([]int, 0)
+		}
+		dispatchedMessages[pchannel] = append(dispatchedMessages[pchannel], msg)
+		indexes[pchannel] = append(indexes[pchannel], idx)
+	}
+	return dispatchedMessages, indexes
+}
+
+// appendToPChannel appends the messages to the specified pchannel.
+func (w *walAccesserImpl) appendToPChannel(ctx context.Context, pchannel string, msgs ...message.MutableMessage) AppendResponses {
+	if len(msgs) == 0 {
+		return newAppendResponseN(0)
+	}
+	resp := newAppendResponseN(len(msgs))
+
+	// get producer of pchannel.
+	p := w.getProducer(pchannel)
+
+	// if only one message here, append it directly, no more goroutine needed.
+	// at most time, there's only one message here.
+	// TODO: only the partition-key with high partition will generate many message in one time on the same pchannel,
+	// we should optimize the message-format, make it into one; but not the goroutine count.
+	if len(msgs) == 1 {
+		msgID, err := p.Produce(ctx, msgs[0])
+		resp.fillResponseAtIdx(AppendResponse{
+			MessageID: msgID,
+			Error:     err,
+		}, 0)
+		return resp
+	}
+
+	// concurrent produce here.
+	wg := sync.WaitGroup{}
+	wg.Add(len(msgs))
+
+	mu := sync.Mutex{}
+	for i, msg := range msgs {
+		i := i
+		msg := msg
+		go func() {
+			defer wg.Done()
+			msgID, err := p.Produce(ctx, msg)
+
+			mu.Lock()
+			resp.fillResponseAtIdx(AppendResponse{
+				MessageID: msgID,
+				Error:     err,
+			}, i)
+			mu.Unlock()
+		}()
+	}
+	wg.Wait()
+	return resp
+}
+
+// createOrGetProducer creates or get a producer.
+// vchannel in same pchannel can share the same producer.
+func (w *walAccesserImpl) getProducer(pchannel string) *producer.ResumableProducer {
+	w.producerMutex.Lock()
+	defer w.producerMutex.Unlock()
+
+	// TODO: A idle producer should be removed maybe?
+	if p, ok := w.producers[pchannel]; ok {
+		return p
+	}
+	p := producer.NewResumableProducer(w.handlerClient.CreateProducer, &producer.ProducerOptions{
+		PChannel: pchannel,
+	})
+	w.producers[pchannel] = p
+	return p
+}

internal/distributed/streaming/internal/consumer/consumer.go

@@ -0,0 +1,38 @@
+package consumer
+
+import (
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+)
+
+var _ ResumableConsumer = (*resumableConsumerImpl)(nil)
+
+// 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
+}
+
+// ResumableConsumer is the interface for consuming message to log service.
+// ResumableConsumer select a right log node to consume automatically.
+// ResumableConsumer will do automatic resume from stream broken and log node re-balance.
+// All error in these package should be marked by streamingservice/errs package.
+type ResumableConsumer interface {
+	// Done returns a channel which will be closed when scanner is finished or closed.
+	Done() <-chan struct{}
+
+	// Error returns the error of the Consumer.
+	Error() error
+
+	// Close the scanner, release the underlying resources.
+	Close()
+}

internal/distributed/streaming/internal/consumer/consumer_impl.go

@@ -0,0 +1,208 @@
+package consumer
+
+import (
+	"context"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"github.com/cockroachdb/errors"
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
+	"github.com/milvus-io/milvus/pkg/log"
+	"github.com/milvus-io/milvus/pkg/metrics"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/syncutil"
+)
+
+var errGracefulShutdown = errors.New("graceful shutdown")
+
+// NewResumableConsumer creates a new resumable consumer.
+func NewResumableConsumer(factory factory, opts *ConsumerOptions) ResumableConsumer {
+	ctx, cancel := context.WithCancel(context.Background())
+	consumer := &resumableConsumerImpl{
+		ctx:            ctx,
+		cancel:         cancel,
+		stopResumingCh: make(chan struct{}),
+		resumingExitCh: make(chan struct{}),
+		logger:         log.With(zap.String("pchannel", opts.PChannel), zap.Any("initialDeliverPolicy", opts.DeliverPolicy)),
+		opts:           opts,
+		mh: &timeTickOrderMessageHandler{
+			inner:                  opts.MessageHandler,
+			lastConfirmedMessageID: nil,
+			lastTimeTick:           0,
+		},
+		factory:    factory,
+		consumeErr: syncutil.NewFuture[error](),
+	}
+	go consumer.resumeLoop()
+	return consumer
+}
+
+// resumableConsumerImpl is a consumer implementation.
+type resumableConsumerImpl struct {
+	ctx            context.Context
+	cancel         context.CancelFunc
+	logger         *log.MLogger
+	stopResumingCh chan struct{}
+	resumingExitCh chan struct{}
+
+	opts       *ConsumerOptions
+	mh         *timeTickOrderMessageHandler
+	factory    factory
+	consumeErr *syncutil.Future[error]
+}
+
+type factory = func(ctx context.Context, opts *handler.ConsumerOptions) (consumer.Consumer, error)
+
+// resumeLoop starts a background task to consume messages from log service to message handler.
+func (rc *resumableConsumerImpl) resumeLoop() {
+	defer func() {
+		// close the message handler.
+		rc.mh.Close()
+		rc.logger.Info("resumable consumer is closed")
+		close(rc.resumingExitCh)
+	}()
+
+	// Use the initialized deliver policy at first running.
+	deliverPolicy := rc.opts.DeliverPolicy
+	deliverFilters := rc.opts.DeliverFilters
+	// consumer need to resume when error occur, so message handler shouldn't close if the internal consumer encounter failure.
+	nopCloseMH := message.NopCloseHandler{
+		Handler: rc.mh,
+	}
+
+	for {
+		// Get last checkpoint sent.
+		// Consume ordering is always time tick order now.
+		if rc.mh.lastConfirmedMessageID != nil {
+			// set the deliver policy start after the last message id.
+			deliverPolicy = options.DeliverPolicyStartAfter(rc.mh.lastConfirmedMessageID)
+			newDeliverFilters := make([]options.DeliverFilter, 0, len(deliverFilters)+1)
+			for _, filter := range deliverFilters {
+				if !options.IsDeliverFilterTimeTick(filter) {
+					newDeliverFilters = append(newDeliverFilters, filter)
+				}
+			}
+			newDeliverFilters = append(newDeliverFilters, options.DeliverFilterTimeTickGT(rc.mh.lastTimeTick))
+			deliverFilters = newDeliverFilters
+		}
+		opts := &handler.ConsumerOptions{
+			PChannel:       rc.opts.PChannel,
+			DeliverPolicy:  deliverPolicy,
+			DeliverFilters: deliverFilters,
+			MessageHandler: nopCloseMH,
+		}
+
+		// Create a new consumer.
+		// Return if context canceled.
+		consumer, err := rc.createNewConsumer(opts)
+		if err != nil {
+			rc.consumeErr.Set(err)
+			return
+		}
+
+		// Wait until the consumer is unavailable or context canceled.
+		if err := rc.waitUntilUnavailable(consumer); err != nil {
+			rc.consumeErr.Set(err)
+			return
+		}
+	}
+}
+
+func (rc *resumableConsumerImpl) createNewConsumer(opts *handler.ConsumerOptions) (consumer.Consumer, error) {
+	// Mark as unavailable.
+	metrics.StreamingServiceClientConsumerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerUnAvailable).Inc()
+	defer metrics.StreamingServiceClientConsumerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerUnAvailable).Dec()
+
+	logger := rc.logger.With(zap.Any("deliverPolicy", opts.DeliverPolicy))
+
+	backoff := backoff.NewExponentialBackOff()
+	backoff.InitialInterval = 100 * time.Millisecond
+	backoff.MaxInterval = 10 * time.Second
+	for {
+		// Create a new consumer.
+		// a underlying stream consumer life time should be equal to the resumable producer.
+		// so ctx of resumable consumer is passed to underlying stream producer creation.
+		consumer, err := rc.factory(rc.ctx, opts)
+		if errors.IsAny(err, context.Canceled, handler.ErrClientClosed) {
+			return nil, err
+		}
+		if err != nil {
+			nextBackoff := backoff.NextBackOff()
+			logger.Warn("create consumer failed, retry...", zap.Error(err), zap.Duration("nextRetryInterval", nextBackoff))
+			time.Sleep(nextBackoff)
+			continue
+		}
+
+		logger.Info("resume on new consumer at new start message id")
+		return consumer, nil
+	}
+}
+
+// waitUntilUnavailable is used to wait until the consumer is unavailable or context canceled.
+func (rc *resumableConsumerImpl) waitUntilUnavailable(consumer handler.Consumer) error {
+	// Mark as available.
+	metrics.StreamingServiceClientConsumerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerAvailable).Inc()
+	defer func() {
+		consumer.Close()
+		if consumer.Error() != nil {
+			rc.logger.Warn("consumer is closed with error", zap.Error(consumer.Error()))
+		}
+		metrics.StreamingServiceClientConsumerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerAvailable).Dec()
+	}()
+
+	select {
+	case <-rc.stopResumingCh:
+		return errGracefulShutdown
+	case <-rc.ctx.Done():
+		return rc.ctx.Err()
+	case <-consumer.Done():
+		rc.logger.Warn("consumer is done or encounter error, try to resume...",
+			zap.Error(consumer.Error()),
+			zap.Any("lastConfirmedMessageID", rc.mh.lastConfirmedMessageID),
+			zap.Uint64("lastTimeTick", rc.mh.lastTimeTick),
+		)
+		return nil
+	}
+}
+
+// gracefulClose graceful close the consumer.
+func (rc *resumableConsumerImpl) gracefulClose() error {
+	close(rc.stopResumingCh)
+	select {
+	case <-rc.resumingExitCh:
+		return nil
+	case <-time.After(50 * time.Millisecond):
+		return context.DeadlineExceeded
+	}
+}
+
+// Close the scanner, release the underlying resources.
+func (rc *resumableConsumerImpl) Close() {
+	if err := rc.gracefulClose(); err != nil {
+		rc.logger.Warn("graceful close a consumer fail, force close is applied")
+	}
+
+	// cancel is always need to be called, even graceful close is success.
+	// force close is applied by cancel context if graceful close is failed.
+	rc.cancel()
+	<-rc.resumingExitCh
+}
+
+// Done returns a channel which will be closed when scanner is finished or closed.
+func (rc *resumableConsumerImpl) Done() <-chan struct{} {
+	return rc.resumingExitCh
+}
+
+// Error returns the error of the Consumer.
+func (rc *resumableConsumerImpl) Error() error {
+	err := rc.consumeErr.Get()
+	if err == nil || errors.Is(err, errGracefulShutdown) {
+		return nil
+	}
+	return err
+}

internal/distributed/streaming/internal/consumer/consumer_test.go

@@ -0,0 +1,71 @@
+package consumer
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/errors"
+
+	"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_consumer"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/consumer"
+	"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/walimpls/impls/walimplstest"
+)
+
+func TestResumableConsumer(t *testing.T) {
+	i := 0
+	c := mock_consumer.NewMockConsumer(t)
+	ch := make(chan struct{})
+	c.EXPECT().Done().Return(ch)
+	c.EXPECT().Error().Return(errors.New("test"))
+	c.EXPECT().Close().Return()
+	rc := NewResumableConsumer(func(ctx context.Context, opts *handler.ConsumerOptions) (consumer.Consumer, error) {
+		if i == 0 {
+			i++
+			opts.MessageHandler.Handle(message.NewImmutableMesasge(
+				walimplstest.NewTestMessageID(123),
+				[]byte("payload"),
+				map[string]string{
+					"key": "value",
+					"_t":  "1",
+					"_tt": message.EncodeUint64(456),
+					"_v":  "1",
+					"_lc": walimplstest.NewTestMessageID(123).Marshal(),
+				}))
+			return c, nil
+		} else if i == 1 {
+			i++
+			return nil, errors.New("test")
+		}
+		newC := mock_consumer.NewMockConsumer(t)
+		newC.EXPECT().Done().Return(make(<-chan struct{}))
+		newC.EXPECT().Error().Return(errors.New("test"))
+		newC.EXPECT().Close().Return()
+		return newC, nil
+	}, &ConsumerOptions{
+		PChannel:      "test",
+		DeliverPolicy: options.DeliverPolicyAll(),
+		DeliverFilters: []options.DeliverFilter{
+			options.DeliverFilterTimeTickGT(1),
+		},
+		MessageHandler: message.ChanMessageHandler(make(chan message.ImmutableMessage, 2)),
+	})
+
+	select {
+	case <-rc.Done():
+		t.Error("should not be done")
+	case <-time.After(100 * time.Millisecond):
+	}
+	close(ch)
+	select {
+	case <-rc.Done():
+		t.Error("should not be done")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	rc.Close()
+	<-rc.Done()
+}

internal/distributed/streaming/internal/consumer/message_handler.go

@@ -0,0 +1,31 @@
+package consumer
+
+import (
+	"github.com/milvus-io/milvus/pkg/metrics"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+)
+
+// timeTickOrderMessageHandler is a message handler that will do metrics and record the last sent message id.
+type timeTickOrderMessageHandler struct {
+	inner                  message.Handler
+	lastConfirmedMessageID message.MessageID
+	lastTimeTick           uint64
+}
+
+func (mh *timeTickOrderMessageHandler) Handle(msg message.ImmutableMessage) {
+	lastConfirmedMessageID := msg.LastConfirmedMessageID()
+	timetick := msg.TimeTick()
+	messageSize := msg.EstimateSize()
+
+	mh.inner.Handle(msg)
+
+	mh.lastConfirmedMessageID = lastConfirmedMessageID
+	mh.lastTimeTick = timetick
+	// Do a metric here.
+	metrics.StreamingServiceClientConsumeBytes.WithLabelValues(paramtable.GetStringNodeID()).Observe(float64(messageSize))
+}
+
+func (mh *timeTickOrderMessageHandler) Close() {
+	mh.inner.Close()
+}

internal/distributed/streaming/internal/errs/error.go

@@ -0,0 +1,11 @@
+package errs
+
+import (
+	"github.com/cockroachdb/errors"
+)
+
+// All error in streamingservice package should be marked by streamingservice/errs package.
+var (
+	ErrClosed   = errors.New("closed")
+	ErrCanceled = errors.New("canceled")
+)

internal/distributed/streaming/internal/producer/producer.go

@@ -0,0 +1,199 @@
+package producer
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"github.com/cockroachdb/errors"
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/errs"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"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/log"
+	"github.com/milvus-io/milvus/pkg/metrics"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/util/syncutil"
+)
+
+var errGracefulShutdown = errors.New("graceful shutdown")
+
+// ProducerOptions is the options for creating a producer.
+type ProducerOptions struct {
+	PChannel string
+}
+
+// NewResumableProducer creates a new producer.
+// Provide an auto resuming producer.
+func NewResumableProducer(f factory, opts *ProducerOptions) *ResumableProducer {
+	ctx, cancel := context.WithCancel(context.Background())
+	p := &ResumableProducer{
+		ctx:            ctx,
+		cancel:         cancel,
+		stopResumingCh: make(chan struct{}),
+		resumingExitCh: make(chan struct{}),
+		lifetime:       lifetime.NewLifetime(lifetime.Working),
+		logger:         log.With(zap.String("pchannel", opts.PChannel)),
+		opts:           opts,
+
+		producer: newProducerWithResumingError(), // lazy initialized.
+		cond:     syncutil.NewContextCond(&sync.Mutex{}),
+		factory:  f,
+	}
+	go p.resumeLoop()
+	return p
+}
+
+// factory is a factory used to create a new underlying streamingnode producer.
+type factory func(ctx context.Context, opts *handler.ProducerOptions) (producer.Producer, error)
+
+// ResumableProducer is implementation for producing message to streaming service.
+// ResumableProducer select a right streaming node to produce automatically.
+// ResumableProducer will do automatic resume from stream broken and streaming node re-balance.
+// All error in these package should be marked by streaming/errs package.
+type ResumableProducer struct {
+	ctx            context.Context
+	cancel         context.CancelFunc
+	resumingExitCh chan struct{}
+	stopResumingCh chan struct{} // Q: why do not use ctx directly?
+	// A: cancel the ctx will cancel the underlying running producer.
+	// Use producer Close is better way to stop producer.
+
+	lifetime lifetime.Lifetime[lifetime.State]
+	logger   *log.MLogger
+	opts     *ProducerOptions
+
+	producer producerWithResumingError
+	cond     *syncutil.ContextCond
+
+	// factory is used to create a new producer.
+	factory factory
+}
+
+// Produce produce a new message to log service.
+func (p *ResumableProducer) Produce(ctx context.Context, msg message.MutableMessage) (msgID message.MessageID, err error) {
+	if p.lifetime.Add(lifetime.IsWorking) != nil {
+		return nil, errors.Wrapf(errs.ErrClosed, "produce on closed producer")
+	}
+	defer p.lifetime.Done()
+
+	for {
+		// get producer.
+		producerHandler, err := p.producer.GetProducerAfterAvailable(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		msgID, err := producerHandler.Produce(ctx, msg)
+		if err == nil {
+			return msgID, nil
+		}
+		// It's ok to stop retry if the error is canceled or deadline exceed.
+		if status.IsCanceled(err) {
+			return nil, errors.Mark(err, errs.ErrCanceled)
+		}
+	}
+}
+
+// resumeLoop is used to resume producer from error.
+func (p *ResumableProducer) resumeLoop() {
+	defer func() {
+		p.logger.Info("stop resuming")
+		close(p.resumingExitCh)
+	}()
+
+	for {
+		producer, err := p.createNewProducer()
+		p.producer.SwapProducer(producer, err)
+		if err != nil {
+			return
+		}
+
+		// Wait until the new producer is unavailable, trigger a new swap operation.
+		if err := p.waitUntilUnavailable(producer); err != nil {
+			p.producer.SwapProducer(nil, err)
+			return
+		}
+	}
+}
+
+// waitUntilUnavailable is used to wait until the producer is unavailable or context canceled.
+func (p *ResumableProducer) waitUntilUnavailable(producer handler.Producer) error {
+	// Mark as available.
+	metrics.StreamingServiceClientProducerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerAvailable).Inc()
+	defer metrics.StreamingServiceClientProducerTotal.WithLabelValues(paramtable.GetStringNodeID(), metrics.StreamingServiceClientProducerAvailable).Dec()
+
+	select {
+	case <-p.stopResumingCh:
+		return errGracefulShutdown
+	case <-p.ctx.Done():
+		return p.ctx.Err()
+	case <-producer.Available():
+		// Wait old producer unavailable, trigger a new resuming operation.
+		p.logger.Warn("producer encounter error, try to resume...")
+		return nil
+	}
+}
+
+// createNewProducer is used to open a new stream producer with backoff.
+func (p *ResumableProducer) createNewProducer() (producer.Producer, error) {
+	backoff := backoff.NewExponentialBackOff()
+	backoff.InitialInterval = 100 * time.Millisecond
+	backoff.MaxInterval = 2 * time.Second
+	for {
+		// Create a new producer.
+		// a underlying stream producer life time should be equal to the resumable producer.
+		// so ctx of resumable producer is passed to underlying stream producer creation.
+		producerHandler, err := p.factory(p.ctx, &handler.ProducerOptions{
+			PChannel: p.opts.PChannel,
+		})
+
+		// Can not resumable:
+		// 1. context canceled: the resumable producer is closed.
+		// 2. ErrClientClosed: the underlying handlerClient is closed.
+		if errors.IsAny(err, context.Canceled, handler.ErrClientClosed) {
+			return nil, err
+		}
+
+		// Otherwise, perform a resuming operation.
+		if err != nil {
+			nextBackoff := backoff.NextBackOff()
+			p.logger.Warn("create producer failed, retry...", zap.Error(err), zap.Duration("nextRetryInterval", nextBackoff))
+			time.Sleep(nextBackoff)
+			continue
+		}
+		return producerHandler, nil
+	}
+}
+
+// gracefulClose graceful close the producer.
+func (p *ResumableProducer) gracefulClose() error {
+	p.lifetime.SetState(lifetime.Stopped)
+	p.lifetime.Wait()
+	// close the stop resuming background to avoid create new producer.
+	close(p.stopResumingCh)
+
+	select {
+	case <-p.resumingExitCh:
+		return nil
+	case <-time.After(50 * time.Millisecond):
+		return context.DeadlineExceeded
+	}
+}
+
+// Close close the producer.
+func (p *ResumableProducer) Close() {
+	if err := p.gracefulClose(); err != nil {
+		p.logger.Warn("graceful close a producer fail, force close is applied")
+	}
+
+	// cancel is always need to be called, even graceful close is success.
+	// force close is applied by cancel context if graceful close is failed.
+	p.cancel()
+	<-p.resumingExitCh
+}

internal/distributed/streaming/internal/producer/producer_resuming.go

@@ -0,0 +1,57 @@
+package producer
+
+import (
+	"context"
+	"sync"
+
+	"github.com/cockroachdb/errors"
+
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/errs"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"github.com/milvus-io/milvus/pkg/util/syncutil"
+)
+
+// newProducerWithResumingError creates a new producer with resuming error.
+func newProducerWithResumingError() producerWithResumingError {
+	return producerWithResumingError{
+		cond: syncutil.NewContextCond(&sync.Mutex{}),
+	}
+}
+
+// producerWithResumingError is a producer that can be resumed.
+type producerWithResumingError struct {
+	cond     *syncutil.ContextCond
+	producer handler.Producer
+	err      error
+}
+
+// GetProducerAfterAvailable gets the producer after it is available.
+func (p *producerWithResumingError) GetProducerAfterAvailable(ctx context.Context) (handler.Producer, error) {
+	p.cond.L.Lock()
+	for p.err == nil && (p.producer == nil || !p.producer.IsAvailable()) {
+		if err := p.cond.Wait(ctx); err != nil {
+			return nil, errors.Mark(err, errs.ErrCanceled)
+		}
+	}
+	err := p.err
+	producer := p.producer
+
+	p.cond.L.Unlock()
+	if err != nil {
+		return nil, errors.Mark(err, errs.ErrClosed)
+	}
+	return producer, nil
+}
+
+// SwapProducer swaps the producer with a new one.
+func (p *producerWithResumingError) SwapProducer(producer handler.Producer, err error) {
+	p.cond.LockAndBroadcast()
+	oldProducer := p.producer
+	p.producer = producer
+	p.err = err
+	p.cond.L.Unlock()
+
+	if oldProducer != nil {
+		oldProducer.Close()
+	}
+}

internal/distributed/streaming/internal/producer/producer_test.go

@@ -0,0 +1,94 @@
+package producer
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/errs"
+	"github.com/milvus-io/milvus/internal/mocks/streamingnode/client/handler/mock_producer"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler/producer"
+	"github.com/milvus-io/milvus/pkg/mocks/streaming/util/mock_message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+)
+
+func TestResumableProducer(t *testing.T) {
+	p := mock_producer.NewMockProducer(t)
+	msgID := mock_message.NewMockMessageID(t)
+	p.EXPECT().Produce(mock.Anything, mock.Anything).Return(msgID, nil)
+	p.EXPECT().Close().Return()
+	ch := make(chan struct{})
+	p.EXPECT().Available().Return(ch)
+	p.EXPECT().IsAvailable().RunAndReturn(func() bool {
+		select {
+		case <-ch:
+			return false
+		default:
+			return true
+		}
+	})
+
+	i := 0
+	ch2 := make(chan struct{})
+	rp := NewResumableProducer(func(ctx context.Context, opts *handler.ProducerOptions) (producer.Producer, error) {
+		if i == 0 {
+			i++
+			return p, nil
+		} else if i == 1 {
+			i++
+			return nil, errors.New("test")
+		} else if i == 2 {
+			p := mock_producer.NewMockProducer(t)
+			msgID := mock_message.NewMockMessageID(t)
+			p.EXPECT().Produce(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, mm message.MutableMessage) (message.MessageID, error) {
+				if ctx.Err() != nil {
+					return nil, ctx.Err()
+				}
+				return msgID, nil
+			})
+			p.EXPECT().Close().Return()
+			p.EXPECT().Available().Return(ch2)
+			p.EXPECT().IsAvailable().RunAndReturn(func() bool {
+				select {
+				case <-ch2:
+					return false
+				default:
+					return true
+				}
+			})
+			i++
+			return p, nil
+		}
+		return nil, handler.ErrClientClosed
+	}, &ProducerOptions{
+		PChannel: "test",
+	})
+
+	msg := mock_message.NewMockMutableMessage(t)
+	id, err := rp.Produce(context.Background(), msg)
+	assert.NotNil(t, id)
+	assert.NoError(t, err)
+	close(ch)
+	id, err = rp.Produce(context.Background(), msg)
+	assert.NotNil(t, id)
+	assert.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	id, err = rp.Produce(ctx, msg)
+	assert.Nil(t, id)
+	assert.Error(t, err)
+	assert.True(t, errors.Is(err, errs.ErrCanceled))
+
+	// Test the underlying handler close.
+	close(ch2)
+	id, err = rp.Produce(context.Background(), msg)
+	assert.Nil(t, id)
+	assert.Error(t, err)
+	assert.True(t, errors.Is(err, errs.ErrClosed))
+	rp.Close()
+}

internal/distributed/streaming/streaming.go

@@ -0,0 +1,66 @@
+package streaming
+
+import (
+	"context"
+
+	clientv3 "go.etcd.io/etcd/client/v3"
+
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+)
+
+var singleton *walAccesserImpl = nil
+
+// Init initializes the wal accesser with the given etcd client.
+// should be called before any other operations.
+func Init(c *clientv3.Client) {
+	singleton = newWALAccesser(c)
+}
+
+// Release releases the resources of the wal accesser.
+func Release() {
+	if singleton != nil {
+		singleton.Close()
+	}
+}
+
+// WAL is the entrance to interact with the milvus write ahead log.
+func WAL() WALAccesser {
+	return singleton
+}
+
+type ReadOption struct {
+	// VChannel is the target vchannel to read.
+	VChannel 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
+}
+
+// Scanner is the interface for reading records from the wal.
+type Scanner interface {
+	// Done returns a channel which will be closed when scanner is finished or closed.
+	Done() <-chan struct{}
+
+	// Error returns the error of the scanner.
+	Error() error
+
+	// Close the scanner, release the underlying resources.
+	Close()
+}
+
+// WALAccesser is the interfaces to interact with the milvus write ahead log.
+type WALAccesser interface {
+	// Append writes a record to the log.
+	// !!! Append didn't promise the order of the message and atomic write.
+	Append(ctx context.Context, msgs ...message.MutableMessage) AppendResponses
+
+	// Read returns a scanner for reading records from the wal.
+	Read(ctx context.Context, opts ReadOption) Scanner
+}

internal/distributed/streaming/streaming_test.go

@@ -0,0 +1,100 @@
+package streaming_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
+	"github.com/milvus-io/milvus/internal/distributed/streaming"
+	kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
+	"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/walimpls/impls/pulsar"
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
+)
+
+const vChannel = "by-dev-rootcoord-dml_4"
+
+func TestMain(m *testing.M) {
+	paramtable.Init()
+	etcd, _ := kvfactory.GetEtcdAndPath()
+	streaming.Init(etcd)
+	defer streaming.Release()
+	m.Run()
+}
+
+func TestStreamingProduce(t *testing.T) {
+	t.Skip()
+	msg, _ := message.NewCreateCollectionMessageBuilderV1().
+		WithHeader(&message.CreateCollectionMessageHeader{
+			CollectionId: 1,
+			PartitionIds: []int64{1, 2, 3},
+		}).
+		WithBody(&msgpb.CreateCollectionRequest{
+			CollectionID: 1,
+		}).
+		WithVChannel(vChannel).
+		BuildMutable()
+	resp := streaming.WAL().Append(context.Background(), msg)
+	fmt.Printf("%+v\n", resp)
+
+	for i := 0; i < 1000; i++ {
+		time.Sleep(time.Millisecond * 1)
+		msg, _ := message.NewInsertMessageBuilderV1().
+			WithHeader(&message.InsertMessageHeader{
+				CollectionId: 1,
+			}).
+			WithBody(&msgpb.InsertRequest{
+				CollectionID: 1,
+			}).
+			WithVChannel(vChannel).
+			BuildMutable()
+		resp := streaming.WAL().Append(context.Background(), msg)
+		fmt.Printf("%+v\n", resp)
+	}
+
+	msg, _ = message.NewDropCollectionMessageBuilderV1().
+		WithHeader(&message.DropCollectionMessageHeader{
+			CollectionId: 1,
+		}).
+		WithBody(&msgpb.DropCollectionRequest{
+			CollectionID: 1,
+		}).
+		WithVChannel(vChannel).
+		BuildMutable()
+	resp = streaming.WAL().Append(context.Background(), msg)
+	fmt.Printf("%+v\n", resp)
+}
+
+func TestStreamingConsume(t *testing.T) {
+	t.Skip()
+	ch := make(message.ChanMessageHandler, 10)
+	s := streaming.WAL().Read(context.Background(), streaming.ReadOption{
+		VChannel:       vChannel,
+		DeliverPolicy:  options.DeliverPolicyAll(),
+		MessageHandler: ch,
+	})
+	defer func() {
+		s.Close()
+	}()
+
+	idx := 0
+	for {
+		time.Sleep(10 * time.Millisecond)
+		select {
+		case msg := <-ch:
+			fmt.Printf("msgID=%+v, tt=%d, lca=%+v, body=%s, idx=%d\n",
+				msg.MessageID(),
+				msg.TimeTick(),
+				msg.LastConfirmedMessageID(),
+				string(msg.Payload()),
+				idx,
+			)
+		case <-time.After(10 * time.Second):
+			return
+		}
+		idx++
+	}
+}

internal/distributed/streaming/wal.go

@@ -0,0 +1,121 @@
+package streaming
+
+import (
+	"context"
+	"sync"
+
+	clientv3 "go.etcd.io/etcd/client/v3"
+
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/consumer"
+	"github.com/milvus-io/milvus/internal/distributed/streaming/internal/producer"
+	"github.com/milvus-io/milvus/internal/streamingcoord/client"
+	"github.com/milvus-io/milvus/internal/streamingnode/client/handler"
+	"github.com/milvus-io/milvus/internal/util/streamingutil/status"
+	"github.com/milvus-io/milvus/pkg/streaming/util/message"
+	"github.com/milvus-io/milvus/pkg/streaming/util/options"
+	"github.com/milvus-io/milvus/pkg/util/funcutil"
+	"github.com/milvus-io/milvus/pkg/util/lifetime"
+)
+
+// newWALAccesser creates a new wal accesser.
+func newWALAccesser(c *clientv3.Client) *walAccesserImpl {
+	// Create a new streaming coord client.
+	streamingCoordClient := client.NewClient(c)
+	// Create a new streamingnode handler client.
+	handlerClient := handler.NewHandlerClient(streamingCoordClient.Assignment())
+	return &walAccesserImpl{
+		lifetime:                       lifetime.NewLifetime(lifetime.Working),
+		streamingCoordAssignmentClient: streamingCoordClient,
+		handlerClient:                  handlerClient,
+		producerMutex:                  sync.Mutex{},
+		producers:                      make(map[string]*producer.ResumableProducer),
+	}
+}
+
+// walAccesserImpl is the implementation of WALAccesser.
+type walAccesserImpl struct {
+	lifetime lifetime.Lifetime[lifetime.State]
+
+	// All services
+	streamingCoordAssignmentClient client.Client
+	handlerClient                  handler.HandlerClient
+
+	producerMutex sync.Mutex
+	producers     map[string]*producer.ResumableProducer
+}
+
+// Append writes a record to the log.
+func (w *walAccesserImpl) Append(ctx context.Context, msgs ...message.MutableMessage) AppendResponses {
+	if err := w.lifetime.Add(lifetime.IsWorking); err != nil {
+		err := status.NewOnShutdownError("wal accesser closed, %s", err.Error())
+		resp := newAppendResponseN(len(msgs))
+		resp.fillAllError(err)
+		return resp
+	}
+	defer w.lifetime.Done()
+
+	// If there is only one message, append it to the corresponding pchannel is ok.
+	if len(msgs) <= 1 {
+		return w.appendToPChannel(ctx, funcutil.ToPhysicalChannel(msgs[0].VChannel()), msgs...)
+	}
+	return w.dispatchByPChannel(ctx, msgs...)
+}
+
+// Read returns a scanner for reading records from the wal.
+func (w *walAccesserImpl) Read(_ context.Context, opts ReadOption) Scanner {
+	if err := w.lifetime.Add(lifetime.IsWorking); err != nil {
+		newErrScanner(status.NewOnShutdownError("wal accesser closed, %s", err.Error()))
+	}
+	defer w.lifetime.Done()
+
+	// TODO: optimize the consumer into pchannel level.
+	pchannel := funcutil.ToPhysicalChannel(opts.VChannel)
+	filters := append(opts.DeliverFilters, options.DeliverFilterVChannel(opts.VChannel))
+	rc := consumer.NewResumableConsumer(w.handlerClient.CreateConsumer, &consumer.ConsumerOptions{
+		PChannel:       pchannel,
+		DeliverPolicy:  opts.DeliverPolicy,
+		DeliverFilters: filters,
+		MessageHandler: opts.MessageHandler,
+	})
+	return rc
+}
+
+// Close closes all the wal accesser.
+func (w *walAccesserImpl) Close() {
+	w.lifetime.SetState(lifetime.Stopped)
+	w.lifetime.Wait()
+
+	w.producerMutex.Lock()
+	for _, p := range w.producers {
+		p.Close()
+	}
+	w.producerMutex.Unlock()
+
+	w.handlerClient.Close()
+	w.streamingCoordAssignmentClient.Close()
+}
+
+// newErrScanner creates a scanner that returns an error.
+func newErrScanner(err error) Scanner {
+	ch := make(chan struct{})
+	return errScanner{
+		closedCh: ch,
+		err:      err,
+	}
+}
+
+type errScanner struct {
+	closedCh chan struct{}
+	err      error
+}
+
+func (s errScanner) Done() <-chan struct{} {
+	return s.closedCh
+}
+
+func (s errScanner) Error() error {
+	return s.err
+}
+
+func (s errScanner) Close() {
+}

internal/distributed/streamingnode/OWNERS

@@ -0,0 +1,5 @@
+reviewers:
+  - chyezh
+
+approvers:
+  - maintainers

internal/mocks/streamingnode/server/mock_wal/mock_WAL.go

@@ -116,6 +116,49 @@ func (_c *MockWAL_AppendAsync_Call) RunAndReturn(run func(context.Context, messa
 	return _c
 }
 
+// Available provides a mock function with given fields:
+func (_m *MockWAL) 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
+}
+
+// MockWAL_Available_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Available'
+type MockWAL_Available_Call struct {
+	*mock.Call
+}
+
+// Available is a helper method to define mock.On call
+func (_e *MockWAL_Expecter) Available() *MockWAL_Available_Call {
+	return &MockWAL_Available_Call{Call: _e.mock.On("Available")}
+}
+
+func (_c *MockWAL_Available_Call) Run(run func()) *MockWAL_Available_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockWAL_Available_Call) Return(_a0 <-chan struct{}) *MockWAL_Available_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockWAL_Available_Call) RunAndReturn(run func() <-chan struct{}) *MockWAL_Available_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Channel provides a mock function with given fields:
 func (_m *MockWAL) Channel() types.PChannelInfo {
 	ret := _m.Called()
@@ -189,6 +232,47 @@ func (_c *MockWAL_Close_Call) RunAndReturn(run func()) *MockWAL_Close_Call {
 	return _c
 }
 
+// IsAvailable provides a mock function with given fields:
+func (_m *MockWAL) 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
+}
+
+// MockWAL_IsAvailable_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsAvailable'
+type MockWAL_IsAvailable_Call struct {
+	*mock.Call
+}
+
+// IsAvailable is a helper method to define mock.On call
+func (_e *MockWAL_Expecter) IsAvailable() *MockWAL_IsAvailable_Call {
+	return &MockWAL_IsAvailable_Call{Call: _e.mock.On("IsAvailable")}
+}
+
+func (_c *MockWAL_IsAvailable_Call) Run(run func()) *MockWAL_IsAvailable_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockWAL_IsAvailable_Call) Return(_a0 bool) *MockWAL_IsAvailable_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockWAL_IsAvailable_Call) RunAndReturn(run func() bool) *MockWAL_IsAvailable_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Read provides a mock function with given fields: ctx, deliverPolicy
 func (_m *MockWAL) Read(ctx context.Context, deliverPolicy wal.ReadOption) (wal.Scanner, error) {
 	ret := _m.Called(ctx, deliverPolicy)

internal/proto/streaming.proto

@@ -279,6 +279,7 @@ message CreateProducerResponse {
     // call at producer level.
 }
 
+// ProduceMessageResponse is the response of the ProduceMessage RPC.
 message ProduceMessageResponse {
     int64 request_id = 1;
     oneof response {

internal/streamingcoord/OWNERS

@@ -0,0 +1,5 @@
+reviewers:
+  - chyezh
+
+approvers:
+  - maintainers

internal/streamingcoord/server/balancer/balancer_impl.go

@@ -140,7 +140,7 @@ func (b *balancerImpl) execute() {
 				// balancer is closed.
 				return
 			}
-			b.logger.Warn("fail to apply balance, start a backoff...")
+			b.logger.Warn("fail to apply balance, start a backoff...", zap.Error(err))
 			balanceTimer.EnableBackoff()
 			continue
 		}

internal/streamingnode/OWNERS

@@ -0,0 +1,5 @@
+reviewers:
+  - chyezh
+
+approvers:
+  - maintainers

internal/streamingnode/client/handler/handler_client.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"time"
 
+	"github.com/cockroachdb/errors"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -23,12 +24,14 @@ import (
 	"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)
+var (
+	_               HandlerClient = (*handlerClientImpl)(nil)
+	ErrClientClosed               = errors.New("handler client is closed")
+)
 
 type (
 	Producer = producer.Producer
@@ -94,15 +97,13 @@ func NewHandlerClient(w types.AssignmentDiscoverWatcher) HandlerClient {
 	})
 	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,
+		lifetime:         lifetime.NewLifetime(lifetime.Working),
+		service:          lazygrpc.WithServiceCreator(conn, streamingpb.NewStreamingNodeHandlerServiceClient),
+		rb:               rb,
+		watcher:          watcher,
+		rebalanceTrigger: w,
+		newProducer:      producer.CreateProducer,
+		newConsumer:      consumer.CreateConsumer,
 	}
 }
 

internal/streamingnode/client/handler/handler_client_impl.go

@@ -19,28 +19,24 @@ import (
 	"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)
+	lifetime         lifetime.Lifetime[lifetime.State]
+	service          lazygrpc.Service[streamingpb.StreamingNodeHandlerServiceClient]
+	rb               resolver.Builder
+	watcher          assignment.Watcher
+	rebalanceTrigger types.AssignmentRebalanceTrigger
+	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")
+		return nil, ErrClientClosed
 	}
 	defer hc.lifetime.Done()
 
@@ -50,7 +46,9 @@ func (hc *handlerClientImpl) CreateProducer(ctx context.Context, opts *ProducerO
 		if err != nil {
 			return nil, err
 		}
-		return hc.createOrGetSharedProducer(ctx, &producer.ProducerOptions{Assignment: assign}, handlerService)
+		return hc.newProducer(ctx, &producer.ProducerOptions{
+			Assignment: assign,
+		}, handlerService)
 	})
 	if err != nil {
 		return nil, err
@@ -61,7 +59,7 @@ func (hc *handlerClientImpl) CreateProducer(ctx context.Context, opts *ProducerO
 // 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")
+		return nil, ErrClientClosed
 	}
 	defer hc.lifetime.Done()
 
@@ -134,60 +132,6 @@ func (hc *handlerClientImpl) waitForNextBackoff(ctx context.Context, pchannel st
 	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)

internal/streamingnode/client/handler/handler_client_test.go

@@ -23,9 +23,7 @@ import (
 	"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) {
@@ -43,8 +41,6 @@ func TestHandlerClient(t *testing.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() {})
@@ -54,13 +50,11 @@ func TestHandlerClient(t *testing.T) {
 
 	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](),
+		lifetime:         lifetime.NewLifetime(lifetime.Working),
+		service:          service,
+		rb:               rb,
+		watcher:          w,
+		rebalanceTrigger: rebalanceTrigger,
 		newProducer: func(ctx context.Context, opts *producer.ProducerOptions, handler streamingpb.StreamingNodeHandlerServiceClient) (Producer, error) {
 			if pK == 0 {
 				pK++
@@ -90,8 +84,6 @@ func TestHandlerClient(t *testing.T) {
 	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)
@@ -122,10 +114,12 @@ func TestHandlerClient(t *testing.T) {
 	handler.Close()
 	producer, err = handler.CreateProducer(ctx, nil)
 	assert.Error(t, err)
+	assert.ErrorIs(t, err, ErrClientClosed)
 	assert.Nil(t, producer)
 
 	consumer, err = handler.CreateConsumer(ctx, nil)
 	assert.Error(t, err)
+	assert.ErrorIs(t, err, ErrClientClosed)
 	assert.Nil(t, consumer)
 }
 

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

@@ -69,6 +69,7 @@ func CreateProducer(
 		pendingRequests:  sync.Map{},
 		requestCh:        make(chan *produceRequest),
 		sendExitCh:       make(chan struct{}),
+		recvExitCh:       make(chan struct{}),
 		finishedCh:       make(chan struct{}),
 	}
 
@@ -104,6 +105,7 @@ type producerImpl struct {
 	pendingRequests sync.Map
 	requestCh       chan *produceRequest
 	sendExitCh      chan struct{}
+	recvExitCh      chan struct{}
 	finishedCh      chan struct{}
 }
 
@@ -143,7 +145,9 @@ func (p *producerImpl) Produce(ctx context.Context, msg message.MutableMessage)
 		return nil, ctx.Err()
 	case p.requestCh <- req:
 	case <-p.sendExitCh:
-		return nil, status.NewInner("producer stream client is closed")
+		return nil, status.NewInner("producer send arm is closed")
+	case <-p.recvExitCh:
+		return nil, status.NewInner("producer recv arm is closed")
 	}
 
 	// Wait for the response from server or context timeout.
@@ -234,21 +238,28 @@ func (p *producerImpl) sendLoop() (err error) {
 		}
 	}()
 
-	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
+	for {
+		select {
+		case <-p.recvExitCh:
+			return errors.New("recv arm of stream closed")
+		case req, ok := <-p.requestCh:
+			if !ok {
+				// all message has been sent, sent close response.
+				return p.grpcStreamClient.SendClose()
+			}
+			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.
@@ -259,6 +270,7 @@ func (p *producerImpl) recvLoop() (err error) {
 			return
 		}
 		p.logger.Info("recv arm of stream closed")
+		close(p.recvExitCh)
 	}()
 
 	for {

internal/streamingnode/client/handler/shared_producer.go

@@ -1,51 +0,0 @@
-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/server/service/handler/producer/produce_server.go

@@ -90,8 +90,22 @@ func (p *ProduceServer) sendLoop() (err error) {
 		}
 		p.logger.Info("send arm of stream closed")
 	}()
+	available := p.wal.Available()
+	var appendWGDoneChan <-chan struct{}
+
 	for {
 		select {
+		case <-available:
+			// If the wal is not available any more, we should stop sending message, and close the server.
+			// appendWGDoneChan make a graceful shutdown for those case.
+			available = nil
+			appendWGDoneChan = p.getWaitAppendChan()
+		case <-appendWGDoneChan:
+			// All pending append request has been finished, we can close the streaming server now.
+			// Recv arm will be closed by context cancel of stream server.
+			// Send an unavailable response to ask client to release resource.
+			p.produceServer.SendClosed()
+			return errors.New("send loop is stopped for close of wal")
 		case resp, ok := <-p.produceMessageCh:
 			if !ok {
 				// all message has been sent, sent close response.
@@ -107,6 +121,15 @@ func (p *ProduceServer) sendLoop() (err error) {
 	}
 }
 
+// getWaitAppendChan returns the channel that can be used to wait for the append operation.
+func (p *ProduceServer) getWaitAppendChan() <-chan struct{} {
+	ch := make(chan struct{})
+	go func() {
+		close(ch)
+	}()
+	return p.wal.Available()
+}
+
 // recvLoop receives the message from client.
 func (p *ProduceServer) recvLoop() (err error) {
 	defer func() {
@@ -142,11 +165,19 @@ func (p *ProduceServer) recvLoop() (err error) {
 
 // handleProduce handles the produce message request.
 func (p *ProduceServer) handleProduce(req *streamingpb.ProduceMessageRequest) {
+	// Stop handling if the wal is not available any more.
+	// The counter of  appendWG will never increased.
+	if !p.wal.IsAvailable() {
+		return
+	}
+
+	p.appendWG.Add(1)
 	p.logger.Debug("recv produce message from client", zap.Int64("requestID", req.RequestId))
 	msg := message.NewMutableMessage(req.GetMessage().GetPayload(), req.GetMessage().GetProperties())
 	if err := p.validateMessage(msg); err != nil {
 		p.logger.Warn("produce message validation failed", zap.Int64("requestID", req.RequestId), zap.Error(err))
 		p.sendProduceResult(req.RequestId, nil, err)
+		p.appendWG.Done()
 		return
 	}
 
@@ -154,7 +185,6 @@ func (p *ProduceServer) handleProduce(req *streamingpb.ProduceMessageRequest) {
 	// Concurrent append request can be executed concurrently.
 	messageSize := msg.EstimateSize()
 	now := time.Now()
-	p.appendWG.Add(1)
 	p.wal.AppendAsync(p.produceServer.Context(), msg, func(id message.MessageID, err error) {
 		defer func() {
 			p.appendWG.Done()

internal/streamingnode/server/service/handler/producer/produce_server_test.go

@@ -91,7 +91,11 @@ func TestProduceSendArm(t *testing.T) {
 		return errors.New("send failure")
 	})
 
+	wal := mock_wal.NewMockWAL(t)
+	wal.EXPECT().Available().Return(make(<-chan struct{}))
+
 	p := &ProduceServer{
+		wal: wal,
 		produceServer: &produceGrpcServerHelper{
 			StreamingNodeHandlerService_ProduceServer: grpcProduceServer,
 		},
@@ -122,6 +126,7 @@ func TestProduceSendArm(t *testing.T) {
 
 	// test send arm failure
 	p = &ProduceServer{
+		wal: wal,
 		produceServer: &produceGrpcServerHelper{
 			StreamingNodeHandlerService_ProduceServer: grpcProduceServer,
 		},
@@ -152,6 +157,7 @@ func TestProduceSendArm(t *testing.T) {
 
 	// test send arm failure
 	p = &ProduceServer{
+		wal: wal,
 		produceServer: &produceGrpcServerHelper{
 			StreamingNodeHandlerService_ProduceServer: grpcProduceServer,
 		},
@@ -191,6 +197,7 @@ func TestProduceServerRecvArm(t *testing.T) {
 		msgID := walimplstest.NewTestMessageID(1)
 		f(msgID, nil)
 	})
+	l.EXPECT().IsAvailable().Return(true)
 
 	p := &ProduceServer{
 		wal: l,

internal/streamingnode/server/wal/adaptor/wal_adaptor.go

@@ -126,6 +126,16 @@ func (w *walAdaptorImpl) Read(ctx context.Context, opts wal.ReadOption) (wal.Sca
 	return s, nil
 }
 
+// IsAvailable returns whether the wal is available.
+func (w *walAdaptorImpl) IsAvailable() bool {
+	return !w.lifetime.IsClosed()
+}
+
+// Available returns a channel that will be closed when the wal is shut down.
+func (w *walAdaptorImpl) Available() <-chan struct{} {
+	return w.lifetime.CloseCh()
+}
+
 // Close overrides Scanner Close function.
 func (w *walAdaptorImpl) Close() {
 	w.lifetime.SetState(lifetime.Stopped)

internal/streamingnode/server/wal/adaptor/wal_test.go

@@ -190,10 +190,12 @@ func (f *testOneWALFramework) testSendCreateCollection(ctx context.Context, w wa
 			CollectionId: 1,
 			PartitionIds: []int64{2},
 		}).
-		WithBody(&msgpb.CreateCollectionRequest{}).BuildMutable()
+		WithBody(&msgpb.CreateCollectionRequest{}).
+		WithVChannel("v1").
+		BuildMutable()
 	assert.NoError(f.t, err)
 
-	msgID, err := w.Append(ctx, createMsg.WithVChannel("v1"))
+	msgID, err := w.Append(ctx, createMsg)
 	assert.NoError(f.t, err)
 	assert.NotNil(f.t, msgID)
 }
@@ -204,10 +206,12 @@ func (f *testOneWALFramework) testSendDropCollection(ctx context.Context, w wal.
 		WithHeader(&message.DropCollectionMessageHeader{
 			CollectionId: 1,
 		}).
-		WithBody(&msgpb.DropCollectionRequest{}).BuildMutable()
+		WithBody(&msgpb.DropCollectionRequest{}).
+		WithVChannel("v1").
+		BuildMutable()
 	assert.NoError(f.t, err)
 
-	msgID, err := w.Append(ctx, dropMsg.WithVChannel("v1"))
+	msgID, err := w.Append(ctx, dropMsg)
 	assert.NoError(f.t, err)
 	assert.NotNil(f.t, msgID)
 }

internal/streamingnode/server/wal/interceptors/segment/manager/seal_queue.go

@@ -182,6 +182,7 @@ func (m *sealQueue) sendFlushMessageIntoWAL(ctx context.Context, collectionID in
 func (m *sealQueue) createNewFlushMessage(collectionID int64, vchannel string, segmentIDs []int64) (message.MutableMessage, error) {
 	// Create a flush message.
 	msg, err := message.NewFlushMessageBuilderV1().
+		WithVChannel(vchannel).
 		WithHeader(&message.FlushMessageHeader{}).
 		WithBody(&message.FlushMessagePayload{
 			CollectionId: collectionID,
@@ -190,6 +191,5 @@ func (m *sealQueue) createNewFlushMessage(collectionID int64, vchannel string, s
 	if err != nil {
 		return nil, errors.Wrap(err, "at create new flush message")
 	}
-	msg.WithVChannel(vchannel)
 	return msg, nil
 }

internal/streamingnode/server/wal/interceptors/timetick/timetick_message.go

@@ -20,7 +20,9 @@ func newTimeTickMsg(ts uint64, sourceID int64) (message.MutableMessage, error) {
 				commonpbutil.WithTimeStamp(ts),
 				commonpbutil.WithSourceID(sourceID),
 			),
-		}).BuildMutable()
+		}).
+		WithBroadcast().
+		BuildMutable()
 	if err != nil {
 		return nil, err
 	}

internal/streamingnode/server/wal/wal.go

@@ -24,6 +24,12 @@ type WAL interface {
 	// Read returns a scanner for reading records from the wal.
 	Read(ctx context.Context, deliverPolicy ReadOption) (Scanner, error)
 
+	// Available return a channel that will be closed when the wal is available.
+	Available() <-chan struct{}
+
+	// IsAvailable returns if the wal is available.
+	IsAvailable() bool
+
 	// Close closes the wal instance.
 	Close()
 }

internal/streamingnode/server/walmanager/manager_impl.go

@@ -3,6 +3,7 @@ package walmanager
 import (
 	"context"
 
+	"github.com/cockroachdb/errors"
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus/internal/streamingnode/server/wal"
@@ -29,7 +30,7 @@ func OpenManager() (Manager, error) {
 // newManager create a wal manager.
 func newManager(opener wal.Opener) Manager {
 	return &managerImpl{
-		lifetime: lifetime.NewLifetime(lifetime.Working),
+		lifetime: lifetime.NewLifetime(managerOpenable | managerRemoveable | managerGetable),
 		wltMap:   typeutil.NewConcurrentMap[string, *walLifetime](),
 		opener:   opener,
 	}
@@ -37,7 +38,7 @@ func newManager(opener wal.Opener) Manager {
 
 // All management operation for a wal will be serialized with order of term.
 type managerImpl struct {
-	lifetime lifetime.Lifetime[lifetime.State]
+	lifetime lifetime.Lifetime[managerState]
 
 	wltMap *typeutil.ConcurrentMap[string, *walLifetime]
 	opener wal.Opener // wal allocator
@@ -46,8 +47,8 @@ type managerImpl struct {
 // Open opens a wal instance for the channel on this Manager.
 func (m *managerImpl) Open(ctx context.Context, channel types.PChannelInfo) (err error) {
 	// reject operation if manager is closing.
-	if m.lifetime.Add(lifetime.IsWorking) != nil {
-		return status.NewOnShutdownError("wal manager is closed")
+	if err := m.lifetime.Add(isOpenable); err != nil {
+		return status.NewOnShutdownError("wal manager is closed, %s", err.Error())
 	}
 	defer func() {
 		m.lifetime.Done()
@@ -64,8 +65,8 @@ func (m *managerImpl) Open(ctx context.Context, channel types.PChannelInfo) (err
 // Remove removes the wal instance for the channel.
 func (m *managerImpl) Remove(ctx context.Context, channel types.PChannelInfo) (err error) {
 	// reject operation if manager is closing.
-	if m.lifetime.Add(lifetime.IsWorking) != nil {
-		return status.NewOnShutdownError("wal manager is closed")
+	if err := m.lifetime.Add(isRemoveable); err != nil {
+		return status.NewOnShutdownError("wal manager is closed, %s", err.Error())
 	}
 	defer func() {
 		m.lifetime.Done()
@@ -83,8 +84,8 @@ func (m *managerImpl) Remove(ctx context.Context, channel types.PChannelInfo) (e
 // Return nil if the wal instance is not found.
 func (m *managerImpl) GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, error) {
 	// reject operation if manager is closing.
-	if m.lifetime.Add(lifetime.IsWorking) != nil {
-		return nil, status.NewOnShutdownError("wal manager is closed")
+	if err := m.lifetime.Add(isGetable); err != nil {
+		return nil, status.NewOnShutdownError("wal manager is closed, %s", err)
 	}
 	defer m.lifetime.Done()
 
@@ -103,8 +104,8 @@ func (m *managerImpl) GetAvailableWAL(channel types.PChannelInfo) (wal.WAL, erro
 // GetAllAvailableChannels returns all available channel info.
 func (m *managerImpl) GetAllAvailableChannels() ([]types.PChannelInfo, error) {
 	// reject operation if manager is closing.
-	if m.lifetime.Add(lifetime.IsWorking) != nil {
-		return nil, status.NewOnShutdownError("wal manager is closed")
+	if err := m.lifetime.Add(isGetable); err != nil {
+		return nil, status.NewOnShutdownError("wal manager is closed, %s", err)
 	}
 	defer m.lifetime.Done()
 
@@ -122,15 +123,16 @@ func (m *managerImpl) GetAllAvailableChannels() ([]types.PChannelInfo, error) {
 
 // Close these manager and release all managed WAL.
 func (m *managerImpl) Close() {
-	m.lifetime.SetState(lifetime.Stopped)
+	m.lifetime.SetState(managerRemoveable)
 	m.lifetime.Wait()
-	m.lifetime.Close()
-
 	// close all underlying walLifetime.
 	m.wltMap.Range(func(channel string, wlt *walLifetime) bool {
 		wlt.Close()
 		return true
 	})
+	m.lifetime.SetState(managerStopped)
+	m.lifetime.Wait()
+	m.lifetime.Close()
 
 	// close all underlying wal instance by allocator if there's resource leak.
 	m.opener.Close()
@@ -151,3 +153,33 @@ func (m *managerImpl) getWALLifetime(channel string) *walLifetime {
 	}
 	return wlt
 }
+
+type managerState int32
+
+const (
+	managerStopped    managerState = 0
+	managerOpenable   managerState = 0x1
+	managerRemoveable managerState = 0x1 << 1
+	managerGetable    managerState = 0x1 << 2
+)
+
+func isGetable(state managerState) error {
+	if state&managerGetable != 0 {
+		return nil
+	}
+	return errors.New("wal manager can not do get operation")
+}
+
+func isRemoveable(state managerState) error {
+	if state&managerRemoveable != 0 {
+		return nil
+	}
+	return errors.New("wal manager can not do remove operation")
+}
+
+func isOpenable(state managerState) error {
+	if state&managerOpenable != 0 {
+		return nil
+	}
+	return errors.New("wal manager can not do open operation")
+}

internal/util/streamingutil/OWNERS

@@ -0,0 +1,5 @@
+reviewers:
+  - chyezh
+
+approvers:
+  - maintainers

pkg/metrics/streaming_service_metrics.go

@@ -37,7 +37,7 @@ var (
 		Name:    "produce_bytes",
 		Help:    "Bytes of produced message",
 		Buckets: bytesBuckets,
-	})
+	}, statusLabelName)
 
 	StreamingServiceClientConsumeBytes = newStreamingServiceClientHistogramVec(prometheus.HistogramOpts{
 		Name:    "consume_bytes",

pkg/mocks/streaming/util/mock_message/mock_MutableMessage.go

@@ -232,6 +232,88 @@ func (_c *MockMutableMessage_Properties_Call) RunAndReturn(run func() message.RP
 	return _c
 }
 
+// TimeTick provides a mock function with given fields:
+func (_m *MockMutableMessage) TimeTick() uint64 {
+	ret := _m.Called()
+
+	var r0 uint64
+	if rf, ok := ret.Get(0).(func() uint64); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(uint64)
+	}
+
+	return r0
+}
+
+// MockMutableMessage_TimeTick_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'TimeTick'
+type MockMutableMessage_TimeTick_Call struct {
+	*mock.Call
+}
+
+// TimeTick is a helper method to define mock.On call
+func (_e *MockMutableMessage_Expecter) TimeTick() *MockMutableMessage_TimeTick_Call {
+	return &MockMutableMessage_TimeTick_Call{Call: _e.mock.On("TimeTick")}
+}
+
+func (_c *MockMutableMessage_TimeTick_Call) Run(run func()) *MockMutableMessage_TimeTick_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockMutableMessage_TimeTick_Call) Return(_a0 uint64) *MockMutableMessage_TimeTick_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockMutableMessage_TimeTick_Call) RunAndReturn(run func() uint64) *MockMutableMessage_TimeTick_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// VChannel provides a mock function with given fields:
+func (_m *MockMutableMessage) VChannel() 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
+}
+
+// MockMutableMessage_VChannel_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'VChannel'
+type MockMutableMessage_VChannel_Call struct {
+	*mock.Call
+}
+
+// VChannel is a helper method to define mock.On call
+func (_e *MockMutableMessage_Expecter) VChannel() *MockMutableMessage_VChannel_Call {
+	return &MockMutableMessage_VChannel_Call{Call: _e.mock.On("VChannel")}
+}
+
+func (_c *MockMutableMessage_VChannel_Call) Run(run func()) *MockMutableMessage_VChannel_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockMutableMessage_VChannel_Call) Return(_a0 string) *MockMutableMessage_VChannel_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockMutableMessage_VChannel_Call) RunAndReturn(run func() string) *MockMutableMessage_VChannel_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Version provides a mock function with given fields:
 func (_m *MockMutableMessage) Version() message.Version {
 	ret := _m.Called()

pkg/streaming/util/message/builder.go

@@ -1,7 +1,6 @@
 package message
 
 import (
-	"fmt"
 	"reflect"
 
 	"github.com/cockroachdb/errors"
@@ -71,6 +70,7 @@ type mutableMesasgeBuilder[H proto.Message, B proto.Message] struct {
 	header     H
 	body       B
 	properties propertiesImpl
+	broadcast  bool
 }
 
 // WithMessageHeader creates a new builder with determined message type.
@@ -85,12 +85,24 @@ func (b *mutableMesasgeBuilder[H, B]) WithBody(body B) *mutableMesasgeBuilder[H,
 	return b
 }
 
+// WithVChannel creates a new builder with virtual channel.
+func (b *mutableMesasgeBuilder[H, B]) WithVChannel(vchannel string) *mutableMesasgeBuilder[H, B] {
+	if b.broadcast {
+		panic("a broadcast message cannot hold vchannel")
+	}
+	b.WithProperty(messageVChannel, vchannel)
+	return b
+}
+
+// WithBroadcast creates a new builder with broadcast property.
+func (b *mutableMesasgeBuilder[H, B]) WithBroadcast() *mutableMesasgeBuilder[H, B] {
+	b.broadcast = true
+	return b
+}
+
 // WithProperty creates a new builder with message property.
 // A key started with '_' is reserved for streaming system, should never used at user of client.
 func (b *mutableMesasgeBuilder[H, B]) WithProperty(key string, val string) *mutableMesasgeBuilder[H, B] {
-	if b.properties.Exist(key) {
-		panic(fmt.Sprintf("message builder already set property field, key = %s", key))
-	}
 	b.properties.Set(key, val)
 	return b
 }
@@ -115,6 +127,9 @@ func (b *mutableMesasgeBuilder[H, B]) BuildMutable() (MutableMessage, error) {
 	if reflect.ValueOf(b.body).IsNil() {
 		panic("message builder not ready for body field")
 	}
+	if !b.broadcast && !b.properties.Exist(messageVChannel) {
+		panic("a non broadcast message builder not ready for vchannel field")
+	}
 
 	// setup header.
 	sp, err := EncodeProto(b.header)

pkg/streaming/util/message/message.go

@@ -52,10 +52,6 @@ type MutableMessage interface {
 	// !!! preserved for streaming system internal usage, don't call it outside of log system.
 	WithTimeTick(tt uint64) MutableMessage
 
-	// WithVChannel sets the virtual channel of current message.
-	// !!! preserved for streaming system internal usage, don't call it outside of log system.
-	WithVChannel(vChannel string) MutableMessage
-
 	// IntoImmutableMessage converts the mutable message to immutable message.
 	IntoImmutableMessage(msgID MessageID) ImmutableMessage
 }
@@ -83,13 +79,7 @@ type ImmutableMessage interface {
 type specializedMutableMessage[H proto.Message, B proto.Message] interface {
 	BasicMessage
 
-	// VChannel returns the vchannel of the message.
-	VChannel() string
-
-	// TimeTick returns the time tick of the message.
-	TimeTick() uint64
-
-	// Header returns the message header.
+	// MessageHeader returns the message header.
 	// Modifications to the returned header will be reflected in the message.
 	Header() H
 

pkg/streaming/util/message/message_builder_test.go

@@ -18,6 +18,7 @@ func TestMessage(t *testing.T) {
 	mutableMessage, err := b.WithHeader(&message.TimeTickMessageHeader{}).
 		WithProperties(map[string]string{"key": "value"}).
 		WithProperty("key2", "value2").
+		WithVChannel("v1").
 		WithBody(&msgpb.TimeTickMsg{}).BuildMutable()
 	assert.NoError(t, err)
 
@@ -31,7 +32,7 @@ func TestMessage(t *testing.T) {
 	assert.Equal(t, "value", v)
 	assert.True(t, ok)
 	assert.Equal(t, message.MessageTypeTimeTick, mutableMessage.MessageType())
-	assert.Equal(t, 30, mutableMessage.EstimateSize())
+	assert.Equal(t, 35, mutableMessage.EstimateSize())
 	mutableMessage.WithTimeTick(123)
 	v, ok = mutableMessage.Properties().Get("_tt")
 	assert.True(t, ok)
@@ -47,7 +48,6 @@ func TestMessage(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, v, "lcMsgID")
 
-	mutableMessage.WithVChannel("v1")
 	v, ok = mutableMessage.Properties().Get("_vc")
 	assert.True(t, ok)
 	assert.Equal(t, "v1", v)

pkg/streaming/util/message/specialized_message_test.go

@@ -12,6 +12,7 @@ import (
 
 func TestAsSpecializedMessage(t *testing.T) {
 	m, err := message.NewInsertMessageBuilderV1().
+		WithVChannel("v1").
 		WithHeader(&message.InsertMessageHeader{
 			CollectionId: 1,
 			Partitions: []*message.PartitionSegmentAssignment{

pkg/streaming/util/message/test_case.go

@@ -91,11 +91,12 @@ func CreateTestInsertMessage(t *testing.T, segmentID int64, totalRows int, timet
 			RowIDs:         rowIDs,
 			Timestamps:     timestamps,
 			NumRows:        uint64(totalRows),
-		}).BuildMutable()
+		}).
+		WithVChannel("v1").
+		BuildMutable()
 	if err != nil {
 		panic(err)
 	}
-	msg.WithVChannel("v1")
 	msg.WithTimeTick(timetick)
 	msg.WithLastConfirmed(messageID)
 	return msg
@@ -122,9 +123,9 @@ func CreateTestCreateCollectionMessage(t *testing.T, collectionID int64, timetic
 	msg, err := NewCreateCollectionMessageBuilderV1().
 		WithHeader(header).
 		WithBody(payload).
+		WithVChannel("v1").
 		BuildMutable()
 	assert.NoError(t, err)
-	msg.WithVChannel("v1")
 	msg.WithTimeTick(timetick)
 	msg.WithLastConfirmed(messageID)
 	return msg
@@ -149,10 +150,11 @@ func CreateTestEmptyInsertMesage(msgID int64, extraProperties map[string]string)
 				MsgID:   msgID,
 			},
 		}).
+		WithVChannel("v1").
 		WithProperties(extraProperties).
 		BuildMutable()
 	if err != nil {
 		panic(err)
 	}
-	return msg.WithVChannel("v1")
+	return msg
 }

pkg/streaming/walimpls/test_framework.go

@@ -248,12 +248,16 @@ func (f *testOneWALImplsFramework) testAppend(ctx context.Context, w WALImpls) (
 		"const": "t",
 		"term":  strconv.FormatInt(int64(f.term), 10),
 	}
-	msg, err := message.NewTimeTickMessageBuilderV1().WithHeader(&message.TimeTickMessageHeader{}).WithBody(&msgpb.TimeTickMsg{
-		Base: &commonpb.MsgBase{
-			MsgType: commonpb.MsgType_TimeTick,
-			MsgID:   int64(f.messageCount - 1),
-		},
-	}).WithProperties(properties).BuildMutable()
+	msg, err := message.NewTimeTickMessageBuilderV1().
+		WithHeader(&message.TimeTickMessageHeader{}).
+		WithBody(&msgpb.TimeTickMsg{
+			Base: &commonpb.MsgBase{
+				MsgType: commonpb.MsgType_TimeTick,
+				MsgID:   int64(f.messageCount - 1),
+			},
+		}).
+		WithVChannel("v1").
+		WithProperties(properties).BuildMutable()
 	assert.NoError(f.t, err)
 
 	id, err := w.Append(ctx, msg)

scripts/run_go_unittest.sh

@@ -170,6 +170,7 @@ function test_streaming()
 go test -race -cover -tags dynamic,test "${MILVUS_DIR}/streamingcoord/..." -failfast -count=1 -ldflags="-r ${RPATH}"
 go test -race -cover -tags dynamic,test "${MILVUS_DIR}/streamingnode/..." -failfast -count=1 -ldflags="-r ${RPATH}"
 go test -race -cover -tags dynamic,test "${MILVUS_DIR}/util/streamingutil/..." -failfast -count=1 -ldflags="-r ${RPATH}"
+go test -race -cover -tags dynamic,test "${MILVUS_DIR}/distributed/streaming/..." -failfast -count=1 -ldflags="-r ${RPATH}"
 pushd pkg
 go test -race -cover -tags dynamic,test "${PKG_DIR}/streaming/..." -failfast -count=1  -ldflags="-r ${RPATH}"
 popd