diff --git a/pkg/hold/pds/scan_broadcaster.go b/pkg/hold/pds/scan_broadcaster.go index f49159e..52e94bb 100644 --- a/pkg/hold/pds/scan_broadcaster.go +++ b/pkg/hold/pds/scan_broadcaster.go @@ -342,13 +342,23 @@ func (sb *ScanBroadcaster) Unsubscribe(sub *ScanSubscriber) { sb.mu.Lock() defer sb.mu.Unlock() + found := false for i, s := range sb.subscribers { if s == sub { sb.subscribers = append(sb.subscribers[:i], sb.subscribers[i+1:]...) + found = true break } } + // A dropped scanner unwinds both handleWriter (on write error) and + // handleReader (in its defer), and each calls Unsubscribe. Everything below + // must happen exactly once: re-running the UPDATE would unassign jobs a + // replacement scanner had already picked up, and closing done twice panics. + if !found { + return + } + // Mark assigned/processing jobs as pending again so they can be re-dispatched. // Including 'processing' handles scanner crashes mid-scan. _, err := sb.db.Exec(` @@ -361,7 +371,10 @@ func (sb *ScanBroadcaster) Unsubscribe(sub *ScanSubscriber) { "error", err) } - close(sub.send) + // Close done, never send. drainPendingJobs runs in its own goroutine + // without holding sb.mu, so closing send here would race an in-flight + // send and panic the process unrecoverably. + close(sub.done) slog.Info("Scanner unsubscribed", "id", sub.id, @@ -411,12 +424,19 @@ func (sb *ScanBroadcaster) dispatchJob(job *ScanJobEvent) { // handleWriter sends jobs to a scanner over its WebSocket connection func (sb *ScanBroadcaster) handleWriter(sub *ScanSubscriber) { - defer func() { - sub.conn.Close() - close(sub.done) - }() + defer sub.conn.Close() + + // done is closed by Unsubscribe, not here — ranging over send would block + // forever now that nothing closes it. handleReader always unsubscribes on + // its way out, so this goroutine is guaranteed to be released. + for { + var job *ScanJobEvent + select { + case <-sub.done: + return + case job = <-sub.send: + } - for job := range sub.send { data, err := json.Marshal(job) if err != nil { slog.Error("Failed to marshal scan job", "seq", job.Seq, "error", err) diff --git a/pkg/hold/pds/scan_broadcaster_test.go b/pkg/hold/pds/scan_broadcaster_test.go new file mode 100644 index 0000000..a6aeae2 --- /dev/null +++ b/pkg/hold/pds/scan_broadcaster_test.go @@ -0,0 +1,180 @@ +package pds + +import ( + "database/sql" + "testing" + "time" +) + +// newTestScanBroadcaster builds a ScanBroadcaster with just a database, no +// background goroutines. Subscriber lifecycle is all that is under test here, +// so the constructor's discovery/dispatch/stale loops (and their s3 and PDS +// dependencies) are deliberately skipped. +func newTestScanBroadcaster(t *testing.T) *ScanBroadcaster { + t.Helper() + + db, err := sql.Open("libsql", "file:"+t.TempDir()+"/scan.db") + if err != nil { + t.Fatalf("open scan db: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + sb := &ScanBroadcaster{ + db: db, + holdDID: "did:web:hold.example.com", + holdEndpoint: "https://hold.example.com", + } + if err := sb.initSchema(); err != nil { + t.Fatalf("initSchema: %v", err) + } + return sb +} + +// newTestScanSubscriber mirrors what Subscribe builds, registered with the +// broadcaster so Unsubscribe finds it. +func newTestScanSubscriber(t *testing.T, sb *ScanBroadcaster, bufSize int) *ScanSubscriber { + t.Helper() + + sub := &ScanSubscriber{ + conn: nil, // no websocket needed; handleWriter is not exercised here + send: make(chan *ScanJobEvent, bufSize), + id: "test-subscriber", + done: make(chan struct{}), + } + + sb.mu.Lock() + sb.subscribers = append(sb.subscribers, sub) + sb.mu.Unlock() + + return sub +} + +func seedPendingJobs(t *testing.T, sb *ScanBroadcaster, n int) { + t.Helper() + + for i := 0; i < n; i++ { + _, err := sb.db.Exec(` + INSERT INTO scan_jobs + (manifest_digest, repository, tag, user_did, user_handle, + hold_did, hold_endpoint, tier, config_json, layers_json, status) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending') + `, "sha256:deadbeef", "repo", "latest", "did:plc:user", "user.example.com", + sb.holdDID, sb.holdEndpoint, "deckhand", "{}", "[]") + if err != nil { + t.Fatalf("seed job %d: %v", i, err) + } + } +} + +// TestScanUnsubscribe_IsIdempotent covers a panic reachable on any scanner that +// dropped mid-write. Unsubscribe used to close sub.send unconditionally, but it +// is called from two places — handleWriter on write error, and handleReader in +// its defer — so a write failure followed by the read side unwinding produced +// "panic: close of closed channel". The slice-removal loop had no guard, so the +// second call fell straight through to the close. +func TestScanUnsubscribe_IsIdempotent(t *testing.T) { + sb := newTestScanBroadcaster(t) + sub := newTestScanSubscriber(t, sb, 4) + + sb.Unsubscribe(sub) + sb.Unsubscribe(sub) // must be a no-op, not a second close + + select { + case <-sub.done: + default: + t.Error("done was not closed by Unsubscribe") + } + + sb.mu.RLock() + n := len(sb.subscribers) + sb.mu.RUnlock() + if n != 0 { + t.Errorf("expected subscriber removed, got %d remaining", n) + } +} + +// TestScanUnsubscribe_UnassignsJobsOnce verifies the idempotency guard protects +// the job-reassignment UPDATE too. Re-running it would unassign jobs that a +// replacement scanner had already been given. +func TestScanUnsubscribe_UnassignsJobsOnce(t *testing.T) { + sb := newTestScanBroadcaster(t) + sub := newTestScanSubscriber(t, sb, 4) + seedPendingJobs(t, sb, 1) + + if _, err := sb.db.Exec(`UPDATE scan_jobs SET status='assigned', assigned_to=?`, sub.id); err != nil { + t.Fatalf("assign: %v", err) + } + + sb.Unsubscribe(sub) + + var status string + if err := sb.db.QueryRow(`SELECT status FROM scan_jobs LIMIT 1`).Scan(&status); err != nil { + t.Fatalf("query: %v", err) + } + if status != "pending" { + t.Errorf("expected job returned to pending, got %q", status) + } + + // Hand the job to a "replacement" scanner, then unsubscribe the dead one + // again. The guard must stop it from stealing the job back. + if _, err := sb.db.Exec(`UPDATE scan_jobs SET status='assigned', assigned_to=?`, "replacement"); err != nil { + t.Fatalf("reassign: %v", err) + } + sb.Unsubscribe(sub) + + if err := sb.db.QueryRow(`SELECT status FROM scan_jobs LIMIT 1`).Scan(&status); err != nil { + t.Fatalf("query: %v", err) + } + if status != "assigned" { + t.Errorf("second Unsubscribe stole the replacement's job, status=%q", status) + } +} + +// TestScanDrainPendingJobs_ConcurrentUnsubscribe is the regression test for +// "panic: send on closed channel" in the scanner path — the same defect as the +// firehose backfill. Subscribe spawns drainPendingJobs in its own goroutine and +// it sends to sub.send without holding sb.mu, so a scanner disconnecting during +// the drain had Unsubscribe close the channel under an in-flight send. +// +// The subscriber uses a 1-slot buffer with no reader so the drain is reliably +// blocked in the send when Unsubscribe fires. +func TestScanDrainPendingJobs_ConcurrentUnsubscribe(t *testing.T) { + for i := 0; i < 25; i++ { + func() { + sb := newTestScanBroadcaster(t) + sub := newTestScanSubscriber(t, sb, 1) + seedPendingJobs(t, sb, 50) + + done := make(chan struct{}) + go func() { + defer close(done) + // Pre-fix this panicked instead of returning. + sb.drainPendingJobs(sub, 0) + }() + + sb.Unsubscribe(sub) + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("drainPendingJobs did not return after Unsubscribe") + } + }() + } +} + +// TestScanDispatchJob_AfterUnsubscribe verifies a removed scanner stops being +// dispatched to and that dispatch does not touch its channel. +func TestScanDispatchJob_AfterUnsubscribe(t *testing.T) { + sb := newTestScanBroadcaster(t) + sub := newTestScanSubscriber(t, sb, 4) + seedPendingJobs(t, sb, 1) + + sb.Unsubscribe(sub) + + sb.dispatchJob(&ScanJobEvent{Type: "job", Seq: 1, Repository: "repo"}) + + if len(sub.send) != 0 { + t.Errorf("unsubscribed scanner received %d jobs", len(sub.send)) + } +}