fix: Remove dispatcher tasks before cleaning up targets

Related to milvus-io#39840 The target could be updated async in previous code. This PR make remove collection from target observer block until all tasks related in dispatchers are removed preventing the metrics being updated after collection released. Signed-off-by: Congqi Xia <[email protected]>
congqixia · Feb 13, 2025 · 79f45dc · 79f45dc
1 parent cb1bf6d
commit 79f45dc
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/internal/querycoordv2/observers/target_observer.go b/internal/querycoordv2/observers/target_observer.go
@@ -219,6 +219,9 @@ func (ob *TargetObserver) schedule(ctx context.Context) {
 				delete(ob.readyNotifiers, req.CollectionID)
 				ob.mut.Unlock()
 
+				ob.loadedDispatcher.RemoveTask(req.CollectionID)
+				ob.loadingDispatcher.RemoveTask(req.CollectionID)
+
 				ob.targetMgr.RemoveCollection(ctx, req.CollectionID)
 				req.Notifier <- nil
 			case ReleasePartition:

diff --git a/internal/querycoordv2/observers/task_dispatcher.go b/internal/querycoordv2/observers/task_dispatcher.go
@@ -33,6 +33,7 @@ type taskDispatcher[K comparable] struct {
 	tasks      *typeutil.ConcurrentMap[K, bool]
 	pool       *conc.Pool[any]
 	notifyCh   chan struct{}
+	removeCh   chan K
 	taskRunner task[K]
 	wg         sync.WaitGroup
 	cancel     context.CancelFunc
@@ -46,6 +47,7 @@ func newTaskDispatcher[K comparable](runner task[K]) *taskDispatcher[K] {
 		tasks:      typeutil.NewConcurrentMap[K, bool](),
 		pool:       conc.NewPool[any](paramtable.Get().QueryCoordCfg.ObserverTaskParallel.GetAsInt()),
 		notifyCh:   make(chan struct{}, 1),
+		removeCh:   make(chan K),
 		taskRunner: runner,
 	}
 }
@@ -81,6 +83,10 @@ func (d *taskDispatcher[K]) AddTask(keys ...K) {
 	}
 }
 
+func (d *taskDispatcher[K]) RemoveTask(key K) {
+	d.removeCh <- key
+}
+
 func (d *taskDispatcher[K]) notify() {
 	select {
 	case d.notifyCh <- struct{}{}:
@@ -105,6 +111,8 @@ func (d *taskDispatcher[K]) schedule(ctx context.Context) {
 				}
 				return true
 			})
+		case key := <-d.removeCh:
+			d.tasks.Remove(key)
 		}
 	}
 }