diff --git a/pkg/hold/pds/scan_broadcaster_ws_test.go b/pkg/hold/pds/scan_broadcaster_ws_test.go new file mode 100644 index 0000000..262e670 --- /dev/null +++ b/pkg/hold/pds/scan_broadcaster_ws_test.go @@ -0,0 +1,125 @@ +package pds + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// newTestScanBroadcaster builds a ScanBroadcaster with only a database, so +// nothing in the existing suite ever reaches handleWriter, handleReader, or the +// Unsubscribe teardown they share — the exact paths 05b856b fixed. These tests +// give the subscriber a real WebSocket so that teardown actually runs. +// +// The invariant under test is stated in Unsubscribe itself: a dropped scanner +// unwinds both goroutines and each calls Unsubscribe, so everything past the +// `found` guard must happen exactly once. Closing `done` twice panics, and +// re-running the requeue UPDATE would unassign jobs a replacement scanner had +// already picked up. + +// wsPair returns the server side of a live WebSocket plus a closer for the +// client side, so a test can drop the connection the way a scanner exiting does. +func wsPair(t *testing.T) (server *websocket.Conn, closeClient func()) { + t.Helper() + + upgrader := websocket.Upgrader{} + accepted := make(chan *websocket.Conn, 1) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + accepted <- conn + })) + t.Cleanup(srv.Close) + + client, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + t.Cleanup(func() { _ = client.Close() }) + + select { + case conn := <-accepted: + return conn, func() { _ = client.Close() } + case <-time.After(5 * time.Second): + t.Fatal("server never completed the upgrade") + return nil, nil + } +} + +func subscriberCount(sb *ScanBroadcaster) int { + sb.mu.Lock() + defer sb.mu.Unlock() + return len(sb.subscribers) +} + +// Unsubscribe must tolerate being called more than once for the same +// subscriber. Without the `found` guard the second call closes an already +// closed channel, which panics and takes the whole hold down. +func TestScanBroadcaster_UnsubscribeIsIdempotent(t *testing.T) { + sb := newTestScanBroadcaster(t) + conn, _ := wsPair(t) + + sub := sb.Subscribe(conn, 0) + if got := subscriberCount(sb); got != 1 { + t.Fatalf("subscriber count after Subscribe = %d, want 1", got) + } + + sb.Unsubscribe(sub) + sb.Unsubscribe(sub) // must not panic + + if got := subscriberCount(sb); got != 0 { + t.Fatalf("subscriber count after Unsubscribe = %d, want 0", got) + } +} + +// handleWriter selects on done rather than ranging over send, so Unsubscribe +// releases it. It closes the connection on the way out, which is what this +// asserts — a writer left parked would leak a goroutine and a socket per +// scanner restart. +func TestScanBroadcaster_WriterExitsOnUnsubscribe(t *testing.T) { + sb := newTestScanBroadcaster(t) + conn, _ := wsPair(t) + + sub := sb.Subscribe(conn, 0) + sb.Unsubscribe(sub) + + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if err := sub.conn.WriteMessage(websocket.TextMessage, []byte("ping")); err != nil { + return // writer ran its deferred Close + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("handleWriter never closed the connection after Unsubscribe") +} + +// The real shape of a scanner going away: the client vanishes, handleReader's +// read fails and it unsubscribes from its defer, handleWriter's write fails and +// it unsubscribes too. Both land on the same subscriber. +func TestScanBroadcaster_DroppedScannerUnwindsOnce(t *testing.T) { + sb := newTestScanBroadcaster(t) + conn, closeClient := wsPair(t) + + sub := sb.Subscribe(conn, 0) + closeClient() + + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if subscriberCount(sb) == 0 { + // Landing here at all means neither goroutine panicked on the way + // out; a double close of done would have taken the process down. + sb.Unsubscribe(sub) // late duplicate, still must be harmless + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("dropped scanner was never unsubscribed") +}