From 87fe506c0ac3c1996c2c0fcf7fc76b4e9e2273dd Mon Sep 17 00:00:00 2001 From: rcourtman Date: Thu, 23 Jul 2026 22:23:30 +0100 Subject: [PATCH] Bound state refresh and recovery ingest ownership --- .../v6/internal/subsystems/api-contracts.md | 10 +- .../v6/internal/subsystems/monitoring.md | 9 +- .../subsystems/performance-and-scalability.md | 7 +- .../internal/subsystems/unified-resources.md | 4 +- internal/monitoring/recovery_ingest.go | 51 +- internal/monitoring/recovery_ingest_test.go | 56 +++ internal/unifiedresources/retention_test.go | 7 +- internal/unifiedresources/store.go | 19 +- internal/unifiedresources/store_test.go | 20 + internal/websocket/hub.go | 468 ++++++++++++------ internal/websocket/hub_more_test.go | 12 + internal/websocket/hub_test.go | 176 +++++++ 12 files changed, 676 insertions(+), 163 deletions(-) diff --git a/docs/release-control/v6/internal/subsystems/api-contracts.md b/docs/release-control/v6/internal/subsystems/api-contracts.md index 22cd35af4..41fe57e52 100644 --- a/docs/release-control/v6/internal/subsystems/api-contracts.md +++ b/docs/release-control/v6/internal/subsystems/api-contracts.md @@ -253,7 +253,15 @@ not a parallel API payload contract. `internal/api/agent_handlers_base.go` and `BroadcastCurrentState` / `BroadcastCurrentStateToTenant`; they must not build or retain full frontend-state payloads at the handler boundary. The WebSocket hub owns tenant-aware state resolution after coalescing, through the same state -getter that backs the canonical `/api/state` payload. +getter that backs the canonical `/api/state` payload. The hub must serialize +whole-state resolution and JSON construction across coalesced broadcasts, +initial-client delivery, and explicit client data requests; cancel delayed or +queued work when its client leaves; and join every state producer before +closing client channels. Reconnect, request, or invalidation churn must not +multiply concurrent clones of that canonical payload or race shutdown sends +against channel closure. Ordinary unregister and slow-client eviction must use +the same synchronized send/close ownership; recovering a send-on-closed panic +is not a substitute for a race-free channel lifecycle. Docker and Podman app-container CPU payloads expose two API facts with different meanings: canonical resource metrics and `/api/metrics-store/history` diff --git a/docs/release-control/v6/internal/subsystems/monitoring.md b/docs/release-control/v6/internal/subsystems/monitoring.md index 7dd113875..a829bcc60 100644 --- a/docs/release-control/v6/internal/subsystems/monitoring.md +++ b/docs/release-control/v6/internal/subsystems/monitoring.md @@ -70,7 +70,10 @@ must be pruned to the retained group set after a completed poll, while preserving cache metadata only for groups still observed or intentionally reused after transient datastore failures. Recovery-point ingestion started by backup polling must be serialized and coalesced so slow store writes cannot retain one -full backup point batch per poll cycle. +full backup point batch per poll cycle. Complete authoritative enumerations +coalesce only within the same provider, ID-prefix, and instance scope; distinct +source scopes and non-reconciling event batches remain FIFO so bounding memory +does not discard independent recovery facts. Removed host-agent reconnect blocks are identity-scoped: matching may use the canonical host ID or token-qualified machine/hostname continuity, but must never block a distinct live host by hostname alone. @@ -345,7 +348,9 @@ data cannot wrap into a fabricated healthy value. worker count rather than datastore cardinality. Per-group cache timestamps must be removed when successful group discovery no longer retains that group, and recovery-store ingestion must be a bounded latest-batch pipeline - rather than an untracked goroutine per poll. + rather than an untracked goroutine per poll. Latest-batch replacement applies + only to complete enumerations with the same provider, ID-prefix, and instance + scope; distinct scopes and event batches must remain independently queued. 12. Add or change agentless availability monitoring only through the poll-provider path. `internal/monitoring/availability_poller.go` owns ICMP, TCP, and HTTP probes, provider health, scheduler task construction, and diff --git a/docs/release-control/v6/internal/subsystems/performance-and-scalability.md b/docs/release-control/v6/internal/subsystems/performance-and-scalability.md index e9514d955..4c3019062 100644 --- a/docs/release-control/v6/internal/subsystems/performance-and-scalability.md +++ b/docs/release-control/v6/internal/subsystems/performance-and-scalability.md @@ -193,7 +193,12 @@ change may globally weaken the Task 03 lifecycle-state idempotency invariant. paths must enqueue coalesced current-state invalidations instead of building full frontend-state payloads at every signal. The hub owns delayed tenant-aware state resolution so superseded signals do not retain large - state snapshots. + state snapshots. Full state resolution and serialization must run through + one cancellable build slot across broadcast, initial-client delivery, and + client-requested refresh delivery: reconnect churn must cancel work for + clients that have already left, and active reconnects or refreshes must not + overlap whole-state clones with each other or with the coalesced broadcast + worker. The operations-loop status endpoint is performance-adjacent because it aggregates fleet, action-audit, workflow-starter, AI-usage, and external-agent activity evidence for every request. Starter and contextual diff --git a/docs/release-control/v6/internal/subsystems/unified-resources.md b/docs/release-control/v6/internal/subsystems/unified-resources.md index 50c033ef9..cc8d19d6a 100644 --- a/docs/release-control/v6/internal/subsystems/unified-resources.md +++ b/docs/release-control/v6/internal/subsystems/unified-resources.md @@ -3990,7 +3990,9 @@ per cycle and grew `resource_changes` without bound (issue #1496, demo outage 2026-07-08). Retention pruning must also enforce a hard row cap on `resource_changes` (`maxResourceChangesRows` in `store.go`) so a pathological writer cannot grow the table unbounded inside the time-based -retention window. +retention window. Store shutdown must close the retention signal exactly once, +join the retention loop, and only then close the single-connection SQLite +pool; a concurrent pruning pass must never outlive its database owner. Action plans in `actions.go` still keep stale-plan protection to the canonical `resourceVersion`, `policyVersion`, and `planHash` fields, so stale execution checks stay in the shared resource action model rather than provider-local diff --git a/internal/monitoring/recovery_ingest.go b/internal/monitoring/recovery_ingest.go index 505a53cbd..84533459f 100644 --- a/internal/monitoring/recovery_ingest.go +++ b/internal/monitoring/recovery_ingest.go @@ -29,6 +29,38 @@ type recoveryIngestBatch struct { reconcile *recoveryReconcileScope } +type recoveryIngestCoalesceKey struct { + provider string + idPrefix string + instance string +} + +func (batch recoveryIngestBatch) coalesceKey() (recoveryIngestCoalesceKey, bool) { + if batch.reconcile == nil { + return recoveryIngestCoalesceKey{}, false + } + return recoveryIngestCoalesceKey{ + provider: batch.reconcile.provider, + idPrefix: batch.reconcile.idPrefix, + instance: batch.reconcile.instance, + }, true +} + +func (m *Monitor) coalescePendingRecoveryIngest(batch recoveryIngestBatch) bool { + key, ok := batch.coalesceKey() + if !ok { + return false + } + for i := len(m.recoveryIngestPending) - 1; i >= 0; i-- { + pendingKey, pendingOK := m.recoveryIngestPending[i].coalesceKey() + if pendingOK && pendingKey == key { + m.recoveryIngestPending[i] = batch + return true + } + } + return false +} + func (m *Monitor) ingestRecoveryPointsAsync(points []recovery.RecoveryPoint) { m.enqueueRecoveryIngest(recoveryIngestBatch{points: points}) } @@ -68,14 +100,27 @@ func (m *Monitor) enqueueRecoveryIngest(batch recoveryIngestBatch) { m.recoveryIngestMu.Lock() if m.recoveryIngestRunning { - // Queue rather than replace: batches come from different sources (PVE - // storage, PBS, snapshots), and dropping one loses a full poll cycle - // for that source. + // A reconcile batch is a complete current enumeration for one source + // scope, so only its newest pending value is authoritative. Distinct + // scopes and non-reconciling event batches remain FIFO; this preserves + // their facts without retaining every superseded full poll result. + if m.coalescePendingRecoveryIngest(batch) { + pending := len(m.recoveryIngestPending) + m.recoveryIngestMu.Unlock() + log.Debug(). + Int("points", len(batch.points)). + Int("provider_observations", len(batch.observations)). + Int("pending_batches", pending). + Msg("Coalesced recovery point ingest behind active batch") + return + } m.recoveryIngestPending = append(m.recoveryIngestPending, batch) + pending := len(m.recoveryIngestPending) m.recoveryIngestMu.Unlock() log.Debug(). Int("points", len(batch.points)). Int("provider_observations", len(batch.observations)). + Int("pending_batches", pending). Msg("Queued recovery point ingest behind active batch") return } diff --git a/internal/monitoring/recovery_ingest_test.go b/internal/monitoring/recovery_ingest_test.go index c8d54944f..df4cdd1d8 100644 --- a/internal/monitoring/recovery_ingest_test.go +++ b/internal/monitoring/recovery_ingest_test.go @@ -267,3 +267,59 @@ func TestIngestRecoveryPointsBestEffortPersistsProviderObservationBeforePointFai ) } } + +func TestRecoveryIngestAcceptedProofCoalescesSameScope(t *testing.T) { + monitor := &Monitor{recoveryIngestRunning: true} + scope := recoveryReconcileScope{ + provider: string(recovery.ProviderProxmoxPBS), + idPrefix: "pbs-backup:", + instance: "pbs-large", + } + + monitor.enqueueRecoveryIngest(recoveryIngestBatch{ + points: []recovery.RecoveryPoint{{ID: "pbs-backup:superseded"}}, + reconcile: &scope, + }) + for i := 1; i < 64; i++ { + monitor.enqueueRecoveryIngest(recoveryIngestBatch{ + points: []recovery.RecoveryPoint{{ID: "pbs-backup:latest"}}, + reconcile: &scope, + }) + } + if got := len(monitor.recoveryIngestPending); got != 1 { + t.Fatalf("pending complete enumerations = %d, want one same-scope latest batch", got) + } + if got := monitor.recoveryIngestPending[0].points[0].ID; got != "pbs-backup:latest" { + t.Fatalf("pending point = %q, want latest complete enumeration", got) + } +} + +func TestRecoveryIngestPendingKeepsDistinctScopesAndEventBatches(t *testing.T) { + monitor := &Monitor{recoveryIngestRunning: true} + pbsScope := recoveryReconcileScope{ + provider: string(recovery.ProviderProxmoxPBS), + idPrefix: "pbs-backup:", + instance: "pbs-a", + } + pveScope := recoveryReconcileScope{ + provider: string(recovery.ProviderProxmoxPVE), + idPrefix: "pve-snapshot:", + instance: "pve-a", + } + + monitor.enqueueRecoveryIngest(recoveryIngestBatch{ + points: []recovery.RecoveryPoint{{ID: "pbs-backup:1"}}, + reconcile: &pbsScope, + }) + monitor.enqueueRecoveryIngest(recoveryIngestBatch{ + points: []recovery.RecoveryPoint{{ID: "pve-snapshot:1"}}, + reconcile: &pveScope, + }) + monitor.enqueueRecoveryIngest(recoveryIngestBatch{ + points: []recovery.RecoveryPoint{{ID: "pve-task:1"}}, + }) + + if got := len(monitor.recoveryIngestPending); got != 3 { + t.Fatalf("pending batches = %d, want two source enumerations and one event batch", got) + } +} diff --git a/internal/unifiedresources/retention_test.go b/internal/unifiedresources/retention_test.go index be502340c..b6650ed81 100644 --- a/internal/unifiedresources/retention_test.go +++ b/internal/unifiedresources/retention_test.go @@ -239,8 +239,11 @@ func TestRetentionLoop_RunsInitialPrune(t *testing.T) { t.Fatalf("RecordChange: %v", err) } - stop := store.startRetentionLoop() - defer close(stop) + stop, done := store.startRetentionLoop() + defer func() { + close(stop) + <-done + }() deadline := time.Now().Add(initialRetentionDelay + 10*time.Second) for time.Now().Before(deadline) { diff --git a/internal/unifiedresources/store.go b/internal/unifiedresources/store.go index dcd4d2d4e..9eb31badc 100644 --- a/internal/unifiedresources/store.go +++ b/internal/unifiedresources/store.go @@ -153,7 +153,9 @@ type SQLiteResourceStore struct { identityPinCache []ResourceIdentityPin identityPinFresh bool - retentionStop chan struct{} + retentionStop chan struct{} + retentionDone chan struct{} + retentionStopOnce sync.Once } const ( @@ -254,7 +256,7 @@ func NewSQLiteResourceStore(dataDir, orgID string) (*SQLiteResourceStore, error) log.Printf("[INFO] unified_resources: database %q recreated successfully after corruption recovery", path) } store.migrateAutoVacuum() - store.retentionStop = store.startRetentionLoop() + store.retentionStop, store.retentionDone = store.startRetentionLoop() return store, nil } @@ -1317,9 +1319,11 @@ func (s *SQLiteResourceStore) reclaimFreePages() { // Without this, resource_changes, action_audits, and related tables grow // without bound and the database file never shrinks (GitHub issue #1496). // Returns a stop channel; closing it signals the goroutine to exit. -func (s *SQLiteResourceStore) startRetentionLoop() chan struct{} { +func (s *SQLiteResourceStore) startRetentionLoop() (chan struct{}, chan struct{}) { stop := make(chan struct{}) + done := make(chan struct{}) go func() { + defer close(done) initial := time.NewTimer(initialRetentionDelay) defer initial.Stop() ticker := time.NewTicker(retentionInterval) @@ -1335,7 +1339,7 @@ func (s *SQLiteResourceStore) startRetentionLoop() chan struct{} { } } }() - return stop + return stop, done } // capResourceChanges deletes the oldest resource_changes rows beyond limit, @@ -1438,7 +1442,12 @@ func (s *SQLiteResourceStore) pruneOldRecords() { func (s *SQLiteResourceStore) Close() error { if s.retentionStop != nil { - close(s.retentionStop) + s.retentionStopOnce.Do(func() { + close(s.retentionStop) + }) + } + if s.retentionDone != nil { + <-s.retentionDone } if s.db != nil { if err := s.db.Close(); err != nil { diff --git a/internal/unifiedresources/store_test.go b/internal/unifiedresources/store_test.go index 6b01fda6a..b6e03b2fe 100644 --- a/internal/unifiedresources/store_test.go +++ b/internal/unifiedresources/store_test.go @@ -3215,3 +3215,23 @@ func TestPendingActionAuditReaderReturnsOldestPendingFirst(t *testing.T) { t.Fatalf("pending actions = %#v, want oldest pending first", got) } } + +func TestSQLiteResourceStoreCloseJoinsRetentionLoopAndIsIdempotent(t *testing.T) { + store, err := NewSQLiteResourceStore(t.TempDir(), "default") + if err != nil { + t.Fatalf("NewSQLiteResourceStore: %v", err) + } + + if err := store.Close(); err != nil { + t.Fatalf("first Close: %v", err) + } + select { + case <-store.retentionDone: + default: + t.Fatal("Close returned before the retention loop exited") + } + + if err := store.Close(); err != nil { + t.Fatalf("second Close: %v", err) + } +} diff --git a/internal/websocket/hub.go b/internal/websocket/hub.go index 518163d3d..69a0bb888 100644 --- a/internal/websocket/hub.go +++ b/internal/websocket/hub.go @@ -26,8 +26,10 @@ const ( // maxWebSocketInboundMessageSize bounds client->server websocket message size. maxWebSocketInboundMessageSize = 64 * 1024 // maxWebSocketOrgIDLength keeps org IDs bounded to prevent oversized header/query abuse. - maxWebSocketOrgIDLength = 64 - websocketHubComponent = "websocket_hub" + maxWebSocketOrgIDLength = 64 + websocketHubComponent = "websocket_hub" + initialWelcomeDelay = 500 * time.Millisecond + initialStateMessageDelay = 100 * time.Millisecond ) // extractPeerIP extracts just the IP part from a RemoteAddr (host:port format) @@ -250,31 +252,48 @@ type Client struct { hub *Hub conn *websocket.Conn send chan []byte + sendMu sync.RWMutex id string orgID string // Organization ID for tenant isolation lastPing time.Time closed atomic.Bool // Set when the client is unregistered; prevents sends to closed channel writeFailures int32 // Consecutive write failures; disconnects after maxWriteFailures + lifecycleDone chan struct{} + lifecycleOnce sync.Once +} + +func (c *Client) initializeLifecycle() { + if c.lifecycleDone == nil { + c.lifecycleDone = make(chan struct{}) + } +} + +func (c *Client) closeLifecycle() { + c.lifecycleOnce.Do(func() { + if c.lifecycleDone != nil { + close(c.lifecycleDone) + } + }) +} + +func (c *Client) closeSend() { + c.sendMu.Lock() + defer c.sendMu.Unlock() + if !c.closed.Swap(true) { + close(c.send) + } } // safeSend attempts to send data to the client's send channel. // Returns false if the client is closed or the channel buffer is full. -// Uses defer/recover to handle the race between close(c.send) and send. func (c *Client) safeSend(data []byte) (sent bool) { - // Early check to avoid most attempts on closed clients. + c.sendMu.RLock() + defer c.sendMu.RUnlock() + if c.closed.Load() { return false } - // Recover from panic if the channel was closed between the check above - // and the send below. This is a defensive pattern to prevent server crashes. - defer func() { - if r := recover(); r != nil { - // Channel was closed concurrently; mark as not sent. - sent = false - } - }() - select { case c.send <- data: return true @@ -413,6 +432,8 @@ type Hub struct { runDone chan struct{} runStartOnce sync.Once runDoneOnce sync.Once + initialStateWG sync.WaitGroup + stateBuildSlot chan struct{} mu sync.RWMutex getState func(orgID string) interface{} // Function to get state for specific tenant allowedOrigins []string // Allowed origins for CORS @@ -420,13 +441,20 @@ type Hub struct { multiTenantChecker MultiTenantChecker // Multi-tenant feature flag and license checker isTrustedProxy func(ip string) bool // Optional: checks if peer IP is a trusted reverse proxy // Broadcast coalescing fields - coalesceWindow time.Duration - coalescePending *Message - coalesceTimer *time.Timer - coalesceMutex sync.Mutex + coalesceWindow time.Duration + coalescePending *Message + coalesceReady *Message + coalesceTimer *time.Timer + coalesceGeneration uint64 + nextGeneration uint64 + coalesceMutex sync.Mutex // Per-tenant coalescing - tenantCoalescePending map[string]*Message - tenantCoalesceTimers map[string]*time.Timer + tenantCoalescePending map[string]*Message + tenantCoalesceReady map[string]*Message + tenantCoalesceTimers map[string]*time.Timer + tenantCoalesceGeneration map[string]uint64 + stateBroadcastWake chan struct{} + stateBroadcastDone chan struct{} } // Message represents a WebSocket message @@ -495,21 +523,26 @@ func (h *Hub) getStateForOrg(orgID string) interface{} { // NewHub creates a new WebSocket hub func NewHub(getState func(orgID string) interface{}) *Hub { return &Hub{ - clients: make(map[*Client]bool), - clientsByTenant: make(map[string]map[*Client]bool), - broadcast: make(chan []byte, 256), - broadcastSeq: make(chan Message, 256), // Buffered sequenced channel - tenantBroadcast: make(chan TenantBroadcast, 256), // Per-tenant broadcasts - register: make(chan *Client), - unregister: make(chan *Client), - stopChan: make(chan struct{}), - runStarted: make(chan struct{}), - runDone: make(chan struct{}), - getState: getState, - allowedOrigins: []string{}, // Default to empty (will be set based on actual host) - coalesceWindow: 100 * time.Millisecond, // Coalesce rapid updates within 100ms - tenantCoalescePending: make(map[string]*Message), - tenantCoalesceTimers: make(map[string]*time.Timer), + clients: make(map[*Client]bool), + clientsByTenant: make(map[string]map[*Client]bool), + broadcast: make(chan []byte, 256), + broadcastSeq: make(chan Message, 256), // Buffered sequenced channel + tenantBroadcast: make(chan TenantBroadcast, 256), // Per-tenant broadcasts + register: make(chan *Client), + unregister: make(chan *Client), + stopChan: make(chan struct{}), + runStarted: make(chan struct{}), + runDone: make(chan struct{}), + stateBuildSlot: make(chan struct{}, 1), + getState: getState, + allowedOrigins: []string{}, // Default to empty (will be set based on actual host) + coalesceWindow: 100 * time.Millisecond, // Coalesce rapid updates within 100ms + tenantCoalescePending: make(map[string]*Message), + tenantCoalesceReady: make(map[string]*Message), + tenantCoalesceTimers: make(map[string]*time.Timer), + tenantCoalesceGeneration: make(map[string]uint64), + stateBroadcastWake: make(chan struct{}, 1), + stateBroadcastDone: make(chan struct{}), } } @@ -525,8 +558,10 @@ func (h *Hub) Run() { defer close(sequencerDone) h.runBroadcastSequencer() }() + go h.runStateBroadcastWorker() defer func() { <-sequencerDone + <-h.stateBroadcastDone h.runDoneOnce.Do(func() { close(h.runDone) }) @@ -539,6 +574,7 @@ func (h *Hub) Run() { for { select { case client := <-h.register: + client.initializeLifecycle() h.mu.Lock() h.clients[client] = true // Also register by tenant if org ID is set @@ -556,67 +592,10 @@ func (h *Hub) Run() { hasGetState := h.hasStateGetter() log.Debug().Bool("hasGetState", hasGetState).Msg("Checking getState function for new client") if hasGetState { - // Add a small delay to ensure client is ready + h.initialStateWG.Add(1) go func() { - log.Debug().Str("client", client.id).Msg("starting initial state goroutine") - time.Sleep(500 * time.Millisecond) - - // First send a small welcome message - welcomeMsg := Message{ - Type: "welcome", - Data: map[string]string{"message": "Connected to Pulse WebSocket", "orgId": client.orgID}, - } - if data, err := json.Marshal(welcomeMsg); err == nil { - // Check if client is still registered before sending (must hold lock) - h.mu.RLock() - _, stillRegistered := h.clients[client] - h.mu.RUnlock() - - if stillRegistered { - log.Info().Str("client", client.id).Msg("sending welcome message") - if client.safeSend(data) { - log.Info().Str("client", client.id).Msg("welcome message sent") - } else { - log.Warn().Str("client", client.id).Msg("failed to send welcome message - client closed or buffer full") - } - } else { - log.Debug().Str("client", client.id).Msg("client disconnected before welcome message") - } - } else { - log.Error().Err(err).Str("client", client.id).Msg("Failed to marshal welcome message") - } - - // Then send the initial state after another delay - time.Sleep(100 * time.Millisecond) - log.Debug().Str("client", client.id).Msg("about to get state") - - // Get the state using tenant-aware getter - stateData := h.getStateForClient(client) - log.Debug().Str("client", client.id).Interface("stateType", fmt.Sprintf("%T", stateData)).Msg("got state for initial message") - - initialMsg := Message{ - Type: "initialState", - Data: sanitizeData(h.prepareStateForBroadcast(stateData)), - } - if data, err := json.Marshal(initialMsg); err == nil { - // Check if client is still registered before sending (must hold lock) - h.mu.RLock() - _, stillRegistered := h.clients[client] - h.mu.RUnlock() - - if stillRegistered { - log.Info().Str("client", client.id).Int("dataLen", len(data)).Int("dataKB", len(data)/1024).Msg("sending initial state to client") - if client.safeSend(data) { - log.Info().Str("client", client.id).Msg("initial state sent successfully") - } else { - log.Warn().Str("client", client.id).Msg("client closed or buffer full, skipping initial state") - } - } else { - log.Debug().Str("client", client.id).Msg("client disconnected before initial state") - } - } else { - log.Error().Err(err).Str("client", client.id).Msg("failed to marshal initial state") - } + defer h.initialStateWG.Done() + h.sendInitialState(client) }() } else { log.Warn(). @@ -644,18 +623,32 @@ func (h *Hub) Run() { case <-h.stopChan: log.Info().Msg("webSocket hub shutting down") - // Close all client connections + // Cancel delayed initial-state work first, then let every producer + // finish before closing client channels. Broadcast workers and + // initial-state delivery otherwise race a final safeSend against + // channel closure during shutdown. h.mu.Lock() for client := range h.clients { - if !client.closed.Swap(true) { - close(client.send) - } + client.closeLifecycle() } for _, tenantClients := range h.clientsByTenant { for client := range tenantClients { - if !client.closed.Swap(true) { - close(client.send) - } + client.closeLifecycle() + } + } + h.mu.Unlock() + + <-sequencerDone + <-h.stateBroadcastDone + h.initialStateWG.Wait() + + h.mu.Lock() + for client := range h.clients { + client.closeSend() + } + for _, tenantClients := range h.clientsByTenant { + for client := range tenantClients { + client.closeSend() } } h.clients = make(map[*Client]bool) @@ -676,6 +669,7 @@ func (h *Hub) Stop() { <-h.runDone default: } + h.initialStateWG.Wait() } func (h *Hub) isStopping() bool { @@ -687,6 +681,140 @@ func (h *Hub) isStopping() bool { } } +func (h *Hub) waitForActiveClient(client *Client, delay time.Duration) bool { + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-timer.C: + h.mu.RLock() + _, registered := h.clients[client] + h.mu.RUnlock() + return registered && !client.closed.Load() && !h.isStopping() + case <-client.lifecycleDone: + return false + case <-h.stopChan: + return false + } +} + +func (h *Hub) sendInitialState(client *Client) { + log.Debug().Str("client", client.id).Msg("starting initial state goroutine") + if !h.waitForActiveClient(client, initialWelcomeDelay) { + log.Debug().Str("client", client.id).Msg("client disconnected before welcome message") + return + } + + welcomeMsg := Message{ + Type: "welcome", + Data: map[string]string{"message": "Connected to Pulse WebSocket", "orgId": client.orgID}, + } + data, err := json.Marshal(welcomeMsg) + if err != nil { + log.Error().Err(err).Str("client", client.id).Msg("Failed to marshal welcome message") + return + } + log.Info().Str("client", client.id).Msg("sending welcome message") + if !client.safeSend(data) { + log.Warn().Str("client", client.id).Msg("failed to send welcome message - client closed or buffer full") + return + } + log.Info().Str("client", client.id).Msg("welcome message sent") + + if !h.waitForActiveClient(client, initialStateMessageDelay) { + log.Debug().Str("client", client.id).Msg("client disconnected before initial state") + return + } + if !h.acquireStateBuildSlot(client) { + log.Debug().Str("client", client.id).Msg("client disconnected while waiting to build initial state") + return + } + defer h.releaseStateBuildSlot() + + log.Debug().Str("client", client.id).Msg("about to get state") + stateData := h.getStateForClient(client) + log.Debug().Str("client", client.id).Interface("stateType", fmt.Sprintf("%T", stateData)).Msg("got state for initial message") + initialMsg := Message{ + Type: "initialState", + Data: sanitizeData(h.prepareStateForBroadcast(stateData)), + } + data, err = json.Marshal(initialMsg) + if err != nil { + log.Error().Err(err).Str("client", client.id).Msg("failed to marshal initial state") + return + } + + h.mu.RLock() + _, stillRegistered := h.clients[client] + h.mu.RUnlock() + if !stillRegistered { + log.Debug().Str("client", client.id).Msg("client disconnected before initial state") + return + } + + log.Info().Str("client", client.id).Int("dataLen", len(data)).Int("dataKB", len(data)/1024).Msg("sending initial state to client") + if client.safeSend(data) { + log.Info().Str("client", client.id).Msg("initial state sent successfully") + } else { + log.Warn().Str("client", client.id).Msg("client closed or buffer full, skipping initial state") + } +} + +func (h *Hub) sendRequestedState(client *Client) { + if client == nil || !h.hasStateGetter() { + return + } + if !h.acquireStateBuildSlot(client) { + log.Debug().Str("client", client.id).Msg("client disconnected while waiting to build requested state") + return + } + defer h.releaseStateBuildSlot() + + stateMsg := Message{ + Type: "rawData", + Data: sanitizeData(h.prepareStateForBroadcast(h.getStateForClient(client))), + } + data, err := json.Marshal(stateMsg) + if err != nil { + log.Error().Err(err).Str("client", client.id).Msg("failed to marshal state for requestData") + return + } + if !client.safeSend(data) { + log.Warn().Str("client", client.id).Msg("Failed to queue requestData state response; client channel closed or full") + } +} + +func (h *Hub) acquireStateBuildSlot(client *Client) bool { + if client == nil { + select { + case h.stateBuildSlot <- struct{}{}: + return true + case <-h.stopChan: + return false + } + } + + select { + case h.stateBuildSlot <- struct{}{}: + h.mu.RLock() + _, registered := h.clients[client] + h.mu.RUnlock() + if registered && !client.closed.Load() && !h.isStopping() { + return true + } + h.releaseStateBuildSlot() + return false + case <-client.lifecycleDone: + return false + case <-h.stopChan: + return false + } +} + +func (h *Hub) releaseStateBuildSlot() { + <-h.stateBuildSlot +} + func (h *Hub) tryRegisterClient(client *Client) bool { if h.isStopping() { return false @@ -891,8 +1019,9 @@ func (h *Hub) removeClientLocked(client *Client) bool { } } - if removed && !client.closed.Swap(true) { - close(client.send) + if removed { + client.closeLifecycle() + client.closeSend() } return removed @@ -919,18 +1048,89 @@ func (h *Hub) dispatchToClients(data []byte, dropLog string) { } } -func (h *Hub) popCoalescedMessage() *Message { +func (h *Hub) markGlobalStateBroadcastReady(generation uint64) { + h.coalesceMutex.Lock() + if generation != h.coalesceGeneration || h.coalescePending == nil { + h.coalesceMutex.Unlock() + return + } + h.coalesceReady = h.coalescePending + h.coalescePending = nil + h.coalesceTimer = nil + h.coalesceMutex.Unlock() + h.wakeStateBroadcastWorker() +} + +func (h *Hub) markTenantStateBroadcastReady(orgID string, generation uint64) { + h.coalesceMutex.Lock() + if generation != h.tenantCoalesceGeneration[orgID] { + h.coalesceMutex.Unlock() + return + } + pending := h.tenantCoalescePending[orgID] + if pending == nil { + h.coalesceMutex.Unlock() + return + } + h.tenantCoalesceReady[orgID] = pending + delete(h.tenantCoalescePending, orgID) + delete(h.tenantCoalesceTimers, orgID) + delete(h.tenantCoalesceGeneration, orgID) + h.coalesceMutex.Unlock() + h.wakeStateBroadcastWorker() +} + +func (h *Hub) wakeStateBroadcastWorker() { + select { + case h.stateBroadcastWake <- struct{}{}: + default: + } +} + +func (h *Hub) popReadyStateBroadcasts() (*Message, map[string]*Message) { h.coalesceMutex.Lock() defer h.coalesceMutex.Unlock() - if h.coalescePending == nil { - return nil - } + global := h.coalesceReady + h.coalesceReady = nil + tenants := h.tenantCoalesceReady + h.tenantCoalesceReady = make(map[string]*Message) + return global, tenants +} - msg := *h.coalescePending - h.coalescePending = nil - h.coalesceTimer = nil - return &msg +func (h *Hub) dispatchStateBroadcast(pending *Message, orgID string) { + if pending == nil || h.isStopping() || !h.hasRecipientsForMessage(*pending, orgID) { + return + } + if !h.acquireStateBuildSlot(nil) { + return + } + defer h.releaseStateBuildSlot() + data, ok := h.marshalBroadcastMessage(*pending, orgID) + if !ok { + return + } + if orgID == "" { + h.dispatchToClients(data, "Client send channel full, dropping coalesced message and closing connection") + return + } + h.dispatchToTenantClients(orgID, data, "Client send channel full, dropping tenant coalesced message and closing connection") +} + +func (h *Hub) runStateBroadcastWorker() { + defer close(h.stateBroadcastDone) + for { + select { + case <-h.stateBroadcastWake: + global, tenants := h.popReadyStateBroadcasts() + h.dispatchStateBroadcast(global, "") + for orgID, pending := range tenants { + h.dispatchStateBroadcast(pending, orgID) + } + case <-h.stopChan: + return + } + } } func (h *Hub) messageHasCurrentStateRequest(msg Message) (stateBroadcastRequest, bool) { @@ -1009,18 +1209,13 @@ func (h *Hub) runBroadcastSequencer() { // Update pending message current := msg h.coalescePending = ¤t + h.nextGeneration++ + generation := h.nextGeneration + h.coalesceGeneration = generation // Set timer to send after coalesce window h.coalesceTimer = time.AfterFunc(h.coalesceWindow, func() { - pending := h.popCoalescedMessage() - if pending != nil { - if !h.hasRecipientsForMessage(*pending, "") { - return - } - if data, ok := h.marshalBroadcastMessage(*pending, ""); ok { - h.dispatchToClients(data, "Client send channel full, dropping coalesced message and closing connection") - } - } + h.markGlobalStateBroadcastReady(generation) }) h.coalesceMutex.Unlock() @@ -1044,24 +1239,14 @@ func (h *Hub) runBroadcastSequencer() { // Update pending message for this tenant msgCopy := tb.Message h.tenantCoalescePending[tb.OrgID] = &msgCopy + h.nextGeneration++ + generation := h.nextGeneration + h.tenantCoalesceGeneration[tb.OrgID] = generation // Set timer to send after coalesce window orgID := tb.OrgID // Capture for closure h.tenantCoalesceTimers[orgID] = time.AfterFunc(h.coalesceWindow, func() { - h.coalesceMutex.Lock() - pending := h.tenantCoalescePending[orgID] - delete(h.tenantCoalescePending, orgID) - delete(h.tenantCoalesceTimers, orgID) - h.coalesceMutex.Unlock() - - if pending != nil { - if !h.hasRecipientsForMessage(*pending, orgID) { - return - } - if data, ok := h.marshalBroadcastMessage(*pending, orgID); ok { - h.dispatchToTenantClients(orgID, data, "Client send channel full, dropping tenant coalesced message and closing connection") - } - } + h.markTenantStateBroadcastReady(orgID, generation) }) h.coalesceMutex.Unlock() @@ -1447,20 +1632,7 @@ func (c *Client) readPump() { log.Error().Err(err).Str("client", c.id).Msg("Failed to marshal pong response") } case "requestData": - // Send current state with lock-safe getter lookup. - if c.hub.hasStateGetter() { - stateMsg := Message{ - Type: "rawData", - Data: sanitizeData(c.hub.prepareStateForBroadcast(c.hub.getStateForClient(c))), - } - if data, err := json.Marshal(stateMsg); err == nil { - if !c.safeSend(data) { - log.Warn().Str("client", c.id).Msg("Failed to queue requestData state response; client channel closed or full") - } - } else { - log.Error().Err(err).Str("client", c.id).Msg("failed to marshal state for requestData") - } - } + c.hub.sendRequestedState(c) default: log.Debug().Str("client", c.id).Str("type", msg.Type).Msg("received WebSocket message") } diff --git a/internal/websocket/hub_more_test.go b/internal/websocket/hub_more_test.go index 6ff944119..31033ba78 100644 --- a/internal/websocket/hub_more_test.go +++ b/internal/websocket/hub_more_test.go @@ -49,6 +49,7 @@ func TestRunBroadcastSequencerImmediate(t *testing.T) { hub.runBroadcastSequencer() close(done) }() + go hub.runStateBroadcastWorker() hub.broadcastSeq <- Message{ Type: "alert", @@ -74,6 +75,11 @@ func TestRunBroadcastSequencerImmediate(t *testing.T) { case <-time.After(200 * time.Millisecond): t.Fatal("broadcast sequencer did not exit") } + select { + case <-hub.stateBroadcastDone: + case <-time.After(200 * time.Millisecond): + t.Fatal("state broadcast worker did not exit") + } } func TestRunBroadcastSequencerCoalescesRawData(t *testing.T) { @@ -94,6 +100,7 @@ func TestRunBroadcastSequencerCoalescesRawData(t *testing.T) { hub.runBroadcastSequencer() close(done) }() + go hub.runStateBroadcastWorker() hub.broadcastSeq <- Message{Type: "rawData", Data: map[string]string{"value": "first"}} hub.broadcastSeq <- Message{Type: "rawData", Data: map[string]string{"value": "second"}} @@ -124,4 +131,9 @@ func TestRunBroadcastSequencerCoalescesRawData(t *testing.T) { case <-time.After(200 * time.Millisecond): t.Fatal("broadcast sequencer did not exit") } + select { + case <-hub.stateBroadcastDone: + case <-time.After(200 * time.Millisecond): + t.Fatal("state broadcast worker did not exit") + } } diff --git a/internal/websocket/hub_test.go b/internal/websocket/hub_test.go index 2e2b8b2a4..6eadbc6fa 100644 --- a/internal/websocket/hub_test.go +++ b/internal/websocket/hub_test.go @@ -4,12 +4,188 @@ import ( "encoding/json" "math" "net/http" + "sync" + "sync/atomic" "testing" "time" "github.com/rcourtman/pulse-go-rewrite/internal/alerts" ) +func TestClientSafeSendAndCloseAreSynchronized(t *testing.T) { + client := &Client{send: make(chan []byte, 1)} + start := make(chan struct{}) + var senders sync.WaitGroup + for i := 0; i < 8; i++ { + senders.Add(1) + go func() { + defer senders.Done() + <-start + for attempt := 0; attempt < 1_000; attempt++ { + client.safeSend([]byte("state")) + } + }() + } + + close(start) + client.closeSend() + senders.Wait() + if client.safeSend([]byte("after-close")) { + t.Fatal("safeSend succeeded after closeSend") + } +} + +func TestHubDisconnectedClientsDoNotBuildInitialState(t *testing.T) { + var stateBuilds atomic.Int64 + hub := NewHub(func(string) interface{} { + stateBuilds.Add(1) + return map[string]interface{}{ + "resources": make([]map[string]interface{}, 4_000), + } + }) + go hub.Run() + t.Cleanup(hub.Stop) + + const reconnects = 32 + for i := 0; i < reconnects; i++ { + client := &Client{ + hub: hub, + id: "disconnected-before-initial-state", + send: make(chan []byte, 1), + } + hub.register <- client + hub.unregister <- client + } + + hub.Stop() + if got := stateBuilds.Load(); got != 0 { + t.Fatalf("built %d initial states for clients that had already disconnected", got) + } +} + +func TestHubCurrentStateBroadcastsStaySerializedUnderChurn(t *testing.T) { + hub := NewHub(nil) + hub.coalesceWindow = time.Millisecond + go hub.Run() + t.Cleanup(hub.Stop) + + client := &Client{ + hub: hub, + id: "state-broadcast-soak", + send: make(chan []byte, 128), + } + hub.register <- client + + var active, maxActive, builds atomic.Int64 + hub.SetStateGetter(func(string) interface{} { + builds.Add(1) + current := active.Add(1) + for { + maximum := maxActive.Load() + if current <= maximum || maxActive.CompareAndSwap(maximum, current) { + break + } + } + time.Sleep(50 * time.Millisecond) + active.Add(-1) + return map[string]string{"status": "ok"} + }) + + for i := 0; i < 50; i++ { + hub.BroadcastCurrentState() + time.Sleep(2 * time.Millisecond) + } + time.Sleep(250 * time.Millisecond) + hub.Stop() + + if got := maxActive.Load(); got != 1 { + t.Fatalf("maximum concurrent current-state builds = %d, want 1", got) + } + if got := builds.Load(); got > 10 { + t.Fatalf("current-state builds = %d for 50 supersedable signals, want at most 10 serialized builds", got) + } +} + +func TestHubInitialAndBroadcastStateBuildsShareSerialization(t *testing.T) { + var active, maxActive atomic.Int64 + hub := NewHub(func(string) interface{} { + current := active.Add(1) + for { + maximum := maxActive.Load() + if current <= maximum || maxActive.CompareAndSwap(maximum, current) { + break + } + } + time.Sleep(20 * time.Millisecond) + active.Add(-1) + return map[string]string{"status": "ok"} + }) + hub.coalesceWindow = time.Millisecond + go hub.Run() + t.Cleanup(hub.Stop) + + const clients = 8 + for i := 0; i < clients; i++ { + hub.register <- &Client{ + hub: hub, + id: "simultaneous-initial-state", + send: make(chan []byte, 8), + } + } + time.AfterFunc(initialWelcomeDelay+initialStateMessageDelay, hub.BroadcastCurrentState) + time.Sleep(initialWelcomeDelay + initialStateMessageDelay + 400*time.Millisecond) + hub.Stop() + + if got := maxActive.Load(); got != 1 { + t.Fatalf("maximum concurrent initial and broadcast state builds = %d, want 1", got) + } +} + +func TestHubRequestedStateBuildsShareSerialization(t *testing.T) { + var active, maxActive atomic.Int64 + hub := NewHub(func(string) interface{} { + current := active.Add(1) + for { + maximum := maxActive.Load() + if current <= maximum || maxActive.CompareAndSwap(maximum, current) { + break + } + } + time.Sleep(20 * time.Millisecond) + active.Add(-1) + return map[string]string{"status": "ok"} + }) + go hub.Run() + t.Cleanup(hub.Stop) + + const clients = 8 + registered := make([]*Client, 0, clients) + for i := 0; i < clients; i++ { + client := &Client{ + hub: hub, + id: "simultaneous-requested-state", + send: make(chan []byte, 8), + } + hub.register <- client + registered = append(registered, client) + } + + var done sync.WaitGroup + done.Add(len(registered)) + for _, client := range registered { + go func() { + defer done.Done() + hub.sendRequestedState(client) + }() + } + done.Wait() + hub.Stop() + + if got := maxActive.Load(); got != 1 { + t.Fatalf("maximum concurrent requested state builds = %d, want 1", got) + } +} + func TestIsValidPrivateOrigin(t *testing.T) { tests := []struct { name string