Refactor filterByTainted

scheduler/reconciler/filters.go

@@ -99,6 +99,171 @@ func (set allocSet) filterOldTerminalAllocs(a ReconcilerState) (remain, ignore a
 	return remain, ignored
 }
 
+// allocCategory represents the classification bucket for an allocation.
+type allocCategory string
+
+const (
+	categoryUntainted     allocCategory = "untainted"
+	categoryMigrate       allocCategory = "migrate"
+	categoryLost          allocCategory = "lost"
+	categoryDisconnecting allocCategory = "disconnecting"
+	categoryReconnecting  allocCategory = "reconnecting"
+	categoryIgnore        allocCategory = "ignore"
+	categoryExpiring      allocCategory = "expiring"
+)
+
+// allocContext holds the per-allocation data used by classification rules.
+type allocContext struct {
+	alloc           *structs.Allocation
+	shouldReconnect bool
+	taintedNode     *structs.Node
+	nodeIsTainted   bool
+	now             time.Time
+}
+
+type classificationRule struct {
+	condition func(allocContext) bool
+	category  allocCategory
+}
+
+// classificationRules are evaluated in order; first match wins.
+var classificationRules = []classificationRule{
+	// Failed allocs that need reconnect and are still desired to run.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.shouldReconnect &&
+				ctx.alloc.DesiredStatus == structs.AllocDesiredStatusRun &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusFailed
+		},
+		category: categoryReconnecting,
+	},
+	// Server-terminal allocs should be ignored when they are not reconnecting.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.alloc.TerminalStatus() && !ctx.shouldReconnect && ctx.alloc.ServerTerminalStatus()
+		},
+		category: categoryIgnore,
+	},
+	// Terminal canaries marked for migration.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.alloc.TerminalStatus() && !ctx.shouldReconnect &&
+				ctx.alloc.DeploymentStatus.IsCanary() && ctx.alloc.DesiredTransition.ShouldMigrate()
+		},
+		category: categoryMigrate,
+	},
+	// Terminal allocs that are not reconnecting are untainted.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.alloc.TerminalStatus() && !ctx.shouldReconnect
+		},
+		category: categoryUntainted,
+	},
+	// Expired allocs are handled before reconnecting/disconnect paths.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.alloc.Expired(ctx.now)
+		},
+		category: categoryExpiring,
+	},
+	// Failed reconnects that server already marked to stop should be ignored.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.shouldReconnect &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusFailed &&
+				ctx.alloc.DesiredStatus == structs.AllocDesiredStatusStop
+		},
+		category: categoryIgnore,
+	},
+	// Disconnected node and alloc already unknown.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode != nil &&
+				ctx.taintedNode.Status == structs.NodeStatusDisconnected &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusUnknown
+		},
+		category: categoryUntainted,
+	},
+	// Disconnected pending allocs should be replaced immediately.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode != nil &&
+				ctx.taintedNode.Status == structs.NodeStatusDisconnected &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusPending
+		},
+		category: categoryLost,
+	},
+	// Disconnected allocs with nil disconnect block or lost_after=0 are lost.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode != nil &&
+				ctx.taintedNode.Status == structs.NodeStatusDisconnected &&
+				ctx.alloc.DisconnectTimeout(ctx.now) == ctx.now
+		},
+		category: categoryLost,
+	},
+	// Remaining disconnected allocs are disconnecting.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode != nil && ctx.taintedNode.Status == structs.NodeStatusDisconnected
+		},
+		category: categoryDisconnecting,
+	},
+	// Non-terminal allocs marked for migration always migrate.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.alloc.DesiredTransition.ShouldMigrate()
+		},
+		category: categoryMigrate,
+	},
+	// Ready or untainted nodes with reconnecting allocs.
+	{
+		condition: func(ctx allocContext) bool {
+			return (!ctx.nodeIsTainted || (ctx.taintedNode != nil && ctx.taintedNode.Status == structs.NodeStatusReady)) &&
+				ctx.shouldReconnect
+		},
+		category: categoryReconnecting,
+	},
+	// Ready or untainted nodes.
+	{
+		condition: func(ctx allocContext) bool {
+			return !ctx.nodeIsTainted || (ctx.taintedNode != nil && ctx.taintedNode.Status == structs.NodeStatusReady)
+		},
+		category: categoryUntainted,
+	},
+	// GC'd (tainted map entry with nil node) nodes are lost.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode == nil
+		},
+		category: categoryLost,
+	},
+	// Terminal node allocations that cannot be replaced.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode.TerminalStatus() &&
+				!ctx.alloc.ReplaceOnDisconnect() &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusUnknown
+		},
+		category: categoryUntainted,
+	},
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode.TerminalStatus() &&
+				!ctx.alloc.ReplaceOnDisconnect() &&
+				ctx.alloc.ClientStatus == structs.AllocClientStatusRunning
+		},
+		category: categoryDisconnecting,
+	},
+	// Remaining terminal node allocs are lost.
+	{
+		condition: func(ctx allocContext) bool {
+			return ctx.taintedNode.TerminalStatus()
+		},
+		category: categoryLost,
+	},
+}
+
 // filterByTainted takes a set of tainted nodes and filters the allocation set
 // into the following groups:
 // 1. Those that exist on untainted nodes
@@ -117,128 +282,50 @@ func (set allocSet) filterByTainted(state ClusterState) (untainted, migrate, los
 	ignore = make(allocSet)
 	expiring = make(allocSet)
 
+	// Create a map for quick category assignment.
+	categoryMap := map[allocCategory]*allocSet{
+		categoryUntainted:     &untainted,
+		categoryMigrate:       &migrate,
+		categoryLost:          &lost,
+		categoryDisconnecting: &disconnecting,
+		categoryReconnecting:  &reconnecting,
+		categoryIgnore:        &ignore,
+		categoryExpiring:      &expiring,
+	}
+
 	for _, alloc := range set {
-		shouldReconnect := false
+		// Build context for classification rules.
+		ctx := allocContext{
+			alloc: alloc,
+			now:   state.Now,
+		}
 
-		// Only compute reconnect for unknown, running, and failed since they
-		// need to go through the reconnect logic.
+		// Compute shouldReconnect for unknown/running/failed allocs.
 		if alloc.ClientStatus == structs.AllocClientStatusUnknown ||
 			alloc.ClientStatus == structs.AllocClientStatusRunning ||
 			alloc.ClientStatus == structs.AllocClientStatusFailed {
-			shouldReconnect = alloc.NeedsToReconnect()
-		}
-
-		// Failed allocs that need to be reconnected must be added to
-		// reconnecting so that they can be handled as a failed reconnect.
-		if shouldReconnect &&
-			alloc.DesiredStatus == structs.AllocDesiredStatusRun &&
-			alloc.ClientStatus == structs.AllocClientStatusFailed {
-			reconnecting[alloc.ID] = alloc
-			continue
-		}
-
-		if alloc.TerminalStatus() && !shouldReconnect {
-			// Server-terminal allocs, if they should not reconnect,
-			// are probably stopped replacements and should be ignored
-			if alloc.ServerTerminalStatus() {
-				ignore[alloc.ID] = alloc
-				continue
-			}
-
-			// Terminal canaries that have been marked for migration need to be
-			// migrated, otherwise we block deployments from progressing by
-			// counting them as running canaries.
-			if alloc.DeploymentStatus.IsCanary() && alloc.DesiredTransition.ShouldMigrate() {
-				migrate[alloc.ID] = alloc
-				continue
-			}
-
-			// Terminal allocs, if not reconnect, are always untainted as they
-			// should never be migrated.
-			untainted[alloc.ID] = alloc
-			continue
-		}
-
-		// The alloc is expired. These end up getting added to "lost" but this
-		// skips having to evaluate the reconnecting logic for these allocs.
-		if alloc.Expired(state.Now) {
-			expiring[alloc.ID] = alloc
-			continue
-		}
-
-		// Ignore failed allocs that need to be reconnected and that have been
-		// marked to stop by the server.
-		if shouldReconnect &&
-			alloc.ClientStatus == structs.AllocClientStatusFailed &&
-			alloc.DesiredStatus == structs.AllocDesiredStatusStop {
-			ignore[alloc.ID] = alloc
-			continue
+			ctx.shouldReconnect = alloc.NeedsToReconnect()
 		}
 
-		taintedNode, nodeIsTainted := state.TaintedNodes[alloc.NodeID]
-		if taintedNode != nil && taintedNode.Status == structs.NodeStatusDisconnected {
-			// ignore allocs already marked unknown, they should already have a followup eval
-			if alloc.ClientStatus == structs.AllocClientStatusUnknown {
-				untainted[alloc.ID] = alloc
-				continue
-			}
-			// If the alloc is pending mark it lost so it is replaced immediately.
-			if alloc.ClientStatus == structs.AllocClientStatusPending {
-				lost[alloc.ID] = alloc
-				continue
+		// Get node taint information, preserving whether key existed so we can
+		// distinguish missing-node from GC'd-node semantics.
+		ctx.taintedNode, ctx.nodeIsTainted = state.TaintedNodes[alloc.NodeID]
+
+		// Apply classification rules in order (first match wins).
+		classified := false
+		for _, rule := range classificationRules {
+			if rule.condition(ctx) {
+				targetSet := categoryMap[rule.category]
+				(*targetSet)[alloc.ID] = alloc
+				classified = true
+				break
 			}
-			// Allocs with nil disconnect blocks or disconnect.lost_after == 0 are lost
-			if alloc.DisconnectTimeout(state.Now) == state.Now {
-				lost[alloc.ID] = alloc
-				continue
-			}
-			disconnecting[alloc.ID] = alloc
-			continue
-		}
-
-		// Non-terminal allocs that should migrate should always migrate
-		if alloc.DesiredTransition.ShouldMigrate() {
-			migrate[alloc.ID] = alloc
-			continue
 		}
 
-		if !nodeIsTainted || (taintedNode != nil && taintedNode.Status == structs.NodeStatusReady) {
-			// Filter allocs on a node that is now re-connected to be resumed.
-			if shouldReconnect {
-				reconnecting[alloc.ID] = alloc
-				continue
-			}
-
-			// Otherwise, Node is untainted so alloc is untainted
+		// Default category if no rule matched.
+		if !classified {
 			untainted[alloc.ID] = alloc
-			continue
-		}
-
-		// Allocs on GC'd (nil) or lost nodes are Lost
-		if taintedNode == nil {
-			lost[alloc.ID] = alloc
-			continue
-		}
-
-		// Allocs on terminal nodes that can't be rescheduled need to be treated
-		// differently than those that can.
-		if taintedNode.TerminalStatus() {
-			if !alloc.ReplaceOnDisconnect() {
-				if alloc.ClientStatus == structs.AllocClientStatusUnknown {
-					untainted[alloc.ID] = alloc
-					continue
-				} else if alloc.ClientStatus == structs.AllocClientStatusRunning {
-					disconnecting[alloc.ID] = alloc
-					continue
-				}
-			}
-
-			lost[alloc.ID] = alloc
-			continue
 		}
-
-		// All other allocs are untainted
-		untainted[alloc.ID] = alloc
 	}
 
 	return