meshcore-analyzer/cmd/ingestor/async_migration_test.go

package main

import (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"log"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// waitForStatus polls AsyncMigrationStatus until it matches `want` or `deadline` passes.
func waitForStatus(t *testing.T, s *Store, name, want string, timeout time.Duration) string {
	t.Helper()
	deadline := time.Now().Add(timeout)
	var status string
	var err error
	for time.Now().Before(deadline) {
		status, err = s.AsyncMigrationStatus(name)
		if err == nil && status == want {
			return status
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("status never reached %q within %s: got %q (err=%v)", want, timeout, status, err)
	return status
}

// TestRunAsyncMigration_PendingThenDone pins the contract for RunAsyncMigration:
//
//  1. After calling, the migration name MUST be queryable in the migrations
//     table with status `pending_async` IMMEDIATELY (no waiting for fn).
//  2. After fn returns, the status MUST transition to `done`.
//  3. RunAsyncMigration MUST return without blocking on fn.
//
// This is the regression test for the recurring "sync migration on large
// table blocks ingestor startup" class (#791, #1483, ...). If this test
// fails the contract is broken — do not relax it; fix the runner.
func TestRunAsyncMigration_PendingThenDone(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()

	started := make(chan struct{})
	release := make(chan struct{})

	const name = "test_async_migration_v1"
	if err := s.RunAsyncMigration(ctx, name, func(ctx context.Context, db *sql.DB) error {
		close(started)
		<-release
		return nil
	}); err != nil {
		t.Fatalf("RunAsyncMigration returned error: %v", err)
	}

	// Wait for the goroutine to actually start before checking status; this
	// proves RunAsyncMigration did not block on fn and that fn is running
	// concurrently.
	select {
	case <-started:
	case <-time.After(2 * time.Second):
		t.Fatal("async migration fn did not start within 2s — RunAsyncMigration may have blocked or never scheduled")
	}

	status, err := s.AsyncMigrationStatus(name)
	if err != nil {
		t.Fatalf("AsyncMigrationStatus while running: %v", err)
	}
	if status != "pending_async" {
		t.Fatalf("status while fn running: got %q, want %q", status, "pending_async")
	}

	close(release)

	// Poll for transition to done.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		status, err = s.AsyncMigrationStatus(name)
		if err == nil && status == "done" {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("status never transitioned to done within 2s: got %q (err=%v)", status, err)
}

// TestRunAsyncMigration_PanicCapture proves that a panic inside fn does NOT
// leak past the recover, AND that the migration row transitions to
// "failed" with the panic message captured — NOT silently to "done".
// Operator visibility into mid-migration crashes is the whole point.
func TestRunAsyncMigration_PanicCapture(t *testing.T) {
	s := newTestStore(t)
	const name = "test_panic_capture_v1"

	if err := s.RunAsyncMigration(context.Background(), name,
		func(ctx context.Context, db *sql.DB) error {
			panic("synthetic boom")
		}); err != nil {
		t.Fatalf("RunAsyncMigration returned error: %v", err)
	}

	s.WaitForAsyncMigrations()

	status, err := s.AsyncMigrationStatus(name)
	if err != nil {
		t.Fatalf("status lookup: %v", err)
	}
	if status != "failed" {
		t.Fatalf("status after panic: got %q, want %q (silent-done would be catastrophic)", status, "failed")
	}

	var errMsg sql.NullString
	if err := s.db.QueryRow(`SELECT error FROM _async_migrations WHERE name = ?`, name).Scan(&errMsg); err != nil {
		t.Fatalf("error column lookup: %v", err)
	}
	if !errMsg.Valid || errMsg.String == "" {
		t.Fatalf("error column empty after panic — operator has no clue what failed")
	}
}

// TestRunAsyncMigration_IdempotentSecondCallNoOps verifies that calling
// RunAsyncMigration a second time with the same name AFTER it has reached
// "done" status does NOT re-run fn. This protects the prod path: ingestor
// restarts must not rebuild already-built indexes.
func TestRunAsyncMigration_IdempotentSecondCallNoOps(t *testing.T) {
	s := newTestStore(t)
	const name = "test_idempotent_v1"

	var calls int32
	fn := func(ctx context.Context, db *sql.DB) error {
		atomic.AddInt32(&calls, 1)
		return nil
	}

	if err := s.RunAsyncMigration(context.Background(), name, fn); err != nil {
		t.Fatalf("first call: %v", err)
	}
	s.WaitForAsyncMigrations()
	waitForStatus(t, s, name, "done", 2*time.Second)

	// Second call must short-circuit; fn must not be invoked again.
	if err := s.RunAsyncMigration(context.Background(), name, fn); err != nil {
		t.Fatalf("second call: %v", err)
	}
	s.WaitForAsyncMigrations()

	if got := atomic.LoadInt32(&calls); got != 1 {
		t.Fatalf("fn invoked %d times, want 1 (done-state row must short-circuit)", got)
	}
}

// Failed and cancelled runs are retried, with stale completion details cleared.
func TestRunAsyncMigration_RestartSafetyInterruptedIsRetried(t *testing.T) {
	for _, previous := range []string{"failed", "cancelled"} {
		t.Run(previous, func(t *testing.T) {
			s := newTestStore(t)
			const name = "test_restart_interrupted_v1"
			if _, err := s.db.Exec(`INSERT INTO _async_migrations (name, status, ended_at, error) VALUES (?, ?, '2026-01-01 00:00:00', 'prior interruption')`, name, previous); err != nil {
				t.Fatalf("seed interrupted row: %v", err)
			}

			ctx, cancel := context.WithCancel(context.Background())
			t.Cleanup(func() { cancel(); s.WaitForAsyncMigrations() })
			release := make(chan struct{})
			var calls int32
			if err := s.RunAsyncMigration(ctx, name, func(ctx context.Context, db *sql.DB) error {
				atomic.AddInt32(&calls, 1)
				select {
				case <-release:
					return nil
				case <-ctx.Done():
					return ctx.Err()
				}
			}); err != nil {
				t.Fatalf("RunAsyncMigration on %s row: %v", previous, err)
			}

			var status string
			var endedAt, errCol sql.NullString
			if err := s.db.QueryRow(`SELECT status, ended_at, error FROM _async_migrations WHERE name = ?`, name).Scan(&status, &endedAt, &errCol); err != nil {
				t.Fatal(err)
			}
			if status != "pending_async" || endedAt.Valid || errCol.Valid {
				t.Fatalf("retry not reset: status=%q ended_at=%v error=%v", status, endedAt, errCol)
			}
			close(release)
			s.WaitForAsyncMigrations()
			if err := s.db.QueryRow(`SELECT status, ended_at, error FROM _async_migrations WHERE name = ?`, name).Scan(&status, &endedAt, &errCol); err != nil {
				t.Fatal(err)
			}
			if status != "done" || !endedAt.Valid || endedAt.String == "" || errCol.Valid || atomic.LoadInt32(&calls) != 1 {
				t.Fatalf("retry did not complete once: status=%q ended_at=%v error=%v calls=%d", status, endedAt, errCol, calls)
			}
		})
	}
}

// TestRunAsyncMigration_RestartSafetyPendingIsRetried simulates the
// ingestor crashing while a migration was still in `pending_async` (the
// goroutine never finished). On next boot the migration MUST be re-picked-up
// — leaving it stuck in pending forever would be a silent prod outage.
func TestRunAsyncMigration_RestartSafetyPendingIsRetried(t *testing.T) {
	s := newTestStore(t)
	const name = "test_restart_pending_v1"

	if err := ensureAsyncMigrationsTable(s.db); err != nil {
		t.Fatalf("ensure table: %v", err)
	}
	if _, err := s.db.Exec(`INSERT INTO _async_migrations (name, status) VALUES (?, 'pending_async')`, name); err != nil {
		t.Fatalf("seed pending row: %v", err)
	}

	var calls int32
	if err := s.RunAsyncMigration(context.Background(), name,
		func(ctx context.Context, db *sql.DB) error {
			atomic.AddInt32(&calls, 1)
			return nil
		}); err != nil {
		t.Fatalf("RunAsyncMigration on pending row: %v", err)
	}
	s.WaitForAsyncMigrations()
	waitForStatus(t, s, name, "done", 2*time.Second)

	if got := atomic.LoadInt32(&calls); got != 1 {
		t.Fatalf("fn invoked %d times, want 1 (pending row must be retried after crash)", got)
	}
}

// Classify the returned error, not ctx.Err(): shutdown must not hide failures.
func TestRunAsyncMigration_CompletionStatus(t *testing.T) {
	for _, tt := range []struct {
		name   string
		result error
		panic  bool
		want   string
	}{
		{"cancelled", context.Canceled, false, "cancelled"},
		{"wrapped cancellation", fmt.Errorf("backfill: %w", context.Canceled), false, "cancelled"},
		{"ordinary error", fmt.Errorf("simulated migration error"), false, "failed"},
		{"deadline exceeded", context.DeadlineExceeded, false, "failed"},
		{"panic", context.Canceled, true, "failed"},
		{"completed", nil, false, "done"},
	} {
		t.Run(tt.name, func(t *testing.T) {
			s := newTestStore(t)
			ctx, cancel := context.WithCancel(context.Background())
			var logs bytes.Buffer
			previousOutput := log.Writer()
			log.SetOutput(&logs)
			t.Cleanup(func() { cancel(); s.WaitForAsyncMigrations(); log.SetOutput(previousOutput) })
			const name = "test_completion_v1"
			if err := s.RunAsyncMigration(ctx, name, func(ctx context.Context, db *sql.DB) error {
				<-ctx.Done()
				if tt.panic {
					panic(tt.result)
				}
				return tt.result
			}); err != nil {
				t.Fatal(err)
			}
			cancel()
			s.WaitForAsyncMigrations()

			var status string
			var endedAt, errCol sql.NullString
			if err := s.db.QueryRow(`SELECT status, ended_at, error FROM _async_migrations WHERE name = ?`, name).Scan(&status, &endedAt, &errCol); err != nil {
				t.Fatal(err)
			}
			if status != tt.want || !endedAt.Valid || endedAt.String == "" {
				t.Fatalf("completion: status=%q want=%q ended_at=%v", status, tt.want, endedAt)
			}
			wantErr := ""
			if tt.result != nil {
				wantErr = tt.result.Error()
			}
			if tt.panic {
				wantErr = "panic: " + wantErr
			}
			if errCol.Valid != (tt.result != nil) || errCol.String != wantErr {
				t.Fatalf("error=%v, want %q", errCol, wantErr)
			}
			if tt.want == "cancelled" && (strings.Contains(logs.String(), "FAILED") || !strings.Contains(logs.String(), "cancelled")) {
				t.Fatalf("cancellation logged as failure: %s", logs.String())
			}
		})
	}
}

// TestRunAsyncMigration_ConcurrentSameNameSerialized validates the
// single-process-instance assumption: ingestor has only one *Store, and
// concurrent RunAsyncMigration(name=X) calls on the SAME *Store must not
// execute fn more than once for a given name. (CoreScope does not support
// multi-ingestor / cluster mode — see MIGRATIONS.md "Concurrency" note —
// so cross-process races are out of scope.)
func TestRunAsyncMigration_ConcurrentSameNameSerialized(t *testing.T) {
	s := newTestStore(t)
	const name = "test_concurrent_serialize_v1"

	var calls int32
	fn := func(ctx context.Context, db *sql.DB) error {
		atomic.AddInt32(&calls, 1)
		time.Sleep(20 * time.Millisecond)
		return nil
	}

	var wg sync.WaitGroup
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			// All concurrent callers use the SAME name. Each is allowed
			// to either no-op (status==done short-circuit) or schedule
			// a re-run; the invariant is "fn never runs more than once
			// concurrently and on second-call-after-done it does not
			// re-execute."
			_ = s.RunAsyncMigration(context.Background(), name, fn)
		}()
	}
	wg.Wait()
	s.WaitForAsyncMigrations()
	waitForStatus(t, s, name, "done", 2*time.Second)

	// The contract per the helper's docstring + Idempotent test is: once
	// status is `done`, subsequent calls short-circuit. Concurrent calls
	// that lose the race to set up the pending_async row may legitimately
	// re-schedule fn (the comment "previous run may have crashed
	// mid-flight" justifies retry on pending_async). The hard bound is
	// "fn runs at most ONCE PER pending->done transition" — for this
	// test we assert fn ran at least once and at most a small bounded
	// number (5 callers, each may have scheduled before any reached done).
	if got := atomic.LoadInt32(&calls); got < 1 || got > 5 {
		t.Fatalf("fn invoked %d times, want 1..5 inclusive (bounded by caller count)", got)
	}
}