POC: speedup single node consolidation by grouping nodes and pools to a single node-type

pkg/controllers/disruption/helpers.go

@@ -94,13 +94,16 @@
 	}
 	pods = append(pods, deletingNodePods...)
 
+	var candidateNodeTypes []string
+	stateNodes, candidateNodeTypes = filterStateNodesByCandidateNodeTypes(stateNodes, candidates)
+
 	var opts []scheduling.Options
 	if options.FromContext(ctx).PreferencePolicy == options.PreferencePolicyIgnore {
 		opts = append(opts, scheduling.IgnorePreferences)
 	}
 	opts = append(opts, scheduling.MinValuesPolicy(options.FromContext(ctx).MinValuesPolicy))
 	scheduler, err := provisioner.NewScheduler(
-		log.IntoContext(ctx, operatorlogging.NopLogger),
+		context.WithValue(log.IntoContext(ctx, operatorlogging.NopLogger), "filterNodePoolsByStateNodeNodeTypes", candidateNodeTypes),
 		pods,
 		stateNodes,
 		opts...,
@@ -169,6 +172,25 @@
 	return len(rhsNames.Intersection(lhsNames)) == len(lhsNames)
 }
 
+// speed up computation by only considering nodes that have the same "node-type" label as the candidates
+func filterStateNodesByCandidateNodeTypes(stateNodes []*state.StateNode, candidates []*Candidate) ([]*state.StateNode, []string) {
+	candidateNodeTypes := lo.Uniq(lo.Map(candidates, func(c *Candidate, _ int) string {
+		if c.Node != nil {
+			return c.Node.Labels["node-type"]
+		}
+		return ""
+	}))
+
+	if lo.Contains(candidateNodeTypes, "") { // something was weird, do not filter, even if some had node-type
+		return stateNodes, candidateNodeTypes
+	}
+
+	stateNodes = lo.Filter(stateNodes, func(n *state.StateNode, _ int) bool {
+		return lo.Contains(candidateNodeTypes, n.Labels()["node-type"])
+	})
+	return stateNodes, candidateNodeTypes
+}
+
 // GetCandidates returns nodes that appear to be currently deprovisionable based off of their nodePool
 func GetCandidates(ctx context.Context, cluster *state.Cluster, kubeClient client.Client, recorder events.Recorder, clk clock.Clock,
 	cloudProvider cloudprovider.CloudProvider, shouldDisrupt CandidateFilter, disruptionClass string, queue *Queue,

pkg/controllers/provisioning/provisioner.go

@@ -255,6 +255,9 @@ func (p *Provisioner) NewScheduler(
 		}
 		return np.DeletionTimestamp.IsZero()
 	})
+	if nodeTypes, ok := ctx.Value("filterNodePoolsByStateNodeNodeTypes").([]string); ok {
+		nodePools = filterNodePoolsByStateNodeNodeTypes(nodePools, nodeTypes)
+	}
 	if len(nodePools) == 0 {
 		return nil, ErrNodePoolsNotFound
 	}
@@ -306,6 +309,17 @@ func (p *Provisioner) NewScheduler(
 	return scheduler.NewScheduler(ctx, p.kubeClient, nodePools, p.cluster, stateNodes, topology, instanceTypes, daemonSetPods, p.recorder, p.clock, volumeReqs, opts...), nil
 }
 
+// speed up computation by only considering pools that have the same "node-type" label as the nodes
+func filterNodePoolsByStateNodeNodeTypes(nodePools []*v1.NodePool, candidateNodeTypes []string) []*v1.NodePool {
+	if lo.Contains(candidateNodeTypes, "") { // something was weird, do not filter
+		return nodePools
+	}
+
+	return lo.Filter(nodePools, func(np *v1.NodePool, _ int) bool {
+		return lo.Contains(candidateNodeTypes, np.Spec.Template.Labels["node-type"])
+	})
+}
+
 func (p *Provisioner) Schedule(ctx context.Context) (scheduler.Results, error) {
 	defer metrics.Measure(scheduler.DurationSeconds, map[string]string{scheduler.ControllerLabel: injection.GetControllerName(ctx)})()
 	start := time.Now()