PR comments

pkg/controllers/state/cluster.go

@@ -172,16 +172,11 @@ func (c *Cluster) DeleteNode(name string) {
 	c.cleanupNode(name)
 }
 
-func (c *Cluster) cleanupNode(name string) {
-	if id, ok := c.nameToProviderID[name]; ok {
-		delete(c.nodes, id)
-		delete(c.nameToProviderID, name)
-		c.RecordConsolidationChange()
-	}
-}
-
 // UpdatePod is called every time the pod is reconciled
 func (c *Cluster) UpdatePod(ctx context.Context, pod *v1.Pod) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
 	var err error
 	if podutils.IsTerminal(pod) {
 		c.updateNodeUsageFromPodCompletion(client.ObjectKeyFromObject(pod))
@@ -195,6 +190,10 @@ func (c *Cluster) UpdatePod(ctx context.Context, pod *v1.Pod) error {
 // DeletePod is called when the pod has been deleted
 func (c *Cluster) DeletePod(podKey types.NamespacedName) {
 	c.antiAffinityPods.Delete(podKey)
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
 	c.updateNodeUsageFromPodCompletion(podKey)
 	c.RecordConsolidationChange()
 }
@@ -228,6 +227,11 @@ func (c *Cluster) Reset() {
 	c.antiAffinityPods = sync.Map{}
 }
 
+// WARNING
+// Everything under this section of code assumes that you have already held a lock when you are calling into these functions
+// and explicitly modifying the cluster state. If you do not hold the cluster state lock before calling any of these helpers
+// you will hit race conditions and data corruption
+
 func (c *Cluster) newStateFromNode(ctx context.Context, node *v1.Node, oldNode *Node) (*Node, error) {
 	if oldNode == nil {
 		oldNode = &Node{
@@ -264,6 +268,14 @@ func (c *Cluster) newStateFromNode(ctx context.Context, node *v1.Node, oldNode *
 	return n, nil
 }
 
+func (c *Cluster) cleanupNode(name string) {
+	if id, ok := c.nameToProviderID[name]; ok {
+		delete(c.nodes, id)
+		delete(c.nameToProviderID, name)
+		c.RecordConsolidationChange()
+	}
+}
+
 func (c *Cluster) populateStartupTaints(ctx context.Context, n *Node) error {
 	if !n.Owned() {
 		return nil
@@ -325,7 +337,7 @@ func (c *Cluster) populateResourceRequests(ctx context.Context, n *Node) error {
 		}
 		c.cleanupOldBindings(pod)
 		n.updateForPod(ctx, pod)
-		c.bindings[client.ObjectKeyFromObject(pod)] = pod.Spec.NodeName // TODO @joinnis: Potentially change this later
+		c.bindings[client.ObjectKeyFromObject(pod)] = pod.Spec.NodeName
 	}
 	return nil
 }
@@ -338,9 +350,6 @@ func (c *Cluster) updateNodeUsageFromPod(ctx context.Context, pod *v1.Pod) error
 		return nil
 	}
 
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	n, ok := c.nodes[c.nameToProviderID[pod.Spec.NodeName]]
 	if !ok {
 		// the node must exist for us to update the resource requests on the node
@@ -353,9 +362,6 @@ func (c *Cluster) updateNodeUsageFromPod(ctx context.Context, pod *v1.Pod) error
 }
 
 func (c *Cluster) updateNodeUsageFromPodCompletion(podKey types.NamespacedName) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	nodeName, bindingKnown := c.bindings[podKey]
 	if !bindingKnown {
 		// we didn't think the pod was bound, so we weren't tracking it and don't need to do anything