diff --git a/cmd/unbounded-net-controller/detail_dispatch.go b/cmd/unbounded-net-controller/detail_dispatch.go index 193d4db48..659ff45e2 100644 --- a/cmd/unbounded-net-controller/detail_dispatch.go +++ b/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() diff --git a/cmd/unbounded-net-controller/node_status.go b/cmd/unbounded-net-controller/node_status.go index eca669423..3dbb37b88 100644 --- a/cmd/unbounded-net-controller/node_status.go +++ b/cmd/unbounded-net-controller/node_status.go @@ -436,6 +436,12 @@ 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 } @@ -443,7 +449,7 @@ func (c *NodeStatusCache) UpdateSource(nodeName, source string) bool { c.mu.Lock() entry, ok := c.entries[nodeName] - if !ok { + if !ok || (expectedSource != "" && entry.Source != expectedSource) { c.mu.Unlock() return false } diff --git a/cmd/unbounded-net-controller/server.go b/cmd/unbounded-net-controller/server.go index 965d6e39f..dace866f1 100644 --- a/cmd/unbounded-net-controller/server.go +++ b/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 { diff --git a/cmd/unbounded-net-controller/ws_source_ownership_test.go b/cmd/unbounded-net-controller/ws_source_ownership_test.go new file mode 100644 index 000000000..ea1e1900a --- /dev/null +++ b/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") + } +}