fix: address review comments on websocket package

- Make Hub.Shutdown() idempotent via sync.Once to prevent double-close panic
- Add Hub.Register()/Unregister() methods that select on hub.done to
  prevent goroutine leaks and deadlocks during shutdown
- Update NewClient to use Hub.Register(), returning error on closed hub
- Update readPump to use Hub.Unregister() instead of raw channel send
- Change NewMessage signature to (Message, error) to surface marshal failures
- Change Message.Bytes() signature to ([]byte, error) for explicit error handling
- Replace time.Sleep in tests with assert.Eventually polling
- Add tests for idempotent shutdown and register-after-shutdown

Author: omattsson

backend/internal/websocket/client.go

@@ -34,16 +34,20 @@ type Client struct {
 // NewClient creates a new Client attached to the given hub and connection,
 // registers it with the hub, and starts the read/write pumps.
 // The caller should not interact with conn after calling NewClient.
-func NewClient(hub *Hub, conn *websocket.Conn) *Client {
+// Returns an error if the hub has already been shut down.
+func NewClient(hub *Hub, conn *websocket.Conn) (*Client, error) {
 	client := &Client{
 		hub:  hub,
 		conn: conn,
 		send: make(chan []byte, sendBufferSize),
 	}
-	hub.register <- client
+	if err := hub.Register(client); err != nil {
+		conn.Close()
+		return nil, err
+	}
 	go client.writePump()
 	go client.readPump()
-	return client
+	return client, nil
 }
 
 // readPump pumps messages from the WebSocket connection to the hub.
@@ -54,7 +58,7 @@ func (c *Client) readPump() {
 		if r := recover(); r != nil {
 			slog.Error("Panic in WebSocket readPump", "recover", r)
 		}
-		c.hub.unregister <- c
+		c.hub.Unregister(c)
 		c.conn.Close()
 	}()
 

backend/internal/websocket/hub.go

@@ -7,6 +7,13 @@ import (
 	"sync"
 )
 
+// ErrHubClosed is returned when attempting to register a client on a shut-down hub.
+var ErrHubClosed = errHubClosed{}
+
+type errHubClosed struct{}
+
+func (errHubClosed) Error() string { return "hub is closed" }
+
 // BroadcastSender is implemented by any type that can broadcast messages
 // to all connected WebSocket clients. Use this interface for decoupled
 // dependency injection (e.g., handlers broadcast events without importing Hub).
@@ -34,6 +41,9 @@ type Hub struct {
 
 	// done signals the Run loop to stop.
 	done chan struct{}
+
+	// shutdownOnce ensures Shutdown is idempotent and safe to call concurrently.
+	shutdownOnce sync.Once
 }
 
 // NewHub creates a new Hub ready to accept clients.
@@ -99,8 +109,31 @@ func (h *Hub) Broadcast(message []byte) {
 }
 
 // Shutdown gracefully stops the hub's Run loop and closes all client connections.
+// It is safe to call multiple times and concurrently.
 func (h *Hub) Shutdown() {
-	close(h.done)
+	h.shutdownOnce.Do(func() {
+		close(h.done)
+	})
+}
+
+// Register safely registers a client with the hub. It returns ErrHubClosed if
+// the hub has been shut down, preventing the caller from blocking forever.
+func (h *Hub) Register(c *Client) error {
+	select {
+	case h.register <- c:
+		return nil
+	case <-h.done:
+		return ErrHubClosed
+	}
+}
+
+// Unregister safely requests client removal from the hub. If the hub has
+// already been shut down the call is a no-op, preventing goroutine leaks.
+func (h *Hub) Unregister(c *Client) {
+	select {
+	case h.unregister <- c:
+	case <-h.done:
+	}
 }
 
 // ClientCount returns the current number of connected clients.

backend/internal/websocket/hub_test.go

@@ -185,7 +185,8 @@ func TestNewMessage(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 
-			msg := NewMessage(tt.msgType, tt.payload)
+			msg, err := NewMessage(tt.msgType, tt.payload)
+			assert.NoError(t, err)
 			assert.Equal(t, tt.wantType, msg.Type)
 			assert.JSONEq(t, tt.wantPayload, string(msg.Payload))
 		})
@@ -195,11 +196,14 @@ func TestNewMessage(t *testing.T) {
 func TestMessage_Bytes(t *testing.T) {
 	t.Parallel()
 
-	msg := NewMessage("item.updated", map[string]int{"id": 1})
-	b := msg.Bytes()
+	msg, err := NewMessage("item.updated", map[string]int{"id": 1})
+	assert.NoError(t, err)
+
+	b, err := msg.Bytes()
+	assert.NoError(t, err)
 
 	var parsed Message
-	err := json.Unmarshal(b, &parsed)
+	err = json.Unmarshal(b, &parsed)
 	assert.NoError(t, err)
 	assert.Equal(t, "item.updated", parsed.Type)
 	assert.JSONEq(t, `{"id":1}`, string(parsed.Payload))

backend/internal/websocket/message.go

@@ -1,6 +1,9 @@
 package websocket
 
-import "encoding/json"
+import (
+	"encoding/json"
+	"fmt"
+)
 
 // Message is the envelope for all WebSocket messages sent to clients.
 type Message struct {
@@ -12,23 +15,23 @@ type Message struct {
 }
 
 // NewMessage creates a Message with the given type and payload.
-// The payload is JSON-marshalled; if marshalling fails the payload is set to null.
-func NewMessage(msgType string, payload interface{}) Message {
+// The payload is JSON-marshalled; an error is returned if marshalling fails.
+func NewMessage(msgType string, payload interface{}) (Message, error) {
 	data, err := json.Marshal(payload)
 	if err != nil {
-		data = []byte("null")
+		return Message{}, fmt.Errorf("marshal payload: %w", err)
 	}
 	return Message{
 		Type:    msgType,
 		Payload: data,
-	}
+	}, nil
 }
 
 // Bytes serialises the Message to JSON bytes suitable for broadcasting.
-func (m Message) Bytes() []byte {
+func (m Message) Bytes() ([]byte, error) {
 	b, err := json.Marshal(m)
 	if err != nil {
-		return []byte(`{"type":"error","payload":null}`)
+		return nil, fmt.Errorf("marshal message: %w", err)
 	}
-	return b
+	return b, nil
 }