Bound state refresh and recovery ingest ownership

This commit is contained in:
rcourtman
2026-07-23 22:23:30 +01:00
parent d618fb8b74
commit 87fe506c0a
12 changed files with 676 additions and 163 deletions
@@ -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`
@@ -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
@@ -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
@@ -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
+48 -3
View File
@@ -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
}
@@ -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)
}
}
+5 -2
View File
@@ -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) {
+14 -5
View File
@@ -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 {
+20
View File
@@ -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)
}
}
+320 -148
View File
@@ -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 = &current
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")
}
+12
View File
@@ -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")
}
}
+176
View File
@@ -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