package snmanager import ( "context" "sync" "github.com/cockroachdb/errors" "go.uber.org/zap" "github.com/milvus-io/milvus/internal/streamingcoord/server/balancer" "github.com/milvus-io/milvus/pkg/v2/log" "github.com/milvus-io/milvus/pkg/v2/streaming/util/types" "github.com/milvus-io/milvus/pkg/v2/util/funcutil" "github.com/milvus-io/milvus/pkg/v2/util/syncutil" "github.com/milvus-io/milvus/pkg/v2/util/typeutil" ) var StaticStreamingNodeManager = newStreamingNodeManager() func newStreamingNodeManager() *StreamingNodeManager { snm := &StreamingNodeManager{ notifier: syncutil.NewAsyncTaskNotifier[struct{}](), balancer: syncutil.NewFuture[balancer.Balancer](), cond: syncutil.NewContextCond(&sync.Mutex{}), latestAssignments: make(map[string]types.PChannelInfoAssigned), streamingNodes: typeutil.NewUniqueSet(), nodeChangedNotifier: syncutil.NewVersionedNotifier(), } go snm.execute() return snm } // StreamingNodeManager is a manager for manage the querynode that embedded into streaming node. // StreamingNodeManager is exclusive with ResourceManager. type StreamingNodeManager struct { notifier *syncutil.AsyncTaskNotifier[struct{}] balancer *syncutil.Future[balancer.Balancer] // The coord is merged after 2.6, so we don't need to make distribution safe. cond *syncutil.ContextCond latestAssignments map[string]types.PChannelInfoAssigned // The latest assignments info got from streaming coord balance module. streamingNodes typeutil.UniqueSet nodeChangedNotifier *syncutil.VersionedNotifier // used to notify that node in streaming node manager has been changed. } // GetWALLocated returns the server id of the node that the wal of the vChannel is located. func (s *StreamingNodeManager) GetWALLocated(vChannel string) int64 { pchannel := funcutil.ToPhysicalChannel(vChannel) var targetServerID int64 s.cond.L.Lock() for { if assignment, ok := s.latestAssignments[pchannel]; ok { targetServerID = assignment.Node.ServerID break } s.cond.Wait(context.Background()) } s.cond.L.Unlock() return targetServerID } // GetStreamingQueryNodeIDs returns the server ids of the streaming query nodes. func (s *StreamingNodeManager) GetStreamingQueryNodeIDs() typeutil.UniqueSet { s.cond.L.Lock() defer s.cond.L.Unlock() return s.streamingNodes.Clone() } // ListenNodeChanged returns a listener for node changed event. func (s *StreamingNodeManager) ListenNodeChanged() *syncutil.VersionedListener { return s.nodeChangedNotifier.Listen(syncutil.VersionedListenAtEarliest) } // SetBalancerReady set the balancer ready for the streaming node manager from streamingcoord initialization. func (s *StreamingNodeManager) SetBalancerReady(b balancer.Balancer) { s.balancer.Set(b) } func (s *StreamingNodeManager) execute() (err error) { defer s.notifier.Finish(struct{}{}) balancer, err := s.balancer.GetWithContext(s.notifier.Context()) if err != nil { return errors.Wrap(err, "failed to wait balancer ready") } for { if err := balancer.WatchChannelAssignments(s.notifier.Context(), func( version typeutil.VersionInt64Pair, relations []types.PChannelInfoAssigned, ) error { s.cond.LockAndBroadcast() s.latestAssignments = make(map[string]types.PChannelInfoAssigned) s.streamingNodes = typeutil.NewUniqueSet() for _, relation := range relations { s.latestAssignments[relation.Channel.Name] = relation s.streamingNodes.Insert(relation.Node.ServerID) } s.nodeChangedNotifier.NotifyAll() log.Info("streaming node manager updated", zap.Any("assignments", s.latestAssignments), zap.Any("streamingNodes", s.streamingNodes)) s.cond.L.Unlock() return nil }); err != nil { return err } } }