fix: custom drain mode should handle parallel drains (#966)

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

Author: XRFXLP

demos/local-slinky-drain-demo/scripts/00-setup.sh

@@ -152,7 +152,7 @@ install_custom_drain_crd() {
 install_nvsentinel() {
     section "Phase 5: Installing NVSentinel"
 
-    local nvsentinel_version="${NVSENTINEL_VERSION:-v0.6.0}"
+    local nvsentinel_version="${NVSENTINEL_VERSION:-v0.10.0}"
 
     log "Installing Prometheus Operator CRDs (for PodMonitor)..."
     kubectl apply --server-side -f https://raw.githubusercontent.com/prometheus-operator/prometheus-operator/v0.68.0/example/prometheus-operator-crd/monitoring.coreos.com_podmonitors.yaml

node-drainer/pkg/evaluator/evaluator.go

@@ -28,10 +28,8 @@ import (
 	"github.com/nvidia/nvsentinel/node-drainer/pkg/config"
 	"github.com/nvidia/nvsentinel/node-drainer/pkg/customdrain"
 	"github.com/nvidia/nvsentinel/node-drainer/pkg/queue"
-	"github.com/nvidia/nvsentinel/store-client/pkg/client"
 	"github.com/nvidia/nvsentinel/store-client/pkg/datastore"
 	"github.com/nvidia/nvsentinel/store-client/pkg/query"
-	"github.com/nvidia/nvsentinel/store-client/pkg/utils"
 )
 
 const (
@@ -89,7 +87,7 @@ func (e *NodeDrainEvaluator) EvaluateEventWithDatabase(ctx context.Context, heal
 	}
 
 	if e.config.CustomDrain.Enabled && e.customDrainClient != nil {
-		return e.evaluateCustomDrain(ctx, healthEvent, database, partialDrainEntity)
+		return e.evaluateCustomDrain(ctx, healthEvent, partialDrainEntity)
 	}
 
 	return e.evaluateUserNamespaceActions(ctx, healthEvent, partialDrainEntity)
@@ -305,19 +303,12 @@ func isTerminalStatus(status model.Status) bool {
 }
 
 func (e *NodeDrainEvaluator) evaluateCustomDrain(ctx context.Context, healthEvent model.HealthEventWithStatus,
-	database queue.DataStore, partialDrainEntity *protos.Entity) (*DrainActionResult, error) {
+	partialDrainEntity *protos.Entity) (*DrainActionResult, error) {
 	nodeName := healthEvent.HealthEvent.NodeName
+	eventID := healthEvent.HealthEvent.Id
 
-	eventID, err := getEventID(ctx, database, nodeName)
-	if err != nil {
-		slog.Error("Failed to get event ID for custom drain",
-			"node", nodeName,
-			"error", err)
-
-		return &DrainActionResult{
-			Action:    ActionWait,
-			WaitDelay: customDrainPollInterval,
-		}, nil
+	if eventID == "" {
+		return nil, fmt.Errorf("health event for node %s is missing Id, cannot generate DrainRequest CR name", nodeName)
 	}
 
 	crName := customdrain.GenerateCRName(nodeName, eventID)
@@ -388,33 +379,6 @@ func (e *NodeDrainEvaluator) evaluateCustomDrain(ctx context.Context, healthEven
 	}, nil
 }
 
-func getEventID(ctx context.Context, database queue.DataStore, nodeName string) (string, error) {
-	opts := &client.FindOneOptions{
-		Sort: map[string]any{"_id": -1},
-	}
-
-	filter := map[string]any{
-		"healthevent.nodename": nodeName,
-	}
-
-	result, err := database.FindDocument(ctx, filter, opts)
-	if err != nil {
-		return "", fmt.Errorf("failed to query database for node %s: %w", nodeName, err)
-	}
-
-	var document map[string]any
-	if err := result.Decode(&document); err != nil {
-		return "", fmt.Errorf("failed to decode health event for node %s: %w", nodeName, err)
-	}
-
-	eventID, err := utils.ExtractDocumentID(document)
-	if err != nil {
-		return "", fmt.Errorf("failed to extract document ID for node %s: %w", nodeName, err)
-	}
-
-	return eventID, nil
-}
-
 /*
 This function determines whether we can skip draining for the current unhealthy HealthEvent. A full drain can be
 skipped if a previous full drain completed against the node after it was most-recently quarantined. Additionally,

node-drainer/pkg/reconciler/reconciler.go

@@ -285,6 +285,12 @@ func (r *Reconciler) ProcessEventGeneric(ctx context.Context,
 
 	eventID := utils.ExtractEventID(event)
 
+	if healthEventWithStatus.HealthEvent != nil && healthEventWithStatus.HealthEvent.Id == "" {
+		if docID, err := utils.ExtractDocumentID(event); err == nil {
+			healthEventWithStatus.HealthEvent.Id = docID
+		}
+	}
+
 	slog.Debug("Processing event", "node", nodeName, "eventID", eventID)
 
 	metrics.TotalEventsReceived.Inc()

node-drainer/pkg/reconciler/reconciler_integration_test.go

@@ -1366,15 +1366,22 @@ func createPod(ctx context.Context, t *testing.T, client kubernetes.Interface, n
 
 type healthEventOptions struct {
 	nodeName          string
+	eventID           string
 	nodeQuarantined   model.Status
 	drainForce        bool
 	recommendedAction protos.RecommendedAction
 	entitiesImpacted  []*protos.Entity
 }
 
 func createHealthEvent(opts healthEventOptions) map[string]interface{} {
+	eventID := opts.nodeName + "-event"
+	if opts.eventID != "" {
+		eventID = opts.eventID
+	}
+
 	healthEvent := &protos.HealthEvent{
 		NodeName:          opts.nodeName,
+		Id:                eventID,
 		CheckName:         "test-check",
 		RecommendedAction: opts.recommendedAction,
 		EntitiesImpacted:  opts.entitiesImpacted,
@@ -1384,9 +1391,6 @@ func createHealthEvent(opts healthEventOptions) map[string]interface{} {
 		healthEvent.DrainOverrides = &protos.BehaviourOverrides{Force: true}
 	}
 
-	eventID := opts.nodeName + "-event"
-	// Return just the fullDocument content, as the event watcher extracts this
-	// from the change stream before passing to the reconciler
 	return map[string]interface{}{
 		"_id":         eventID,
 		"healthevent": healthEvent,
@@ -1859,6 +1863,100 @@ func TestReconciler_CustomDrainHappyPath(t *testing.T) {
 	}, 10*time.Second, 500*time.Millisecond, "CR should be deleted after drain completes")
 }
 
+func TestReconciler_CustomDrainMultipleEventsOnSameNode(t *testing.T) {
+	customDrainCfg := config.CustomDrainConfig{
+		Enabled:               true,
+		TemplateMountPath:     "../customdrain/testdata",
+		TemplateFileName:      "drain-template.yaml",
+		Namespace:             "default",
+		ApiGroup:              "drain.example.com",
+		Version:               "v1alpha1",
+		Kind:                  "DrainRequest",
+		StatusConditionType:   "Complete",
+		StatusConditionStatus: "True",
+		Timeout:               config.Duration{Duration: 30 * time.Minute},
+	}
+
+	setup := setupCustomDrainTest(t, customDrainCfg)
+
+	nodeName := "multi-event-node"
+	createNode(setup.ctx, t, setup.client, nodeName)
+	createNamespace(setup.ctx, t, setup.client, "default")
+
+	gvr := schema.GroupVersionResource{
+		Group:    "drain.example.com",
+		Version:  "v1alpha1",
+		Resource: "drainrequests",
+	}
+
+	eventA := createHealthEvent(healthEventOptions{
+		nodeName:        nodeName,
+		eventID:         "event-aaa",
+		nodeQuarantined: model.Quarantined,
+	})
+	eventB := createHealthEvent(healthEventOptions{
+		nodeName:        nodeName,
+		eventID:         "event-bbb",
+		nodeQuarantined: model.Quarantined,
+	})
+
+	setup.mockDB.storeEvent(nodeName, eventB)
+
+	// Process Event A — should create its own DrainRequest CR
+	err := setup.reconciler.ProcessEventGeneric(setup.ctx, eventA, setup.mockDB, setup.healthEventStore, nodeName)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "waiting for custom drain CR to complete")
+
+	crNameA := customdrain.GenerateCRName(nodeName, "event-aaa")
+	crNameB := customdrain.GenerateCRName(nodeName, "event-bbb")
+
+	_, err = setup.dynamicClient.Resource(gvr).Namespace("default").Get(setup.ctx, crNameA, metav1.GetOptions{})
+	require.NoError(t, err, "DrainRequest for Event A should exist")
+
+	// Process Event B — should create a separate DrainRequest CR
+	err = setup.reconciler.ProcessEventGeneric(setup.ctx, eventB, setup.mockDB, setup.healthEventStore, nodeName)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "waiting for custom drain CR to complete")
+
+	_, err = setup.dynamicClient.Resource(gvr).Namespace("default").Get(setup.ctx, crNameB, metav1.GetOptions{})
+	require.NoError(t, err, "DrainRequest for Event B should exist")
+
+	// Retry Event A while CR is incomplete — should poll Event A's CR, not Event B's
+	err = setup.reconciler.ProcessEventGeneric(setup.ctx, eventA, setup.mockDB, setup.healthEventStore, nodeName)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "waiting for retry delay",
+		"Event A should be waiting on its own CR, not trying to create a new one")
+
+	// Complete Event A's DrainRequest
+	crA, err := setup.dynamicClient.Resource(gvr).Namespace("default").Get(setup.ctx, crNameA, metav1.GetOptions{})
+	require.NoError(t, err)
+	err = unstructured.SetNestedField(crA.Object, []any{
+		map[string]any{
+			"type":               "Complete",
+			"status":             "True",
+			"lastTransitionTime": time.Now().Format(time.RFC3339),
+			"reason":             "DrainSucceeded",
+			"message":            "All pods drained",
+		},
+	}, "status", "conditions")
+	require.NoError(t, err)
+	_, err = setup.dynamicClient.Resource(gvr).Namespace("default").UpdateStatus(setup.ctx, crA, metav1.UpdateOptions{})
+	require.NoError(t, err)
+
+	// Process Event A again — should detect completion and mark AlreadyDrained
+	err = setup.reconciler.ProcessEventGeneric(setup.ctx, eventA, setup.mockDB, setup.healthEventStore, nodeName)
+	require.NoError(t, err, "Event A should complete successfully after its CR is done")
+
+	require.Eventually(t, func() bool {
+		_, err := setup.dynamicClient.Resource(gvr).Namespace("default").Get(setup.ctx, crNameA, metav1.GetOptions{})
+		return errors.IsNotFound(err)
+	}, 10*time.Second, 500*time.Millisecond, "Event A's CR should be cleaned up")
+
+	// Event B's CR should still exist independently
+	_, err = setup.dynamicClient.Resource(gvr).Namespace("default").Get(setup.ctx, crNameB, metav1.GetOptions{})
+	require.NoError(t, err, "Event B's CR should still exist — it hasn't been completed yet")
+}
+
 func TestReconciler_CustomDrainCRDNotFound(t *testing.T) {
 	customDrainCfg := config.CustomDrainConfig{
 		Enabled:               true,