chore: stop an interrupted repack from slowing every write

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_<oid> 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 == "" {

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

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