mirror of https://github.com/milvus-io/milvus.git
109 lines
2.8 KiB
Go
109 lines
2.8 KiB
Go
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/v2/log"
|
|
"github.com/milvus-io/milvus/pkg/v2/streaming/util/types"
|
|
"github.com/milvus-io/milvus/pkg/v2/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()
|
|
}
|