net-controller: preserve fresh status during WebSocket replacement

cmd/unbounded-net-controller/detail_dispatch.go

@@ -15,6 +15,15 @@ type nodeWSConnection struct {
 	send   func(context.Context, statusv1alpha1.DetailRequest) error
 }
 
+func (h *healthState) markNodeWSStale(nodeName string, connection *nodeWSConnection, source string) {
+	h.nodeWSMu.Lock()
+	defer h.nodeWSMu.Unlock()
+
+	if connection != nil && h.nodeWSRegistry[nodeName] == connection {
+		h.statusCache.UpdateSourceIf(nodeName, source, "stale-cache")
+	}
+}
+
 func (h *healthState) setNodeWSDetailSender(nodeName string, connection *nodeWSConnection, send func(context.Context, statusv1alpha1.DetailRequest) error) {
 	h.nodeWSMu.Lock()
 

cmd/unbounded-net-controller/node_status.go

@@ -436,14 +436,20 @@ func (c *NodeStatusCache) Delete(nodeName string) {
 // UpdateSource updates the cached status source for a node without changing
 // the cached payload or ReceivedAt timestamp.
 func (c *NodeStatusCache) UpdateSource(nodeName, source string) bool {
+	return c.UpdateSourceIf(nodeName, "", source)
+}
+
+// UpdateSourceIf changes the source only if the expected transport still owns it.
+// An empty expected source preserves the unconditional UpdateSource behavior.
+func (c *NodeStatusCache) UpdateSourceIf(nodeName, expectedSource, source string) bool {
 	if source == "" {
 		return false
 	}
 
 	c.mu.Lock()
 
 	entry, ok := c.entries[nodeName]
-	if !ok {
+	if !ok || (expectedSource != "" && entry.Source != expectedSource) {
 		c.mu.Unlock()
 		return false
 	}

cmd/unbounded-net-controller/server.go

@@ -667,7 +667,10 @@ func registerPushHandlers(mux *http.ServeMux, health *healthState, webhookServer
 			defer wsCancel()
 
 			var registration *nodeWSConnection
-			defer func() { health.unregisterNodeWS(lastWSNodeName, registration) }()
+			defer func() {
+				health.markNodeWSStale(lastWSNodeName, registration, source)
+				health.unregisterNodeWS(lastWSNodeName, registration)
+			}()
 
 			writeGate := make(chan struct{}, 1)
 
@@ -795,7 +798,7 @@ func registerPushHandlers(mux *http.ServeMux, health *healthState, webhookServer
 					return
 				case readErr := <-errCh:
 					if lastWSNodeName != "" {
-						health.statusCache.UpdateSource(lastWSNodeName, "stale-cache")
+						health.markNodeWSStale(lastWSNodeName, registration, source)
 					}
 					// Log graceful close frames and expected disconnections at
 					// V(4) to reduce noise during rolling restarts.
@@ -818,7 +821,7 @@ func registerPushHandlers(mux *http.ServeMux, health *healthState, webhookServer
 				case frame, ok := <-recvCh:
 					if !ok {
 						if lastWSNodeName != "" {
-							health.statusCache.UpdateSource(lastWSNodeName, "stale-cache")
+							health.markNodeWSStale(lastWSNodeName, registration, source)
 						}
 
 						return
@@ -948,7 +951,7 @@ func registerPushHandlers(mux *http.ServeMux, health *healthState, webhookServer
 							klog.V(2).Infof("Node WebSocket keepalive closing connection after reaching failure threshold (source=%s, node=%s, failures=%d, threshold=%d)", source, nodeNameLog, keepaliveFailures, health.statusWSKeepaliveFailureCount)
 
 							if lastWSNodeName != "" {
-								health.statusCache.UpdateSource(lastWSNodeName, "stale-cache")
+								health.markNodeWSStale(lastWSNodeName, registration, source)
 							}
 
 							if closeErr := conn.Close(websocket.StatusGoingAway, "keepalive failure threshold reached"); closeErr != nil {

cmd/unbounded-net-controller/ws_source_ownership_test.go

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+import "testing"
+
+func TestWebSocketTeardownPreservesNewerStatusSource(t *testing.T) {
+	health := &healthState{statusCache: NewNodeStatusCache()}
+	status := NodeStatusResponse{NodeInfo: NodeInfo{Name: "node"}}
+	health.statusCache.StoreFull("node", status, "ws")
+	old := health.registerNodeWS("node", func() {})
+	current := health.registerNodeWS("node", func() {})
+
+	health.markNodeWSStale("node", old, "ws")
+
+	if cached, _ := health.statusCache.Get("node"); cached.Source != "ws" {
+		t.Fatal("old connection teardown marked the replacement stale")
+	}
+
+	health.statusCache.StoreFull("node", status, "push")
+	health.markNodeWSStale("node", current, "ws")
+
+	if cached, _ := health.statusCache.Get("node"); cached.Source != "push" {
+		t.Fatal("WebSocket teardown overwrote a newer HTTP publication")
+	}
+
+	health.statusCache.StoreFull("node", status, "ws")
+	before := health.statusCache.GetAll()
+	health.markNodeWSStale("node", current, "ws")
+
+	if cached, _ := health.statusCache.Get("node"); cached.Source != "stale-cache" || before["node"].Source != "ws" {
+		t.Fatal("current teardown lost its stale signal or mutated an old snapshot")
+	}
+}