diff --git a/pkg/rest/socket_controller.go b/pkg/rest/socket_controller.go index fc3f2e5..fc8c1a6 100644 --- a/pkg/rest/socket_controller.go +++ b/pkg/rest/socket_controller.go @@ -1,7 +1,9 @@ package rest import ( + "errors" "net/http" + "sync" "time" "github.com/gorilla/websocket" @@ -43,9 +45,11 @@ var upgrader = websocket.Upgrader{ type msgListener[T any] struct { hub *msghub.Hub c chan T - mailbox string // Name of mailbox to monitor, "" == all mailboxes + done chan struct{} // Closed by Close; stops event delivery and the write pump + mailbox string // Name of mailbox to monitor, "" == all mailboxes toStored func(msg event.MessageMetadata) T toDeleted func(mailbox string, id string) T + closeOnce sync.Once } // newMsgListener creates a listener and registers it with the hub. The optional mailbox @@ -59,6 +63,7 @@ func newMsgListener[T any]( ml := &msgListener[T]{ hub: hub, c: make(chan T, socketChanLen), + done: make(chan struct{}), mailbox: mailbox, toStored: toStored, toDeleted: toDeleted, @@ -72,13 +77,22 @@ func (ml *msgListener[T]) watches(mailbox string) bool { return ml.mailbox == "" || ml.mailbox == mailbox } -// Receive handles an incoming message. +// errListenerClosed is returned by Receive/Delete after Close; the hub removes listeners +// that return errors, providing deregistration even if RemoveListener has not run yet. +var errListenerClosed = errors.New("monitor listener closed") + +// Receive handles an incoming message. If the listener is closed while its event queue is +// full, the enqueue is abandoned rather than blocking the hub actor forever. func (ml *msgListener[T]) Receive(msg event.MessageMetadata) error { if !ml.watches(msg.Mailbox) { return nil } - ml.c <- ml.toStored(msg) - return nil + select { + case ml.c <- ml.toStored(msg): + return nil + case <-ml.done: + return errListenerClosed + } } // Delete handles a deleted message. @@ -86,19 +100,30 @@ func (ml *msgListener[T]) Delete(mailbox string, id string) error { if ml.toDeleted == nil || !ml.watches(mailbox) { return nil } - ml.c <- ml.toDeleted(mailbox, id) - return nil + select { + case ml.c <- ml.toDeleted(mailbox, id): + return nil + case <-ml.done: + return errListenerClosed + } } -// Close removes the listener registration and closes the event queue. +// Close removes the listener registration and stops event delivery. Safe for concurrent +// use: the read and write pumps both invoke it when the connection drops. The done channel +// is closed before the removal is queued: a hub actor blocked in this listener's Receive or +// Delete is released first, so it keeps draining the hub op queue and the removal can be +// queued even when that queue is full. The event queue is not closed: RemoveListener +// returns once the op is queued, not processed, so a concurrent hub-actor send into the +// queue could otherwise panic on the close; the done channel ends delivery and unblocks +// any such send. func (ml *msgListener[T]) Close() { - select { - case <-ml.c: - // Already closed - default: + ml.closeOnce.Do(func() { + // Release a hub actor blocked in Receive or Delete before queueing the removal: + // with the removal queued first, a full hub op queue would block it here forever, + // never reaching the close below -- the only operation able to release the actor. + close(ml.done) ml.hub.RemoveListener(ml) - close(ml.c) - } + }) } // socketReadPump makes sure the websocket client is still connected, discarding any messages @@ -140,9 +165,9 @@ func socketReadPump(conn *websocket.Conn, closeFn func()) { } } -// socketWritePump relays events from c to the peer as JSON messages until c is closed, sending -// pings while idle to keep the connection alive. closeFn is invoked on exit. -func socketWritePump[T any](c chan T, conn *websocket.Conn, closeFn func()) { +// socketWritePump relays events from c to the peer as JSON messages until done is closed, +// sending pings while idle to keep the connection alive. closeFn is invoked on exit. +func socketWritePump[T any](c chan T, done <-chan struct{}, conn *websocket.Conn, closeFn func()) { slog := log.With().Str("module", "rest").Str("proto", "WebSocket"). Str("remote", conn.RemoteAddr().String()).Logger() @@ -168,6 +193,10 @@ func socketWritePump[T any](c chan T, conn *websocket.Conn, closeFn func()) { // Write failed return } + case <-done: + // Listener closed, exit + _ = conn.WriteMessage(websocket.CloseMessage, []byte{}) + return case <-ticker.C: // Send ping if err := conn.SetWriteDeadline(time.Now().Add(socketWriteWait)); err != nil { @@ -205,7 +234,7 @@ func serveMonitor[T any]( Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket") // Create, register listener; then interact with conn. ml := newListener(hub, mailbox) - go socketWritePump(ml.c, conn, ml.Close) + go socketWritePump(ml.c, ml.done, conn, ml.Close) socketReadPump(conn, ml.Close) return nil } diff --git a/pkg/rest/socket_controller_test.go b/pkg/rest/socket_controller_test.go index cb1100e..23e39c2 100644 --- a/pkg/rest/socket_controller_test.go +++ b/pkg/rest/socket_controller_test.go @@ -6,6 +6,7 @@ import ( "errors" "net/http/httptest" "net/mail" + "strconv" "strings" "testing" "time" @@ -399,12 +400,279 @@ func TestMsgListenerNilDeletedDropsDeletes(t *testing.T) { func TestMsgListenerCloseIdempotent(t *testing.T) { hub := startMsgHub(t) - ml := newMsgListener(hub, "", - func(msg event.MessageMetadata) string { return "stored" }, nil) + ml := newMsgListener(hub, "watched", + func(msg event.MessageMetadata) string { return "stored" }, + func(mailbox, id string) string { return "deleted" }) ml.Close() - ml.Close() // second close must not double-close the event channel + ml.Close() // second close must not panic or double-close the done channel - _, ok := <-ml.c - assert.False(t, ok, "event channel should be closed") + // Close is exactly-once: the done channel ends delivery and the write pump. + select { + case <-ml.done: + default: + t.Fatal("done channel should be closed after Close") + } + + // Closed listeners either deliver into the never-again-read queue or report closed; + // both outcomes are terminal, as the hub-actor send races the done channel. + err := ml.Receive(testMetadata("watched", "0001")) + if err != nil { + require.ErrorIs(t, err, errListenerClosed) + } + err = ml.Delete("watched", "0001") + if err != nil { + require.ErrorIs(t, err, errListenerClosed) + } +} + +func TestMsgListenerCloseFullQueue(t *testing.T) { + // A disconnecting client may leave a full event queue behind; Close must still + // deregister the listener, or the hub actor blocks on its next Receive into the full + // queue forever, wedging all monitor dispatch. Regression: the old check-then-act + // select mistook the buffered events for a closed-channel signal and skipped + // RemoveListener entirely. + hub := startMsgHub(t) + ml := newMsgListener(hub, "watched", + func(msg event.MessageMetadata) string { return "stored:" + msg.ID }, nil) + + // Fill the event queue (capacity 100) with no reader draining it. + for i := range 100 { + hub.Dispatch(testMetadata("watched", strconv.Itoa(i))) + } + hub.Sync() + + ml.Close() + hub.Sync() + + // The backlog remains queued, but the listener is deregistered: dispatch must not + // block the hub, and must not be delivered. Pre-fix, the no-op Close left the + // listener registered, so the dispatch was delivered (and the next dispatch would + // have wedged the actor inside Receive). + backlog, _ := drainEvents(ml.c) + assert.Equal(t, 100, backlog, "backlog should remain queued after Close") + + hub.Dispatch(testMetadata("watched", "after-close")) + hub.Sync() + + count, _ := drainEvents(ml.c) + assert.Equal(t, 0, count, "closed listener must not receive further events") +} + +// drainEvents reads buffered events from c until it is empty; it reports how many events +// were drained and whether the channel was found closed. A receive from a closed channel +// succeeds immediately, so a plain select-default loop would never terminate on one. +func drainEvents[T any](c chan T) (count int, closed bool) { + for { + select { + case _, ok := <-c: + if !ok { + return count, true + } + count++ + default: + return count, false + } + } +} + +// blockingListener stalls the hub actor inside Receive until released. +type blockingListener struct { + release chan struct{} +} + +func (l *blockingListener) Receive(msg event.MessageMetadata) error { + <-l.release + return nil +} + +func (l *blockingListener) Delete(mailbox string, id string) error { + return nil +} + +func TestMsgListenerCloseConcurrent(t *testing.T) { + // The read and write pumps both call Close when the connection drops, so Close must be + // safe to call concurrently. Regression: the check-then-act select let both callers + // through the default branch and panicked with "close of closed channel" on the second + // close. + // + // Deterministic reproduction: a blocking listener stalls the hub actor mid-dispatch, + // and enough further dispatches are queued to keep the hub op queue (capacity 100) + // full. Every concurrent Close then blocks inside RemoveListener -- already past its + // check of the listener state -- before any of them can proceed. Releasing the actor + // lets it process the pending RemoveListeners, resuming all the blocked closers at + // once. + hub := startMsgHub(t) + blocker := &blockingListener{release: make(chan struct{})} + hub.AddListener(blocker) + ml := newMsgListener(hub, "watched", + func(msg event.MessageMetadata) string { return "stored:" + msg.ID }, + func(mailbox, id string) string { return "deleted:" + id }) + + const queued = 101 // 1 dispatch stalls the actor; 100 fill the hub op queue + for i := range queued { + hub.Dispatch(testMetadata("watched", strconv.Itoa(i))) + } + + // The stalled dispatch may have delivered 0 or 1 events to ml (listener iteration + // order is random); drain whatever arrived so the queue is empty when the closers + // check it. + drainEvents(ml.c) + + const callers = 8 + done := make(chan struct{}, callers) // ensure capacity so we do not block population + for range callers { + go func() { + ml.Close() + done <- struct{}{} + }() + } + // Give every closer time to block inside RemoveListener before the actor resumes; + // otherwise their ops interleave with the queued dispatches and the reproduction + // becomes timing-dependent. + time.Sleep(50 * time.Millisecond) + + close(blocker.release) + for range callers { + <-done + } + + // Events racing in before the hub processed the RemoveListeners may sit in the queue, + // but once it has (Sync), the listener must be deregistered: later dispatches are + // neither delivered to it nor blocking the hub. + hub.Sync() + drainEvents(ml.c) + hub.Dispatch(testMetadata("watched", "after-close")) + hub.Sync() + count, _ := drainEvents(ml.c) + assert.Equal(t, 0, count, "closed listener must not receive further events") +} + +func TestMsgListenerCloseUnblocksHub(t *testing.T) { + // Close must unblock a hub actor stuck sending into the listener's full event queue: + // a slow monitor client disconnecting with a backlog of events would otherwise wedge + // every hub dispatch. (With the original close(ml.c) approach, this + // send-on-closed-channel panicked inside the actor instead.) + hub := startMsgHub(t) + blocker := &blockingListener{release: make(chan struct{})} + hub.AddListener(blocker) + ml := newMsgListener(hub, "watched", + func(msg event.MessageMetadata) string { return "stored:" + msg.ID }, nil) + + // Stall the actor, then queue enough dispatches to fill both the hub op queue and the + // listener's event queue once deliveries resume. + const queued = 101 // 1 dispatch stalls the actor; 100 fill the hub op queue + for i := range queued { + hub.Dispatch(testMetadata("watched", strconv.Itoa(i))) + } + drainEvents(ml.c) // drop the 0 or 1 events delivered before the actor stalled + close(blocker.release) + + // Once the queued dispatches fill the event queue, the next dispatch blocks the actor + // inside Receive; run it from its own goroutine so the test can proceed to Close. + dispatched := make(chan struct{}, 1) + go func() { + hub.Dispatch(testMetadata("watched", "one-too-many")) + dispatched <- struct{}{} + }() + + ml.Close() + + // The blocked send must abort and the hub must keep processing. + synced := make(chan struct{}, 1) + go func() { + hub.Sync() + synced <- struct{}{} + }() + select { + case <-synced: + case <-time.After(5 * time.Second): + t.Fatal("hub actor still wedged after Close") + } + select { + case <-dispatched: + case <-time.After(5 * time.Second): + t.Fatal("dispatch still blocked after Close") + } +} + +func TestMsgListenerCloseFullHubOpQueue(t *testing.T) { + // Regression (PR #623 review): with the hub actor blocked inside this listener's + // Receive (full event queue) and the hub op queue (capacity 100) also full, + // RemoveListener cannot enqueue. Close must therefore close done before queueing the + // removal: with the removal queued first, execution never reached close(ml.done) -- the + // only operation able to release the actor -- and Close deadlocked under the same + // slow-client condition while concurrent dispatchers kept the hub queue saturated. + hub := startMsgHub(t) + blocker := &blockingListener{release: make(chan struct{})} + hub.AddListener(blocker) + ml := newMsgListener(hub, "watched", + func(msg event.MessageMetadata) string { return "stored:" + msg.ID }, nil) + + // Stall the actor inside the blocker, then fill the hub op queue behind it: the first + // dispatch wedges the actor, the next 100 fill the queue. The stalled dispatch may + // have delivered 0 or 1 events to ml (listener iteration order is random); drain + // whatever arrived so the counts below are stable. + for i := range 101 { + hub.Dispatch(testMetadata("watched", strconv.Itoa(i))) + } + drainEvents(ml.c) + + // Once the blocker is released, a dispatcher goroutine keeps the op queue saturated: + // every slot the actor frees is refilled, so by the time the event queue (capacity 100) + // fills and the actor wedges inside Receive, the op queue is full again and the + // dispatcher is blocked on its next enqueue. 150 refills exceed the ~102 slots freed + // until the wedge, leaving spare dispatches to top the queue up. + dispatched := make(chan struct{}) + go func() { + defer close(dispatched) + for i := range 150 { + hub.Dispatch(testMetadata("watched", strconv.Itoa(100+i))) + } + }() + close(blocker.release) + + // Wait for the wedge: the event queue is full and nothing drains it. The sleep lets + // the dispatcher refill the one slot the wedging op vacated and block itself. + require.Eventually(t, func() bool { + return len(ml.c) == socketChanLen + }, monitorTestTimeout, time.Millisecond, "event queue should fill and wedge the hub actor") + time.Sleep(50 * time.Millisecond) + + // Close must return despite the wedged actor and the saturated op queue. + closed := make(chan struct{}) + go func() { + ml.Close() + close(closed) + }() + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatal("Close deadlocked: hub actor blocked in Receive with a full op queue") + } + + // The released actor must drain the saturated op queue: the dispatcher completes and + // the hub answers Sync again. + select { + case <-dispatched: + case <-time.After(5 * time.Second): + t.Fatal("dispatches still blocked after Close") + } + synced := make(chan struct{}, 1) + go func() { + hub.Sync() + synced <- struct{}{} + }() + select { + case <-synced: + case <-time.After(5 * time.Second): + t.Fatal("hub actor still wedged after Close") + } + + // The listener is deregistered: later dispatches are not delivered to it. + drainEvents(ml.c) + hub.Dispatch(testMetadata("watched", "after-close")) + hub.Sync() + count, _ := drainEvents(ml.c) + assert.Equal(t, 0, count, "closed listener must not receive further events") }