package pgstore import ( "context" "testing" "time" ) // runWithAttempt is one live run of one funded account, with attempt 1 carrying a named unit. func runWithAttempt(t *testing.T, unit string) (*Store, context.Context, string, time.Time) { t.Helper() s, ctx := testDB(t) now := fundedAccount(t, s, ctx, "u1", "10") bookID, err := s.AddBook(ctx, NewBook{OwnerID: "u1", Title: "b", SourceLang: "zh", TargetLang: "ru", ChapterCount: 10, Workdir: "/srv/books/b", Now: now}) if err != nil { t.Fatal(err) } started, err := s.StartRun(ctx, StartRunInput{UserID: "u1", BookID: bookID, OrderedChapters: 1, Ceiling: 1_000_000, Now: now}, 0, nil) if err != nil { t.Fatal(err) } // StartRun opens attempt 1 without a unit name — the row exists before anything is spawned. The // test owns the numbering from here, so it names that attempt itself. if _, err := s.pool.Exec(ctx, `update run_attempts set unit_name = $2 where run_id = $1`, started.ID, unit); err != nil { t.Fatal(err) } return s, ctx, started.ID, now } func liveUnitOf(t *testing.T, s *Store, ctx context.Context, runID string) string { t.Helper() var unit string if err := s.pool.QueryRow(ctx, ` select coalesce(unit_name, '') from run_attempts where run_id = $1 and ended_at is null order by attempt_no desc limit 1`, runID).Scan(&unit); err != nil { t.Fatal(err) } return unit } // ⛔ THE UNIT THIS ANSWERS WITH IS THE ONE THE CALLER SIGNALS, so a stale name is a signal sent to a // dead unit while the run's new process keeps working — and the user holds a `202` that says the stop // was accepted. This is that race, run against a real database rather than argued. // // The mechanism is READ COMMITTED, and it is why one statement could not do this. As a single UPDATE // the call blocks on the run row's lock, which the reconciler's restart holds while it ends one // attempt and opens the next; when the lock is released the UPDATE's own WHERE is re-checked against // the new row, so the intent lands correctly — but a SUBQUERY in RETURNING is still evaluated on the // snapshot the statement began with, which predates the restart. // // ⚠ MEASURED IN BOTH DIRECTIONS BEFORE THIS WAS WRITTEN, which is why it is a test and not a // precaution: with the single-statement form this fixture answered `unit-OLD` while the live attempt // was `unit-NEW`; with the locking read it answers `unit-NEW`. // // The waiting time is PRINTED and asserted, because it is what proves the contended path was taken at // all: a call that did not block never met the race, and «it answered the right unit» would then be // true of a test that measured nothing. func TestAStopAnswersTheUnitThatIsLiveWhenTheIntentLands(t *testing.T) { const hold = 400 * time.Millisecond s, ctx, runID, now := runWithAttempt(t, "unit-OLD") tx, err := s.pool.Begin(ctx) if err != nil { t.Fatal(err) } defer func() { _ = tx.Rollback(ctx) }() // What a restart does: it takes the run row's lock, ends the live attempt and opens the next one. for _, q := range []struct { sql string args []any }{ {`update runs set revision = revision + 1 where id = $1`, []any{runID}}, {`update run_attempts set ended_at = $2 where run_id = $1 and attempt_no = 1`, []any{runID, now}}, {`insert into run_attempts (run_id, attempt_no, unit_name, started_at) values ($1, 2, 'unit-NEW', $2)`, []any{runID, now}}, } { if _, err := tx.Exec(ctx, q.sql, q.args...); err != nil { t.Fatal(err) } } type answer struct { unit string waited time.Duration err error } done := make(chan answer, 1) go func() { // A context of its own: this call must block on the row lock, not on the test's clock. start := time.Now() _, unit, err := s.RequestStop(context.Background(), "u1", runID, now.Add(time.Second)) done <- answer{unit, time.Since(start), err} }() time.Sleep(hold) // let the stop reach the lock and wait there if err := tx.Commit(ctx); err != nil { t.Fatal(err) } got := <-done if got.err != nil { t.Fatal(got.err) } live := liveUnitOf(t, s, ctx, runID) t.Logf("a restart committed %q while a stop waited %.2fs for the run's lock; the stop answered %q", live, got.waited.Seconds(), got.unit) if got.waited < hold { t.Fatalf("the stop returned after %.2fs and the restart held the lock for %s: it never waited on the "+ "lock, so this fixture did not reproduce the race and proves nothing", got.waited.Seconds(), hold) } if live != "unit-NEW" { t.Fatalf("the live attempt after the restart is %q, want unit-NEW: the fixture did not restart anything", live) } if got.unit != live { t.Errorf("the stop answered %q while the live unit is %q: the caller would signal a unit that is already "+ "dead, the run's new process would keep working, and the user would hold a 202 saying it was accepted", got.unit, live) } // The intent itself must still be there — the fix is about WHICH unit, not about whether the stop // was recorded. var at *time.Time if err := s.pool.QueryRow(ctx, `select stop_requested_at from runs where id = $1`, runID).Scan(&at); err != nil { t.Fatal(err) } if at == nil { t.Error("the stop answered a unit but recorded no intent: the sweep would restart the run it just stopped") } } // The uncontended path must not have changed: no restart, no lock to wait for, and the answer is the // unit that is there. Without this the test above could pass over a call that always answered the // newest attempt for the wrong reason. func TestAnUncontendedStopStillAnswersItsUnit(t *testing.T) { s, ctx, runID, now := runWithAttempt(t, "unit-ONLY") run, unit, err := s.RequestStop(ctx, "u1", runID, now) if err != nil { t.Fatal(err) } if unit != "unit-ONLY" || run.ID != runID { t.Errorf("an ordinary stop answered run %q unit %q", run.ID, unit) } }