feat(websocketspec): fire BeforeDisconnect/AfterDisconnect hooks on close

Hooks fire from Connection.Close() once per connection. ConnectionManager
Shutdown now closes connections outside its lock so hooks can call back
into the manager. Update the single-transaction audit plan status.
This commit is contained in:
Hein
2026-10-01 10:53:21 +02:00
parent ae2b0a4ef4
commit 982c90bfdd
4 changed files with 115 additions and 13 deletions
+44 -3
View File
@@ -193,12 +193,17 @@ func (cm *ConnectionManager) GetConnection(id string) (*Connection, bool) {
func (cm *ConnectionManager) Shutdown() {
cm.cancel()
// Close all connections
cm.mu.Lock()
// Close all connections outside the lock: disconnect hooks may call back
// into the manager (Count, GetConnection).
cm.mu.RLock()
conns := make([]*Connection, 0, len(cm.connections))
for _, conn := range cm.connections {
conns = append(conns, conn)
}
cm.mu.RUnlock()
for _, conn := range conns {
conn.Close()
}
cm.mu.Unlock()
}
// ReadPump reads messages from the WebSocket connection
@@ -300,6 +305,14 @@ func (c *Connection) SendJSON(v interface{}) error {
// Close closes the connection
func (c *Connection) Close() {
c.closedOnce.Do(func() {
// Disconnect hooks fire once per connection, however it closes. They
// are best-effort: a failure is logged and never blocks the close.
hookCtx := c.disconnectHookContext()
c.logDisconnectHook(BeforeDisconnect, c.executeHook(func(r *HookRegistry) error { return r.Execute(BeforeDisconnect, hookCtx) }))
defer func() {
c.logDisconnectHook(AfterDisconnect, c.executeHook(func(r *HookRegistry) error { return r.Execute(AfterDisconnect, hookCtx) }))
}()
if c.cancel != nil {
c.cancel()
}
@@ -321,6 +334,34 @@ func (c *Connection) Close() {
})
}
// disconnectHookContext builds the hook context for the disconnect hooks. The
// context is detached from the connection's cancellation so hooks that touch
// the database still run after the connection context is cancelled.
func (c *Connection) disconnectHookContext() *HookContext {
parent := c.ctx
if parent == nil {
parent = context.Background()
}
return &HookContext{
Context: context.WithoutCancel(parent),
Handler: c.handler,
Connection: c,
}
}
func (c *Connection) executeHook(run func(*HookRegistry) error) error {
if c.handler == nil || c.handler.hooks == nil {
return nil
}
return run(c.handler.hooks)
}
func (c *Connection) logDisconnectHook(hook HookType, err error) {
if err != nil {
logger.Warn("[WebSocketSpec] %s hook failed for connection %s: %v", hook, c.ID, err)
}
}
// AddSubscription adds a subscription to this connection
func (c *Connection) AddSubscription(sub *Subscription) {
c.mu.Lock()
+54
View File
@@ -594,3 +594,57 @@ func TestConnectionManager_CompleteLifecycle(t *testing.T) {
_, exists = cm.GetConnection("conn-1")
assert.False(t, exists)
}
func TestConnection_Close_FiresDisconnectHooksOnce(t *testing.T) {
h := &Handler{hooks: NewHookRegistry()}
var order []string
var ctxErrAfter error
h.hooks.Register(BeforeDisconnect, func(hc *HookContext) error {
order = append(order, "before")
assert.NoError(t, hc.Context.Err())
return nil
})
h.hooks.Register(AfterDisconnect, func(hc *HookContext) error {
order = append(order, "after")
ctxErrAfter = hc.Context.Err()
return nil
})
conn := createTestConnection("c1")
conn.handler = h
conn.Close()
conn.Close() // second close must not re-fire
assert.Equal(t, []string{"before", "after"}, order)
assert.NoError(t, ctxErrAfter, "hook context must survive the connection cancel")
}
func TestConnection_Close_HookErrorDoesNotBlockClose(t *testing.T) {
h := &Handler{hooks: NewHookRegistry()}
h.hooks.Register(BeforeDisconnect, func(*HookContext) error { return assert.AnError })
conn := createTestConnection("c2")
conn.handler = h
conn.Close()
assert.Error(t, conn.ctx.Err(), "connection must still be cancelled")
}
func TestConnectionManager_Shutdown_HookMayCallManager(t *testing.T) {
cm := NewConnectionManager(context.Background())
h := &Handler{hooks: NewHookRegistry(), connManager: cm}
h.hooks.Register(BeforeDisconnect, func(*HookContext) error { cm.Count(); return nil })
conn := createTestConnection("c3")
conn.handler = h
cm.connections[conn.ID] = conn
done := make(chan struct{})
go func() { cm.Shutdown(); close(done) }()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Shutdown deadlocked on a hook calling the manager")
}
}