diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go index 45868a0e8..6f62df0d6 100644 --- a/cmd/config/config_env.go +++ b/cmd/config/config_env.go @@ -86,6 +86,7 @@ func bindEnvVars() { viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS") viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS") viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS") + viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS") viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES") viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES") viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING") @@ -378,6 +379,7 @@ func parseSchemaSnapshotConfig(pgurl string) (*snapshotbuilder.SchemaSnapshotCon if err != nil { return nil, err } + return &snapshotbuilder.SchemaSnapshotConfig{ DumpRestore: &pgdumprestore.Config{ SourcePGURL: pgurl, @@ -393,6 +395,7 @@ func parseSchemaSnapshotConfig(pgurl string) (*snapshotbuilder.SchemaSnapshotCon ExcludedSecurityLabels: viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS"), RefreshMaterializedViews: viper.GetBool("PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS"), IndexConstraintSessionSettings: viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS"), + IndexRestoreWorkers: viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS"), IncludeObjectTypes: viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES"), ExcludeObjectTypes: viper.GetStringSlice("PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES"), }, diff --git a/cmd/config/config_env_test.go b/cmd/config/config_env_test.go index 2bfa53c66..570a78889 100644 --- a/cmd/config/config_env_test.go +++ b/cmd/config/config_env_test.go @@ -6,6 +6,8 @@ import ( "os" "testing" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -53,6 +55,7 @@ func Test_EnvVarsToStreamConfig(t *testing.T) { os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_NO_OWNER", "true") os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_MODE", "full") os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS", "maintenance_work_mem=4GB max_parallel_maintenance_workers=4 synchronous_commit=off statement_timeout=0 lock_timeout=0") + os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS", "4") os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING", "true") os.Setenv("PGSTREAM_KAFKA_READER_SERVERS", "localhost:9092") @@ -130,3 +133,29 @@ func Test_EnvVarsToOtelConfig(t *testing.T) { validateTestOtelConfig(t, otelConfig) } + +// Test_EnvVarsToSchemaSnapshotConfig_IndexRestoreWorkers resolves the setting +// from the process environment on a viper with nothing else in it. The +// surrounding tests share viper's global config map with the ones that load +// test_config.env, so a variable that is never bound with BindEnv still +// resolves there, out of the file, and a missing binding goes unnoticed. +func Test_EnvVarsToSchemaSnapshotConfig_IndexRestoreWorkers(t *testing.T) { + reset := func() { + viper.Reset() + bindEnvVars() + } + reset() + t.Cleanup(reset) + + t.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS", "4") + cfg, err := parseSchemaSnapshotConfig("postgresql://user:password@localhost:5432/mydatabase") + require.NoError(t, err) + require.Equal(t, uint(4), cfg.DumpRestore.IndexRestoreWorkers) + + // unset, the generator applies its own default + os.Unsetenv("PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS") + reset() + cfg, err = parseSchemaSnapshotConfig("postgresql://user:password@localhost:5432/mydatabase") + require.NoError(t, err) + require.Equal(t, uint(0), cfg.DumpRestore.IndexRestoreWorkers) +} diff --git a/cmd/config/config_yaml.go b/cmd/config/config_yaml.go index a9bbcf50d..386bfd778 100644 --- a/cmd/config/config_yaml.go +++ b/cmd/config/config_yaml.go @@ -138,6 +138,7 @@ type PgDumpPgRestoreConfig struct { ExcludedSecurityLabels []string `mapstructure:"excluded_security_labels" yaml:"excluded_security_labels"` RefreshMaterializedViews bool `mapstructure:"refresh_materialized_views" yaml:"refresh_materialized_views"` IndexConstraintSessionSettings []string `mapstructure:"index_constraint_session_settings" yaml:"index_constraint_session_settings"` + IndexRestoreWorkers uint `mapstructure:"index_restore_workers" yaml:"index_restore_workers"` IncludeObjectTypes []string `mapstructure:"include_object_types" yaml:"include_object_types"` ExcludeObjectTypes []string `mapstructure:"exclude_object_types" yaml:"exclude_object_types"` } @@ -657,6 +658,7 @@ func (c *YAMLConfig) parseSchemaSnapshotConfig() (*snapshotbuilder.SchemaSnapsho streamSchemaCfg.DumpRestore.ExcludedSecurityLabels = schemaSnapshotCfg.PgDumpPgRestore.ExcludedSecurityLabels streamSchemaCfg.DumpRestore.RefreshMaterializedViews = schemaSnapshotCfg.PgDumpPgRestore.RefreshMaterializedViews streamSchemaCfg.DumpRestore.IndexConstraintSessionSettings = schemaSnapshotCfg.PgDumpPgRestore.IndexConstraintSessionSettings + streamSchemaCfg.DumpRestore.IndexRestoreWorkers = schemaSnapshotCfg.PgDumpPgRestore.IndexRestoreWorkers streamSchemaCfg.DumpRestore.IncludeObjectTypes = schemaSnapshotCfg.PgDumpPgRestore.IncludeObjectTypes streamSchemaCfg.DumpRestore.ExcludeObjectTypes = schemaSnapshotCfg.PgDumpPgRestore.ExcludeObjectTypes diff --git a/cmd/config/helper_test.go b/cmd/config/helper_test.go index 8eab0fdc6..1e83976eb 100644 --- a/cmd/config/helper_test.go +++ b/cmd/config/helper_test.go @@ -90,6 +90,7 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) { "statement_timeout=0", "lock_timeout=0", }, + IndexRestoreWorkers: 4, }, }, Recorder: &builder.SnapshotRecorderConfig{ diff --git a/cmd/config/test/test_config.env b/cmd/config/test/test_config.env index 6f7ec7ed7..921168c33 100644 --- a/cmd/config/test/test_config.env +++ b/cmd/config/test/test_config.env @@ -30,6 +30,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_NO_PRIVILEGES=true PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS="anon" PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS=true PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS="maintenance_work_mem=4GB max_parallel_maintenance_workers=4 synchronous_commit=off statement_timeout=0 lock_timeout=0" +PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS=4 PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING=true # Kafka diff --git a/cmd/config/test/test_config.yaml b/cmd/config/test/test_config.yaml index d576f0643..ef33968b0 100644 --- a/cmd/config/test/test_config.yaml +++ b/cmd/config/test/test_config.yaml @@ -34,6 +34,7 @@ source: - synchronous_commit=off - statement_timeout=0 - lock_timeout=0 + index_restore_workers: 4 # number of standalone indexes restored concurrently disable_progress_tracking: true # whether to disable progress tracking for the snapshot replication: # when mode is replication or snapshot_and_replication replication_slot: "pgstream_mydatabase_slot" diff --git a/config_template.yaml b/config_template.yaml index 103a39254..16cc23248 100644 --- a/config_template.yaml +++ b/config_template.yaml @@ -50,6 +50,11 @@ source: # - statement_timeout=0 # - lock_timeout=0 # - synchronous_commit=off # faster restore, but a target crash right after the restore can lose the final index/constraint commits + # Number of standalone CREATE INDEX/CREATE UNIQUE INDEX statements restored concurrently against the target. Statements that depend on an index + # (constraints, comments, partition attachments, REPLICA IDENTITY, CLUSTER) always wait for every index to be created first. + # Each worker holds one target connection and applies the index_constraint_session_settings above to it, so the target needs this many free + # connection slots and up to workers x maintenance_work_mem of memory. Defaults to 1 (sequential), maximum 32. + index_restore_workers: 1 dump_file: pg_dump.sql # name of the file where the contents of the schema pg_dump command and output will be written for debugging purposes. # Granular object type filtering for schema snapshots. Only one of include_object_types or exclude_object_types can be set. # Available categories: tables, sequences, types, indexes, constraints, functions, views, materialized_views, triggers, event_triggers, policies, rules, comments, extensions, collations, text_search diff --git a/docs/configuration.md b/docs/configuration.md index e84e86012..a0f3fc279 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -54,6 +54,11 @@ source: # - statement_timeout=0 # - lock_timeout=0 # - synchronous_commit=off # faster restore, but a target crash right after the restore can lose the final index/constraint commits + # Number of standalone CREATE INDEX/CREATE UNIQUE INDEX statements restored concurrently against the target. Statements that depend on an index + # (constraints, comments, partition attachments, REPLICA IDENTITY, CLUSTER) always wait for every index to be created first. + # Each worker holds one target connection and applies the index_constraint_session_settings above to it, so the target needs this many free + # connection slots and up to workers x maintenance_work_mem of memory. Defaults to 1 (sequential), maximum 32. + index_restore_workers: 1 dump_file: pg_dump.sql # name of the file where the contents of the schema pg_dump command and output will be written for debugging purposes. # Granular object type filtering for schema snapshots. Only one of include_object_types or exclude_object_types can be set. # Available categories: tables, sequences, types, indexes, constraints, functions, views, materialized_views, triggers, event_triggers, policies, rules, comments, extensions, collations, text_search @@ -248,6 +253,7 @@ Here's a list of all the environment variables that can be used to configure the | PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDED_SECURITY_LABELS | [] | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of providers whose security labels will be excluded. | | PGSTREAM_POSTGRES_SNAPSHOT_REFRESH_MATERIALIZED_VIEWS | False | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, whether to refresh materialized views (REFRESH MATERIALIZED VIEW ... WITH DATA) after the table data has been restored. | | PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS | [] | No | Space-separated PostgreSQL `name=value` session settings applied only while restoring indexes and constraints, for example `maintenance_work_mem=4GB max_parallel_maintenance_workers=4`. Each setting must be a whitespace-free `name=value` pair; invalid entries fail at startup. Unset or empty preserves existing behavior. | +| PGSTREAM_POSTGRES_SNAPSHOT_INDEX_RESTORE_WORKERS | 1 | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, number of standalone `CREATE INDEX`/`CREATE UNIQUE INDEX` statements restored concurrently against the target. Statements that depend on an index (constraints added `USING INDEX`, comments, partition attachments, `REPLICA IDENTITY`, `CLUSTER`) always wait for every index to be created first. Each worker holds one target connection and applies `PGSTREAM_POSTGRES_SNAPSHOT_INDEX_CONSTRAINT_SESSION_SETTINGS` to it, so the target needs this many free connection slots and up to workers × `maintenance_work_mem` of memory. Defaults to 1 (sequential); values above 32 are rejected at startup. | | PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES | [] | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of object type categories to include in the schema snapshot. Everything else is excluded. Mutually exclusive with `PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES`. See [object type filtering](#object-type-filtering). | | PGSTREAM_POSTGRES_SNAPSHOT_EXCLUDE_OBJECT_TYPES | [] | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, list of object type categories to exclude from the schema snapshot. Mutually exclusive with `PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_OBJECT_TYPES`. See [object type filtering](#object-type-filtering). | | PGSTREAM_POSTGRES_SNAPSHOT_ROLE | "" | No | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, role name to be used to create the dump. | diff --git a/internal/postgres/pg_restore.go b/internal/postgres/pg_restore.go index b46175fae..4c5fa2fba 100644 --- a/internal/postgres/pg_restore.go +++ b/internal/postgres/pg_restore.go @@ -426,20 +426,58 @@ func NewPGRestoreErrors(errs ...error) *PGRestoreErrors { return pgrestoreErrs } +// MergePGRestoreErrors combines the errors of several restores into a single +// PGRestoreErrors that keeps each error in the bucket it was originally +// classified into, so that HasCriticalErrors and IsRetryable describe the set +// as a whole. It returns nil when none of the restores failed. +// +// Concurrent restores need this: classifying their combined failure by +// whichever error happened to arrive first would make the retry decision +// depend on scheduling, and a merged error built with NewPGRestoreErrors would +// demote every nested restore error to critical, since a PGRestoreErrors +// matches none of the classifications addError checks for. +func MergePGRestoreErrors(errs ...error) error { + merged := &PGRestoreErrors{} + for _, err := range errs { + if err == nil { + continue + } + restoreErrs := &PGRestoreErrors{} + if errors.As(err, &restoreErrs) { + merged.ignoredErrs = append(merged.ignoredErrs, restoreErrs.ignoredErrs...) + merged.criticalErrs = append(merged.criticalErrs, restoreErrs.criticalErrs...) + merged.retryableErrs = append(merged.retryableErrs, restoreErrs.retryableErrs...) + continue + } + merged.addError(err) + } + + if !merged.HasErrors() { + return nil + } + return merged +} + func (e PGRestoreErrors) Error() string { if !e.HasErrors() { return "" } + return errors.Join(e.Unwrap()...).Error() +} + +func (e PGRestoreErrors) HasErrors() bool { + return len(e.criticalErrs) > 0 || len(e.retryableErrs) > 0 || len(e.ignoredErrs) > 0 +} +// Unwrap exposes the individual restore errors, so that a caller holding the +// collection can still interrogate it with errors.Is and errors.As rather than +// by matching on the joined message. +func (e PGRestoreErrors) Unwrap() []error { all := make([]error, 0, len(e.criticalErrs)+len(e.retryableErrs)+len(e.ignoredErrs)) all = append(all, e.criticalErrs...) all = append(all, e.retryableErrs...) all = append(all, e.ignoredErrs...) - return errors.Join(all...).Error() -} - -func (e PGRestoreErrors) HasErrors() bool { - return len(e.criticalErrs) > 0 || len(e.retryableErrs) > 0 || len(e.ignoredErrs) > 0 + return all } // HasCriticalErrors reports whether the restore hit an error that must not be diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/exotic_dump_shapes_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/exotic_dump_shapes_test.go new file mode 100644 index 000000000..f64a9fe79 --- /dev/null +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/exotic_dump_shapes_test.go @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 + +package pgdumprestore + +import ( + "os" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestParseDump_exoticShapes runs a real pg_dump --schema-only output through +// the parser and the block partitioning that the concurrent index restore +// depends on. The fixture is captured verbatim from postgres and covers the +// shapes whose classification is not obvious from the code: an exclusion +// constraint, a partial index with a multi-condition WHERE clause, a long +// expression index, a constraint whose index is reused (UNIQUE), CLUSTER ON, +// a table partition attachment and index partition attachments. +// +// Two properties matter for the parallel restore and neither is visible from +// the parser alone: a statement must land in the group that can run it, and no +// statement may be dropped or merged into another on the way there. +func TestParseDump_exoticShapes(t *testing.T) { + t.Parallel() + + dump, err := os.ReadFile("testdata/exotic_schema_dump.sql") + require.NoError(t, err) + + sg := SnapshotGenerator{objectTypeFilter: &objectTypeFilter{}} + parsed := sg.parseDump(dump) + connectBlocks, indexBlocks, otherBlocks := partitionDumpBlocks(parsed.indicesAndConstraints, isIndexStatement) + require.Empty(t, connectBlocks) + + // pg_dump emits every CREATE INDEX on a single line, however long the + // expression or WHERE clause, so each one becomes its own block and can be + // restored on its own connection + require.ElementsMatch(t, []string{ + "CREATE INDEX parts_val_idx ON ONLY public.parts USING btree (val);", + "CREATE INDEX parts_0_val_idx ON public.parts_0 USING btree (val);", + "CREATE INDEX parts_1_val_idx ON public.parts_1 USING btree (val);", + "CREATE INDEX rooms_expr_idx ON public.rooms USING btree (lower(name), upper(email), ((val * (2)::numeric)), COALESCE(name, email, 'a-fairly-long-default-value'::text));", + "CREATE INDEX rooms_partial_idx ON public.rooms USING btree (name) WHERE ((val > (100)::numeric) AND (name IS NOT NULL) AND (email IS NOT NULL));", + }, indexBlocks) + + // everything that builds its own index or references one has to wait for + // the index phase, and keeps its relative order + require.Equal(t, []string{ + "ALTER TABLE ONLY public.rooms\n ADD CONSTRAINT rooms_email_key UNIQUE (email);", + "ALTER TABLE public.rooms CLUSTER ON rooms_email_key;", + "ALTER TABLE ONLY public.rooms\n ADD CONSTRAINT rooms_no_overlap EXCLUDE USING gist (room WITH =, during WITH &&);", + "COMMENT ON INDEX public.rooms_partial_idx IS 'partial';", + "ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_0_val_idx;", + "ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_1_val_idx;", + }, otherBlocks) + + // a table partition attachment is not an index dependency and stays in the + // schema dump, which is restored before the data + require.Contains(t, string(parsed.filtered), "ALTER TABLE ONLY public.parts ATTACH PARTITION public.parts_0 FOR VALUES FROM (0) TO (10);") + + // no index or constraint statement may be lost or glued to its neighbour + // on the way into a group. Statements outside that domain are excluded: + // the parser deliberately drops some of them (ownership of excluded roles, + // legacy PL/pgSQL handlers, filtered security labels). + restored := string(parsed.filtered) + string(parsed.indicesAndConstraints) + string(parsed.views) + for _, line := range strings.Split(string(dump), "\n") { + line = strings.TrimSpace(line) + if !isIndexStatement(line) && !strings.HasPrefix(line, "ADD CONSTRAINT") && + !isAttachPartitionIndexStatement(line) && !strings.HasPrefix(line, "COMMENT ON INDEX") && + !isClusterOnAlterTable(line) { + continue + } + require.Contains(t, restored, line, "statement dropped by the dump parsing") + } +} + +// TestParseDump_singleLineAddConstraint covers the shape the block splitting +// would be most sensitive to: an ALTER TABLE ... ADD CONSTRAINT on one line. +// Real pg_dump wraps these onto two lines (see the fixture above), so this +// exercises the defensive branch rather than an observed output, and pins that +// consecutive constraints stay separate blocks instead of merging into one. +func TestParseDump_singleLineAddConstraint(t *testing.T) { + t.Parallel() + + sg := SnapshotGenerator{objectTypeFilter: &objectTypeFilter{}} + parsed := sg.parseDump([]byte( + "ALTER TABLE public.a ADD CONSTRAINT a_check CHECK (v > 0);\n" + + "ALTER TABLE public.b ADD CONSTRAINT b_check CHECK (v > 0);\n" + + "CREATE INDEX a_idx ON public.a USING btree (v);\n")) + + _, indexBlocks, otherBlocks := partitionDumpBlocks(parsed.indicesAndConstraints, isIndexStatement) + require.Equal(t, []string{"CREATE INDEX a_idx ON public.a USING btree (v);"}, indexBlocks) + require.Equal(t, []string{ + "ALTER TABLE public.a ADD CONSTRAINT a_check CHECK (v > 0);", + "ALTER TABLE public.b ADD CONSTRAINT b_check CHECK (v > 0);", + }, otherBlocks) +} diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker.go index 96a039af8..f4df5ba9e 100644 --- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker.go +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker.go @@ -17,13 +17,15 @@ import ( // pg restore, such as index creation, and displays progress bars. type snapshotTracker struct { conn pglib.Querier - progressBars *synclib.Map[string, progress.Bar] + progressBars *synclib.Map[int, progress.Bar] barBuilder func(total int64, description, unit string) progress.Bar clock clockwork.Clock } // indexCreationRow representation of a row from pg_stat_progress_create_index type indexCreationRow struct { + // PID of the backend building the index. + PID int // Table on which the index is being created. Table string // OID of the index being created or reindexed. During a non-concurrent CREATE INDEX, this is 0. @@ -48,7 +50,7 @@ func newSnapshotTracker(ctx context.Context, pgurl string) (*snapshotTracker, er } return &snapshotTracker{ conn: connPool, - progressBars: synclib.NewMap[string, progress.Bar](), + progressBars: synclib.NewMap[int, progress.Bar](), clock: clockwork.NewRealClock(), barBuilder: progress.NewBar, }, nil @@ -61,8 +63,8 @@ func (st *snapshotTracker) trackIndexesCreation(ctx context.Context) { for { select { case <-ctx.Done(): - for table := range st.progressBars.GetMap() { - st.markProgressBarCompleted(table) + for pid := range st.progressBars.GetMap() { + st.markProgressBarCompleted(pid) } return case <-ticker.Chan(): @@ -71,7 +73,7 @@ func (st *snapshotTracker) trackIndexesCreation(ctx context.Context) { continue } - for table, row := range rowMap { + for pid, row := range rowMap { // skip initialization phase where total is 0 if row.TuplesTotal == 0 { continue @@ -80,39 +82,37 @@ func (st *snapshotTracker) trackIndexesCreation(ctx context.Context) { // We can't use the index oid in the row to uniquely identify // the index being tracked since it is not set for CREATE INDEX // which is the command the restore produces. Instead we use the - // table name. - // - // There can only be one index being created per table at a - // time, so we can track progress bars by table name. When the - // number of tuples done is lower than the previous recorded - // value, it means a new index is being created for the same - // table and the previous one can be marked as completed. - existingBar, found := st.progressBars.Get(table) + // pid of the backend building it: a backend runs one statement + // at a time, so a single bar tracks it. When the number of + // tuples done drops below the previous value, the backend has + // moved on to the next index and the previous one is complete. + // Keying by table instead would collapse the bars of indices + // that concurrent restores build on the same table. + existingBar, found := st.progressBars.Get(pid) switch { case found && row.TuplesDone >= existingBar.Current(): existingBar.SetCurrent(row.TuplesDone) continue case found && row.TuplesDone < existingBar.Current(): // if we're setting a lower current value, it's likely that - // a new index creation has started on the same table. So + // a new index creation has started on the same backend. So // complete the old bar and create a new one. - st.markProgressBarCompleted(table) + st.markProgressBarCompleted(pid) fallthrough default: // Create new progress bar for the index being created if not // found in the bar map bar := st.barBuilder(row.TuplesTotal, st.barDescription(row.Table), "tuples") - st.progressBars.Set(row.Table, bar) + st.progressBars.Set(pid, bar) bar.SetCurrent(row.TuplesDone) } } - // when the rows no longer return an existing table index being tracked, - // it means the index creation is done and we can mark it as - // complete. - for table := range st.progressBars.GetMap() { - if _, found := rowMap[table]; !found { - st.markProgressBarCompleted(table) + // when the rows no longer return an index being tracked, it means + // the index creation is done and we can mark it as complete. + for pid := range st.progressBars.GetMap() { + if _, found := rowMap[pid]; !found { + st.markProgressBarCompleted(pid) } } } @@ -120,22 +120,22 @@ func (st *snapshotTracker) trackIndexesCreation(ctx context.Context) { } // https://www.postgresql.org/docs/current/progress-reporting.html#CREATE-INDEX-PROGRESS-REPORTING -const createIndexProgressQuery = `SELECT relid::regclass AS table,index_relid::regclass AS index, phase, tuples_done, tuples_total, command FROM pg_stat_progress_create_index;` +const createIndexProgressQuery = `SELECT pid, relid::regclass AS table,index_relid::regclass AS index, phase, tuples_done, tuples_total, command FROM pg_stat_progress_create_index;` -func (st *snapshotTracker) getCreateIndexProgressRows(ctx context.Context) (map[string]indexCreationRow, error) { +func (st *snapshotTracker) getCreateIndexProgressRows(ctx context.Context) (map[int]indexCreationRow, error) { rows, err := st.conn.Query(ctx, createIndexProgressQuery) if err != nil { return nil, err } defer rows.Close() - result := map[string]indexCreationRow{} + result := map[int]indexCreationRow{} for rows.Next() { var row indexCreationRow - if err := rows.Scan(&row.Table, &row.Index, &row.Phase, &row.TuplesDone, &row.TuplesTotal, &row.Command); err != nil { + if err := rows.Scan(&row.PID, &row.Table, &row.Index, &row.Phase, &row.TuplesDone, &row.TuplesTotal, &row.Command); err != nil { return nil, err } - result[row.Table] = row + result[row.PID] = row } if err := rows.Err(); err != nil { return nil, err @@ -143,12 +143,12 @@ func (st *snapshotTracker) getCreateIndexProgressRows(ctx context.Context) (map[ return result, nil } -func (st *snapshotTracker) markProgressBarCompleted(name string) { - bar, found := st.progressBars.Get(name) +func (st *snapshotTracker) markProgressBarCompleted(pid int) { + bar, found := st.progressBars.Get(pid) if found { bar.Close() } - st.progressBars.Delete(name) + st.progressBars.Delete(pid) } func (st *snapshotTracker) barDescription(table string) string { diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker_test.go index d7fa10ba8..d2a781ead 100644 --- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker_test.go +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/pg_snapshot_tracker_test.go @@ -19,10 +19,12 @@ import ( synclib "github.com/xataio/pgstream/internal/sync" ) +const testPID = 4242 + func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { t.Parallel() - testRows := func(retTuplesDone, retTuplesTotal int64, tableName ...string) *pglibmocks.Rows { + testRowsForPID := func(pidValue int, retTuplesDone, retTuplesTotal int64, tableName ...string) *pglibmocks.Rows { testTable := "test_table" if len(tableName) > 0 { testTable = tableName[0] @@ -31,20 +33,23 @@ func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { return &pglibmocks.Rows{ NextFn: func(i uint) bool { return i == 1 }, ScanFn: func(i uint, dest ...any) error { - require.Len(t, dest, 6) - table, ok := dest[0].(*string) + require.Len(t, dest, 7) + pid, ok := dest[0].(*int) + require.True(t, ok) + table, ok := dest[1].(*string) require.True(t, ok) - index, ok := dest[1].(*string) + index, ok := dest[2].(*string) require.True(t, ok) - phase, ok := dest[2].(*string) + phase, ok := dest[3].(*string) require.True(t, ok) - tuplesDone, ok := dest[3].(*int64) + tuplesDone, ok := dest[4].(*int64) require.True(t, ok) - tuplesTotal, ok := dest[4].(*int64) + tuplesTotal, ok := dest[5].(*int64) require.True(t, ok) - command, ok := dest[5].(*string) + command, ok := dest[6].(*string) require.True(t, ok) + *pid = pidValue *table = testTable *index = "test_index" *phase = "index creation" @@ -59,6 +64,10 @@ func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { } } + testRows := func(retTuplesDone, retTuplesTotal int64, tableName ...string) *pglibmocks.Rows { + return testRowsForPID(testPID, retTuplesDone, retTuplesTotal, tableName...) + } + tests := []struct { name string querier func(doneChan chan struct{}) pglib.Querier @@ -247,7 +256,7 @@ func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { wantMarkCompletedCalls: 2, }, { - name: "ok - create a new bar for a different table, completing previous one", + name: "ok - create a new bar for a different backend, completing previous one", querier: func(_ chan struct{}) pglib.Querier { return &pglibmocks.Querier{ QueryFn: func(ctx context.Context, i uint, sql string, args ...any) (pglib.Rows, error) { @@ -255,7 +264,7 @@ func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { case 1: return testRows(100, 100), nil case 2: - return testRows(100, 100, "another_table"), nil + return testRowsForPID(testPID+1, 100, 100, "another_table"), nil default: t.Fatalf("unexpected Query call %d", i) return nil, nil @@ -320,7 +329,7 @@ func TestSnapshotTracker_trackIndexesCreation(t *testing.T) { st := snapshotTracker{ conn: tc.querier(doneChan), - progressBars: synclib.NewMap[string, progress.Bar](), + progressBars: synclib.NewMap[int, progress.Bar](), barBuilder: barBuilder, clock: fakeClock, } diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go index 1b5a7db41..b7e0363aa 100644 --- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator.go @@ -24,6 +24,7 @@ import ( "github.com/xataio/pgstream/pkg/snapshot" "github.com/xataio/pgstream/pkg/snapshot/generator" "github.com/xataio/pgstream/pkg/wal/processor" + "golang.org/x/sync/errgroup" ) // SnapshotGenerator generates postgres schema snapshots using pg_dump and @@ -64,6 +65,10 @@ type SnapshotGenerator struct { // created by an earlier attempt fail with "already exists", which the // restore ignores. indexRestoreBackoff backoff.Provider + // indexRestoreWorkers controls how many standalone CREATE INDEX/CREATE + // UNIQUE INDEX statements are restored concurrently. Defaults to 1 + // (sequential, unchanged behaviour). + indexRestoreWorkers uint } type snapshotProgressTracker interface { @@ -104,6 +109,17 @@ type Config struct { // exponential backoff with an initial interval of 1s, a max interval of // 1min and 3 retries. IndexRestoreRetries backoff.Config + // IndexRestoreWorkers is the number of standalone CREATE INDEX/CREATE + // UNIQUE INDEX statements restored concurrently against the target + // database. Statements that depend on an index existing (constraints + // added USING INDEX, comments, partition attachments, REPLICA IDENTITY + // and CLUSTER) are always restored afterwards, once every index has been + // created. Defaults to 1, which preserves the existing sequential + // behaviour. Each worker holds a target connection and applies + // IndexConstraintSessionSettings to it, so both the connections and the + // memory those settings reserve scale with this value. Capped at + // maxIndexRestoreWorkers. + IndexRestoreWorkers uint // IncludeObjectTypes is a list of object type categories to include in the // schema snapshot. Only one of IncludeObjectTypes or ExcludeObjectTypes // can be set. @@ -128,8 +144,22 @@ const ( defaultIndexRestoreRetryInitialInterval = time.Second defaultIndexRestoreRetryMaxInterval = time.Minute defaultIndexRestoreMaxRetries = 3 + defaultIndexRestoreWorkers = 1 + // maxIndexRestoreWorkers keeps IndexRestoreWorkers from exhausting the + // target's connection slots or memory, and catches a negative value that + // wrapped around into a huge uint in the config decoder. + maxIndexRestoreWorkers = 32 ) +var errTooManyIndexRestoreWorkers = errors.New("index restore workers above the maximum") + +func (c *Config) indexRestoreWorkers() uint { + if c.IndexRestoreWorkers > 0 { + return c.IndexRestoreWorkers + } + return defaultIndexRestoreWorkers +} + func (c *Config) indexRestoreBackoffConfig() *backoff.Config { if c.IndexRestoreRetries.DisableRetries { return &backoff.Config{} @@ -185,6 +215,10 @@ func NewSnapshotGenerator(ctx context.Context, c *Config, opts ...Option) (*Snap return nil, err } + if workers := c.indexRestoreWorkers(); workers > maxIndexRestoreWorkers { + return nil, fmt.Errorf("%w: %d configured, maximum is %d", errTooManyIndexRestoreWorkers, workers, maxIndexRestoreWorkers) + } + sourceConnPool, err := pglib.NewConnPool(ctx, c.SourcePGURL) if err != nil { return nil, err @@ -211,6 +245,7 @@ func NewSnapshotGenerator(ctx context.Context, c *Config, opts ...Option) (*Snap indexConstraintSessionSettings: c.IndexConstraintSessionSettings, objectTypeFilter: objTypeFilter, indexRestoreBackoff: backoff.NewProvider(c.indexRestoreBackoffConfig()), + indexRestoreWorkers: c.indexRestoreWorkers(), } for _, opt := range opts { @@ -264,6 +299,11 @@ func WithRestoreToWAL(processor processor.Processor) Option { return func(sg *SnapshotGenerator) { sg.pgRestoreFn = newPGSnapshotWALRestore(processor, sg.sourceQuerier).restoreToWAL sg.indexConstraintSessionSettings = nil + // this path converts the dump into WAL events for a non-postgres + // target instead of running index builds against a target database, + // so concurrency buys nothing and would emit the DDL events out of + // order + sg.indexRestoreWorkers = defaultIndexRestoreWorkers } } @@ -484,7 +524,8 @@ func (s *SnapshotGenerator) restoreIndicesAndConstraints(ctx context.Context, du "retries": retries, "backoff": d, }) - }) + }, + ) if err == nil { return nil } @@ -513,7 +554,116 @@ func (s *SnapshotGenerator) restoreIndices(ctx context.Context, opts pglib.PGRes if s.snapshotTracker != nil { return s.restoreIndicesWithTracking(ctx, opts, dump) } - return s.restoreDumpWithOptions(ctx, opts, dump) + return s.restoreIndexDump(ctx, opts, dump) +} + +// restoreIndexDump restores standalone CREATE INDEX/CREATE UNIQUE INDEX +// statements concurrently across up to indexRestoreWorkers connections, since +// they have no dependencies on each other. Every other statement in dump may +// depend on an index existing (a constraint added USING INDEX, a comment on +// the index, a partition attachment, REPLICA IDENTITY or CLUSTER), so it is +// only restored once all indexes above have been created, preserving its +// original relative order. +func (s *SnapshotGenerator) restoreIndexDump(ctx context.Context, opts pglib.PGRestoreOptions, dump []byte) error { + if s.indexRestoreWorkers <= 1 { + return s.restoreDumpWithOptions(ctx, opts, dump) + } + + connectBlocks, indexBlocks, otherBlocks := partitionDumpBlocks(dump, isIndexStatement) + if len(indexBlocks) <= 1 { + s.logger.Debug("restoring indices sequentially, too few index statements to parallelise", + loglib.Fields{"index_statements": len(indexBlocks)}) + return s.restoreDumpWithOptions(ctx, opts, dump) + } + + s.logger.Info("restoring indices concurrently", loglib.Fields{ + "index_restore_workers": s.indexRestoreWorkers, + "index_statements": len(indexBlocks), + "dependent_statements": len(otherBlocks), + }) + indexErr := s.restoreIndexBlocksInParallel(ctx, opts, connectBlocks, indexBlocks) + + // the dependent statements are restored even when an index failed. The + // single restore this replaces applied all of them and reported the + // failure, and the ones whose index is missing fail with "does not + // exist", which the restore already ignores. Skipping them would leave + // the target with data but without its constraints, foreign keys and + // replica identities, and nothing short of a re-snapshot to repair it. + if ctx.Err() != nil { + return errors.Join(indexErr, ctx.Err()) + } + + otherErr := s.restoreDumpWithOptions(ctx, opts, joinDumpBlocks(connectBlocks, otherBlocks)) + return pglib.MergePGRestoreErrors(indexErr, otherErr) +} + +// restoreIndexBlocksInParallel restores each of indexBlocks on its own restore +// call, prefixed with any connect statements, running up to +// indexRestoreWorkers of them at a time. +// +// A failing index deliberately does not cancel the restores still in flight. +// Cancelling kills psql mid-CREATE INDEX, which the server only notices once +// the statement completes, so the index it was building is committed anyway +// and collides with the retry that restoreIndicesAndConstraints issues moments +// later: the reissued statement waits on the orphan's uncommitted catalog row +// and then fails with a duplicate key error that is classified as permanent, +// turning a transient failure into a lost snapshot. Every block is attempted +// instead, and the failures are merged so the retry decision is taken over the +// whole wave rather than over whichever error happened to arrive first. +func (s *SnapshotGenerator) restoreIndexBlocksInParallel(ctx context.Context, opts pglib.PGRestoreOptions, connectBlocks, indexBlocks []string) error { + eg := errgroup.Group{} + // the conversion is safe: the worker count is capped at + // maxIndexRestoreWorkers when the generator is built + eg.SetLimit(int(s.indexRestoreWorkers)) + + mutex := sync.Mutex{} + restoreErrs := make([]error, 0, len(indexBlocks)) + + for _, block := range indexBlocks { + if ctx.Err() != nil { + break + } + eg.Go(func() error { + err := s.restoreDumpWithOptions(ctx, opts, joinDumpBlocks(connectBlocks, []string{block})) + if err != nil { + s.logger.Warn(err, "restoring index", loglib.Fields{"statement": block}) + mutex.Lock() + restoreErrs = append(restoreErrs, err) + mutex.Unlock() + } + // the failure is collected rather than returned, so that the + // group keeps the remaining indices going + return nil + }) + } + _ = eg.Wait() + + return pglib.MergePGRestoreErrors(restoreErrs...) +} + +// partitionDumpBlocks splits a dump into the blank line separated statement +// blocks pg_dump emits, and sorts them into the connect statements, the blocks +// matching the predicate, and the rest, preserving the relative order of the +// last two groups. +// +// The connect statements are kept apart because they select the database the +// statements after them apply to, so every subset of the blocks has to be +// prefixed with them (see joinDumpBlocks). +func partitionDumpBlocks(dump []byte, matches func(block string) bool) (connectBlocks, matchingBlocks, otherBlocks []string) { + for _, block := range strings.Split(string(dump), "\n\n") { + block = strings.TrimSpace(block) + switch { + case block == "": + continue + case strings.Contains(block, `\connect`): + connectBlocks = append(connectBlocks, block) + case matches(block): + matchingBlocks = append(matchingBlocks, block) + default: + otherBlocks = append(otherBlocks, block) + } + } + return connectBlocks, matchingBlocks, otherBlocks } // asRetryError marks any error that rerunning the dump cannot resolve as @@ -921,6 +1071,10 @@ func (s *SnapshotGenerator) parseDump(d []byte) *dump { indicesAndConstraints.WriteString("\n\n") case strings.HasPrefix(line, "ALTER TABLE") && strings.Contains(line, "ADD CONSTRAINT"): indicesAndConstraints.WriteString(line) + // the separator keeps this statement from being glued to the next + // one, which would hide the statement that follows it from the + // block splitting the restore does + indicesAndConstraints.WriteString("\n\n") case strings.HasPrefix(line, "ALTER TABLE") && isClusterOnAlterTable(line): indicesAndConstraints.WriteString(line) indicesAndConstraints.WriteString("\n\n") @@ -1295,7 +1449,7 @@ func (s *SnapshotGenerator) restoreIndicesWithTracking(ctx context.Context, opts defer wg.Done() s.snapshotTracker.trackIndexesCreation(trackingCtx) }() - err := s.restoreDumpWithOptions(ctx, opts, dump) + err := s.restoreIndexDump(ctx, opts, dump) // wait for the tracking to finish once the restore is done cancel() wg.Wait() diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go index 3c6ad960e..66d30cc65 100644 --- a/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/snapshot_pg_dump_restore_generator_test.go @@ -7,7 +7,10 @@ import ( "errors" "fmt" "os" + "slices" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -2599,3 +2602,443 @@ func TestSnapshotGenerator_restoreIndicesAndConstraints_noBackoffProvider(t *tes require.Error(t, err) require.Equal(t, 1, attempts) } + +func TestPartitionDumpBlocks(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + dump []byte + wantConnectBlocks []string + wantIndexBlocks []string + wantOtherBlocks []string + }{ + { + name: "empty dump", + }, + { + name: "mixed statements", + dump: []byte("\\connect testdb\n\n" + + "CREATE INDEX idx_a ON public.a USING btree (id);\n\n" + + "CREATE UNIQUE INDEX idx_b ON public.b USING btree (id);\n\n" + + "ALTER TABLE public.a ADD CONSTRAINT a_pkey PRIMARY KEY USING INDEX idx_a;\n\n" + + "COMMENT ON INDEX idx_b IS 'test';\n\n"), + wantConnectBlocks: []string{`\connect testdb`}, + wantIndexBlocks: []string{ + "CREATE INDEX idx_a ON public.a USING btree (id);", + "CREATE UNIQUE INDEX idx_b ON public.b USING btree (id);", + }, + wantOtherBlocks: []string{ + "ALTER TABLE public.a ADD CONSTRAINT a_pkey PRIMARY KEY USING INDEX idx_a;", + "COMMENT ON INDEX idx_b IS 'test';", + }, + }, + { + // ATTACH PARTITION depends on both the parent and the child index + // already existing, so it must not be treated as an independent + // index statement. + name: "attach partition is not an index statement", + dump: []byte("CREATE INDEX events_partition_id_idx ON ONLY public.events USING btree (partition_id);\n\n" + + "CREATE INDEX events_0_partition_id_idx ON public.events_0 USING btree (partition_id);\n\n" + + "ALTER INDEX public.events_partition_id_idx ATTACH PARTITION public.events_0_partition_id_idx;\n\n"), + wantIndexBlocks: []string{ + "CREATE INDEX events_partition_id_idx ON ONLY public.events USING btree (partition_id);", + "CREATE INDEX events_0_partition_id_idx ON public.events_0 USING btree (partition_id);", + }, + wantOtherBlocks: []string{ + "ALTER INDEX public.events_partition_id_idx ATTACH PARTITION public.events_0_partition_id_idx;", + }, + }, + { + // pg_dump splits an ALTER TABLE ONLY constraint over two lines, + // which parseDump keeps together in a single block + name: "multi line constraint stays one block", + dump: []byte("CREATE INDEX idx_a ON public.a USING btree (id);\n\n" + + "ALTER TABLE ONLY public.a\n ADD CONSTRAINT a_fkey FOREIGN KEY (b_id) REFERENCES public.b(id);\n\n"), + wantIndexBlocks: []string{"CREATE INDEX idx_a ON public.a USING btree (id);"}, + wantOtherBlocks: []string{"ALTER TABLE ONLY public.a\n ADD CONSTRAINT a_fkey FOREIGN KEY (b_id) REFERENCES public.b(id);"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + gotConnect, gotIndex, gotOther := partitionDumpBlocks(tc.dump, isIndexStatement) + require.Equal(t, tc.wantConnectBlocks, gotConnect) + require.Equal(t, tc.wantIndexBlocks, gotIndex) + require.Equal(t, tc.wantOtherBlocks, gotOther) + }) + } +} + +const ( + testConnectBlock = `\connect testdb` + testIndexA = "CREATE INDEX idx_a ON public.a USING btree (id);" + testIndexB = "CREATE UNIQUE INDEX idx_b ON public.b USING btree (id);" + testDependent1 = "ALTER TABLE public.a ADD CONSTRAINT a_pkey PRIMARY KEY USING INDEX idx_a;" + testDependent2 = "COMMENT ON INDEX idx_b IS 'test';" +) + +func testIndexDump(blocks ...string) []byte { + return []byte(strings.Join(blocks, "\n\n") + "\n\n") +} + +// restoreRecorder records the dumps a restore was called with, so that the +// assertions can run on the test goroutine: testify's require calls FailNow, +// which must not be reached from the restore worker goroutines. +type restoreRecorder struct { + mutex sync.Mutex + dumps []string +} + +func (r *restoreRecorder) record(dump []byte) { + r.mutex.Lock() + defer r.mutex.Unlock() + r.dumps = append(r.dumps, string(dump)) +} + +func (r *restoreRecorder) recorded() []string { + r.mutex.Lock() + defer r.mutex.Unlock() + return slices.Clone(r.dumps) +} + +func TestSnapshotGenerator_restoreIndexDump(t *testing.T) { + t.Parallel() + + dump := testIndexDump(testConnectBlock, testIndexA, testIndexB, testDependent1, testDependent2) + indexACall := string(testIndexDump(testConnectBlock, testIndexA)) + indexBCall := string(testIndexDump(testConnectBlock, testIndexB)) + dependentsCall := string(testIndexDump(testConnectBlock, testDependent1, testDependent2)) + + transientErr := func() error { + return pglib.NewPGRestoreErrors(&pglib.ErrTransientFailure{Details: "psql: error: connection to server was lost"}) + } + + t.Run("no workers configured restores the whole dump in one call", func(t *testing.T) { + t.Parallel() + + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + require.NoError(t, sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump)) + require.Equal(t, []string{string(dump)}, recorder.recorded()) + }) + + t.Run("a single index statement falls back to one sequential call", func(t *testing.T) { + t.Parallel() + + singleIndexDump := testIndexDump(testIndexA, testDependent1) + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 4, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + require.NoError(t, sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, singleIndexDump)) + require.Equal(t, []string{string(singleIndexDump)}, recorder.recorded()) + }) + + t.Run("each index is restored on its own, and the dependent statements after all of them", func(t *testing.T) { + t.Parallel() + + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + require.NoError(t, sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump)) + + recorded := recorder.recorded() + require.Len(t, recorded, 3) + // the indices are restored one statement per call, in any order, + // each carrying the connect statements + require.ElementsMatch(t, []string{indexACall, indexBCall}, recorded[:2]) + // the dependent statements are restored last, in one call, in their + // original relative order + require.Equal(t, dependentsCall, recorded[2]) + }) + + t.Run("more workers than index statements", func(t *testing.T) { + t.Parallel() + + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 8, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + require.NoError(t, sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump)) + + recorded := recorder.recorded() + require.Len(t, recorded, 3) + require.ElementsMatch(t, []string{indexACall, indexBCall}, recorded[:2]) + require.Equal(t, dependentsCall, recorded[2]) + }) + + t.Run("indices are restored concurrently", func(t *testing.T) { + t.Parallel() + + // both index restores block until the other one has started, so the + // test times out rather than passing if the restores are serialised + started := make(chan struct{}, 2) + release := make(chan struct{}) + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + if !isIndexStatement(strings.TrimPrefix(string(gotDump), testConnectBlock+"\n\n")) { + return "", nil + } + started <- struct{}{} + <-release + return "", nil + }, + } + + done := make(chan error, 1) + go func() { done <- sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump) }() + + for range 2 { + select { + case <-started: + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for both index restores to start: they are not running concurrently") + } + } + close(release) + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for the restore to finish") + } + }) + + t.Run("a failing index does not stop the other indices or the dependent statements", func(t *testing.T) { + t.Parallel() + + errTest := errors.New("oh noes") + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + if string(gotDump) == indexACall { + return "", errTest + } + return "", nil + }, + } + + err := sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump) + require.ErrorIs(t, err, errTest) + + // cancelling the sibling restores would orphan the index builds they + // have already started on the server, and skipping the dependent + // statements would leave the target without its constraints + recorded := recorder.recorded() + require.Len(t, recorded, 3) + require.ElementsMatch(t, []string{indexACall, indexBCall}, recorded[:2]) + require.Equal(t, dependentsCall, recorded[2]) + }) + + t.Run("a wave of transient failures stays retryable", func(t *testing.T) { + t.Parallel() + + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + if isIndexStatement(strings.TrimPrefix(string(gotDump), testConnectBlock+"\n\n")) { + return "", transientErr() + } + return "", nil + }, + } + + err := sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump) + require.Error(t, err) + + // the merged error decides the retry for the whole wave, rather than + // whichever worker's error arrived first + restoreErrs := &pglib.PGRestoreErrors{} + require.ErrorAs(t, err, &restoreErrs) + require.True(t, restoreErrs.IsRetryable()) + require.Len(t, restoreErrs.GetRetryableErrors(), 2) + }) + + t.Run("a critical failure alongside a transient one is not retryable", func(t *testing.T) { + t.Parallel() + + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + switch string(gotDump) { + case indexACall: + return "", transientErr() + case indexBCall: + return "", pglib.NewPGRestoreErrors(errors.New("oh noes")) + default: + return "", nil + } + }, + } + + err := sg.restoreIndexDump(context.Background(), pglib.PGRestoreOptions{}, dump) + require.Error(t, err) + + restoreErrs := &pglib.PGRestoreErrors{} + require.ErrorAs(t, err, &restoreErrs) + require.False(t, restoreErrs.IsRetryable()) + require.True(t, restoreErrs.HasCriticalErrors()) + }) + + t.Run("a cancelled context stops the restore before the dependent statements", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + recorder := &restoreRecorder{} + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + err := sg.restoreIndexDump(ctx, pglib.PGRestoreOptions{}, dump) + require.ErrorIs(t, err, context.Canceled) + require.Empty(t, recorder.recorded()) + }) + + t.Run("the progress tracking path restores every block too", func(t *testing.T) { + t.Parallel() + + recorder := &restoreRecorder{} + tracked := make(chan struct{}) + sg := &SnapshotGenerator{ + logger: log.NewNoopLogger(), + indexRestoreWorkers: 2, + snapshotTracker: &mockSnapshotTracker{ + trackIndexesCreationFn: func(ctx context.Context) { + close(tracked) + <-ctx.Done() + }, + }, + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + recorder.record(gotDump) + return "", nil + }, + } + + require.NoError(t, sg.restoreIndices(context.Background(), pglib.PGRestoreOptions{}, dump)) + <-tracked + + recorded := recorder.recorded() + require.Len(t, recorded, 3) + require.ElementsMatch(t, []string{indexACall, indexBCall}, recorded[:2]) + require.Equal(t, dependentsCall, recorded[2]) + }) +} + +func TestConfig_indexRestoreWorkers(t *testing.T) { + t.Parallel() + + require.Equal(t, uint(1), (&Config{}).indexRestoreWorkers()) + require.Equal(t, uint(3), (&Config{IndexRestoreWorkers: 3}).indexRestoreWorkers()) +} + +func TestNewSnapshotGenerator_indexRestoreWorkersValidation(t *testing.T) { + t.Parallel() + + // the worker count is validated before any connection is opened, so an + // unusable value fails at startup instead of at the index restore, hours + // into a snapshot + _, err := NewSnapshotGenerator(context.Background(), &Config{ + SourcePGURL: "postgres://not-used", + IndexRestoreWorkers: maxIndexRestoreWorkers + 1, + }) + require.ErrorIs(t, err, errTooManyIndexRestoreWorkers) + + // a negative value that wrapped around on its way through the config + // decoder is caught by the same cap + _, err = NewSnapshotGenerator(context.Background(), &Config{ + SourcePGURL: "postgres://not-used", + IndexRestoreWorkers: ^uint(0), + }) + require.ErrorIs(t, err, errTooManyIndexRestoreWorkers) +} + +// TestSnapshotGenerator_restoreIndicesAndConstraints_parallel covers the +// interaction between the concurrent index restore and the retry wrapper: a +// transient failure in any worker has to be classified as retryable for the +// whole wave, and the rerun has to succeed once the objects the first attempt +// did create come back as "already exists". +func TestSnapshotGenerator_restoreIndicesAndConstraints_parallel(t *testing.T) { + t.Parallel() + + dump := testIndexDump(testIndexA, testIndexB, testDependent1) + indexACall := string(testIndexDump(testIndexA)) + indexBCall := string(testIndexDump(testIndexB)) + + attempts := atomic.Int32{} + sg := SnapshotGenerator{ + targetURL: "target-url", + logger: log.NewNoopLogger(), + optionGenerator: &optionGenerator{targetURL: "target-url"}, + indexRestoreWorkers: 2, + indexRestoreBackoff: backoff.NewProvider(&backoff.Config{ + Constant: &backoff.ConstantConfig{ + Interval: time.Millisecond, + MaxRetries: 2, + }, + }), + pgRestoreFn: func(_ context.Context, _ pglib.PGRestoreOptions, gotDump []byte) (string, error) { + got := string(gotDump) + if got != indexACall && got != indexBCall { + return "", nil + } + // the first attempt loses its connection on both indices, the + // second finds idx_a already there from the first one + if attempts.Add(1) <= 2 { + return "", pglib.NewPGRestoreErrors(&pglib.ErrTransientFailure{ + Details: "psql: error: connection to server was lost", + }) + } + return "", pglib.NewPGRestoreErrors(&pglib.ErrRelationAlreadyExists{ + Details: `ERROR: relation "idx_a" already exists`, + }) + }, + } + + require.NoError(t, sg.restoreIndicesAndConstraints(context.Background(), dump, &snapshot.Snapshot{})) + // both indices on the first attempt, both again on the retry + require.Equal(t, int32(4), attempts.Load()) +} diff --git a/pkg/snapshot/generator/postgres/schema/pgdumprestore/testdata/exotic_schema_dump.sql b/pkg/snapshot/generator/postgres/schema/pgdumprestore/testdata/exotic_schema_dump.sql new file mode 100644 index 000000000..5a0dfda86 --- /dev/null +++ b/pkg/snapshot/generator/postgres/schema/pgdumprestore/testdata/exotic_schema_dump.sql @@ -0,0 +1,186 @@ +-- +-- PostgreSQL database dump +-- + +\restrict ddneeN6Oxl0ecekB9eBT4Gf3pl0KK8idsoGfccvrHpbmW08qS8rI99Mb8UlV0pK + + +SET statement_timeout = 0; +SET lock_timeout = 0; +SET idle_in_transaction_session_timeout = 0; +SET transaction_timeout = 0; +SET client_encoding = 'UTF8'; +SET standard_conforming_strings = on; +SELECT pg_catalog.set_config('search_path', '', false); +SET check_function_bodies = false; +SET xmloption = content; +SET client_min_messages = warning; +SET row_security = off; + +-- +-- Name: btree_gist; Type: EXTENSION; Schema: -; Owner: - +-- + +CREATE EXTENSION IF NOT EXISTS btree_gist WITH SCHEMA public; + + +-- +-- Name: EXTENSION btree_gist; Type: COMMENT; Schema: -; Owner: +-- + +COMMENT ON EXTENSION btree_gist IS 'support for indexing common datatypes in GiST'; + + +SET default_tablespace = ''; + +-- +-- Name: parts; Type: TABLE; Schema: public; Owner: postgres +-- + +CREATE TABLE public.parts ( + id integer NOT NULL, + val text +) +PARTITION BY RANGE (id); + + +ALTER TABLE public.parts OWNER TO postgres; + +SET default_table_access_method = heap; + +-- +-- Name: parts_0; Type: TABLE; Schema: public; Owner: postgres +-- + +CREATE TABLE public.parts_0 ( + id integer NOT NULL, + val text +); + + +ALTER TABLE public.parts_0 OWNER TO postgres; + +-- +-- Name: parts_1; Type: TABLE; Schema: public; Owner: postgres +-- + +CREATE TABLE public.parts_1 ( + id integer NOT NULL, + val text +); + + +ALTER TABLE public.parts_1 OWNER TO postgres; + +-- +-- Name: rooms; Type: TABLE; Schema: public; Owner: postgres +-- + +CREATE TABLE public.rooms ( + id integer NOT NULL, + room integer, + during tsrange, + name text, + email text, + val numeric, + CONSTRAINT rooms_room_positive CHECK ((room >= 0)), + CONSTRAINT rooms_val_positive CHECK ((val >= (0)::numeric)) +); + + +ALTER TABLE public.rooms OWNER TO postgres; + +-- +-- Name: parts_0; Type: TABLE ATTACH; Schema: public; Owner: postgres +-- + +ALTER TABLE ONLY public.parts ATTACH PARTITION public.parts_0 FOR VALUES FROM (0) TO (10); + + +-- +-- Name: parts_1; Type: TABLE ATTACH; Schema: public; Owner: postgres +-- + +ALTER TABLE ONLY public.parts ATTACH PARTITION public.parts_1 FOR VALUES FROM (10) TO (20); + + +-- +-- Name: rooms rooms_email_key; Type: CONSTRAINT; Schema: public; Owner: postgres +-- + +ALTER TABLE ONLY public.rooms + ADD CONSTRAINT rooms_email_key UNIQUE (email); + +ALTER TABLE public.rooms CLUSTER ON rooms_email_key; + + +-- +-- Name: rooms rooms_no_overlap; Type: CONSTRAINT; Schema: public; Owner: postgres +-- + +ALTER TABLE ONLY public.rooms + ADD CONSTRAINT rooms_no_overlap EXCLUDE USING gist (room WITH =, during WITH &&); + + +-- +-- Name: parts_val_idx; Type: INDEX; Schema: public; Owner: postgres +-- + +CREATE INDEX parts_val_idx ON ONLY public.parts USING btree (val); + + +-- +-- Name: parts_0_val_idx; Type: INDEX; Schema: public; Owner: postgres +-- + +CREATE INDEX parts_0_val_idx ON public.parts_0 USING btree (val); + + +-- +-- Name: parts_1_val_idx; Type: INDEX; Schema: public; Owner: postgres +-- + +CREATE INDEX parts_1_val_idx ON public.parts_1 USING btree (val); + + +-- +-- Name: rooms_expr_idx; Type: INDEX; Schema: public; Owner: postgres +-- + +CREATE INDEX rooms_expr_idx ON public.rooms USING btree (lower(name), upper(email), ((val * (2)::numeric)), COALESCE(name, email, 'a-fairly-long-default-value'::text)); + + +-- +-- Name: rooms_partial_idx; Type: INDEX; Schema: public; Owner: postgres +-- + +CREATE INDEX rooms_partial_idx ON public.rooms USING btree (name) WHERE ((val > (100)::numeric) AND (name IS NOT NULL) AND (email IS NOT NULL)); + + +-- +-- Name: INDEX rooms_partial_idx; Type: COMMENT; Schema: public; Owner: postgres +-- + +COMMENT ON INDEX public.rooms_partial_idx IS 'partial'; + + +-- +-- Name: parts_0_val_idx; Type: INDEX ATTACH; Schema: public; Owner: postgres +-- + +ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_0_val_idx; + + +-- +-- Name: parts_1_val_idx; Type: INDEX ATTACH; Schema: public; Owner: postgres +-- + +ALTER INDEX public.parts_val_idx ATTACH PARTITION public.parts_1_val_idx; + + +-- +-- PostgreSQL database dump complete +-- + +\unrestrict ddneeN6Oxl0ecekB9eBT4Gf3pl0KK8idsoGfccvrHpbmW08qS8rI99Mb8UlV0pK + diff --git a/pkg/stream/integration/snapshot_pg_integration_test.go b/pkg/stream/integration/snapshot_pg_integration_test.go index d9717e47a..d762581b2 100644 --- a/pkg/stream/integration/snapshot_pg_integration_test.go +++ b/pkg/stream/integration/snapshot_pg_integration_test.go @@ -624,6 +624,113 @@ func Test_SnapshotToPostgres_MaterializedViewRefresh(t *testing.T) { require.Contains(t, indexes, indexName) } +// Test_SnapshotToPostgres_ParallelIndexRestore verifies that, when the schema +// snapshot restores indexes concurrently (IndexRestoreWorkers > 1), every +// standalone index is still created and everything that depends on an index +// existing is restored correctly afterwards: a constraint added USING INDEX, +// a comment on an index, and a partitioned index attached to its parent via +// ALTER INDEX ... ATTACH PARTITION. +func Test_SnapshotToPostgres_ParallelIndexRestore(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + var snapshotPGURL string + pgcleanup, err := testcontainers.SetupPostgresContainer(context.Background(), &snapshotPGURL, testcontainers.Postgres14, "config/postgresql.conf") + require.NoError(t, err) + defer pgcleanup() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + suffix := time.Now().UnixNano() + testTable := fmt.Sprintf("parallel_idx_%d", suffix) + uniqueIndex := fmt.Sprintf("parallel_idx_ui_%d", suffix) + otherIndex1 := fmt.Sprintf("parallel_idx_i2_%d", suffix) + otherIndex2 := fmt.Sprintf("parallel_idx_i3_%d", suffix) + exprIndex := fmt.Sprintf("parallel_idx_expr_%d", suffix) + uniqueConstraint := fmt.Sprintf("parallel_idx_uc_%d", suffix) + indexComment := "parallel restore test comment" + + partitionedTable := fmt.Sprintf("parallel_part_%d", suffix) + partitionedIndex := fmt.Sprintf("parallel_part_idx_%d", suffix) + + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE TABLE %s( + id serial PRIMARY KEY, + val text NOT NULL, + val2 text NOT NULL, + val3 integer NOT NULL + )`, testTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `INSERT INTO %s(val, val2, val3) VALUES ('a', 'x', 1), ('b', 'y', 2)`, testTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE UNIQUE INDEX %s ON %s(val)`, uniqueIndex, testTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE INDEX %s ON %s(val2)`, otherIndex1, testTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE INDEX %s ON %s(val3)`, otherIndex2, testTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE INDEX %s ON %s(lower(val))`, exprIndex, testTable)) + // depends on the unique index above already existing. This also renames + // the index to the constraint name. + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `ALTER TABLE %s ADD CONSTRAINT %s UNIQUE USING INDEX %s`, testTable, uniqueConstraint, uniqueIndex)) + // depends on the index above already existing + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `COMMENT ON INDEX %s IS '%s'`, otherIndex1, indexComment)) + + // a partitioned index: pg_dump represents it as a parent CREATE INDEX ON + // ONLY, one CREATE INDEX per partition, and an ALTER INDEX ... ATTACH + // PARTITION per partition that depends on both indexes already existing. + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE TABLE %s(id integer NOT NULL, val text) PARTITION BY RANGE (id)`, partitionedTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE TABLE %s_0 PARTITION OF %s FOR VALUES FROM (0) TO (100)`, partitionedTable, partitionedTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE TABLE %s_1 PARTITION OF %s FOR VALUES FROM (100) TO (200)`, partitionedTable, partitionedTable)) + execQueryWithURL(t, ctx, snapshotPGURL, fmt.Sprintf( + `CREATE INDEX %s ON %s(val)`, partitionedIndex, partitionedTable)) + + cfg := &stream.Config{ + Listener: testPostgresListenerCfgWithSnapshot(snapshotPGURL, targetPGURL, []string{"public.*"}), + Processor: testPostgresProcessorCfg(), + } + cfg.Listener.Postgres.Snapshot.Schema.DumpRestore.IndexRestoreWorkers = 4 + require.NoError(t, stream.Snapshot(ctx, testLogger(), cfg, nil)) + + targetConn, err := pglib.NewConn(ctx, targetPGURL) + require.NoError(t, err) + defer targetConn.Close(ctx) + + indexes := getTableIndexes(t, ctx, targetConn, "public", testTable) + // the unique index was renamed to the constraint name by ADD CONSTRAINT + // ... UNIQUE USING INDEX + require.Contains(t, indexes, uniqueConstraint) + require.Contains(t, indexes, otherIndex1) + require.Contains(t, indexes, otherIndex2) + require.Contains(t, indexes, exprIndex) + + var constraintType string + err = targetConn.QueryRow(ctx, []any{&constraintType}, ` + SELECT contype FROM pg_constraint WHERE conname = $1 + `, uniqueConstraint) + require.NoError(t, err) + require.Equal(t, "u", constraintType) + + var gotComment string + err = targetConn.QueryRow(ctx, []any{&gotComment}, fmt.Sprintf( + `SELECT obj_description('%s'::regclass, 'pg_class')`, otherIndex1)) + require.NoError(t, err) + require.Equal(t, indexComment, gotComment) + + var attachedPartitions int + err = targetConn.QueryRow(ctx, []any{&attachedPartitions}, fmt.Sprintf( + `SELECT count(*) FROM pg_inherits WHERE inhparent = '%s'::regclass`, partitionedIndex)) + require.NoError(t, err) + require.Equal(t, 2, attachedPartitions) +} + // Test_SnapshotToPostgres_SkipsLegacyPLPGSQLHandlers verifies that legacy // public PL/pgSQL handler functions from a source dump are not restored as // ordinary user functions, while normal user-defined functions are preserved.