From e8f8c4100eb63d0d5e948b70bdb2f7704d7b641e Mon Sep 17 00:00:00 2001 From: Michael Buntarman Date: Wed, 30 Sep 2026 06:38:24 +0700 Subject: [PATCH] chore: stop an interrupted repack from slowing every write --- extensions/tn_vacuum/mechanism_repack.go | 78 +++++++++- .../tn_vacuum/mechanism_repack_pg_test.go | 137 ++++++++++++++++++ extensions/tn_vacuum/mechanism_repack_test.go | 72 ++++++++- 3 files changed, 285 insertions(+), 2 deletions(-) create mode 100644 extensions/tn_vacuum/mechanism_repack_pg_test.go diff --git a/extensions/tn_vacuum/mechanism_repack.go b/extensions/tn_vacuum/mechanism_repack.go index f1eac0cc..9adf1428 100644 --- a/extensions/tn_vacuum/mechanism_repack.go +++ b/extensions/tn_vacuum/mechanism_repack.go @@ -16,14 +16,25 @@ import ( var ErrPgRepackUnavailable = errors.New("pg_repack binary not found in PATH") +// pgRepackStopGrace is how long a cancelled pg_repack gets to drop its trigger +// and log table before it is killed. pg_repack cleans up on SIGINT only; a +// SIGKILL leaves both behind. +const pgRepackStopGrace = 30 * time.Second + +// leftoverLockTimeout bounds how long clearing an interrupted run may wait for a +// table lock, since block execution queues behind that wait. +const leftoverLockTimeout = "5s" + type pgRepackMechanism struct { logger log.Logger binaryPath string db DBConnConfig + // clearLeftovers runs before every pg_repack; tests replace it. + clearLeftovers func(ctx context.Context, db DBConnConfig, logger log.Logger) error } func NewPgRepackMechanism() Mechanism { - return &pgRepackMechanism{} + return &pgRepackMechanism{clearLeftovers: clearInterruptedRepack} } func (m *pgRepackMechanism) Name() string { return "pg_repack" } @@ -62,6 +73,16 @@ func (m *pgRepackMechanism) Run(ctx context.Context, req RunRequest) (*RunReport return nil, fmt.Errorf("pg_repack requires database name") } + if m.clearLeftovers != nil { + if err := m.clearLeftovers(ctx, db, m.logger); err != nil { + report.Duration = time.Since(startTime) + report.Status = StatusFailed + report.Error = err.Error() + m.logger.Warn("pg_repack skipped: could not clear an interrupted run", "error", err) + return report, err + } + } + args := []string{fmt.Sprintf("--dbname=%s", db.Database), "--all"} if db.Host != "" { args = append(args, fmt.Sprintf("--host=%s", db.Host)) @@ -80,6 +101,10 @@ func (m *pgRepackMechanism) Run(ctx context.Context, req RunRequest) (*RunReport args = append(args, "--no-order") cmd := exec.CommandContext(ctx, m.binaryPath, args...) + // On cancellation send SIGINT, which pg_repack handles by dropping its trigger + // and log table, and kill it only if it has not exited after the grace period. + cmd.Cancel = func() error { return cmd.Process.Signal(os.Interrupt) } + cmd.WaitDelay = pgRepackStopGrace env := os.Environ() if db.Password != "" { env = append(env, fmt.Sprintf("PGPASSWORD=%s", db.Password)) @@ -168,6 +193,57 @@ func ensurePgRepackExtension(ctx context.Context, db DBConnConfig, logger log.Lo return nil } +// leftoverCountSQL counts what a pg_repack run leaves when it is killed before it +// can clean up: repack_trigger on a source table, and the tables it creates in the +// repack schema. A fresh pg_repack extension owns no tables there. +const leftoverCountSQL = `SELECT + (SELECT count(*) FROM pg_trigger WHERE tgname = 'repack_trigger' AND NOT tgisinternal), + (SELECT count(*) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = 'repack' AND c.relkind = 'r')` + +// clearInterruptedRepack removes the trigger and log tables an interrupted +// pg_repack run leaves behind. The trigger copies every write to its table into a +// repack.log_ table until someone drops it, and a later --all run repacks +// those log tables as well, so the triggers chain and one write becomes many. +// Dropping the extension with CASCADE removes all of it, and the extension is +// created again in the same transaction. Nothing in the repack schema is +// consensus state. +func clearInterruptedRepack(ctx context.Context, db DBConnConfig, logger log.Logger) error { + conn, err := pgx.Connect(ctx, buildConnString(db)) + if err != nil { + return fmt.Errorf("connect to clear pg_repack leftovers: %w", err) + } + defer conn.Close(ctx) + + var triggers, tables int64 + if err := conn.QueryRow(ctx, leftoverCountSQL).Scan(&triggers, &tables); err != nil { + return fmt.Errorf("count pg_repack leftovers: %w", err) + } + if triggers == 0 && tables == 0 { + return nil + } + logger.Warn("clearing what an interrupted pg_repack run left behind", "triggers", triggers, "tables", tables) + + tx, err := conn.Begin(ctx) + if err != nil { + return fmt.Errorf("begin clearing pg_repack leftovers: %w", err) + } + defer func() { _ = tx.Rollback(ctx) }() + for _, stmt := range []string{ + "SET LOCAL lock_timeout = '" + leftoverLockTimeout + "'", + "DROP EXTENSION IF EXISTS pg_repack CASCADE", + "CREATE EXTENSION pg_repack", + } { + if _, err := tx.Exec(ctx, stmt); err != nil { + return fmt.Errorf("clear pg_repack leftovers: %w", err) + } + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("commit clearing pg_repack leftovers: %w", err) + } + return nil +} + func buildConnString(db DBConnConfig) string { host := db.Host if host == "" { diff --git a/extensions/tn_vacuum/mechanism_repack_pg_test.go b/extensions/tn_vacuum/mechanism_repack_pg_test.go new file mode 100644 index 00000000..73becca7 --- /dev/null +++ b/extensions/tn_vacuum/mechanism_repack_pg_test.go @@ -0,0 +1,137 @@ +//go:build kwiltest + +package tn_vacuum + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/stretchr/testify/require" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" + "github.com/trufnetwork/kwil-db/core/log" +) + +// startRepackPostgres starts the Postgres image nodes run, which ships the +// pg_repack extension. +func startRepackPostgres(t *testing.T) DBConnConfig { + t.Helper() + ctx := context.Background() + container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: testcontainers.ContainerRequest{ + Image: "ghcr.io/trufnetwork/kwil-postgres:16.8-1", + ExposedPorts: []string{"5432/tcp"}, + Env: map[string]string{ + "POSTGRES_DB": "kwil", + "POSTGRES_USER": "kwil", + "POSTGRES_PASSWORD": "kwil", + }, + WaitingFor: wait.ForLog("database system is ready to accept connections"). + WithOccurrence(2). + WithStartupTimeout(60 * time.Second), + }, + Started: true, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = container.Terminate(context.Background()) }) + + host, err := container.Host(ctx) + require.NoError(t, err) + port, err := container.MappedPort(ctx, "5432/tcp") + require.NoError(t, err) + return DBConnConfig{Host: host, Port: port.Port(), User: "kwil", Password: "kwil", Database: "kwil"} +} + +func countOne(t *testing.T, conn *pgx.Conn, query string) int64 { + t.Helper() + var n int64 + require.NoError(t, conn.QueryRow(context.Background(), query).Scan(&n)) + return n +} + +// leaveInterruptedRepack runs what pg_repack runs on a table before it copies +// it, taken from pg_repack's own repack.tables view, and stops there, as a +// killed run does: a primary key type, a log table, and repack_trigger. +func leaveInterruptedRepack(t *testing.T, conn *pgx.Conn, table string) { + t.Helper() + ctx := context.Background() + var pktype, logTable, trigger string + require.NoError(t, conn.QueryRow(ctx, + "SELECT create_pktype, create_log, create_trigger FROM repack.tables WHERE relname = $1", table, + ).Scan(&pktype, &logTable, &trigger)) + for _, stmt := range []string{pktype, logTable, trigger} { + _, err := conn.Exec(ctx, stmt) + require.NoError(t, err) + } +} + +func TestClearInterruptedRepack(t *testing.T) { + ctx := context.Background() + db := startRepackPostgres(t) + conn, err := pgx.Connect(ctx, buildConnString(db)) + require.NoError(t, err) + defer conn.Close(ctx) + + for _, stmt := range []string{ + "CREATE EXTENSION pg_repack", + "CREATE SCHEMA main", + "CREATE TABLE main.events (id INT PRIMARY KEY, v TEXT)", + "INSERT INTO main.events SELECT g, 'r' || g FROM generate_series(1, 3) g", + } { + _, err := conn.Exec(ctx, stmt) + require.NoError(t, err) + } + + // Nothing to clear: the extension is left as it is. + var extOID uint32 + require.NoError(t, conn.QueryRow(ctx, "SELECT oid FROM pg_extension WHERE extname = 'pg_repack'").Scan(&extOID)) + require.NoError(t, clearInterruptedRepack(ctx, db, log.DiscardLogger)) + require.EqualValues(t, extOID, countOne(t, conn, "SELECT oid::bigint FROM pg_extension WHERE extname = 'pg_repack'"), + "with nothing left behind, the extension must not be dropped") + + // A run interrupted on main.events, then a later --all run interrupted on + // that run's log table: the triggers chain. + leaveInterruptedRepack(t, conn, "main.events") + var eventsOID uint32 + require.NoError(t, conn.QueryRow(ctx, "SELECT 'main.events'::regclass::oid").Scan(&eventsOID)) + firstLog := fmt.Sprintf("repack.log_%d", eventsOID) + leaveInterruptedRepack(t, conn, firstLog) + var firstLogOID uint32 + require.NoError(t, conn.QueryRow(ctx, "SELECT $1::regclass::oid", firstLog).Scan(&firstLogOID)) + secondLog := fmt.Sprintf("repack.log_%d", firstLogOID) + + _, err = conn.Exec(ctx, "INSERT INTO main.events VALUES (4, 'r4')") + require.NoError(t, err) + require.EqualValues(t, 1, countOne(t, conn, "SELECT count(*) FROM "+firstLog), "one write lands in the first log table") + require.EqualValues(t, 1, countOne(t, conn, "SELECT count(*) FROM "+secondLog), "and again in the second") + + require.NoError(t, clearInterruptedRepack(ctx, db, log.DiscardLogger)) + + require.EqualValues(t, 0, countOne(t, conn, + "SELECT count(*) FROM pg_trigger WHERE tgname = 'repack_trigger' AND NOT tgisinternal")) + require.EqualValues(t, 0, countOne(t, conn, + "SELECT count(*) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = 'repack' AND c.relkind = 'r'")) + require.EqualValues(t, 1, countOne(t, conn, "SELECT count(*) FROM pg_extension WHERE extname = 'pg_repack'"), + "the extension is created again") + require.EqualValues(t, 4, countOne(t, conn, "SELECT count(*) FROM main.events"), "the table's own rows stay") + _, err = conn.Exec(ctx, "INSERT INTO main.events VALUES (5, 'r5')") + require.NoError(t, err, "writes work once the trigger is gone") + + // A run can also die after dropping its trigger but before dropping its log + // table. A table left in the repack schema on its own is still cleared. + var pktype, logTable string + require.NoError(t, conn.QueryRow(ctx, + "SELECT create_pktype, create_log FROM repack.tables WHERE relname = 'main.events'", + ).Scan(&pktype, &logTable)) + for _, stmt := range []string{pktype, logTable} { + _, err := conn.Exec(ctx, stmt) + require.NoError(t, err) + } + require.NoError(t, clearInterruptedRepack(ctx, db, log.DiscardLogger)) + require.EqualValues(t, 0, countOne(t, conn, + "SELECT count(*) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = 'repack' AND c.relkind = 'r'"), + "a log table without its trigger is cleared too") +} diff --git a/extensions/tn_vacuum/mechanism_repack_test.go b/extensions/tn_vacuum/mechanism_repack_test.go index 3892cb28..1e97ae28 100644 --- a/extensions/tn_vacuum/mechanism_repack_test.go +++ b/extensions/tn_vacuum/mechanism_repack_test.go @@ -1,6 +1,17 @@ package tn_vacuum -import "testing" +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + "github.com/trufnetwork/kwil-db/core/log" +) func TestCountRepackedTables(t *testing.T) { tests := []struct { @@ -77,3 +88,62 @@ func TestDetectPgRepackSoftFailure(t *testing.T) { }) } } + +// fakePgRepack writes a stand-in for the pg_repack binary. It records that it +// started, then waits; on SIGINT it records that too, the way pg_repack drops +// its trigger and log table when interrupted. +func fakePgRepack(t *testing.T) (binary, started, interrupted string) { + t.Helper() + dir := t.TempDir() + binary = filepath.Join(dir, "pg_repack") + started = filepath.Join(dir, "started") + interrupted = filepath.Join(dir, "interrupted") + script := fmt.Sprintf("#!/bin/sh\ntrap 'touch %s; exit 1' INT\ntouch %s\nwhile :; do sleep 0.05; done\n", interrupted, started) + require.NoError(t, os.WriteFile(binary, []byte(script), 0o755)) + return binary, started, interrupted +} + +func fileExists(path string) bool { + _, err := os.Stat(path) + return err == nil +} + +func TestRunInterruptsPgRepackOnCancel(t *testing.T) { + binary, started, interrupted := fakePgRepack(t) + m := &pgRepackMechanism{ + logger: log.DiscardLogger, + binaryPath: binary, + clearLeftovers: func(context.Context, DBConnConfig, log.Logger) error { return nil }, + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := m.Run(ctx, RunRequest{DB: DBConnConfig{Database: "kwild"}}) + done <- err + }() + require.Eventually(t, func() bool { return fileExists(started) }, 5*time.Second, 10*time.Millisecond) + + cancel() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("Run did not return after its context was cancelled") + } + require.True(t, fileExists(interrupted), "a cancelled pg_repack must get SIGINT, the signal it cleans up on") +} + +func TestRunSkipsPgRepackWhenLeftoversCannotBeCleared(t *testing.T) { + binary, started, _ := fakePgRepack(t) + clearErr := errors.New("lock timeout") + m := &pgRepackMechanism{ + logger: log.DiscardLogger, + binaryPath: binary, + clearLeftovers: func(context.Context, DBConnConfig, log.Logger) error { return clearErr }, + } + + report, err := m.Run(context.Background(), RunRequest{DB: DBConnConfig{Database: "kwild"}}) + require.ErrorIs(t, err, clearErr) + require.Equal(t, StatusFailed, report.Status) + require.False(t, fileExists(started), "pg_repack must not run on top of an interrupted run's trigger and log tables") +}