Repository: hashicorp/nomad
File: scheduler/reconciler/reconcile_node.go (Go)

Code before:
// computeCanaryNodes is a helper function that, given required task groups, a
// mapping of nodes to their allocs and a map of eligible nodes, outputs a
// map[nodeID] -> map[TG] -> bool which indicates which TGs this node is a
// canary for, and a map[TG] -> int to indicate how many total canaries are to
// be placed for a TG.
func computeCanaryNodes(required map[string]*structs.TaskGroup, nodeAllocs map[string][]*structs.Allocation,
	eligibleNodes map[string]*structs.Node) (map[string]map[string]bool, map[string]int) {

	canaryNodes := map[string]map[string]bool{}
	eligibleNodesList := slices.Collect(maps.Values(eligibleNodes))
	canariesPerTG := map[string]int{}

	for _, tg := range required {
		if tg.Update.IsEmpty() || tg.Update.Canary == 0 {
			continue
		}

		// round up to the nearest integer
		numberOfCanaryNodes := int(math.Ceil(float64(tg.Update.Canary) * float64(len(eligibleNodes)) / 100))
		canariesPerTG[tg.Name] = numberOfCanaryNodes

		// check if there are any allocations on any nodes that are/were
		// canaries.
		for nodeID, allocs := range nodeAllocs {
			for _, a := range allocs {
				if a.DeploymentStatus == nil || a.DeploymentStatus.Canary == false {
					continue
				}
				if a.TaskGroup == tg.Name {
					canaryNodes[nodeID] = map[string]bool{}
					canaryNodes[nodeID][tg.Name] = true

					// this node should no longer be considered when searching
					// for canary nodes
					numberOfCanaryNodes -= 1
					eligibleNodesList = slices.DeleteFunc(
						eligibleNodesList,
						func(n *structs.Node) bool { return n.ID == nodeID },
					)
				}
			}
		}

		for i, n := range eligibleNodesList {
			canaryNodes[n.ID] = map[string]bool{}

			if i > numberOfCanaryNodes-1 {
				break
			}

			canaryNodes[n.ID][tg.Name] = true

Review comment (suggestion):
I think we can also skip this if `nr.Deployment == nil` or `nr.Deployment.ID != a.DeploymentID` to skip any straggling canaries from a previous incomplete deployment.

Code after:
// computeCanaryNodes is a helper function that, given required task groups,
// mappings of nodes to their live allocs and terminal allocs, and a map of
// eligible nodes, outputs a map[nodeID] -> map[TG] -> bool which indicates
// which TGs this node is a canary for, and a map[TG] -> int to indicate how
// many total canaries are to be placed for a TG.
func (nr *NodeReconciler) computeCanaryNodes(required map[string]*structs.TaskGroup,
	liveAllocs map[string][]*structs.Allocation, terminalAllocs structs.TerminalByNodeByName,
	eligibleNodes map[string]*structs.Node) (map[string]map[string]bool, map[string]int) {

	canaryNodes := map[string]map[string]bool{}
	eligibleNodesList := slices.Collect(maps.Values(eligibleNodes))
	canariesPerTG := map[string]int{}

	for _, tg := range required {
		if tg.Update.IsEmpty() || tg.Update.Canary == 0 {
			continue
		}

		// round up to the nearest integer
		numberOfCanaryNodes := int(math.Ceil(float64(tg.Update.Canary) * float64(len(eligibleNodes)) / 100))
		canariesPerTG[tg.Name] = numberOfCanaryNodes

		// check if there are any live allocations on any nodes that are/were
		// canaries.
		for nodeID, allocs := range liveAllocs {
			for _, a := range allocs {
				eligibleNodesList, numberOfCanaryNodes = nr.findOldCanaryNodes(
					eligibleNodesList, numberOfCanaryNodes, a, tg, canaryNodes, nodeID)
			}
		}

		// check if there are any terminal allocations that were canaries
		for nodeID, terminalAlloc := range terminalAllocs {
			for _, a := range terminalAlloc {
				eligibleNodesList, numberOfCanaryNodes = nr.findOldCanaryNodes(
					eligibleNodesList, numberOfCanaryNodes, a, tg, canaryNodes, nodeID)
			}
		}

		for i, n := range eligibleNodesList {
			if i > numberOfCanaryNodes-1 {
				break
			}

			if _, ok := canaryNodes[n.ID]; !ok {
				canaryNodes[n.ID] = map[string]bool{}
			}

			canaryNodes[n.ID][tg.Name] = true
		}
	}