Add logic for a compute-domain daemon to remove its IP on shutdown

Signed-off-by: Kevin Klues <[email protected]>

Author: klueska

cmd/compute-domain-daemon/computedomain.go

@@ -113,7 +113,19 @@ func (m *ComputeDomainManager) Start(ctx context.Context) (rerr error) {
 }
 
 // Stop stops the compute domain manager.
+//
+//nolint:contextcheck
 func (m *ComputeDomainManager) Stop() error {
+	// Create a new context for cleanup operations since the original context might be cancelled
+	cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cleanupCancel()
+
+	// Attempt to remove this node from the ComputeDomain status before shutting down
+	// Don't return error here as we still want to proceed with shutdown
+	if err := m.removeNodeFromComputeDomain(cleanupCtx); err != nil {
+		klog.Errorf("Failed to remove node from ComputeDomain during shutdown: %v", err)
+	}
+
 	if m.cancelContext != nil {
 		m.cancelContext()
 	}
@@ -226,6 +238,44 @@ func (m *ComputeDomainManager) GetNodesUpdateChan() chan []*nvapi.ComputeDomainN
 	return m.updatedNodesChan
 }
 
+// removeNodeFromComputeDomain removes the current node's entry from the ComputeDomain status.
+func (m *ComputeDomainManager) removeNodeFromComputeDomain(ctx context.Context) error {
+	objs := m.informer.GetIndexer().List()
+	if len(objs) == 0 {
+		klog.Infof("No ComputeDomain objects found in informer cache during cleanup")
+		return nil
+	}
+
+	cd, ok := objs[0].(*nvapi.ComputeDomain)
+	if !ok {
+		return fmt.Errorf("failed to cast object to ComputeDomain")
+	}
+
+	newCD := cd.DeepCopy()
+	var updatedNodes []*nvapi.ComputeDomainNode
+
+	for _, node := range newCD.Status.Nodes {
+		if node.Name != m.config.nodeName {
+			updatedNodes = append(updatedNodes, node)
+		} else {
+			klog.Infof("Removing node %s (%s) from ComputeDomain %s/%s during shutdown", node.Name, node.IPAddress, cd.Namespace, cd.Name)
+		}
+	}
+
+	if len(updatedNodes) == len(newCD.Status.Nodes) {
+		klog.Infof("Node %s not found in ComputeDomain %s/%s during cleanup", m.config.nodeName, cd.Namespace, cd.Name)
+		return nil
+	}
+
+	newCD.Status.Nodes = updatedNodes
+	if _, err := m.config.clientsets.Nvidia.ResourceV1beta1().ComputeDomains(newCD.Namespace).UpdateStatus(ctx, newCD, metav1.UpdateOptions{}); err != nil {
+		return fmt.Errorf("error removing node from ComputeDomain status: %w", err)
+	}
+
+	klog.Infof("Successfully removed node %s from ComputeDomain %s/%s", m.config.nodeName, cd.Namespace, cd.Name)
+	return nil
+}
+
 func getIPSet(nodeInfos []*nvapi.ComputeDomainNode) IPSet {
 	set := make(IPSet)
 	for _, n := range nodeInfos {