From 16623efd46ed2797f6d2ebb1be1d23c236242fe3 Mon Sep 17 00:00:00 2001 From: Rory Shanks Date: Tue, 8 Sep 2026 10:55:56 +0200 Subject: [PATCH] Add state migrations and reduce Postgres scheduling overhead --- README.md | 10 +- cmd/partforge/main.go | 140 +--- docs/deployment.md | 2 + docs/development.md | 9 + docs/postgres.md | 59 +- docs/setup.md | 3 + e2e/run.sh | 3 + internal/rewrite/processor.go | 70 +- internal/rewrite/processor_test.go | 88 +- internal/state/maintenance.go | 171 ++++ internal/state/migrations.go | 143 ++++ internal/state/postgres.go | 882 +++++++------------- internal/state/postgres_integration_test.go | 652 +++++++++++++++ internal/state/postgres_test.go | 333 -------- 14 files changed, 1464 insertions(+), 1101 deletions(-) create mode 100644 internal/state/maintenance.go create mode 100644 internal/state/migrations.go create mode 100644 internal/state/postgres_integration_test.go diff --git a/README.md b/README.md index 237f0fc..2703837 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,15 @@ on source node parts to S3; local ClickHouse; destinati ## Getting started -Two SQL files define your migration; everything else is mechanical. Write them, then run four commands. +Initialize or upgrade the Postgres state schema before running uploads or workers: + +```sh +partforge migrate -postgres-url="$POSTGRES_URL" +``` + +For an existing deployment, stop workers while migrating. See [Postgres migrations](docs/postgres.md#migrate-before-starting-workers). + +Two SQL files define your ClickHouse rewrite. Write them, then run the pipeline commands below. ### 1. Destination schema — `dest.sql` diff --git a/cmd/partforge/main.go b/cmd/partforge/main.go index 02c30f8..231f279 100644 --- a/cmd/partforge/main.go +++ b/cmd/partforge/main.go @@ -71,6 +71,7 @@ type commandHelp struct { } var commandHelps = []commandHelp{ + {Name: "migrate", Usage: "[flags]", Summary: "Apply pending Postgres state schema migrations.", Details: "Run once before starting workers. Adopts existing state tables. Stop workers while migrating; schema changes and backfills run in one transaction."}, { Name: "upload-backup", Usage: "[flags]", @@ -266,6 +267,8 @@ func run() error { defer stop() switch os.Args[1] { + case "migrate": + return runMigrate(ctx, os.Args[2:]) case "upload-backup": return runUploadBackup(ctx, os.Args[2:]) case "upload-freeze": @@ -1725,6 +1728,7 @@ func runWorker(ctx context.Context, args []string) error { if roleSettings.Compact { didCompactWork, err := runWorkerCompaction(ctx, workerCompactionConfig{ StateStore: stateStore, + Once: *once, WorkerID: resolvedWorkerID, WorkDir: *workDir, ClickHouseURL: *clickHouseURL, @@ -2018,6 +2022,7 @@ func createWorkerRunDirs(workDir string) (workerRunDirs, error) { } type workerCompactionConfig struct { + Once bool StateStore *state.Store WorkerID string WorkDir string @@ -2068,38 +2073,17 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool, return false, err } } - if cfg.CompactLeaseStaleAfter > 0 { - now := time.Now().UTC() - released, err := cfg.StateStore.ReleaseStaleCompactingParts(ctx, now, cfg.CompactLeaseStaleAfter) - if err != nil { - if ctx.Err() != nil { - slog.Info("worker shutdown requested while releasing stale compact work", "stage", "shutdown") - return false, nil - } - return false, err - } - if released > 0 { - slog.Warn("released stale compacting parts", "stage", "release_stale_compact", "worker_id", cfg.WorkerID, "released", released, "stale_after", cfg.CompactLeaseStaleAfter) - } + finalized, err := cfg.StateStore.MaintainCompaction(ctx, cfg.CompactWindow, cfg.CompactLeaseStaleAfter, time.Now().UTC(), cfg.Once) + if err != nil { + return false, err } - finalization := compactFinalizationResult{} - if cfg.CompactWindow > 0 { - var err error - finalization, err = finalizeCompactReadyJobs(ctx, cfg.StateStore, cfg.CompactWindow, time.Now().UTC()) - if err != nil { - return false, err - } - if finalization.Finalized > 0 { - slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalization.Finalized) - return true, nil - } - if len(finalization.ExpiredJobIDs) > 0 { - slog.Info("skipping compact claims for jobs past compact window", "stage", "claim_compact", "worker_id", cfg.WorkerID, "jobs", len(finalization.ExpiredJobIDs)) - } + if finalized > 0 { + slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalized) + return true, nil } slog.Info("claiming compact-ready batch", "stage", "claim_compact", "worker_id", cfg.WorkerID) batch, err := cfg.StateStore.ClaimNextCompactBatch(ctx, cfg.WorkerID, time.Now().UTC(), state.CompactClaimOptions{ - ExcludedJobIDs: finalization.ExpiredJobIDs, + CompactWindow: cfg.CompactWindow, }) if err != nil { if ctx.Err() != nil { @@ -2109,14 +2093,6 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool, return false, err } if batch == nil { - finalization, err := finalizeCompactReadyJobs(ctx, cfg.StateStore, cfg.CompactWindow, time.Now().UTC()) - if err != nil { - return false, err - } - if finalization.Finalized > 0 { - slog.Info("finalized compact-ready artifacts", "stage", "finalize_compact", "artifacts", finalization.Finalized) - return true, nil - } return false, nil } currentBatch := func() state.CompactBatch { @@ -2174,11 +2150,7 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool, DestinationSchema: batch.Parts[0].DestinationSchema, Inputs: compactInputs(batch.Parts), } - jobParts, err := cfg.StateStore.ListJobParts(ctx, batch.JobID) - if err != nil { - return true, markCompactBatchFailed(fmt.Errorf("list job parts for compact deadline: %w", err)) - } - compactDeadline, err := compactBatchDeadline(jobParts, cfg.CompactWindow, time.Now().UTC()) + compactDeadline, err := cfg.StateStore.CompactDeadline(ctx, batch.JobID, cfg.CompactWindow) if err != nil { return true, markCompactBatchFailed(err) } @@ -2191,13 +2163,13 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool, } slog.Info("compact window expired before compact batch started; released", "stage", "compact_window_expired", "job_id", batch.JobID, "output_part_id", outputPartID) finalizeCtx, finalizeCancel := workerStateUpdateContext() - finalization, finalizeErr := finalizeCompactReadyJob(finalizeCtx, cfg.StateStore, batch.JobID, cfg.CompactWindow, time.Now().UTC()) + finalized, finalizeErr := cfg.StateStore.FinalizeCompactReadyJob(finalizeCtx, batch.JobID, cfg.CompactWindow, time.Now().UTC()) finalizeCancel() if finalizeErr != nil { return true, finalizeErr } - if finalization.Finalized > 0 { - slog.Info("finalized compact-ready artifacts after compact window expiration", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalization.Finalized) + if finalized > 0 { + slog.Info("finalized compact-ready artifacts after compact window expiration", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalized) } return true, nil } @@ -2274,14 +2246,14 @@ func runWorkerCompaction(ctx context.Context, cfg workerCompactionConfig) (bool, return true, nil } finalizeCtx, finalizeCancel := workerStateUpdateContext() - finalization, finalizeErr := finalizeCompactReadyJob(finalizeCtx, cfg.StateStore, batch.JobID, cfg.CompactWindow, time.Now().UTC()) + finalized, finalizeErr := cfg.StateStore.FinalizeCompactReadyJob(finalizeCtx, batch.JobID, cfg.CompactWindow, time.Now().UTC()) finalizeCancel() if finalizeErr != nil { cleanupCompactNow() return true, finalizeErr } - if finalization.Finalized > 0 { - slog.Info("finalized compact-ready artifacts after no-reduction compaction", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalization.Finalized) + if finalized > 0 { + slog.Info("finalized compact-ready artifacts after no-reduction compaction", "stage", "finalize_compact", "job_id", batch.JobID, "artifacts", finalized) } cleanupCompactNow() return true, nil @@ -2591,59 +2563,6 @@ func compactClaimSplayMax(compactWindow time.Duration) time.Duration { return 250 * time.Millisecond } -type compactFinalizationResult struct { - Finalized int - ExpiredJobIDs map[string]struct{} -} - -func finalizeCompactReadyJobs(ctx context.Context, store *state.Store, compactWindow time.Duration, now time.Time) (compactFinalizationResult, error) { - result := compactFinalizationResult{ExpiredJobIDs: map[string]struct{}{}} - jobIDs, err := store.ListJobIDsByStatus(ctx, state.StatusCompactReady) - if err != nil { - return result, err - } - for _, jobID := range jobIDs { - jobResult, err := finalizeCompactReadyJob(ctx, store, jobID, compactWindow, now) - if err != nil { - return result, err - } - result.Finalized += jobResult.Finalized - for expiredJobID := range jobResult.ExpiredJobIDs { - result.ExpiredJobIDs[expiredJobID] = struct{}{} - } - } - return result, nil -} - -func finalizeCompactReadyJob(ctx context.Context, store *state.Store, jobID string, compactWindow time.Duration, now time.Time) (compactFinalizationResult, error) { - result := compactFinalizationResult{ExpiredJobIDs: map[string]struct{}{}} - parts, err := store.ListJobParts(ctx, jobID) - if err != nil { - return result, err - } - expired, err := compactWindowExpired(parts, compactWindow, now) - if err != nil { - return result, err - } - if expired { - result.ExpiredJobIDs[jobID] = struct{}{} - } - compactReady, ok, err := finalizableCompactReadyParts(parts, compactWindow, now) - if err != nil { - return result, err - } - if !ok { - return result, nil - } - for _, part := range compactReady { - if err := store.MarkCompactReadyFinished(ctx, part, now); err != nil { - return result, err - } - result.Finalized++ - } - return result, nil -} - func finalizableCompactReadyParts(parts []state.Part, compactWindow time.Duration, now time.Time) ([]state.Part, bool, error) { compactReady := compactReadyParts(parts) if len(compactReady) == 0 { @@ -2862,6 +2781,27 @@ func runImportFinished(ctx context.Context, args []string) error { }) } +func runMigrate(ctx context.Context, args []string) error { + fs := newCommandFlagSet("migrate") + configPath := fs.String("config", defaultConfigPath, "JSON config file path; CLI flags override config values") + stateTable := fs.String("state-table", defaultStateTable, "Postgres table used for PartForge state") + region := fs.String("aws-region", "", "AWS region for Postgres IAM auth") + postgresURL := fs.String("postgres-url", "", "Postgres state store connection URL") + postgresIAMAuth := fs.Bool("postgres-iam-auth", false, "use AWS IAM authentication for the Postgres state store") + if err := parseFlags(fs, args); err != nil { + return err + } + if err := applyConfigDefaults(fs, *configPath, "migrate"); err != nil { + return err + } + applied, err := state.Migrate(ctx, state.Config{Region: *region, Endpoint: *postgresURL, IAMAuth: *postgresIAMAuth, Table: *stateTable}) + if err != nil { + return err + } + slog.Info("state schema is up to date", "migrations_applied", applied, "state_table", *stateTable) + return nil +} + func runListJobs(ctx context.Context, args []string) error { fs := newCommandFlagSet("list-jobs") var ( diff --git a/docs/deployment.md b/docs/deployment.md index 0a4fac9..10c6a25 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -4,6 +4,8 @@ The worker image is published on every push to `main` to `ghcr.io//partfo The image is a single Ubuntu container with `clickhouse-server`, `clickhouse-client`, `s5cmd`, and the Go binary. Its entrypoint is the binary and the default command is `worker`. It **runs as root** (so it can write its work directory on root-owned host mounts) and starts a local `clickhouse server` child process for each claimed part. +Before starting a new worker release, stop the existing workers and run `partforge migrate` once with that release's image and the same Postgres/state-table settings. The migration backfills and indexes the existing table under table locks; restart workers after it succeeds. See [the migration procedure](postgres.md#migrate-before-starting-workers). + ## Recommended: workers on ECS with an IAM task role Run the workers as an ECS service and give the task an **IAM role** scoped to the S3 bucket and the RDS/Aurora PostgreSQL database user. This is the recommended setup: diff --git a/docs/development.md b/docs/development.md index 5713056..103961b 100644 --- a/docs/development.md +++ b/docs/development.md @@ -16,6 +16,15 @@ go test ./... The e2e script stands up LocalStack, Postgres, and a ClickHouse container, builds the worker image, and runs the full pipeline against `e2e/sql/`, diffing the result against `e2e/expected.tsv`. It builds the image each run; set `PARTFORGE_E2E_SKIP_BUILD=1` to reuse an existing `partforge-worker:latest`. +PostgreSQL integration checks exercise fresh and hand-created schemas, migration rollback and concurrent migration runs, concurrent work claims, ownership, maintenance, and 20,000-row query plans. With local compose Postgres running: + +```sh +PARTFORGE_TEST_POSTGRES_URL='postgres://partforge:partforge@localhost:15432/partforge?sslmode=disable' \ + go test ./internal/state -run Postgres -count=1 +``` + +These tests create and remove isolated schemas. They skip when the connection variable is unset. The e2e script also runs `migrate` twice before uploading, checking both initial setup and a no-op rerun. + ## Build ```sh diff --git a/docs/postgres.md b/docs/postgres.md index e7246aa..03bcaf2 100644 --- a/docs/postgres.md +++ b/docs/postgres.md @@ -20,30 +20,47 @@ For local compose runs: Default table name is `partforge_state`; override with `-state-table` or `state_table`. -## Schema +## Migrate before starting workers -The app creates the table and indexes on startup: - -```sql -CREATE TABLE IF NOT EXISTS partforge_state ( - job_id text NOT NULL, - part_id text NOT NULL, - status text NOT NULL, - worker_id text NOT NULL DEFAULT '', - created_at text NOT NULL, - updated_at text NOT NULL, - data jsonb NOT NULL, - PRIMARY KEY (job_id, part_id) -); - -CREATE INDEX IF NOT EXISTS partforge_state_status_idx - ON partforge_state (status, created_at, job_id, part_id); - -CREATE INDEX IF NOT EXISTS partforge_state_job_status_idx - ON partforge_state (job_id, status, part_id); +```sh +partforge migrate -postgres-url="$POSTGRES_URL" ``` -The scalar columns support claims and job/status scans. The full part record is stored in `data`. +Use the same `-config`, `-state-table`, `-postgres-iam-auth`, and `-aws-region` settings as the workers. A schema-qualified state table is supported. The database role running migrations must own the existing state table and have permission to create tables and indexes in its schema. + +For an existing deployment: + +1. Stop workers and pause uploads and state-changing admin commands. +2. Run `partforge migrate` with the new binary against the existing state table. +3. Start the new workers after the command succeeds. + +Migrations backfill columns and build indexes in one transaction, holding table locks. Plan a maintenance window; this is not an online migration command. Existing rows, statuses, progress, and lineage are retained. The first migration adopts the original hand-created table with `CREATE TABLE IF NOT EXISTS`; do not recreate or empty it. + +The command records numbered versions in `_migrations`. An advisory transaction lock serializes concurrent migration commands, and a failure rolls back the pending changes and their version records together. Re-running a successful migration command is a no-op. Unexpected versions or gaps fail explicitly. Workers check the version at startup and tell you to run `partforge migrate` if an upgrade is needed; they no longer execute schema DDL. + +Migrations are compiled into the binary in [internal/state/migrations.go](../internal/state/migrations.go). Append new SQL migrations to the list; never edit, reorder, or remove an already released migration. There is no automatic downgrade command. + +## Schema and scheduling + +The original primary key `(job_id, part_id)`, state columns, and full `data` JSON record remain. Additional scalar columns support scheduling without indexing the mutable JSON itself: + +| Columns | Purpose | +|---|---| +| `source_artifact_bytes` | Ordered partial index for largest-first READY claims | +| `compact_bytes`, `compact_eligible` | Ordered partial index for eligible compact claims | +| `compact_normalized` | Identify artifacts containing one physical part | +| `compact_stale_at` | Indexed stale-compaction lookup, preserving the existing earlier-of-heartbeat-and-claim timeout | +| `original_compact_ready_at` | Indexed job compact deadline lookup | + +Application writes update these projections in the same SQL statement as the JSON or status change. No triggers are installed. Use PartForge commands for state changes; direct SQL must maintain the corresponding derived columns too. The full Part record remains in `data`; no job metadata or progress table is required. + +Compactors claim the largest eligible unlocked artifact. Job, destination, and explicit partition filters still apply; partitions with an active compactor are no longer deprioritized. + +A single `_maintenance` record reserves compaction maintenance once every ten seconds across the worker fleet. `worker -once` requests an immediate pass but still skips a pass already in progress. Only the reservation holder scans job summaries and expires stale work. The reservation commits with the work, and active maintenance is skipped by other workers without waiting. Workers sharing a state table should use the same compact-window and stale-timeout settings. New finalizable output can wait up to the next maintenance pass; a long maintenance run or database contention can extend that delay. + +Rewrite progress combines live query counters and stage timing in one periodic heartbeat. Rewrite and compact progress updates use conditional SQL patches that check worker ownership. A failed rewrite heartbeat cancels processing and surfaces the error. + +Connection pool configuration is unchanged. ## IAM Auth diff --git a/docs/setup.md b/docs/setup.md index 5f1555d..e5f39c6 100644 --- a/docs/setup.md +++ b/docs/setup.md @@ -14,6 +14,9 @@ Requirements, configuration, and how to run the four stages by hand. For the hig ```sh docker compose up -d localstack postgres +docker compose build worker +docker compose run --rm worker migrate \ + -postgres-url='postgres://partforge:partforge@postgres:5432/partforge?sslmode=disable' ``` This creates the `partforge` S3 bucket in LocalStack and starts a local Postgres database. Point commands at them with: diff --git a/e2e/run.sh b/e2e/run.sh index 61d6823..99e7a4a 100755 --- a/e2e/run.sh +++ b/e2e/run.sh @@ -114,6 +114,9 @@ for _ in $(seq 1 60); do done docker compose exec -T postgres pg_isready -U partforge -d partforge >/dev/null +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker migrate -postgres-url="$POSTGRES_URL" +CLICKHOUSE_DATA_DIR="$DATA_DIR" docker compose run --rm worker migrate -postgres-url="$POSTGRES_URL" + docker compose exec -T clickhouse clickhouse-client --multiquery < e2e/sql/setup_and_freeze.sql docker compose exec -T clickhouse clickhouse-client --query \ diff --git a/internal/rewrite/processor.go b/internal/rewrite/processor.go index ebcc91c..7d7b9fe 100644 --- a/internal/rewrite/processor.go +++ b/internal/rewrite/processor.go @@ -86,6 +86,7 @@ func StageOrder() []string { } type Processor struct { + progressTracker *rewriteStageTracker S3Copy s3copy.Copier ClickHouse chhttp.Client WorkDir string @@ -208,6 +209,7 @@ type workerTableInfo struct { } type rewriteStageTracker struct { + queryProgress *metrics.QueryProgress mu sync.Mutex reportMu sync.Mutex startedAt time.Time @@ -306,6 +308,7 @@ func (p Processor) ProcessPart(ctx context.Context, item WorkItem) (result Proce progressManifest := manifest.Manifest{JobID: item.JobID, PartID: item.PartID} stageTracker := newRewriteStageTracker(startedAt, stageProcessPart) + p.progressTracker = stageTracker defer p.recorder().ClearStageProgress(progressManifest) heartbeat, err := p.startProgressHeartbeat(ctx, progressManifest, stageTracker) if err != nil { @@ -928,7 +931,7 @@ func (p Processor) runInsertSelect(ctx context.Context, m manifest.Manifest, att recorder := p.recorder() progress := metrics.QueryProgress{} defer recorder.ClearCurrentProgress(m) - lastProgressReport := time.Time{} + p.recordQueryProgress(metrics.QueryProgress{}) ticker := time.NewTicker(time.Second) defer ticker.Stop() @@ -944,13 +947,12 @@ func (p Processor) runInsertSelect(ctx context.Context, m manifest.Manifest, att } if found { recorder.ObserveProgress(m, progress, finalProgress) - if err := p.reportProgress(ctx, m, ProgressSnapshot{QueryProgress: &finalProgress}); err != nil { + if err := p.reportFinalQueryProgress(ctx, m, finalProgress); err != nil { return err } } return nil case <-ticker.C: - now := time.Now() current, found, err := p.queryProgress(ctx, queryID) if err != nil { cancel() @@ -960,14 +962,7 @@ func (p Processor) runInsertSelect(ctx context.Context, m manifest.Manifest, att if found { recorder.ObserveProgress(m, progress, current) progress = current - if shouldReportProgress(p.ProgressInterval, lastProgressReport, now) { - if err := p.reportProgress(ctx, m, ProgressSnapshot{QueryProgress: ¤t}); err != nil { - cancel() - <-errCh - return err - } - lastProgressReport = now - } + p.recordQueryProgress(current) } case <-ctx.Done(): cancel() @@ -977,6 +972,36 @@ func (p Processor) runInsertSelect(ctx context.Context, m manifest.Manifest, att } } +// Query polling records the latest counters; the existing stage heartbeat is +// the only periodic writer. Stage changes and the final query result still flush. +func (p Processor) recordQueryProgress(progress metrics.QueryProgress) { + if p.progressTracker == nil { + return + } + p.progressTracker.mu.Lock() + defer p.progressTracker.mu.Unlock() + p.progressTracker.queryProgress = &progress +} + +func (t *rewriteStageTracker) currentQueryProgress() *metrics.QueryProgress { + t.mu.Lock() + defer t.mu.Unlock() + if t.queryProgress == nil { + return nil + } + copy := *t.queryProgress + return © +} + +func (p Processor) reportFinalQueryProgress(ctx context.Context, m manifest.Manifest, progress metrics.QueryProgress) error { + if p.progressTracker != nil { + p.progressTracker.reportMu.Lock() + defer p.progressTracker.reportMu.Unlock() + } + p.recordQueryProgress(progress) + return p.reportProgress(ctx, m, ProgressSnapshot{QueryProgress: &progress}) +} + func (p Processor) reportProgress(ctx context.Context, m manifest.Manifest, snapshot ProgressSnapshot) error { if p.ReportProgress == nil { return nil @@ -985,6 +1010,7 @@ func (p Processor) reportProgress(ctx context.Context, m manifest.Manifest, snap } type progressHeartbeat struct { + err error ctx context.Context cancel context.CancelFunc done chan struct{} @@ -999,7 +1025,8 @@ func (p Processor) startProgressHeartbeat(ctx context.Context, m manifest.Manife heartbeat.ctx, heartbeat.cancel = context.WithCancel(ctx) heartbeat.done = make(chan struct{}) if err := p.reportStageSnapshot(heartbeat.ctx, m, tracker); err != nil { - slog.Warn("progress heartbeat update failed; continuing", "job_id", m.JobID, "part_id", m.PartID, "error", err) + heartbeat.cancel() + return nil, err } go func() { @@ -1013,7 +1040,9 @@ func (p Processor) startProgressHeartbeat(ctx context.Context, m manifest.Manife if heartbeat.ctx.Err() != nil { return } - slog.Warn("progress heartbeat update failed; continuing", "job_id", m.JobID, "part_id", m.PartID, "error", err) + heartbeat.err = err + heartbeat.cancel() + return } case <-heartbeat.ctx.Done(): return @@ -1038,7 +1067,7 @@ func (p Processor) reportStageProgress(ctx context.Context, m manifest.Manifest, TotalElapsed: progress.TotalElapsed, CompletedStageDurations: progress.CompletedStageDurations, }) - return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress}) + return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress, QueryProgress: tracker.currentQueryProgress()}) } func (p Processor) reportStageComplete(ctx context.Context, m manifest.Manifest, tracker *rewriteStageTracker, stage string) error { @@ -1056,7 +1085,7 @@ func (p Processor) reportStageComplete(ctx context.Context, m manifest.Manifest, TotalElapsed: progress.TotalElapsed, CompletedStageDurations: progress.CompletedStageDurations, }) - return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress}) + return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress, QueryProgress: tracker.currentQueryProgress()}) } func (p Processor) reportStageSnapshot(ctx context.Context, m manifest.Manifest, tracker *rewriteStageTracker) error { @@ -1073,7 +1102,7 @@ func (p Processor) reportStageSnapshot(ctx context.Context, m manifest.Manifest, TotalElapsed: progress.TotalElapsed, CompletedStageDurations: progress.CompletedStageDurations, }) - return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress}) + return p.reportProgress(ctx, m, ProgressSnapshot{StageProgress: &progress, QueryProgress: tracker.currentQueryProgress()}) } func (h *progressHeartbeat) Context() context.Context { @@ -1086,14 +1115,7 @@ func (h *progressHeartbeat) Stop() error { } h.cancel() <-h.done - return nil -} - -func shouldReportProgress(interval time.Duration, last time.Time, now time.Time) bool { - if interval <= 0 { - return false - } - return last.IsZero() || !now.Before(last.Add(interval)) + return h.err } func resetDestinationTable(ctx context.Context, ch chhttp.Client, m manifest.Manifest, destDDL string) error { diff --git a/internal/rewrite/processor_test.go b/internal/rewrite/processor_test.go index aa7589d..af0c9c6 100644 --- a/internal/rewrite/processor_test.go +++ b/internal/rewrite/processor_test.go @@ -17,6 +17,7 @@ import ( "github.com/PostHog/partforge/internal/chhttp" "github.com/PostHog/partforge/internal/freeze" "github.com/PostHog/partforge/internal/manifest" + "github.com/PostHog/partforge/internal/metrics" "github.com/PostHog/partforge/internal/s3copy" ) @@ -967,19 +968,33 @@ func TestInsertSelectRetryBackoff(t *testing.T) { } } -func TestShouldReportProgress(t *testing.T) { - now := time.Unix(100, 0) - if shouldReportProgress(0, time.Time{}, now) { - t.Fatal("expected disabled interval to skip progress report") +func TestQueryProgressSharesStageHeartbeat(t *testing.T) { + tracker := newRewriteStageTracker(time.Now(), stageInsertSelect) + var snapshots []ProgressSnapshot + p := Processor{progressTracker: tracker, ReportProgress: func(_ context.Context, _ manifest.Manifest, snapshot ProgressSnapshot) error { + snapshots = append(snapshots, snapshot) + return nil + }} + p.recordQueryProgress(metrics.QueryProgress{ReadRows: 10}) + p.recordQueryProgress(metrics.QueryProgress{ReadRows: 20}) + if len(snapshots) != 0 { + t.Fatal("query polling wrote progress") + } + m := manifest.Manifest{JobID: "job", PartID: "part"} + if err := p.reportStageSnapshot(context.Background(), m, tracker); err != nil { + t.Fatal(err) + } + if len(snapshots) != 1 || snapshots[0].QueryProgress.ReadRows != 20 || snapshots[0].StageProgress.Stage != stageInsertSelect { + t.Fatalf("combined snapshots: %+v", snapshots) } - if !shouldReportProgress(15*time.Second, time.Time{}, now) { - t.Fatal("expected first progress report") + if err := p.reportFinalQueryProgress(context.Background(), m, metrics.QueryProgress{ReadRows: 30}); err != nil { + t.Fatal(err) } - if shouldReportProgress(15*time.Second, now.Add(-14*time.Second), now) { - t.Fatal("expected interval gate to skip report") + if err := p.reportStageSnapshot(context.Background(), m, tracker); err != nil { + t.Fatal(err) } - if !shouldReportProgress(15*time.Second, now.Add(-15*time.Second), now) { - t.Fatal("expected interval gate to allow report") + if len(snapshots) != 3 || snapshots[1].QueryProgress.ReadRows != 30 || snapshots[2].QueryProgress.ReadRows != 30 { + t.Fatalf("final snapshots: %+v", snapshots) } } @@ -1076,44 +1091,33 @@ func TestProgressHeartbeatDisabled(t *testing.T) { } } -func TestProgressHeartbeatReportFailureContinues(t *testing.T) { - reportErr := errors.New("progress update failed") - attempts := make(chan struct{}, 8) +func TestProgressHeartbeatReportFailureCancelsProcessing(t *testing.T) { + reportErr := errors.New("lost ownership") reports := 0 - processor := Processor{ - ProgressInterval: time.Millisecond, - ReportProgress: func(ctx context.Context, m manifest.Manifest, snapshot ProgressSnapshot) error { - reports++ - attempts <- struct{}{} - if reports == 2 { - return reportErr - } - return nil - }, - } - - tracker := newRewriteStageTracker(time.Now(), stageProcessPart) - heartbeat, err := processor.startProgressHeartbeat(context.Background(), manifest.Manifest{JobID: "job-1", PartID: "part-1"}, tracker) + processor := Processor{ProgressInterval: time.Millisecond, ReportProgress: func(context.Context, manifest.Manifest, ProgressSnapshot) error { + reports++ + if reports > 1 { + return reportErr + } + return nil + }} + heartbeat, err := processor.startProgressHeartbeat(context.Background(), manifest.Manifest{}, newRewriteStageTracker(time.Now(), stageProcessPart)) if err != nil { t.Fatal(err) } - defer func() { - if err := heartbeat.Stop(); err != nil { - t.Fatal(err) - } - }() - - for i := 0; i < 4; i++ { - select { - case <-attempts: - case <-time.After(time.Second): - t.Fatal("timed out waiting for heartbeat report") - } - } select { case <-heartbeat.Context().Done(): - t.Fatal("heartbeat context was canceled by report failure") - default: + case <-time.After(time.Second): + t.Fatal("processing was not canceled") + } + if err := heartbeat.Stop(); !errors.Is(err, reportErr) { + t.Fatalf("stop error = %v", err) + } + if reports != 2 { + t.Fatalf("reports = %d", reports) + } + if _, err := processor.startProgressHeartbeat(context.Background(), manifest.Manifest{}, newRewriteStageTracker(time.Now(), stageProcessPart)); !errors.Is(err, reportErr) { + t.Fatalf("initial report error = %v", err) } } diff --git a/internal/state/maintenance.go b/internal/state/maintenance.go new file mode 100644 index 0000000..80a1292 --- /dev/null +++ b/internal/state/maintenance.go @@ -0,0 +1,171 @@ +package state + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" +) + +// MaintainCompaction runs at most once per ten seconds across looping workers. +// A one-shot worker requests a pass immediately, still skipping active maintenance. +// The reservation and work commit together, so failure does not consume the run. +func (s *Store) MaintainCompaction(ctx context.Context, window, staleAfter time.Duration, now time.Time, once bool) (int, error) { + tx, err := s.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer tx.Rollback(ctx) + tag, err := tx.Exec(ctx, `UPDATE `+s.relatedSQL("maintenance")+` SET next_run_at = $1 + WHERE id IN (SELECT id FROM `+s.relatedSQL("maintenance")+` WHERE next_run_at <= $2 OR $3::boolean FOR UPDATE SKIP LOCKED)`, now.Add(10*time.Second), now, once) + if err != nil { + return 0, err + } + if tag.RowsAffected() == 0 { + return 0, nil + } + if staleAfter > 0 { + if _, err := s.releaseStaleCompactingPartsTx(ctx, tx, now, staleAfter); err != nil { + return 0, err + } + } + finalized, err := s.finalizeCompactReadyTx(ctx, tx, "", window, now) + if err != nil { + return 0, err + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return finalized, nil +} + +func (s *Store) releaseStaleCompactingPartsTx(ctx context.Context, tx pgx.Tx, now time.Time, staleAfter time.Duration) (int, error) { + rows, err := tx.Query(ctx, `SELECT data FROM `+s.tableSQL+` WHERE status = 'COMPACTING' AND compact_stale_at <= $1 FOR UPDATE SKIP LOCKED`, now.Add(-staleAfter)) + if err != nil { + return 0, err + } + parts, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (Part, error) { + var data []byte + if err := row.Scan(&data); err != nil { + return Part{}, err + } + return partFromJSON(data) + }) + if err != nil { + return 0, err + } + for _, part := range parts { + if part.CompactReadyAt == "" { + part.CompactReadyAt = compactReadyAtForRelease(part, now) + } + setStatus(&part, StatusCompactReady, now) + part.WorkerID = "" + part.CompactingAt = "" + part.Error = "" + part.CompactCooldownUntil = "" + clearCompactProgress(&part) + if err := s.savePartTx(ctx, tx, part); err != nil { + return 0, err + } + } + return len(parts), nil +} + +func (s *Store) CompactDeadline(ctx context.Context, jobID string, window time.Duration) (time.Time, error) { + if window <= 0 { + return time.Time{}, nil + } + var readyAt time.Time + err := s.pool.QueryRow(ctx, `SELECT original_compact_ready_at FROM `+s.tableSQL+` WHERE job_id = $1 AND original_compact_ready_at IS NOT NULL ORDER BY original_compact_ready_at DESC LIMIT 1`, jobID).Scan(&readyAt) + if errors.Is(err, pgx.ErrNoRows) { + return time.Time{}, errors.New("no original compact-ready timestamp found") + } + if err != nil { + return time.Time{}, err + } + return readyAt.Add(window), nil +} + +func (s *Store) FinalizeCompactReadyJob(ctx context.Context, jobID string, window time.Duration, now time.Time) (int, error) { + if jobID == "" { + return 0, errors.New("job id is required") + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer tx.Rollback(ctx) + n, err := s.finalizeCompactReadyTx(ctx, tx, jobID, window, now) + if err != nil { + return 0, err + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return n, nil +} + +func (s *Store) finalizeCompactReadyTx(ctx context.Context, tx pgx.Tx, jobID string, window time.Duration, now time.Time) (int, error) { + type summary struct { + jobID string + readyAt *time.Time + blocked, eligible, normalized bool + } + query := `SELECT job_id, max(original_compact_ready_at), + bool_or(status IN ('READY', 'IN_PROGRESS', 'COMPACTING', 'FAILED')), + bool_or(status = 'COMPACT_READY' AND compact_eligible), + bool_or(status = 'COMPACT_READY' AND compact_normalized) + FROM ` + s.tableSQL + args := []any{} + if jobID != "" { + query += ` WHERE job_id = $1` + args = append(args, jobID) + } + query += ` GROUP BY job_id HAVING bool_or(status = 'COMPACT_READY')` + rows, err := tx.Query(ctx, query, args...) + if err != nil { + return 0, err + } + jobs, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (summary, error) { + var j summary + err := row.Scan(&j.jobID, &j.readyAt, &j.blocked, &j.eligible, &j.normalized) + return j, err + }) + if err != nil { + return 0, err + } + finalized := 0 + for _, j := range jobs { + // A normalized artifact can finish even while other work remains active. + normalizedOnly := j.normalized + if !normalizedOnly { + if j.blocked { + continue + } + if window > 0 { + if j.readyAt == nil { + return 0, fmt.Errorf("job %s: no original compact-ready timestamp found", j.jobID) + } + if now.Before(j.readyAt.Add(window)) { + continue + } + } else if j.eligible { + continue + } // Leave useful work for claimers when the window is disabled. + } + tag, err := tx.Exec(ctx, `WITH candidates AS ( + SELECT job_id, part_id FROM `+s.tableSQL+` + WHERE job_id = $1 AND status = 'COMPACT_READY' AND (NOT $2::boolean OR compact_normalized) + FOR UPDATE SKIP LOCKED) + UPDATE `+s.tableSQL+` p SET status = 'FINISHED', updated_at = $3, compact_stale_at = NULL, + data = (p.data - 'error' - 'compact_cooldown_until') || jsonb_build_object('status', 'FINISHED', 'updated_at', $3::text, 'finished_at', $3::text) + FROM candidates c WHERE p.job_id = c.job_id AND p.part_id = c.part_id`, j.jobID, normalizedOnly, formatTime(now)) + if err != nil { + return 0, err + } + finalized += int(tag.RowsAffected()) + } + return finalized, nil +} diff --git a/internal/state/migrations.go b/internal/state/migrations.go new file mode 100644 index 0000000..58672d8 --- /dev/null +++ b/internal/state/migrations.go @@ -0,0 +1,143 @@ +package state + +import ( + "context" + "errors" + "fmt" + "strings" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +// Migrations are append-only. Each state table has its own version ledger. +func (s *Store) migrations() []string { + return []string{ + fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( + job_id text NOT NULL, part_id text NOT NULL, status text NOT NULL, + worker_id text NOT NULL DEFAULT '', created_at text NOT NULL, + updated_at text NOT NULL, data jsonb NOT NULL, PRIMARY KEY (job_id, part_id)); + CREATE INDEX IF NOT EXISTS %s ON %s (status, created_at, job_id, part_id); + CREATE INDEX IF NOT EXISTS %s ON %s (job_id, status, part_id);`, + s.tableSQL, s.statusIndexSQL, s.tableSQL, s.jobStatusIndexSQL, s.tableSQL), + fmt.Sprintf(`ALTER TABLE %[1]s ADD COLUMN source_artifact_bytes numeric(20,0) NOT NULL DEFAULT 0; + UPDATE %[1]s SET source_artifact_bytes = COALESCE((data->>'source_artifact_bytes')::numeric, 0); + ALTER TABLE %[1]s ADD CHECK (source_artifact_bytes >= 0); + CREATE INDEX %[2]s ON %[1]s (source_artifact_bytes DESC, created_at, job_id, part_id) WHERE status = 'READY';`, + s.tableSQL, s.indexSQL("ready_priority_idx")), + fmt.Sprintf(`ALTER TABLE %[1]s + ADD COLUMN compact_bytes numeric(20,0) NOT NULL DEFAULT 0, + ADD COLUMN compact_eligible boolean NOT NULL DEFAULT false, + ADD COLUMN compact_normalized boolean NOT NULL DEFAULT false, + ADD COLUMN compact_stale_at timestamptz, + ADD COLUMN original_compact_ready_at timestamptz; + UPDATE %[1]s SET + original_compact_ready_at = CASE WHEN COALESCE((data->>'compact_generation')::int, 0) <= 0 AND jsonb_array_length(COALESCE(NULLIF(data->'compact_input_part_ids', 'null'::jsonb), '[]'::jsonb)) = 0 THEN NULLIF(data->>'compact_ready_at', '')::timestamptz END, + compact_bytes = COALESCE((data->>'destination_active_part_bytes')::numeric, 0), + compact_eligible = + COALESCE(btrim(data->>'destination_database'), '') <> '' AND + COALESCE(btrim(data->>'destination_table'), '') <> '' AND + COALESCE(btrim(data->>'destination_schema'), '') <> '' AND + COALESCE((data->>'destination_active_part_count')::numeric, 0) > 0 AND + EXISTS (SELECT FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 1), + compact_normalized = COALESCE((data->>'destination_active_part_count')::numeric, 0) = 1 AND + (SELECT count(*) = 1 AND COALESCE(bool_and(p.value::numeric = 1), false) + FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 0), + compact_stale_at = CASE WHEN status = 'COMPACTING' THEN + LEAST(NULLIF(updated_at, '')::timestamptz, NULLIF(data->>'compacting_at', '')::timestamptz) END; + ALTER TABLE %[1]s ADD CHECK (status <> 'COMPACTING' OR compact_stale_at IS NOT NULL); + CREATE INDEX %[2]s ON %[1]s (compact_bytes DESC, created_at, job_id, part_id) WHERE status = 'COMPACT_READY' AND compact_eligible; + CREATE INDEX %[3]s ON %[1]s (compact_stale_at, job_id, part_id) WHERE status = 'COMPACTING'; + CREATE INDEX %[4]s ON %[1]s (job_id, original_compact_ready_at DESC) WHERE original_compact_ready_at IS NOT NULL; + CREATE TABLE %[5]s (id boolean PRIMARY KEY CHECK (id), next_run_at timestamptz NOT NULL); + INSERT INTO %[5]s VALUES (true, '-infinity');`, + s.tableSQL, s.indexSQL("compact_priority_idx"), s.indexSQL("compact_stale_idx"), s.indexSQL("compact_deadline_idx"), s.relatedSQL("maintenance")), + } +} + +func (s *Store) indexSQL(suffix string) string { return quoteIndexName(s.tableName, suffix) } + +func (s *Store) relatedSQL(suffix string) string { + parts := strings.Split(s.tableName, ".") + name := quoteIndexName(parts[len(parts)-1], suffix) + if len(parts) > 1 { + return pgx.Identifier(parts[:len(parts)-1]).Sanitize() + "." + name + } + return name +} + +func (s *Store) schemaVersion(ctx context.Context, q interface { + Query(context.Context, string, ...any) (pgx.Rows, error) +}) (int, error) { + rows, err := q.Query(ctx, `SELECT version FROM `+s.relatedSQL("migrations")+` ORDER BY version`) + if err != nil { + return 0, err + } + defer rows.Close() + version := 0 + for rows.Next() { + var next int + if err := rows.Scan(&next); err != nil { + return 0, err + } + if next != version+1 || next > len(s.migrations()) { + return 0, fmt.Errorf("unsupported migration history: version %d after %d", next, version) + } + version = next + } + return version, rows.Err() +} + +func (s *Store) checkSchema(ctx context.Context) error { + version, err := s.schemaVersion(ctx, s.pool) + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "42P01" { + return errors.New("state schema is not migrated; run partforge migrate with the same Postgres URL and state table") + } + if err != nil { + return fmt.Errorf("check state schema: %w", err) + } + if version != len(s.migrations()) { + return fmt.Errorf("state schema is at version %d, expected %d; run partforge migrate", version, len(s.migrations())) + } + return nil +} + +// Migrate adopts the original hand-created schema and applies pending upgrades. +// Stop workers for migration: backfills and index builds hold table locks. +func Migrate(ctx context.Context, cfg Config) (int, error) { + s, err := openStore(ctx, cfg) + if err != nil { + return 0, err + } + defer s.pool.Close() + tx, err := s.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer tx.Rollback(ctx) + // ponytail: serialize migrations per database; per-table locks if parallel migrations are ever needed. + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('partforge:migrate', 0))`); err != nil { + return 0, err + } + if _, err := tx.Exec(ctx, `CREATE TABLE IF NOT EXISTS `+s.relatedSQL("migrations")+` (version integer PRIMARY KEY, applied_at timestamptz NOT NULL DEFAULT now())`); err != nil { + return 0, err + } + version, err := s.schemaVersion(ctx, tx) + if err != nil { + return 0, err + } + migrations := s.migrations() + for i := version; i < len(migrations); i++ { + if _, err := tx.Exec(ctx, migrations[i]); err != nil { + return 0, fmt.Errorf("migration %d: %w", i+1, err) + } + if _, err := tx.Exec(ctx, `INSERT INTO `+s.relatedSQL("migrations")+` (version) VALUES ($1)`, i+1); err != nil { + return 0, err + } + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return len(migrations) - version, nil +} diff --git a/internal/state/postgres.go b/internal/state/postgres.go index 4c82ebe..cff978c 100644 --- a/internal/state/postgres.go +++ b/internal/state/postgres.go @@ -5,8 +5,8 @@ import ( "encoding/json" "errors" "fmt" + "math/big" "net" - "sort" "strconv" "strings" "time" @@ -14,6 +14,7 @@ import ( "github.com/aws/aws-sdk-go-v2/config" rdsauth "github.com/aws/aws-sdk-go-v2/feature/rds/auth" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgxpool" ) @@ -57,6 +58,7 @@ type Config struct { } type Store struct { + tableName string pool *pgxpool.Pool tableSQL string statusIndexSQL string @@ -174,6 +176,7 @@ func clonePartitionCounts(counts map[string]uint64) map[string]uint64 { } type CompactClaimOptions struct { + CompactWindow time.Duration ExcludedJobIDs map[string]struct{} JobID string Bucket string @@ -209,6 +212,18 @@ type RewriteStageProgress struct { } func New(ctx context.Context, cfg Config) (*Store, error) { + store, err := openStore(ctx, cfg) + if err != nil { + return nil, err + } + if err := store.checkSchema(ctx); err != nil { + store.pool.Close() + return nil, err + } + return store, nil +} + +func openStore(ctx context.Context, cfg Config) (*Store, error) { if strings.TrimSpace(cfg.Table) == "" { cfg.Table = defaultStateTable } @@ -236,14 +251,11 @@ func New(ctx context.Context, cfg Config) (*Store, error) { } store := &Store{ pool: pool, + tableName: strings.TrimSpace(cfg.Table), tableSQL: tableSQL, statusIndexSQL: statusIndexSQL, jobStatusIndexSQL: jobStatusIndexSQL, } - if err := store.ensureSchema(ctx); err != nil { - pool.Close() - return nil, err - } return store, nil } @@ -309,29 +321,6 @@ func quoteIndexName(table, suffix string) string { return pgx.Identifier{base + "_" + suffix}.Sanitize() } -func (s *Store) ensureSchema(ctx context.Context) error { - statements := []string{ - fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s ( - job_id text NOT NULL, - part_id text NOT NULL, - status text NOT NULL, - worker_id text NOT NULL DEFAULT '', - created_at text NOT NULL, - updated_at text NOT NULL, - data jsonb NOT NULL, - PRIMARY KEY (job_id, part_id) - )`, s.tableSQL), - fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON %s (status, created_at, job_id, part_id)`, s.statusIndexSQL, s.tableSQL), - fmt.Sprintf(`CREATE INDEX IF NOT EXISTS %s ON %s (job_id, status, part_id)`, s.jobStatusIndexSQL, s.tableSQL), - } - for _, statement := range statements { - if _, err := s.pool.Exec(ctx, statement); err != nil { - return fmt.Errorf("ensure postgres state schema: %w", err) - } - } - return nil -} - func NewPart(jobID, partID, bucket, sourceKey, finishedKey string, now time.Time) Part { createdAt := formatTime(now) return Part{ @@ -399,15 +388,73 @@ func partFromJSON(data []byte) (Part, error) { return part, nil } -func (s *Store) savePartTx(ctx context.Context, tx pgx.Tx, part Part) error { +const partColumns = "job_id, part_id, status, worker_id, created_at, updated_at, data, source_artifact_bytes, compact_bytes, compact_eligible, compact_normalized, compact_stale_at, original_compact_ready_at" + +// Full writes derive scheduling values once and persist them alongside the JSON. +func partWriteValues(part Part) ([]any, error) { data, err := partJSON(part) + if err != nil { + return nil, err + } + partitions, fragmented, single := 0, false, true + for id, count := range part.DestinationActivePartitionCounts { + if strings.TrimSpace(id) == "" || count == 0 { + continue + } + partitions++ + fragmented = fragmented || count > 1 + single = single && count == 1 + } + eligible := strings.TrimSpace(part.DestinationDatabase) != "" && strings.TrimSpace(part.DestinationTable) != "" && strings.TrimSpace(part.DestinationSchema) != "" && part.DestinationActivePartCount > 0 && fragmented + normalized := part.DestinationActivePartCount == 1 && partitions == 1 && single + var staleAt, originalReadyAt *time.Time + if part.Status == StatusCompacting { + for _, value := range []string{part.UpdatedAt, part.CompactingAt} { + if strings.TrimSpace(value) == "" { + continue + } + parsed, err := time.Parse(time.RFC3339Nano, value) + if err != nil { + return nil, fmt.Errorf("compact time for %s/%s: %w", part.JobID, part.PartID, err) + } + if staleAt == nil || parsed.Before(*staleAt) { + staleAt = &parsed + } + } + if staleAt == nil { + return nil, fmt.Errorf("compacting part %s/%s has no updated_at or compacting_at", part.JobID, part.PartID) + } + } + if !isGeneratedCompactPart(part) && strings.TrimSpace(part.CompactReadyAt) != "" { + parsed, err := time.Parse(time.RFC3339Nano, part.CompactReadyAt) + if err != nil { + return nil, fmt.Errorf("compact-ready time for %s/%s: %w", part.JobID, part.PartID, err) + } + originalReadyAt = &parsed + } + return []any{part.JobID, part.PartID, string(part.Status), part.WorkerID, part.CreatedAt, part.UpdatedAt, data, + pgtype.Numeric{Int: new(big.Int).SetUint64(part.SourceArtifactBytes), Valid: true}, + pgtype.Numeric{Int: new(big.Int).SetUint64(part.DestinationActivePartBytes), Valid: true}, + eligible, normalized, staleAt, originalReadyAt}, nil +} + +func (s *Store) insertPartTx(ctx context.Context, tx pgx.Tx, part Part) error { + values, err := partWriteValues(part) + if err != nil { + return err + } + _, err = tx.Exec(ctx, `INSERT INTO `+s.tableSQL+` (`+partColumns+`) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13)`, values...) + return err +} + +func (s *Store) savePartTx(ctx context.Context, tx pgx.Tx, part Part) error { + values, err := partWriteValues(part) if err != nil { return err } - tag, err := tx.Exec(ctx, - `UPDATE `+s.tableSQL+` SET status = $1, worker_id = $2, created_at = $3, updated_at = $4, data = $5 WHERE job_id = $6 AND part_id = $7`, - string(part.Status), part.WorkerID, part.CreatedAt, part.UpdatedAt, data, part.JobID, part.PartID, - ) + tag, err := tx.Exec(ctx, `UPDATE `+s.tableSQL+` SET status=$3, worker_id=$4, created_at=$5, updated_at=$6, data=$7, + source_artifact_bytes=$8, compact_bytes=$9, compact_eligible=$10, compact_normalized=$11, + compact_stale_at=$12, original_compact_ready_at=$13 WHERE job_id=$1 AND part_id=$2`, values...) if err != nil { return err } @@ -508,10 +555,6 @@ func (s *Store) CreatePart(ctx context.Context, part Part) error { if err := validatePart(part); err != nil { return err } - data, err := partJSON(part) - if err != nil { - return err - } tx, err := s.pool.Begin(ctx) if err != nil { return err @@ -529,10 +572,7 @@ func (s *Store) CreatePart(ctx context.Context, part Part) error { return fmt.Errorf("source part reference for %s/%s does not match source artifact %s/%s", part.JobID, part.PartID, source.JobID, source.PartID) } } - _, err = tx.Exec(ctx, - `INSERT INTO `+s.tableSQL+` (job_id, part_id, status, worker_id, created_at, updated_at, data) VALUES ($1, $2, $3, $4, $5, $6, $7)`, - part.JobID, part.PartID, string(part.Status), part.WorkerID, part.CreatedAt, part.UpdatedAt, data, - ) + err = s.insertPartTx(ctx, tx, part) if err != nil { return fmt.Errorf("create state item for %s/%s: %w", part.JobID, part.PartID, err) } @@ -580,6 +620,10 @@ func (s *Store) MarkCompactReady(ctx context.Context, part Part, workerID, finis return nil } +func (s *Store) readyClaimQuery() string { + return `SELECT data FROM ` + s.tableSQL + ` WHERE status = 'READY' ORDER BY source_artifact_bytes DESC, created_at, job_id, part_id LIMIT 1 FOR UPDATE SKIP LOCKED` +} + func (s *Store) ClaimNextReady(ctx context.Context, workerID string, now time.Time) (*Part, error) { if strings.TrimSpace(workerID) == "" { return nil, errors.New("worker id is required") @@ -591,7 +635,7 @@ func (s *Store) ClaimNextReady(ctx context.Context, workerID string, now time.Ti defer tx.Rollback(ctx) var data []byte - err = tx.QueryRow(ctx, `SELECT data FROM `+s.tableSQL+` WHERE status = $1 ORDER BY COALESCE((data->>'source_artifact_bytes')::bigint, 0) DESC, created_at, job_id, part_id LIMIT 1 FOR UPDATE SKIP LOCKED`, string(StatusReady)).Scan(&data) + err = tx.QueryRow(ctx, s.readyClaimQuery()).Scan(&data) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } @@ -616,72 +660,77 @@ func (s *Store) ClaimNextCompactBatch(ctx context.Context, workerID string, now if strings.TrimSpace(workerID) == "" { return nil, errors.New("worker id is required") } - candidates, err := s.listPartsByStatusIndex(ctx, StatusCompactReady) + tx, err := s.pool.Begin(ctx) if err != nil { - return nil, fmt.Errorf("query compact-ready parts: %w", err) + return nil, err } - if len(candidates) == 0 { + defer tx.Rollback(ctx) + query, args := s.compactClaimQuery(opts, now) + var data []byte + err = tx.QueryRow(ctx, query, args...).Scan(&data) + if errors.Is(err, pgx.ErrNoRows) { return nil, nil } - compacting, err := s.listPartsByStatusIndex(ctx, StatusCompacting) if err != nil { - return nil, fmt.Errorf("query compacting parts: %w", err) + return nil, fmt.Errorf("claim compact-ready part: %w", err) } - - groups := compactCandidateGroups(candidates, compacting, opts) - for _, selected := range compactCandidateSelections(groups, opts) { - claimed, err := s.claimCompactParts(ctx, selected, workerID, now) - if IsConditionalCheckFailed(err) { - continue - } - if err != nil { - return nil, err - } - batch, err := compactBatchFromParts(claimed) - if err != nil { - _ = s.ReleaseCompactBatch(ctx, CompactBatch{JobID: claimed[0].JobID, Parts: claimed}, workerID, now) - return nil, err - } - return batch, nil - } - return nil, nil -} - -func compactCandidateSelections(groups []compactGroup, opts CompactClaimOptions) [][]Part { - selections := make([][]Part, 0, len(groups)) - for _, group := range groups { - if selected := selectCompactBatchParts(group, opts); len(selected) > 0 { - selections = append(selections, selected) - } + part, err := partFromJSON(data) + if err != nil { + return nil, err } - sort.SliceStable(selections, func(i, j int) bool { - return selections[i][0].DestinationActivePartBytes > selections[j][0].DestinationActivePartBytes - }) - return selections -} - -func (s *Store) listPartsByStatusIndex(ctx context.Context, status Status) ([]Part, error) { - rows, err := s.pool.Query(ctx, `SELECT data FROM `+s.tableSQL+` WHERE status = $1 ORDER BY created_at, job_id, part_id`, string(status)) + setStatus(&part, StatusCompacting, now) + part.CompactingAt = formatTime(now) + part.WorkerID = workerID + part.Error = "" + part.CompactCooldownUntil = "" + batch, err := compactBatchFromParts([]Part{part}) if err != nil { return nil, err } - defer rows.Close() - var parts []Part - for rows.Next() { - var data []byte - if err := rows.Scan(&data); err != nil { - return nil, err + if err := s.savePartTx(ctx, tx, part); err != nil { + return nil, err + } + if err := tx.Commit(ctx); err != nil { + return nil, err + } + return batch, nil +} + +// Claim the largest eligible unlocked artifact using the compact queue index. +func (s *Store) compactClaimQuery(opts CompactClaimOptions, now time.Time) (string, []any) { + args := []any{} + bind := func(v any) string { args = append(args, v); return fmt.Sprintf("$%d", len(args)) } + filters := []string{"p.status = 'COMPACT_READY'", "p.compact_eligible"} + for _, filter := range []struct{ field, value string }{ + {"p.job_id", opts.JobID}, {"p.data->>'bucket'", opts.Bucket}, + {"p.data->>'destination_database'", opts.DestinationDatabase}, + {"p.data->>'destination_table'", opts.DestinationTable}, + {"p.data->>'destination_schema'", opts.DestinationSchema}, + } { + if filter.value != "" { + filters = append(filters, filter.field+" = "+bind(filter.value)) } - part, err := partFromJSON(data) - if err != nil { - return nil, err + } + if len(opts.ExcludedJobIDs) > 0 { + ids := make([]string, 0, len(opts.ExcludedJobIDs)) + for id := range opts.ExcludedJobIDs { + ids = append(ids, id) } - parts = append(parts, part) + filters = append(filters, "NOT (p.job_id = ANY("+bind(ids)+"::text[]))") } - if err := rows.Err(); err != nil { - return nil, err + required := "" + if len(opts.RequiredPartitionIDs) > 0 { + required = " AND partition.key = ANY(" + bind(opts.RequiredPartitionIDs) + "::text[])" + filters = append(filters, "EXISTS (SELECT FROM jsonb_each_text(p.data->'destination_active_partition_counts') partition WHERE btrim(partition.key) <> '' AND partition.value::numeric > 0"+required+")") } - return parts, nil + from := s.tableSQL + " p" + if opts.CompactWindow > 0 { + cutoff := bind(now.Add(-opts.CompactWindow)) + // A lateral join lets Postgres memoize the deadline by job while walking the queue. + from += " LEFT JOIN LATERAL (SELECT original_compact_ready_at FROM " + s.tableSQL + " original WHERE original.job_id = p.job_id AND original_compact_ready_at IS NOT NULL ORDER BY original_compact_ready_at DESC LIMIT 1) deadline ON true" + filters = append(filters, "COALESCE(deadline.original_compact_ready_at > "+cutoff+", true)") + } + return "SELECT p.data FROM " + from + " WHERE " + strings.Join(filters, " AND ") + " ORDER BY p.compact_bytes DESC, p.created_at, p.job_id, p.part_id LIMIT 1 FOR UPDATE OF p SKIP LOCKED", args } func (s *Store) ReleaseCompactBatch(ctx context.Context, batch CompactBatch, workerID string, now time.Time) error { @@ -760,28 +809,33 @@ func (s *Store) HeartbeatCompactBatch(ctx context.Context, batch CompactBatch, w if strings.TrimSpace(workerID) == "" { return false, errors.New("worker id is required") } - finalizeRequested := false + requested := false for _, part := range batch.Parts { - updated, err := s.updatePart(ctx, part.JobID, part.PartID, func(current Part) bool { - return compactOwnedOrUnownedReady(current, workerID) - }, func(current *Part) error { - setStatus(current, StatusCompacting, now) - if strings.TrimSpace(current.CompactingAt) == "" { - current.CompactingAt = formatTime(now) - } - current.WorkerID = workerID - current.Error = "" - current.CompactCooldownUntil = "" - return nil - }) + finalize, err := s.updateCompactProgress(ctx, part, workerID, []byte(`{}`), now) if err != nil { return false, fmt.Errorf("heartbeat compacting part %s/%s: %w", part.JobID, part.PartID, err) } - if strings.TrimSpace(updated.CompactFinalizeRequestedAt) != "" { - finalizeRequested = true - } + requested = requested || finalize + } + return requested, nil +} + +// One conditional statement preserves ownership and finalization requests while +// patching only progress fields, rather than reading and rewriting the full part. +func (s *Store) updateCompactProgress(ctx context.Context, part Part, workerID string, patch []byte, now time.Time) (bool, error) { + var requested bool + err := s.pool.QueryRow(ctx, `UPDATE `+s.tableSQL+` SET status = 'COMPACTING', worker_id = $1, updated_at = $2, + compact_stale_at = LEAST($2::text::timestamptz, COALESCE(NULLIF(btrim(data->>'compacting_at'), '')::timestamptz, $2::text::timestamptz)), + data = (data - 'error' - 'compact_cooldown_until') || $3::jsonb || + jsonb_build_object('status', 'COMPACTING', 'worker_id', $1::text, 'updated_at', $2::text, + 'compacting_at', COALESCE(NULLIF(btrim(data->>'compacting_at'), ''), $2::text)) + WHERE job_id = $4 AND part_id = $5 AND + ((status = 'COMPACTING' AND worker_id = $1) OR (status = 'COMPACT_READY' AND btrim(worker_id) = '')) + RETURNING COALESCE(btrim(data->>'compact_finalize_requested_at'), '') <> ''`, workerID, formatTime(now), patch, part.JobID, part.PartID).Scan(&requested) + if errors.Is(err, pgx.ErrNoRows) { + return false, &conditionalCheckFailedError{message: fmt.Sprintf("part %s/%s did not match expected state", part.JobID, part.PartID)} } - return finalizeRequested, nil + return requested, err } func (s *Store) RequestCompactFinalization(ctx context.Context, part Part, now time.Time) error { @@ -817,34 +871,21 @@ func (s *Store) UpdateCompactProgress(ctx context.Context, batch CompactBatch, o if progress.MergeProgress < 0 { return fmt.Errorf("compact merge progress must be non-negative, got %f", progress.MergeProgress) } + patch, err := json.Marshal(map[string]any{ + "compact_progress_at": formatTime(now), "compact_output_part_id": outputPartID, + "compact_input_part_count": inputStats.Count, "compact_input_rows": inputStats.Rows, "compact_input_bytes": inputStats.Bytes, + "compact_output_part_count": outputStats.Count, "compact_output_rows": outputStats.Rows, "compact_output_bytes": outputStats.Bytes, + "compact_stage": strings.TrimSpace(progress.Stage), "compact_active_merges": progress.ActiveMerges, "compact_merge_progress": progress.MergeProgress, + }) + if err != nil { + return err + } for _, part := range batch.Parts { - _, err := s.updatePart(ctx, part.JobID, part.PartID, func(current Part) bool { - return compactOwnedOrUnownedReady(current, workerID) - }, func(current *Part) error { - setStatus(current, StatusCompacting, now) - if strings.TrimSpace(current.CompactingAt) == "" { - current.CompactingAt = formatTime(now) - } - current.WorkerID = workerID - current.CompactProgressAt = formatTime(now) - current.CompactOutputPartID = outputPartID - current.CompactInputPartCount = inputStats.Count - current.CompactInputRows = inputStats.Rows - current.CompactInputBytes = inputStats.Bytes - current.CompactOutputPartCount = outputStats.Count - current.CompactOutputRows = outputStats.Rows - current.CompactOutputBytes = outputStats.Bytes - current.CompactStage = strings.TrimSpace(progress.Stage) - current.CompactActiveMerges = progress.ActiveMerges - current.CompactMergeProgress = progress.MergeProgress - current.Error = "" - current.CompactCooldownUntil = "" - return nil - }) - if err != nil { + if _, err := s.updateCompactProgress(ctx, part, workerID, patch, now); err != nil { return fmt.Errorf("update compact progress for %s/%s: %w", part.JobID, part.PartID, err) } } + return nil } @@ -852,53 +893,19 @@ func (s *Store) ReleaseStaleCompactingParts(ctx context.Context, now time.Time, if staleAfter <= 0 { return 0, fmt.Errorf("compact stale timeout must be greater than zero, got %s", staleAfter) } - parts, err := s.listPartsByStatusIndex(ctx, StatusCompacting) + tx, err := s.pool.Begin(ctx) if err != nil { - return 0, fmt.Errorf("query compacting parts: %w", err) - } - cutoff := now.Add(-staleAfter) - released := 0 - for _, part := range parts { - staleAt, err := compactStaleTime(part) - if err != nil { - return released, err - } - if staleAt.After(cutoff) { - continue - } - ok, err := s.releaseStaleCompactingPart(ctx, part, now) - if err != nil { - return released, err - } - if ok { - released++ - } - } - return released, nil -} - -func (s *Store) releaseStaleCompactingPart(ctx context.Context, part Part, now time.Time) (bool, error) { - _, err := s.updatePart(ctx, part.JobID, part.PartID, func(current Part) bool { - return current.Status == StatusCompacting && current.UpdatedAt == part.UpdatedAt - }, func(current *Part) error { - setStatus(current, StatusCompactReady, now) - if strings.TrimSpace(current.CompactReadyAt) == "" { - current.CompactReadyAt = compactReadyAtForRelease(part, now) - } - current.WorkerID = "" - current.CompactingAt = "" - current.Error = "" - current.CompactCooldownUntil = "" - clearCompactProgress(current) - return nil - }) - if IsConditionalCheckFailed(err) { - return false, nil + return 0, err } + defer tx.Rollback(ctx) + n, err := s.releaseStaleCompactingPartsTx(ctx, tx, now, staleAfter) if err != nil { - return false, fmt.Errorf("release stale compacting part %s/%s: %w", part.JobID, part.PartID, err) + return 0, err + } + if err := tx.Commit(ctx); err != nil { + return 0, err } - return true, nil + return n, nil } func (s *Store) CompleteCompaction(ctx context.Context, batch CompactBatch, output Part, workerID string, now time.Time) error { @@ -927,16 +934,10 @@ func (s *Store) CompleteCompaction(ctx context.Context, batch CompactBatch, outp } defer tx.Rollback(ctx) - outputData, err := partJSON(output) - if err != nil { - return err - } - if _, err := tx.Exec(ctx, - `INSERT INTO `+s.tableSQL+` (job_id, part_id, status, worker_id, created_at, updated_at, data) VALUES ($1, $2, $3, $4, $5, $6, $7)`, - output.JobID, output.PartID, string(output.Status), output.WorkerID, output.CreatedAt, output.UpdatedAt, outputData, - ); err != nil { + if err := s.insertPartTx(ctx, tx, output); err != nil { return fmt.Errorf("complete compaction for %s/%s: %w", batch.JobID, output.PartID, err) } + for _, part := range batch.Parts { current, err := s.readPartTx(ctx, tx, part.JobID, part.PartID) if err != nil { @@ -979,144 +980,6 @@ func (s *Store) MarkCompactReadyFinished(ctx context.Context, part Part, now tim return nil } -type compactGroup struct { - key string - parts []Part - compactingPartitionIDs []string -} - -func compactCandidateGroups(parts, compacting []Part, opts CompactClaimOptions) []compactGroup { - groupsByKey := map[string][]Part{} - var order []string - for _, part := range parts { - if strings.TrimSpace(part.DestinationDatabase) == "" || - strings.TrimSpace(part.DestinationTable) == "" || - strings.TrimSpace(part.DestinationSchema) == "" || - part.DestinationActivePartCount == 0 || - len(part.DestinationActivePartitionCounts) == 0 || - !matchesCompactClaimOptions(part, opts) { - continue - } - key := compactGroupKey(part) - if _, ok := groupsByKey[key]; !ok { - order = append(order, key) - } - groupsByKey[key] = append(groupsByKey[key], part) - } - compactingPartitionsByKey := compactingPartitionIDsByGroup(compacting) - groups := make([]compactGroup, 0, len(order)) - for _, key := range order { - groupParts := groupsByKey[key] - sort.SliceStable(groupParts, func(i, j int) bool { - if groupParts[i].CompactGeneration != groupParts[j].CompactGeneration { - return groupParts[i].CompactGeneration < groupParts[j].CompactGeneration - } - if groupParts[i].UpdatedAt != groupParts[j].UpdatedAt { - return groupParts[i].UpdatedAt < groupParts[j].UpdatedAt - } - return groupParts[i].PartID < groupParts[j].PartID - }) - groups = append(groups, compactGroup{ - key: key, - parts: groupParts, - compactingPartitionIDs: compactingPartitionsByKey[key], - }) - } - return groups -} - -func compactGroupKey(part Part) string { - return strings.Join([]string{part.JobID, part.Bucket, part.DestinationDatabase, part.DestinationTable, part.DestinationSchema}, "\x00") -} - -func compactingPartitionIDsByGroup(parts []Part) map[string][]string { - sets := map[string]map[string]struct{}{} - for _, part := range parts { - if part.Status != StatusCompacting || - strings.TrimSpace(part.DestinationDatabase) == "" || - strings.TrimSpace(part.DestinationTable) == "" || - strings.TrimSpace(part.DestinationSchema) == "" { - continue - } - key := compactGroupKey(part) - if _, ok := sets[key]; !ok { - sets[key] = map[string]struct{}{} - } - for _, partitionID := range partPartitionIDs(part) { - sets[key][partitionID] = struct{}{} - } - } - out := make(map[string][]string, len(sets)) - for key, set := range sets { - partitionIDs := make([]string, 0, len(set)) - for partitionID := range set { - partitionIDs = append(partitionIDs, partitionID) - } - sort.Strings(partitionIDs) - out[key] = partitionIDs - } - return out -} - -func matchesCompactClaimOptions(part Part, opts CompactClaimOptions) bool { - if _, excluded := opts.ExcludedJobIDs[part.JobID]; excluded { - return false - } - if opts.JobID != "" && part.JobID != opts.JobID { - return false - } - if opts.Bucket != "" && part.Bucket != opts.Bucket { - return false - } - if opts.DestinationDatabase != "" && part.DestinationDatabase != opts.DestinationDatabase { - return false - } - if opts.DestinationTable != "" && part.DestinationTable != opts.DestinationTable { - return false - } - if opts.DestinationSchema != "" && part.DestinationSchema != opts.DestinationSchema { - return false - } - if len(opts.RequiredPartitionIDs) > 0 && !partOverlapsRequiredPartitions(part, opts.RequiredPartitionIDs) { - return false - } - return true -} - -func compactHeartbeatTime(part Part) (time.Time, error) { - for _, value := range []string{part.UpdatedAt, part.CompactingAt} { - if strings.TrimSpace(value) == "" { - continue - } - t, err := time.Parse(timeFormat, value) - if err != nil { - return time.Time{}, fmt.Errorf("parse compact heartbeat time for part %s/%s: %w", part.JobID, part.PartID, err) - } - return t, nil - } - return time.Time{}, fmt.Errorf("compacting part %s/%s has no updated_at or compacting_at", part.JobID, part.PartID) -} - -func compactStaleTime(part Part) (time.Time, error) { - var staleAt time.Time - for _, value := range []string{part.UpdatedAt, part.CompactingAt} { - if strings.TrimSpace(value) == "" { - continue - } - t, err := time.Parse(timeFormat, value) - if err != nil { - return time.Time{}, fmt.Errorf("parse compact stale time for part %s/%s: %w", part.JobID, part.PartID, err) - } - if staleAt.IsZero() || t.Before(staleAt) { - staleAt = t - } - } - if staleAt.IsZero() { - return time.Time{}, fmt.Errorf("compacting part %s/%s has no updated_at or compacting_at", part.JobID, part.PartID) - } - return staleAt, nil -} - func compactReadyAtForRelease(part Part, now time.Time) string { for _, value := range []string{part.CompactReadyAt, part.ProgressUpdatedAt, part.UpdatedAt, part.CompactingAt} { if strings.TrimSpace(value) != "" { @@ -1126,154 +989,6 @@ func compactReadyAtForRelease(part Part, now time.Time) string { return formatTime(now) } -func selectCompactBatchParts(group compactGroup, opts CompactClaimOptions) []Part { - partitions := orderedCandidatePartitions(group.parts, opts.RequiredPartitionIDs) - preferredPartitions := partitionsWithout(partitions, group.compactingPartitionIDs) - if part, ok := selectFragmentedCompactPart(group.parts, preferredPartitions); ok { - return []Part{part} - } - fallbackPartitions := partitionsWithout(partitions, preferredPartitions) - if part, ok := selectFragmentedCompactPart(group.parts, fallbackPartitions); ok { - return []Part{part} - } - return nil -} - -func selectFragmentedCompactPart(parts []Part, partitions []string) (Part, bool) { - var selected Part - found := false - for _, part := range parts { - eligible := false - for _, partitionID := range partitions { - if part.DestinationActivePartitionCounts[partitionID] > 0 { - eligible = true - break - } - } - if !eligible { - continue - } - fragmented := false - for _, count := range part.DestinationActivePartitionCounts { - fragmented = fragmented || count > 1 - } - if fragmented && (!found || part.DestinationActivePartBytes > selected.DestinationActivePartBytes) { - selected = part - found = true - } - } - return selected, found -} - -func partitionsWithout(partitions, excluded []string) []string { - excludedSet := partitionSet(excluded) - if len(excludedSet) == 0 { - return append([]string(nil), partitions...) - } - out := make([]string, 0, len(partitions)) - for _, partitionID := range partitions { - if _, ok := excludedSet[partitionID]; ok { - continue - } - out = append(out, partitionID) - } - return out -} - -func orderedCandidatePartitions(parts []Part, required []string) []string { - requiredSet := partitionSet(required) - seen := map[string]struct{}{} - var partitions []string - for _, part := range parts { - ids := partPartitionIDs(part) - for _, partitionID := range ids { - if len(requiredSet) > 0 { - if _, ok := requiredSet[partitionID]; !ok { - continue - } - } - if _, ok := seen[partitionID]; ok { - continue - } - seen[partitionID] = struct{}{} - partitions = append(partitions, partitionID) - } - } - return partitions -} - -func partitionSet(partitionIDs []string) map[string]struct{} { - out := map[string]struct{}{} - for _, partitionID := range partitionIDs { - if strings.TrimSpace(partitionID) == "" { - continue - } - out[partitionID] = struct{}{} - } - return out -} - -func partOverlapsRequiredPartitions(part Part, required []string) bool { - for partitionID := range partitionSet(required) { - if part.DestinationActivePartitionCounts[partitionID] > 0 { - return true - } - } - return false -} - -func partPartitionIDs(part Part) []string { - ids := make([]string, 0, len(part.DestinationActivePartitionCounts)) - for partitionID, count := range part.DestinationActivePartitionCounts { - if strings.TrimSpace(partitionID) == "" || count == 0 { - continue - } - ids = append(ids, partitionID) - } - sort.Strings(ids) - return ids -} - -func (s *Store) claimCompactParts(ctx context.Context, parts []Part, workerID string, now time.Time) ([]Part, error) { - if err := validateCompactBatchParts(parts); err != nil { - return nil, err - } - for _, part := range parts { - if part.Status != StatusCompactReady { - return nil, fmt.Errorf("compact batch part %s/%s is %s, expected %s", part.JobID, part.PartID, part.Status, StatusCompactReady) - } - } - tx, err := s.pool.Begin(ctx) - if err != nil { - return nil, err - } - defer tx.Rollback(ctx) - - claimed := make([]Part, 0, len(parts)) - for _, part := range parts { - claimedPart, err := s.readPartTx(ctx, tx, part.JobID, part.PartID) - if err != nil { - return nil, fmt.Errorf("claim compact-ready part %s/%s: %w", part.JobID, part.PartID, err) - } - if claimedPart.Status != StatusCompactReady { - return nil, fmt.Errorf("claim compact-ready part %s/%s: %w", part.JobID, part.PartID, &conditionalCheckFailedError{}) - } - setStatus(&claimedPart, StatusCompacting, now) - claimedPart.CompactingAt = formatTime(now) - claimedPart.WorkerID = workerID - claimedPart.Error = "" - claimedPart.CompactCooldownUntil = "" - if err := s.savePartTx(ctx, tx, claimedPart); err != nil { - return nil, fmt.Errorf("claim compact-ready part %s/%s: %w", part.JobID, part.PartID, err) - } - claimed = append(claimed, claimedPart) - } - if err := tx.Commit(ctx); err != nil { - return nil, err - } - return claimed, nil -} - func compactBatchFromParts(parts []Part) (*CompactBatch, error) { if len(parts) == 0 { return nil, nil @@ -1421,40 +1136,59 @@ func (s *Store) UpdateRewriteProgress(ctx context.Context, jobID, partID, worker if strings.TrimSpace(workerID) == "" { return errors.New("worker id is required") } - _, err := s.updatePart(ctx, jobID, partID, func(current Part) bool { - return current.Status == StatusInProgress && current.WorkerID == workerID - }, func(current *Part) error { - current.UpdatedAt = formatTime(now) - current.ProgressUpdatedAt = formatTime(now) - if progress.QueryProgress != nil { - current.ReadRows = progress.QueryProgress.ReadRows - current.ReadBytes = progress.QueryProgress.ReadBytes - current.TotalRowsApprox = progress.QueryProgress.TotalRowsApprox - current.WrittenRows = progress.QueryProgress.WrittenRows - current.WrittenBytes = progress.QueryProgress.WrittenBytes - } - if progress.SourceActivePartStats != nil { - current.SourceActivePartCount = progress.SourceActivePartStats.Count - current.SourceActivePartRows = progress.SourceActivePartStats.Rows - current.SourceActivePartBytes = progress.SourceActivePartStats.Bytes - } - if progress.DestinationActivePartStats != nil { - current.DestinationActivePartCount = progress.DestinationActivePartStats.Count - current.DestinationActivePartRows = progress.DestinationActivePartStats.Rows - current.DestinationActivePartBytes = progress.DestinationActivePartStats.Bytes - } - if progress.DestinationFailedMerges != nil { - current.DestinationFailedMerges = *progress.DestinationFailedMerges - } - if progress.StageProgress != nil { - current.RewriteStage = progress.StageProgress.Stage - current.RewriteStageStartedAt = formatTime(progress.StageProgress.StageStartedAt) - current.RewriteStageElapsedMs = progress.StageProgress.StageElapsedMs - current.RewriteTotalElapsedMs = progress.StageProgress.TotalElapsedMs - current.RewriteStageDurationsMs = progress.StageProgress.CompletedStageDurationsMs - } - return nil - }) + patch := map[string]any{} + patch["updated_at"] = formatTime(now) + patch["progress_updated_at"] = formatTime(now) + if progress.QueryProgress != nil { + patch["read_rows"] = progress.QueryProgress.ReadRows + patch["read_bytes"] = progress.QueryProgress.ReadBytes + patch["total_rows_approx"] = progress.QueryProgress.TotalRowsApprox + patch["written_rows"] = progress.QueryProgress.WrittenRows + patch["written_bytes"] = progress.QueryProgress.WrittenBytes + } + if progress.SourceActivePartStats != nil { + patch["source_active_part_count"] = progress.SourceActivePartStats.Count + patch["source_active_part_rows"] = progress.SourceActivePartStats.Rows + patch["source_active_part_bytes"] = progress.SourceActivePartStats.Bytes + } + if progress.DestinationActivePartStats != nil { + patch["destination_active_part_count"] = progress.DestinationActivePartStats.Count + patch["destination_active_part_rows"] = progress.DestinationActivePartStats.Rows + patch["destination_active_part_bytes"] = progress.DestinationActivePartStats.Bytes + } + if progress.DestinationFailedMerges != nil { + patch["destination_failed_merges"] = *progress.DestinationFailedMerges + } + if progress.StageProgress != nil { + patch["rewrite_stage"] = progress.StageProgress.Stage + patch["rewrite_stage_started_at"] = formatTime(progress.StageProgress.StageStartedAt) + patch["rewrite_stage_elapsed_ms"] = progress.StageProgress.StageElapsedMs + patch["rewrite_total_elapsed_ms"] = progress.StageProgress.TotalElapsedMs + patch["rewrite_stage_durations_ms"] = progress.StageProgress.CompletedStageDurationsMs + } + data, err := json.Marshal(patch) + if err != nil { + return err + } + updates := `updated_at = $1, data = data || $2::jsonb` + args := []any{formatTime(now), data, jobID, partID, workerID} + if stats := progress.DestinationActivePartStats; stats != nil { + updates += `, compact_bytes = $6, compact_eligible = $7::boolean AND + COALESCE(btrim(data->>'destination_database'), '') <> '' AND + COALESCE(btrim(data->>'destination_table'), '') <> '' AND + COALESCE(btrim(data->>'destination_schema'), '') <> '' AND + EXISTS (SELECT FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 1), + compact_normalized = $8::boolean AND + (SELECT count(*) = 1 AND COALESCE(bool_and(p.value::numeric = 1), false) + FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 0)` + args = append(args, pgtype.Numeric{Int: new(big.Int).SetUint64(stats.Bytes), Valid: true}, stats.Count > 0, stats.Count == 1) + } + tag, err := s.pool.Exec(ctx, `UPDATE `+s.tableSQL+` SET `+updates+` WHERE job_id = $3 AND part_id = $4 AND status = 'IN_PROGRESS' AND worker_id = $5`, args...) + + if err == nil && tag.RowsAffected() != 1 { + err = &conditionalCheckFailedError{message: fmt.Sprintf("part %s/%s did not match expected state", jobID, partID)} + } + if err != nil { return fmt.Errorf("update rewrite progress for %s/%s: %w", jobID, partID, err) } @@ -1470,79 +1204,83 @@ func (s *Store) ListJobs(ctx context.Context) ([]Job, error) { } func (s *Store) ListJobIDsByStatus(ctx context.Context, statuses ...Status) ([]string, error) { - jobs, err := s.ListJobsByStatus(ctx, statuses...) + values := make([]string, 0, len(statuses)) + for _, status := range statuses { + if strings.TrimSpace(string(status)) == "" { + return nil, errors.New("status is required") + } + values = append(values, string(status)) + } + rows, err := s.pool.Query(ctx, `SELECT DISTINCT job_id FROM `+s.tableSQL+` WHERE status = ANY($1::text[]) ORDER BY job_id`, values) if err != nil { return nil, err } - jobIDs := make([]string, 0, len(jobs)) - for _, job := range jobs { - jobIDs = append(jobIDs, job.JobID) - } - return jobIDs, nil + return pgx.CollectRows(rows, pgx.RowTo[string]) } func (s *Store) ListJobsByStatus(ctx context.Context, statuses ...Status) ([]Job, error) { - jobsByID := map[string]Job{} - jobPartitionsByID := map[string]map[string]struct{}{} - queried := map[Status]struct{}{} + values := make([]string, 0, len(statuses)) for _, status := range statuses { if strings.TrimSpace(string(status)) == "" { return nil, errors.New("status is required") } - if _, ok := queried[status]; ok { - continue + values = append(values, string(status)) + } + rows, err := s.pool.Query(ctx, `WITH selected AS MATERIALIZED ( + SELECT job_id, status, created_at, updated_at, COALESCE(data->>'job_name', '') AS name, + COALESCE((data->>'destination_active_part_count')::numeric, 0) AS part_count, + COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb) AS partitions + FROM `+s.tableSQL+` WHERE status = ANY($1::text[]) + ), partition_counts AS ( + SELECT job_id, count(DISTINCT p.key) AS count FROM selected, + LATERAL jsonb_each_text(partitions) p + WHERE status <> 'SUPERSEDED' AND btrim(p.key) <> '' AND p.value::numeric > 0 GROUP BY job_id + ) SELECT s.job_id, s.status, count(*), COALESCE(min(NULLIF(s.name, '')), ''), max(s.name), + min(s.created_at), max(s.updated_at), sum(CASE WHEN s.status <> 'SUPERSEDED' THEN part_count ELSE 0 END)::text, + COALESCE(max(p.count), 0) + FROM selected s LEFT JOIN partition_counts p USING (job_id) + GROUP BY s.job_id, s.status ORDER BY s.job_id, s.status`, values) + if err != nil { + return nil, err + } + defer rows.Close() + var jobs []Job + for rows.Next() { + var jobID, minName, maxName, createdAt, updatedAt, activeCount string + var status Status + var count, partitions int + if err := rows.Scan(&jobID, &status, &count, &minName, &maxName, &createdAt, &updatedAt, &activeCount, &partitions); err != nil { + return nil, err } - queried[status] = struct{}{} - parts, err := s.listPartsByStatusIndex(ctx, status) + if minName != maxName { + return nil, fmt.Errorf("job %s has conflicting job_name values %q and %q", jobID, minName, maxName) + } + if len(jobs) == 0 || jobs[len(jobs)-1].JobID != jobID { + jobs = append(jobs, Job{JobID: jobID, Counts: map[Status]int{}, SubmittedAt: createdAt}) + } + job := &jobs[len(jobs)-1] + if job.Name != "" && minName != "" && job.Name != minName { + return nil, fmt.Errorf("job %s has conflicting job_name values %q and %q", jobID, job.Name, minName) + } + if minName != "" { + job.Name = minName + } + n, err := strconv.ParseUint(activeCount, 10, 64) if err != nil { - return nil, fmt.Errorf("query job ids for status %s: %w", status, err) + return nil, err } - for _, part := range parts { - if part.JobID == "" { - continue - } - existing := jobsByID[part.JobID] - if existing.JobID == "" { - existing = Job{JobID: part.JobID, Name: part.JobName, Counts: map[Status]int{}} - } - if existing.Name == "" && part.JobName != "" { - existing.Name = part.JobName - } - if existing.Name != "" && part.JobName != "" && existing.Name != part.JobName { - return nil, fmt.Errorf("job %s has conflicting job_name values %q and %q", part.JobID, existing.Name, part.JobName) - } - existing.Total++ - existing.Counts[status]++ - if status != StatusSuperseded { - existing.DestinationActivePartCount += part.DestinationActivePartCount - if jobPartitionsByID[part.JobID] == nil { - jobPartitionsByID[part.JobID] = map[string]struct{}{} - } - for partitionID, count := range part.DestinationActivePartitionCounts { - if strings.TrimSpace(partitionID) != "" && count > 0 { - jobPartitionsByID[part.JobID][partitionID] = struct{}{} - } - } - existing.DestinationPartitionCount = len(jobPartitionsByID[part.JobID]) - } - if part.CreatedAt != "" && (existing.SubmittedAt == "" || part.CreatedAt < existing.SubmittedAt) { - existing.SubmittedAt = part.CreatedAt - } - if part.UpdatedAt != "" && part.UpdatedAt > existing.UpdatedAt { - existing.UpdatedAt = part.UpdatedAt - } - jobsByID[part.JobID] = existing + job.Total += count + job.Counts[status] = count + job.DestinationActivePartCount += n + job.DestinationPartitionCount = partitions + if createdAt < job.SubmittedAt { + job.SubmittedAt = createdAt + } + if updatedAt > job.UpdatedAt { + job.UpdatedAt = updatedAt } } - - jobs := make([]Job, 0, len(jobsByID)) - for _, job := range jobsByID { - jobs = append(jobs, job) - } - sort.Slice(jobs, func(i, j int) bool { - return jobs[i].JobID < jobs[j].JobID - }) - return jobs, nil + return jobs, rows.Err() } func (s *Store) ListJobParts(ctx context.Context, jobID string) ([]Part, error) { @@ -1569,9 +1307,6 @@ func (s *Store) ListJobParts(ctx context.Context, jobID string) ([]Part, error) if err := rows.Err(); err != nil { return nil, err } - sort.Slice(parts, func(i, j int) bool { - return parts[i].PartID < parts[j].PartID - }) return parts, nil } @@ -1946,19 +1681,6 @@ func validateOriginalResetPart(part Part) error { return nil } -func (s *Store) claimPart(ctx context.Context, part Part, workerID string, now time.Time) (*Part, error) { - claimed, err := s.updatePart(ctx, part.JobID, part.PartID, func(current Part) bool { - return current.Status == StatusReady - }, func(current *Part) error { - claimPartInMemory(current, workerID, now) - return nil - }) - if err != nil { - return nil, fmt.Errorf("claim state item for %s/%s: %w", part.JobID, part.PartID, err) - } - return &claimed, nil -} - func claimPartInMemory(part *Part, workerID string, now time.Time) { setStatus(part, StatusInProgress, now) part.StartedAt = formatTime(now) diff --git a/internal/state/postgres_integration_test.go b/internal/state/postgres_integration_test.go new file mode 100644 index 0000000..cdcd41a --- /dev/null +++ b/internal/state/postgres_integration_test.go @@ -0,0 +1,652 @@ +package state + +import ( + "context" + "fmt" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5" +) + +func postgresTestConfig(t testing.TB) Config { + t.Helper() + url := os.Getenv("PARTFORGE_TEST_POSTGRES_URL") + if url == "" { + t.Skip("set PARTFORGE_TEST_POSTGRES_URL to run PostgreSQL integration checks") + } + ctx := context.Background() + conn, err := pgx.Connect(ctx, url) + if err != nil { + t.Fatal(err) + } + schema := fmt.Sprintf("partforge_test_%d", time.Now().UnixNano()) + if _, err := conn.Exec(ctx, "CREATE SCHEMA "+pgx.Identifier{schema}.Sanitize()); err != nil { + conn.Close(ctx) + t.Fatal(err) + } + t.Cleanup(func() { + if _, err := conn.Exec(ctx, "DROP SCHEMA "+pgx.Identifier{schema}.Sanitize()+" CASCADE"); err != nil { + t.Error(err) + } + conn.Close(ctx) + }) + return Config{Endpoint: url, Table: schema + ".parts"} +} + +func postgresTestStore(t testing.TB) *Store { + t.Helper() + cfg := postgresTestConfig(t) + if _, err := Migrate(context.Background(), cfg); err != nil { + t.Fatal(err) + } + s, err := New(context.Background(), cfg) + if err != nil { + t.Fatal(err) + } + t.Cleanup(s.pool.Close) + t.Cleanup(func() { assertSchedulingColumns(t, s) }) + return s +} + +func seedPostgresParts(t testing.TB, s *Store, parts []Part) { + t.Helper() + rows := make([][]any, 0, len(parts)) + for _, part := range parts { + values, err := partWriteValues(part) + if err != nil { + t.Fatal(err) + } + rows = append(rows, values) + } + if _, err := s.pool.CopyFrom(context.Background(), pgx.Identifier(strings.Split(s.tableName, ".")), strings.Split(partColumns, ", "), pgx.CopyFromRows(rows)); err != nil { + t.Fatal(err) + } +} + +func TestPostgresMigrations(t *testing.T) { + ctx := context.Background() + cfg := postgresTestConfig(t) + if _, err := New(ctx, cfg); err == nil || !strings.Contains(err.Error(), "migrate") { + t.Fatalf("unmigrated schema error = %v", err) + } + const runners = 4 + var wg sync.WaitGroup + applied := make(chan int, runners) + errs := make(chan error, runners) + for i := 0; i < runners; i++ { + wg.Add(1) + go func() { defer wg.Done(); n, err := Migrate(ctx, cfg); applied <- n; errs <- err }() + } + wg.Wait() + close(applied) + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + total := 0 + for n := range applied { + total += n + } + s, err := New(ctx, cfg) + if err != nil { + t.Fatal(err) + } + defer s.pool.Close() + if total != len(s.migrations()) { + t.Fatalf("applied %d migrations", total) + } + if n, err := Migrate(ctx, cfg); err != nil || n != 0 { + t.Fatalf("repeat migration = %d, %v", n, err) + } + // Qualified and search-path-qualified names address the same migration ledger. + alias := cfg + name := strings.Split(cfg.Table, ".") + alias.Table = name[1] + alias.Endpoint += "&search_path=" + name[0] + if n, err := Migrate(ctx, alias); err != nil || n != 0 { + t.Fatalf("aliased migration = %d, %v", n, err) + } + aliasedStore, err := New(ctx, alias) + if err != nil { + t.Fatal(err) + } + aliasedStore.pool.Close() + + if _, err := s.pool.Exec(ctx, `INSERT INTO `+s.relatedSQL("migrations")+` (version) VALUES (99)`); err != nil { + t.Fatal(err) + } + if _, err := New(ctx, cfg); err == nil { + t.Fatal("accepted unsupported schema version") + } + if _, err := Migrate(ctx, cfg); err == nil { + t.Fatal("migrated unsupported history") + } +} + +func TestPostgresLegacyMigrationIsAtomic(t *testing.T) { + ctx := context.Background() + cfg := postgresTestConfig(t) + s, err := openStore(ctx, cfg) + if err != nil { + t.Fatal(err) + } + defer s.pool.Close() + if _, err := s.pool.Exec(ctx, s.migrations()[0]); err != nil { + t.Fatal(err) + } + part := NewPart("legacy", "part", "bucket", "source", "finished", time.Now()) + part.SourceArtifactBytes = 12345 + legacyData, err := partJSON(part) + if err != nil { + t.Fatal(err) + } + if _, err := s.pool.Exec(ctx, `INSERT INTO `+s.tableSQL+` (job_id,part_id,status,worker_id,created_at,updated_at,data) VALUES ($1,$2,$3,$4,$5,$6,$7)`, part.JobID, part.PartID, string(part.Status), part.WorkerID, part.CreatedAt, part.UpdatedAt, legacyData); err != nil { + t.Fatal(err) + } + if _, err := s.pool.Exec(ctx, `UPDATE `+s.tableSQL+` SET data = data || '{"source_artifact_bytes":"invalid"}'::jsonb`); err != nil { + t.Fatal(err) + } + if _, err := Migrate(ctx, cfg); err == nil { + t.Fatal("migration accepted invalid size") + } + var added bool + if err := s.pool.QueryRow(ctx, `SELECT EXISTS(SELECT FROM pg_attribute WHERE attrelid=$1::regclass AND attname='source_artifact_bytes')`, s.tableSQL).Scan(&added); err != nil { + t.Fatal(err) + } + if added { + t.Fatal("failed migration left schema changes") + } + data, err := partJSON(part) + if err != nil { + t.Fatal(err) + } + if _, err := s.pool.Exec(ctx, `UPDATE `+s.tableSQL+` SET data=$1::jsonb || '{"compact_input_part_ids":null,"destination_active_partition_counts":null}'::jsonb`, data); err != nil { + t.Fatal(err) + } + if _, err := Migrate(ctx, cfg); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + var size string + var count int + if err := s.pool.QueryRow(ctx, `SELECT source_artifact_bytes::text, count(*) OVER() FROM `+s.tableSQL).Scan(&size, &count); err != nil { + t.Fatal(err) + } + if size != "12345" || count != 1 { + t.Fatalf("legacy data changed: size=%s count=%d", size, count) + } + if _, err := s.ListJobs(ctx); err != nil { + t.Fatalf("legacy optional null fields: %v", err) + } + +} + +func TestPostgresClaimsAndProgress(t *testing.T) { + ctx := context.Background() + s := postgresTestStore(t) + now := time.Now().UTC() + parts := make([]Part, 64) + for i := range parts { + parts[i] = NewPart(fmt.Sprintf("job-%d", i%10), fmt.Sprintf("part-%03d", i), "bucket", "source", "finished", now) + parts[i].SourceArtifactBytes = uint64(i + 1) + } + seedPostgresParts(t, s, parts) + // A locked highest-priority part must not stall another claimant. + tx, err := s.pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, `SELECT FROM `+s.tableSQL+` WHERE part_id='part-063' FOR UPDATE`); err != nil { + t.Fatal(err) + } + timed, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + first, err := s.ClaimNextReady(timed, "first", now) + if err != nil { + t.Fatal(err) + } + if first == nil || first.PartID != "part-062" { + t.Fatalf("first claim = %+v", first) + } + if err := tx.Rollback(ctx); err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + claimed := make(chan *Part, 63) + errs := make(chan error, 63) + for i := 0; i < 63; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + p, err := s.ClaimNextReady(ctx, fmt.Sprint(i), now) + claimed <- p + errs <- err + }(i) + } + wg.Wait() + close(claimed) + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + seen := map[string]bool{first.PartID: true} + for p := range claimed { + if p == nil || seen[p.PartID] { + t.Fatalf("duplicate or missing claim: %+v", p) + } + seen[p.PartID] = true + } + if len(seen) != 64 { + t.Fatalf("claimed %d parts", len(seen)) + } + if p, err := s.ClaimNextReady(ctx, "empty", now); err != nil || p != nil { + t.Fatalf("empty queue: %+v %v", p, err) + } + if err := s.UpdateRewriteProgress(ctx, first.JobID, first.PartID, "wrong", RewriteProgress{}, now); !IsConditionalCheckFailed(err) { + t.Fatalf("ownership error = %v", err) + } + if err := s.UpdateRewriteProgress(ctx, first.JobID, first.PartID, "first", RewriteProgress{QueryProgress: &QueryProgress{ReadRows: 42}, SourceActivePartStats: &PartStats{Count: 3}}, now); err != nil { + t.Fatal(err) + } + if err := s.UpdateRewriteProgress(ctx, first.JobID, first.PartID, "first", RewriteProgress{QueryProgress: &QueryProgress{}}, now); err != nil { + t.Fatal(err) + } + current, err := s.ListJobParts(ctx, first.JobID) + if err != nil { + t.Fatal(err) + } + for _, p := range current { + if p.PartID == first.PartID && (p.ReadRows != 0 || p.SourceActivePartCount != 3 || p.SourceArtifactBytes != 63) { + t.Fatalf("progress patch lost fields: %+v", p) + } + } +} + +func TestPostgresCompactScheduling(t *testing.T) { + ctx := context.Background() + s := postgresTestStore(t) + now := time.Now().UTC() + parts := make([]Part, 20) + for i := range parts { + p := NewPart("job", fmt.Sprintf("part-%02d", i), "bucket", "source", "finished", now) + p.Status = StatusCompactReady + p.CompactReadyAt = formatTime(now) + p.DestinationDatabase = "db" + p.DestinationTable = "table" + p.DestinationSchema = "schema" + p.DestinationActivePartCount = 2 + p.DestinationActivePartBytes = uint64(i + 1) + p.DestinationActivePartitionCounts = map[string]uint64{"p": 2} + parts[i] = p + } + seedPostgresParts(t, s, parts) + tx, err := s.pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, `SELECT FROM `+s.tableSQL+` WHERE part_id='part-19' FOR UPDATE`); err != nil { + t.Fatal(err) + } + timed, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + batch, err := s.ClaimNextCompactBatch(timed, "worker", now, CompactClaimOptions{CompactWindow: time.Hour}) + if err != nil { + t.Fatal(err) + } + if batch == nil || batch.Parts[0].PartID != "part-18" { + t.Fatalf("compact claim: %+v", batch) + } + if err := tx.Rollback(ctx); err != nil { + t.Fatal(err) + } + // All contenders originally picked a single part in this one destination group. + var wg sync.WaitGroup + claimed := make(chan *CompactBatch, 19) + errs := make(chan error, 19) + for i := 0; i < 19; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + b, err := s.ClaimNextCompactBatch(ctx, fmt.Sprint(i), now, CompactClaimOptions{}) + claimed <- b + errs <- err + }(i) + } + wg.Wait() + close(claimed) + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + seen := map[string]bool{batch.Parts[0].PartID: true} + for b := range claimed { + if b == nil || seen[b.Parts[0].PartID] { + t.Fatalf("duplicate or missing compact claim: %+v", b) + } + seen[b.Parts[0].PartID] = true + } + if n, err := s.FinalizeCompactReadyJob(ctx, "job", time.Hour, now); err != nil || n != 0 { + t.Fatalf("active job finalized: %d %v", n, err) + } + if n, err := s.ReleaseStaleCompactingParts(ctx, now.Add(2*time.Hour), time.Hour); err != nil || n != 20 { + t.Fatalf("stale release: %d %v", n, err) + } + if b, err := s.ClaimNextCompactBatch(ctx, "late", now.Add(2*time.Hour), CompactClaimOptions{CompactWindow: time.Hour}); err != nil || b != nil { + t.Fatalf("claimed expired work: %+v %v", b, err) + } + if n, err := s.MaintainCompaction(ctx, time.Hour, time.Hour, now.Add(2*time.Hour), false); err != nil || n != 20 { + t.Fatalf("finalized: %d %v", n, err) + } + // The shared cadence must skip a second scan even when new normalized work arrives. + normalized := parts[0] + normalized.JobID = "new" + normalized.PartID = "normalized" + normalized.DestinationActivePartCount = 1 + normalized.DestinationActivePartitionCounts = map[string]uint64{"p": 1} + seedPostgresParts(t, s, []Part{normalized}) + if n, err := s.MaintainCompaction(ctx, time.Hour, time.Hour, now.Add(2*time.Hour), false); err != nil || n != 0 { + t.Fatalf("maintenance ran twice: %d %v", n, err) + } + if n, err := s.MaintainCompaction(ctx, time.Hour, time.Hour, now.Add(2*time.Hour), true); err != nil || n != 1 { + t.Fatalf("normalized finalization: %d %v", n, err) + } +} + +func TestPostgresCompactOptionsAndSummaries(t *testing.T) { + ctx := context.Background() + s := postgresTestStore(t) + now := time.Now().UTC() + makePart := func(job, id, partition string, size uint64, status Status) Part { + p := NewPart(job, id, "bucket", "source", "finished", now) + p.Status = status + p.JobName = "name" + p.CompactReadyAt = formatTime(now) + p.DestinationDatabase = "db" + p.DestinationTable = "table" + p.DestinationSchema = "schema" + p.DestinationActivePartCount = 2 + p.DestinationActivePartBytes = size + p.DestinationActivePartitionCounts = map[string]uint64{partition: 2} + if status == StatusCompacting { + p.WorkerID = "busy" + p.CompactingAt = formatTime(now) + } + return p + } + parts := []Part{makePart("a", "busy", "busy", 1, StatusCompacting), makePart("a", "large", "busy", 1000, StatusCompactReady), makePart("a", "idle", "idle", 100, StatusCompactReady), makePart("b", "other", "busy", 200, StatusCompactReady)} + seedPostgresParts(t, s, parts) + // Partition filters remain explicit; busy partitions no longer change priority. + for _, test := range []struct { + opts CompactClaimOptions + want string + }{ + {CompactClaimOptions{}, "large"}, {CompactClaimOptions{JobID: "a"}, "large"}, + {CompactClaimOptions{RequiredPartitionIDs: []string{"busy"}}, "large"}, + {CompactClaimOptions{RequiredPartitionIDs: []string{"idle"}}, "idle"}, + {CompactClaimOptions{JobID: "b"}, "other"}, + {CompactClaimOptions{ExcludedJobIDs: map[string]struct{}{"a": {}}}, "other"}, + {CompactClaimOptions{Bucket: "other"}, ""}, {CompactClaimOptions{DestinationDatabase: "other"}, ""}, + {CompactClaimOptions{DestinationTable: "other"}, ""}, {CompactClaimOptions{DestinationSchema: "other"}, ""}, + } { + opts := test.opts + batch, err := s.ClaimNextCompactBatch(ctx, "test", now, opts) + if err != nil { + t.Fatal(err) + } + if test.want == "" { + if batch != nil { + t.Fatalf("unexpected claim for %+v: %+v", opts, batch) + } + continue + } + if batch == nil || batch.Parts[0].PartID != test.want { + t.Fatalf("claim for %+v = %+v, want %s", opts, batch, test.want) + } + if err := s.RequestCompactFinalization(ctx, batch.Parts[0], now); err != nil { + t.Fatal(err) + } + if err := s.UpdateCompactProgress(ctx, *batch, "output", "test", PartStats{Count: 2}, PartStats{Count: 1}, CompactProgress{Stage: "merging"}, now); err != nil { + t.Fatal(err) + } + if _, err := s.HeartbeatCompactBatch(ctx, *batch, "wrong", now); !IsConditionalCheckFailed(err) { + t.Fatalf("compact ownership error = %v", err) + } + if requested, err := s.HeartbeatCompactBatch(ctx, *batch, "test", now); err != nil || !requested { + t.Fatalf("finalize request lost: %v %v", requested, err) + } + if err := s.ReleaseCompactBatch(ctx, *batch, "test", now); err != nil { + t.Fatal(err) + } + } + ids, err := s.ListJobIDsByStatus(ctx, StatusCompactReady, StatusCompactReady) + if err != nil { + t.Fatal(err) + } + if strings.Join(ids, ",") != "a,b" { + t.Fatalf("job IDs = %v", ids) + } + jobs, err := s.ListJobs(ctx) + if err != nil { + t.Fatal(err) + } + if len(jobs) != 2 || jobs[0].Total != 3 || jobs[0].Counts[StatusCompactReady] != 2 || jobs[0].Name != "name" || jobs[0].DestinationPartitionCount != 2 || jobs[0].DestinationActivePartCount != 6 { + t.Fatalf("job summaries = %+v", jobs) + } + if _, err := s.pool.Exec(ctx, `UPDATE `+s.tableSQL+` SET data=data || '{"job_name":"conflict"}'::jsonb WHERE part_id='large'`); err != nil { + t.Fatal(err) + } + if _, err := s.ListJobs(ctx); err == nil { + t.Fatal("accepted conflicting job names") + } +} + +func TestPostgresClaimPlans(t *testing.T) { + ctx := context.Background() + s := postgresTestStore(t) + now := time.Now().UTC() + parts := make([]Part, 20000) + for i := range parts { + p := NewPart(fmt.Sprintf("job-%d", i%10), fmt.Sprintf("part-%05d", i), "bucket", "source", "finished", now) + p.SourceArtifactBytes = uint64(i + 1) + p.DestinationActivePartBytes = uint64(i + 1) + p.DestinationDatabase = "db" + p.DestinationTable = "table" + p.DestinationSchema = strings.Repeat("column String, ", 80) + p.DestinationActivePartCount = 2 + p.DestinationActivePartitionCounts = map[string]uint64{"p": 2} + p.CompactReadyAt = formatTime(now) + if i%2 == 0 { + p.Status = StatusCompactReady + } + parts[i] = p + } + seedPostgresParts(t, s, parts) + if _, err := s.pool.Exec(ctx, "ANALYZE "+s.tableSQL); err != nil { + t.Fatal(err) + } + compact, args := s.compactClaimQuery(CompactClaimOptions{CompactWindow: time.Hour}, now) + expired, expiredArgs := s.compactClaimQuery(CompactClaimOptions{CompactWindow: time.Hour}, now.Add(2*time.Hour)) + for _, test := range []struct { + name, query, index string + args []any + }{ + {"rewrite", s.readyClaimQuery(), "ready_priority_idx", nil}, + {"compact", compact, "compact_priority_idx", args}, + {"expired compact", expired, "compact_priority_idx", expiredArgs}, + } { + rows, err := s.pool.Query(ctx, "EXPLAIN (ANALYZE, BUFFERS) "+test.query, test.args...) + if err != nil { + t.Fatal(err) + } + plan, err := pgx.CollectRows(rows, pgx.RowTo[string]) + if err != nil { + t.Fatal(err) + } + text := strings.Join(plan, "\n") + t.Log(test.name + "\n" + text) + if !strings.Contains(text, "Index Scan using "+strings.Trim(s.indexSQL(test.index), `"`)) || strings.Contains(text, "Sort Key:") { + t.Fatalf("%s claim lost ordered index scan:\n%s", test.name, text) + } + } +} + +// Check application writes against values derived independently from the JSON. +func assertSchedulingColumns(t testing.TB, s *Store) { + t.Helper() + var mismatches int + err := s.pool.QueryRow(context.Background(), `SELECT count(*) FROM `+s.tableSQL+` WHERE + source_artifact_bytes IS DISTINCT FROM COALESCE((data->>'source_artifact_bytes')::numeric, 0) OR + compact_bytes IS DISTINCT FROM COALESCE((data->>'destination_active_part_bytes')::numeric, 0) OR + compact_eligible IS DISTINCT FROM ( + COALESCE(btrim(data->>'destination_database'), '') <> '' AND + COALESCE(btrim(data->>'destination_table'), '') <> '' AND + COALESCE(btrim(data->>'destination_schema'), '') <> '' AND + COALESCE((data->>'destination_active_part_count')::numeric, 0) > 0 AND + EXISTS (SELECT FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 1)) OR + compact_normalized IS DISTINCT FROM ( + COALESCE((data->>'destination_active_part_count')::numeric, 0) = 1 AND + (SELECT count(*) = 1 AND COALESCE(bool_and(p.value::numeric = 1), false) + FROM jsonb_each_text(COALESCE(NULLIF(data->'destination_active_partition_counts', 'null'::jsonb), '{}'::jsonb)) p WHERE btrim(p.key) <> '' AND p.value::numeric > 0)) OR + compact_stale_at IS DISTINCT FROM (CASE WHEN status = 'COMPACTING' THEN LEAST(NULLIF(updated_at, '')::timestamptz, NULLIF(data->>'compacting_at', '')::timestamptz) END) OR + original_compact_ready_at IS DISTINCT FROM (CASE WHEN COALESCE((data->>'compact_generation')::int, 0) <= 0 AND jsonb_array_length(COALESCE(NULLIF(data->'compact_input_part_ids', 'null'::jsonb), '[]'::jsonb)) = 0 THEN NULLIF(data->>'compact_ready_at', '')::timestamptz END)`).Scan(&mismatches) + if err != nil { + t.Fatal(err) + } + if mismatches != 0 { + t.Fatalf("%d rows have stale scheduling columns", mismatches) + } +} + +func TestPostgresApplicationSchedulingColumns(t *testing.T) { + ctx := context.Background() + s := postgresTestStore(t) + var triggers int + if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM pg_trigger WHERE tgrelid=$1::regclass AND NOT tgisinternal`, s.tableSQL).Scan(&triggers); err != nil { + t.Fatal(err) + } + if triggers != 0 { + t.Fatalf("found %d triggers", triggers) + } + now := time.Now().UTC().Truncate(time.Microsecond) + part := NewPart("job", "part", "bucket", "source", "finished", now) + part.SourceArtifactBytes = ^uint64(0) + if err := s.CreatePart(ctx, part); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + claimed, err := s.ClaimNextReady(ctx, "worker", now) + if err != nil { + t.Fatal(err) + } + if claimed == nil { + t.Fatal("no ready claim") + } + // Exercise partial updates with retained destination metadata and partitions. + if _, err := s.updatePart(ctx, part.JobID, part.PartID, nil, func(p *Part) error { + p.DestinationDatabase, p.DestinationTable, p.DestinationSchema = "db", "table", "schema" + p.DestinationActivePartitionCounts = map[string]uint64{"p": 2} + return nil + }); err != nil { + t.Fatal(err) + } + for _, stats := range []*PartStats{{Count: 2, Bytes: ^uint64(0)}, nil, {}} { + if err := s.UpdateRewriteProgress(ctx, part.JobID, part.PartID, "worker", RewriteProgress{DestinationActivePartStats: stats}, now); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + } + if err := s.MarkCompactReady(ctx, part, "worker", "finished", "db", "table", "schema", PartStats{Count: 2, Bytes: 99}, map[string]uint64{"p": 2}, now); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + batch, err := s.ClaimNextCompactBatch(ctx, "worker", now, CompactClaimOptions{}) + if err != nil { + t.Fatal(err) + } + if batch == nil { + t.Fatal("no compact claim") + } + assertSchedulingColumns(t, s) + if _, err := s.HeartbeatCompactBatch(ctx, *batch, "worker", now.Add(time.Minute)); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + if err := s.ReleaseCompactBatch(ctx, *batch, "worker", now.Add(2*time.Minute)); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + // A late heartbeat can reacquire an unowned ready part. + if _, err := s.HeartbeatCompactBatch(ctx, *batch, "worker", now.Add(3*time.Minute)); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + output := batch.Parts[0] + output.PartID, output.Status, output.WorkerID = "output", StatusCompactReady, "" + output.CompactingAt = "" + output.CompactGeneration = 1 + output.CompactInputPartIDs = []string{part.PartID} + output.DestinationActivePartCount = 1 + output.DestinationActivePartitionCounts = map[string]uint64{"p": 1} + if err := s.CompleteCompaction(ctx, *batch, output, "worker", now.Add(4*time.Minute)); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) + if n, err := s.FinalizeCompactReadyJob(ctx, part.JobID, time.Hour, now.Add(5*time.Minute)); err != nil || n != 1 { + t.Fatalf("finalized %d: %v", n, err) + } + assertSchedulingColumns(t, s) +} + +func TestPostgresSchedulingBackfill(t *testing.T) { + ctx := context.Background() + cfg := postgresTestConfig(t) + s, err := openStore(ctx, cfg) + if err != nil { + t.Fatal(err) + } + defer s.pool.Close() + if _, err := s.pool.Exec(ctx, s.migrations()[0]); err != nil { + t.Fatal(err) + } + now := time.Now().UTC().Truncate(time.Microsecond) + for i, status := range []Status{StatusReady, StatusCompactReady, StatusCompacting, StatusCompactReady} { + p := NewPart("legacy", fmt.Sprint(i), "bucket", "source", "finished", now) + p.Status, p.SourceArtifactBytes = status, ^uint64(0) + p.DestinationDatabase, p.DestinationTable, p.DestinationSchema = "db", "table", "schema" + p.DestinationActivePartCount, p.DestinationActivePartBytes = 2, ^uint64(0) + p.DestinationActivePartitionCounts = map[string]uint64{"p": 2, " ": 9, "empty": 0} + p.CompactReadyAt = formatTime(now.Add(-time.Hour)) + if status == StatusCompacting { + p.CompactingAt = formatTime(now.Add(-time.Minute)) + } + if i == 3 { + p.CompactGeneration = 1 + p.CompactInputPartIDs = []string{"input"} + p.DestinationActivePartCount = 1 + p.DestinationActivePartitionCounts = map[string]uint64{"p": 1} + } + data, err := partJSON(p) + if err != nil { + t.Fatal(err) + } + if _, err := s.pool.Exec(ctx, `INSERT INTO `+s.tableSQL+` (job_id,part_id,status,worker_id,created_at,updated_at,data) VALUES ($1,$2,$3,$4,$5,$6,$7)`, p.JobID, p.PartID, string(p.Status), p.WorkerID, p.CreatedAt, p.UpdatedAt, data); err != nil { + t.Fatal(err) + } + } + if _, err := Migrate(ctx, cfg); err != nil { + t.Fatal(err) + } + assertSchedulingColumns(t, s) +} diff --git a/internal/state/postgres_test.go b/internal/state/postgres_test.go index e6aeb5b..e50eec7 100644 --- a/internal/state/postgres_test.go +++ b/internal/state/postgres_test.go @@ -60,269 +60,6 @@ func TestFailedRetryTarget(t *testing.T) { } } -func TestSelectCompactBatchPartsAllowsSingleMultiPartArtifact(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{parts: []Part{ - { - PartID: "part-1", - DestinationActivePartCount: 4, - DestinationActivePartBytes: 1024, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 4, - }, - }, - }}, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "part-1" { - t.Fatalf("selected = %+v, want part-1", selected) - } -} - -func TestSelectCompactBatchPartsAllowsOversizedSingleMultiPartArtifact(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{parts: []Part{ - { - PartID: "part-1", - DestinationActivePartCount: 4, - DestinationActivePartBytes: 4096, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 4, - }, - }, - }}, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "part-1" { - t.Fatalf("selected = %+v, want oversized part-1", selected) - } -} - -func TestSelectCompactBatchPartsNormalizesFragmentedArtifactAlone(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{parts: []Part{ - { - PartID: "fragmented", - DestinationActivePartCount: 3, - DestinationActivePartBytes: 300, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 3, - }, - }, - { - PartID: "normalized", - DestinationActivePartCount: 1, - DestinationActivePartBytes: 100, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 1, - }, - }, - }}, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "fragmented" { - t.Fatalf("selected = %+v, want fragmented artifact alone", selected) - } -} - -func TestSelectCompactBatchPartsNormalizesIdlePartitionFirst(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{ - parts: []Part{ - { - PartID: "busy", - DestinationActivePartCount: 2, - DestinationActivePartBytes: 1000, - DestinationActivePartitionCounts: map[string]uint64{ - "busy": 2, - }, - }, - { - PartID: "idle", - DestinationActivePartCount: 2, - DestinationActivePartBytes: 100, - DestinationActivePartitionCounts: map[string]uint64{ - "idle": 2, - }, - }, - }, - compactingPartitionIDs: []string{"busy"}, - }, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "idle" { - t.Fatalf("selected = %+v, want idle fragmented artifact alone", selected) - } -} - -func TestSelectCompactBatchPartsChoosesLargestEligibleArtifact(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{parts: []Part{ - { - PartID: "small", - DestinationActivePartBytes: 100, - DestinationActivePartitionCounts: map[string]uint64{ - "partition": 2, - }, - }, - { - PartID: "large", - DestinationActivePartBytes: 1000, - DestinationActivePartitionCounts: map[string]uint64{ - "partition": 2, - }, - }, - }}, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "large" { - t.Fatalf("selected = %+v, want largest eligible artifact", selected) - } -} - -func TestCompactCandidateSelectionsOrdersLargestAcrossGroups(t *testing.T) { - selections := compactCandidateSelections([]compactGroup{ - {parts: []Part{{PartID: "small", DestinationActivePartBytes: 100, DestinationActivePartitionCounts: map[string]uint64{"partition": 2}}}}, - {parts: []Part{{PartID: "large", DestinationActivePartBytes: 1000, DestinationActivePartitionCounts: map[string]uint64{"partition": 2}}}}, - }, CompactClaimOptions{}) - - if len(selections) != 2 || selections[0][0].PartID != "large" { - t.Fatalf("selections = %+v, want largest group candidate first", selections) - } -} - -func TestSelectCompactBatchPartsDoesNotCombineFragmentedBusyPartitionThroughIdleOverlap(t *testing.T) { - selected := selectCompactBatchParts(compactGroup{ - parts: []Part{ - { - PartID: "fragmented", - DestinationActivePartCount: 3, - DestinationActivePartitionCounts: map[string]uint64{ - "busy": 2, - "idle": 1, - }, - }, - { - PartID: "normalized", - DestinationActivePartCount: 1, - DestinationActivePartitionCounts: map[string]uint64{ - "idle": 1, - }, - }, - }, - compactingPartitionIDs: []string{"busy"}, - }, CompactClaimOptions{}) - - if len(selected) != 1 || selected[0].PartID != "fragmented" { - t.Fatalf("selected = %+v, want fragmented artifact alone", selected) - } -} - -func TestSelectCompactBatchPartsIgnoresCooldownField(t *testing.T) { - now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) - selected := selectCompactBatchParts(compactGroup{ - parts: []Part{ - { - PartID: "part-cooldown", - DestinationActivePartCount: 2, - DestinationActivePartBytes: 1024, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 2, - }, - CompactCooldownUntil: formatTime(now.Add(time.Hour)), - }, - }, - }, CompactClaimOptions{}) - if len(selected) != 1 || selected[0].PartID != "part-cooldown" { - t.Fatalf("selected = %+v, want cooldown field ignored", selected) - } -} - -func TestSelectCompactBatchPartsDoesNotCombineNormalizedArtifacts(t *testing.T) { - now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) - selected := selectCompactBatchParts(compactGroup{ - parts: []Part{ - { - PartID: "part-fresh", - DestinationActivePartCount: 1, - DestinationActivePartBytes: 100, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 1, - }, - }, - { - PartID: "part-cooldown", - DestinationActivePartCount: 1, - DestinationActivePartBytes: 100, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 1, - }, - CompactCooldownUntil: formatTime(now.Add(time.Hour)), - }, - }, - }, CompactClaimOptions{}) - - if len(selected) != 0 { - t.Fatalf("selected = %+v, want no multi-artifact batch", selected) - } -} - -func TestCompactCandidateGroupsIncludesRowsWithCooldownField(t *testing.T) { - now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) - groups := compactCandidateGroups([]Part{ - { - JobID: "job-1", - PartID: "part-cooldown", - Bucket: "bucket", - DestinationDatabase: "db", - DestinationTable: "table", - DestinationSchema: "schema", - DestinationActivePartCount: 2, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 2, - }, - CompactCooldownUntil: formatTime(now.Add(time.Hour)), - }, - { - JobID: "job-1", - PartID: "part-ready", - Bucket: "bucket", - DestinationDatabase: "db", - DestinationTable: "table", - DestinationSchema: "schema", - DestinationActivePartCount: 2, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 2, - }, - }, - }, nil, CompactClaimOptions{}) - if len(groups) != 1 || len(groups[0].parts) != 2 || groups[0].parts[0].PartID != "part-cooldown" || groups[0].parts[1].PartID != "part-ready" { - t.Fatalf("groups = %+v, want cooldown and ready parts", groups) - } -} - -func TestCompactCandidateGroupsSkipsExcludedJobs(t *testing.T) { - groups := compactCandidateGroups([]Part{ - { - JobID: "job-1", - PartID: "part-1", - Bucket: "bucket", - DestinationDatabase: "db", - DestinationTable: "table", - DestinationSchema: "schema", - DestinationActivePartCount: 2, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 2, - }, - }, - { - JobID: "job-2", - PartID: "part-2", - Bucket: "bucket", - DestinationDatabase: "db", - DestinationTable: "table", - DestinationSchema: "schema", - DestinationActivePartCount: 2, - DestinationActivePartitionCounts: map[string]uint64{ - "202606": 2, - }, - }, - }, nil, CompactClaimOptions{ExcludedJobIDs: map[string]struct{}{"job-1": {}}}) - if len(groups) != 1 || len(groups[0].parts) != 1 || groups[0].parts[0].JobID != "job-2" { - t.Fatalf("groups = %+v, want only non-excluded job-2", groups) - } -} - func TestValidatePartRejectsPartialSourceRef(t *testing.T) { part := NewPart("job-1", "part-1", "bucket", "source/part-1", "finished/part-1", time.Now().UTC()) part.SourceJobID = "job-source" @@ -366,41 +103,6 @@ func TestValidatePartRejectsSelfSourceRef(t *testing.T) { } } -func TestCompactCandidateGroupsSeparateJobs(t *testing.T) { - candidates := []Part{ - compactBatchTestPart("job-a", "part-a1", StatusCompactReady), - compactBatchTestPart("job-b", "part-b1", StatusCompactReady), - compactBatchTestPart("job-a", "part-a2", StatusCompactReady), - compactBatchTestPart("job-b", "part-b2", StatusCompactReady), - } - compacting := []Part{ - compactBatchTestPart("job-a", "part-a-busy", StatusCompacting), - } - - groups := compactCandidateGroups(candidates, compacting, CompactClaimOptions{}) - if len(groups) != 2 { - t.Fatalf("groups = %+v, want one group per job", groups) - } - for _, group := range groups { - if len(group.parts) != 2 { - t.Fatalf("group = %+v, want two parts for one job", group) - } - jobID := group.parts[0].JobID - switch jobID { - case "job-a": - if len(group.compactingPartitionIDs) != 1 || group.compactingPartitionIDs[0] != "partition-a" { - t.Fatalf("job-a compacting partitions = %v, want partition-a", group.compactingPartitionIDs) - } - case "job-b": - if len(group.compactingPartitionIDs) != 0 { - t.Fatalf("job-b compacting partitions = %v, want none", group.compactingPartitionIDs) - } - default: - t.Fatalf("unexpected job group %s", jobID) - } - } -} - func TestCompactBatchFromPartsRejectsMixedJobs(t *testing.T) { _, err := compactBatchFromParts([]Part{ compactBatchTestPart("job-a", "part-a", StatusCompacting), @@ -493,41 +195,6 @@ func TestCompactReadyAtForReleaseBackfillsExistingRowsFromProgress(t *testing.T) } } -func TestCompactHeartbeatTimeUsesUpdatedAt(t *testing.T) { - now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) - part := Part{ - JobID: "job-1", - PartID: "part-1", - UpdatedAt: formatTime(now), - CompactingAt: formatTime(now.Add(-time.Hour)), - } - got, err := compactHeartbeatTime(part) - if err != nil { - t.Fatal(err) - } - if !got.Equal(now) { - t.Fatalf("compactHeartbeatTime = %s, want %s", got, now) - } -} - -func TestCompactStaleTimeUsesOldestLeaseTimestamp(t *testing.T) { - now := time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC) - part := Part{ - JobID: "job-1", - PartID: "part-1", - UpdatedAt: formatTime(now), - CompactingAt: formatTime(now.Add(-2 * time.Hour)), - } - got, err := compactStaleTime(part) - if err != nil { - t.Fatal(err) - } - want := now.Add(-2 * time.Hour) - if !got.Equal(want) { - t.Fatalf("compactStaleTime = %s, want %s", got, want) - } -} - func compactBatchTestPart(jobID, partID string, status Status) Part { now := formatTime(time.Date(2026, 6, 23, 12, 0, 0, 0, time.UTC)) return Part{