chore: import upstream snapshot with attribution

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

@@ -0,0 +1,22 @@
+package assignment
+
+import (
+	"context"
+
+	"github.com/milvus-io/milvus/pkg/v3/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,106 @@
+package assignment
+
+import (
+	"context"
+	"sync"
+
+	"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/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/v3/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() {
+	mlog.Info(w.ctx, "assignment watcher start")
+	var err error
+	defer func() {
+		// error can be ignored here, so use info level log here.
+		mlog.Info(w.ctx, "assignment watcher close", mlog.Err(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/v3/streaming/util/types"
+	"github.com/milvus-io/milvus/pkg/v3/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)
+}