textmachine/platform/internal/pgstore/stoprace_test.go

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)
	}
}