feat: use priority queue for node drainer worker pipeline (#1341)

Signed-off-by: Ajay Mishra <ajmishra@nvidia.com>

Author: XRFXLP

…signs/040-node-drainer-priority-queue.md …signs/041-node-drainer-priority-queue.mddocs/designs/040-node-drainer-priority-queue.md renamed to docs/designs/041-node-drainer-priority-queue.md docs/designs/040-node-drainer-priority-queue.md renamed to docs/designs/041-node-drainer-priority-queue.md

@@ -1,4 +1,4 @@
-# ADR-040: Node Drainer — Priority Queue
+# ADR-041: Node Drainer — Priority Queue
 
 ## Context
 

node-drainer/pkg/metrics/metrics.go

@@ -119,6 +119,15 @@ var (
 		},
 	)
 
+	// QueueItemsAssigned tracks priority decisions made when events enter the ready queue
+	QueueItemsAssigned = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "node_drainer_queue_items_assigned_total",
+			Help: "Total number of ready queue items assigned by priority and reason.",
+		},
+		[]string{"priority", "reason"},
+	)
+
 	// CustomDrainCRDNotFound tracks failures when custom drain CRD is not found
 	CustomDrainCRDNotFound = promauto.NewCounterVec(
 		prometheus.CounterOpts{

node-drainer/pkg/queue/priority_queue.go

@@ -0,0 +1,180 @@
+// Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package queue
+
+import (
+	"log/slog"
+	"sync"
+
+	"k8s.io/client-go/util/workqueue"
+
+	"github.com/nvidia/nvsentinel/node-drainer/pkg/metrics"
+)
+
+type queuePriority string
+type nodePriorityStateKind int
+
+const (
+	queuePriorityHigh queuePriority = "high"
+	queuePriorityLow  queuePriority = "low"
+
+	priorityReasonNodeNotYetDraining     = "node_not_yet_draining"
+	priorityReasonNodeAlreadyDraining    = "node_already_draining"
+	priorityReasonNodeHighPriorityQueued = "node_high_priority_queued"
+)
+
+const (
+	nodePriorityStateNone nodePriorityStateKind = iota
+	nodePriorityStateDraining
+	nodePriorityStateRepresented
+)
+
+// nodePriorityState tracks the node-level state needed to classify new ready
+// queue items without changing the event lifecycle itself.
+type nodePriorityState struct {
+	mu sync.Mutex
+
+	nodes map[string]nodePriorityStateEntry
+}
+
+type nodePriorityStateEntry struct {
+	kind              nodePriorityStateKind
+	representativeKey string // only meaningful when kind == nodePriorityStateRepresented
+}
+
+// nodeEventPriorityQueue is the ready queue used under Kubernetes' rate
+// limiting workqueue. It preserves FIFO order within each priority lane.
+type nodeEventPriorityQueue struct {
+	state *nodePriorityState
+	high  []NodeEvent
+	low   []NodeEvent
+}
+
+var _ workqueue.Queue[NodeEvent] = (*nodeEventPriorityQueue)(nil)
+
+func newNodePriorityState() *nodePriorityState {
+	return &nodePriorityState{
+		nodes: make(map[string]nodePriorityStateEntry),
+	}
+}
+
+func newNodeEventPriorityQueue(state *nodePriorityState) *nodeEventPriorityQueue {
+	return &nodeEventPriorityQueue{
+		state: state,
+	}
+}
+
+func (q *nodeEventPriorityQueue) Touch(NodeEvent) {}
+
+// Push assigns priority when an item becomes ready. Only the first queued item
+// for a not-yet-draining node gets the high-priority lane.
+func (q *nodeEventPriorityQueue) Push(item NodeEvent) {
+	priority, reason := q.state.classifyForEnqueue(item)
+	if priority == queuePriorityHigh {
+		q.high = append(q.high, item)
+	} else {
+		q.low = append(q.low, item)
+	}
+
+	metrics.QueueItemsAssigned.WithLabelValues(string(priority), reason).Inc()
+	slog.Debug("Assigned node-drainer queue priority",
+		"node", item.NodeName,
+		"eventID", item.EventID,
+		"priority", priority,
+		"reason", reason)
+}
+
+func (q *nodeEventPriorityQueue) Len() int {
+	return len(q.high) + len(q.low)
+}
+
+// Pop always drains high-priority representatives before duplicate or
+// already-draining work in the low-priority lane.
+func (q *nodeEventPriorityQueue) Pop() NodeEvent {
+	if len(q.high) > 0 {
+		return popNodeEvent(&q.high)
+	}
+
+	return popNodeEvent(&q.low)
+}
+
+// classifyForEnqueue decides whether an item can still improve the time to
+// first draining transition for its node.
+func (s *nodePriorityState) classifyForEnqueue(item NodeEvent) (queuePriority, string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	entry := s.nodes[item.NodeName]
+	if entry.kind == nodePriorityStateDraining {
+		return queuePriorityLow, priorityReasonNodeAlreadyDraining
+	}
+
+	if entry.kind == nodePriorityStateRepresented {
+		return queuePriorityLow, priorityReasonNodeHighPriorityQueued
+	}
+
+	s.nodes[item.NodeName] = nodePriorityStateEntry{
+		kind:              nodePriorityStateRepresented,
+		representativeKey: representativeKey(item),
+	}
+
+	return queuePriorityHigh, priorityReasonNodeNotYetDraining
+}
+
+// releaseRepresentative allows a node to receive another high-priority item
+// after its current high-priority representative has left the ready queue.
+func (s *nodePriorityState) releaseRepresentative(item NodeEvent) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	entry := s.nodes[item.NodeName]
+	if entry.kind != nodePriorityStateRepresented || entry.representativeKey != representativeKey(item) {
+		return
+	}
+
+	delete(s.nodes, item.NodeName)
+}
+
+// markNodeDraining moves future work for this node to low priority until the
+// draining label is removed or replaced by a terminal state.
+func (s *nodePriorityState) markNodeDraining(nodeName string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.nodes[nodeName] = nodePriorityStateEntry{kind: nodePriorityStateDraining}
+}
+
+// clearNodeDraining lets the node receive high-priority work again after it
+// leaves the draining state.
+func (s *nodePriorityState) clearNodeDraining(nodeName string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.nodes[nodeName].kind == nodePriorityStateDraining {
+		delete(s.nodes, nodeName)
+	}
+}
+
+func popNodeEvent(items *[]NodeEvent) NodeEvent {
+	item := (*items)[0]
+	(*items)[0] = NodeEvent{}
+	*items = (*items)[1:]
+
+	return item
+}
+
+func representativeKey(item NodeEvent) string {
+	return item.NodeName + ":" + item.EventID
+}

node-drainer/pkg/queue/queue.go

@@ -29,11 +29,23 @@ import (
 )
 
 func NewEventQueueManager() EventQueueManager {
+	priorityState := newNodePriorityState()
+	baseQueue := workqueue.NewTypedWithConfig(workqueue.TypedQueueConfig[NodeEvent]{
+		Queue: newNodeEventPriorityQueue(priorityState),
+	})
+	delayingQueue := workqueue.NewTypedDelayingQueueWithConfig(workqueue.TypedDelayingQueueConfig[NodeEvent]{
+		Queue: baseQueue,
+	})
+
 	mgr := &eventQueueManager{
-		queue: workqueue.NewTypedRateLimitingQueue(
+		queue: workqueue.NewTypedRateLimitingQueueWithConfig(
 			workqueue.NewTypedItemExponentialFailureRateLimiter[NodeEvent](10*time.Second, 2*time.Minute),
+			workqueue.TypedRateLimitingQueueConfig[NodeEvent]{
+				DelayingQueue: delayingQueue,
+			},
 		),
-		shutdown: make(chan struct{}),
+		shutdown:      make(chan struct{}),
+		priorityState: priorityState,
 	}
 
 	return mgr
@@ -45,6 +57,14 @@ func (m *eventQueueManager) SetDataStoreEventProcessor(processor DataStoreEventP
 	m.dataStoreEventProcessor = processor
 }
 
+func (m *eventQueueManager) MarkNodeDraining(nodeName string) {
+	m.priorityState.markNodeDraining(nodeName)
+}
+
+func (m *eventQueueManager) ClearNodeDraining(nodeName string) {
+	m.priorityState.clearNodeDraining(nodeName)
+}
+
 // EnqueueEventGeneric enqueues an event using the new database-agnostic interface.
 // Only the document ID is stored in the queue; the full event is fetched from the
 // database lazily when the worker processes the item, keeping queue memory minimal.

node-drainer/pkg/queue/queue_test.go

@@ -343,6 +343,122 @@ func TestWorkqueueDeduplication_DifferentNodes(t *testing.T) {
 	queueImpl.queue.Done(item2)
 }
 
+func TestPriorityQueue_GroupedFloodPrioritizesUnrepresentedNodes(t *testing.T) {
+	mgr := NewEventQueueManager()
+	defer mgr.Shutdown()
+
+	ctx := context.Background()
+	mockDB := &mockDataStore{}
+	mockHealthEventStore := &MockHealthEventStore{}
+
+	node1Event1 := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439011",
+		"nodeName": "node-1",
+	}
+	node1Event2 := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439012",
+		"nodeName": "node-1",
+	}
+	node2Event1 := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439013",
+		"nodeName": "node-2",
+	}
+
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-1", node1Event1, mockDB, mockHealthEventStore, node1Event1["_id"]))
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-1", node1Event2, mockDB, mockHealthEventStore, node1Event2["_id"]))
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-2", node2Event1, mockDB, mockHealthEventStore, node2Event1["_id"]))
+
+	queueImpl := mgr.(*eventQueueManager)
+
+	item1, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439011", item1.EventID)
+
+	item2, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439013", item2.EventID, "unrepresented node gets the high-priority lane")
+
+	item3, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439012", item3.EventID, "duplicate work for represented node remains low priority")
+
+	queueImpl.queue.Done(item1)
+	queueImpl.queue.Done(item2)
+	queueImpl.queue.Done(item3)
+}
+
+func TestPriorityQueue_DrainingNodesStayLowPriorityUntilCleared(t *testing.T) {
+	mgr := NewEventQueueManager()
+	defer mgr.Shutdown()
+
+	ctx := context.Background()
+	mockDB := &mockDataStore{}
+	mockHealthEventStore := &MockHealthEventStore{}
+
+	queueImpl := mgr.(*eventQueueManager)
+	queueImpl.MarkNodeDraining("node-1")
+
+	node1DrainingEvent := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439021",
+		"nodeName": "node-1",
+	}
+	node2Event := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439022",
+		"nodeName": "node-2",
+	}
+	node1AfterClearEvent := datastore.Event{
+		"_id":      "507f1f77bcf86cd799439023",
+		"nodeName": "node-1",
+	}
+
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-1", node1DrainingEvent, mockDB, mockHealthEventStore, node1DrainingEvent["_id"]))
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-2", node2Event, mockDB, mockHealthEventStore, node2Event["_id"]))
+
+	item1, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439022", item1.EventID, "not-yet-draining node should run before draining duplicate work")
+
+	queueImpl.ClearNodeDraining("node-1")
+	require.NoError(t, mgr.EnqueueEventGeneric(ctx, "node-1", node1AfterClearEvent, mockDB, mockHealthEventStore, node1AfterClearEvent["_id"]))
+
+	item2, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439023", item2.EventID, "cleared node can receive high priority again")
+
+	item3, shutdown := queueImpl.queue.Get()
+	require.False(t, shutdown)
+	assert.Equal(t, "507f1f77bcf86cd799439021", item3.EventID)
+
+	queueImpl.queue.Done(item1)
+	queueImpl.queue.Done(item2)
+	queueImpl.queue.Done(item3)
+}
+
+func TestPriorityQueue_NonComparableDocumentID_DoesNotPanic(t *testing.T) {
+	state := newNodePriorityState()
+	queueImpl := newNodeEventPriorityQueue(state)
+	item := NodeEvent{
+		NodeName:   "node-1",
+		EventID:    "event-1",
+		DocumentID: map[string]string{"id": "non-comparable"},
+	}
+
+	require.NotPanics(t, func() {
+		queueImpl.Push(item)
+		entry := state.nodes[item.NodeName]
+		require.Equal(t, nodePriorityStateRepresented, entry.kind)
+		require.Equal(t, representativeKey(item), entry.representativeKey)
+
+		got := queueImpl.Pop()
+		assert.Equal(t, item.NodeName, got.NodeName)
+		assert.Equal(t, item.EventID, got.EventID)
+
+		state.releaseRepresentative(got)
+		_, represented := state.nodes[item.NodeName]
+		assert.False(t, represented)
+	})
+}
+
 // Mock DataStore for testing
 type mockDataStore struct{}
 

node-drainer/pkg/queue/types.go

@@ -56,6 +56,8 @@ type EventQueueManager interface {
 
 	Start(ctx context.Context)
 	Shutdown()
+	MarkNodeDraining(nodeName string)
+	ClearNodeDraining(nodeName string)
 
 	// New database-agnostic method
 	SetDataStoreEventProcessor(processor DataStoreEventProcessor)
@@ -68,4 +70,5 @@ type eventQueueManager struct {
 	dataStoreEventProcessor DataStoreEventProcessor // New database-agnostic processor
 	shutdown                chan struct{}
 	sessions                sync.Map // EventID -> *DrainSession
+	priorityState           *nodePriorityState
 }

node-drainer/pkg/queue/worker.go

@@ -73,7 +73,10 @@ func (m *eventQueueManager) processNextWorkItem(ctx context.Context) bool {
 		return false
 	}
 
-	defer m.queue.Done(nodeEvent)
+	defer func() {
+		m.priorityState.releaseRepresentative(nodeEvent)
+		m.queue.Done(nodeEvent)
+	}()
 
 	if nodeEvent.Database == nil || nodeEvent.HealthEventStore == nil {
 		slog.ErrorContext(ctx, "NodeEvent missing database or health event store, dropping",