diff --git a/pkg/snapshot/generator/postgres/data/ctid_table_reader.go b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go new file mode 100644 index 00000000..3f2367ce --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + "github.com/jackc/pgx/v5/pgtype" + pglib "github.com/xataio/pgstream/internal/postgres" + loglib "github.com/xataio/pgstream/pkg/log" + "github.com/xataio/pgstream/pkg/snapshot" + "github.com/xataio/pgstream/pkg/wal/processor" + "golang.org/x/sync/errgroup" +) + +// ctidReader reads a schema's tables by ranging over a stable transaction +// snapshot using the ctid. The transaction snapshot is exported once per schema +// in beginSchema and imported by every page range transaction, which lets the +// reader parallelise the work across page ranges while keeping a consistent +// view of each table. +type ctidReader struct { + conn pglib.Querier + logger loglib.Logger + adapter *adapter + processor processor.Processor + tableWorkers uint + batchBytes uint64 + + // progress is shared by value with the snapshot generator; the underlying + // bars map is shared by reference. + progress progressTracker +} + +// beginSchema opens the transaction that exports the shared transaction +// snapshot and keeps it open for the duration of fn, so that every readTable +// call can import it. The snapshot is only exported when the schema has at least +// one ctid table to read. +func (r *ctidReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error { + // use a transaction snapshot to ensure the table rows can be parallelised. + // The transaction snapshot is available for use only until the end of the + // transaction that exported it. + // https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION + return r.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error { + session := &readSession{} + if len(st.tables) > 0 { + snapshotID, err := exportSnapshot(ctx, tx) + if err != nil { + return snapshot.NewSchemaErrors(st.schema, err) + } + session.snapshotID = snapshotID + } + + return fn(ctx, session) + }, snapshotTxOptions()) +} + +func (r *ctidReader) readTable(ctx context.Context, session *readSession, table *table) error { + snapshotID := session.snapshotID + tableInfo, err := r.getTableInfo(ctx, table.schema, table.name, snapshotID) + if err != nil { + return err + } + if tableInfo.isEmpty() { + return nil + } + table.rowSize = tableInfo.avgRowBytes + + // If one page range fails, we abort the entire table snapshot. The + // snapshot relies on the transaction snapshot id to ensure all workers + // have the same table view, which allows us to use the ctid to + // parallelise the work. + rangeChan := make(chan pageRange, tableInfo.pageCount) + errGroup, ctx := errgroup.WithContext(ctx) + for i := uint(0); i < r.tableWorkers; i++ { + errGroup.Go(func() error { + return r.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan) + }) + } + + // page count returned by postgres starts at 0, so we need to include it + // when creating the page ranges. + for start := uint(0); start <= uint(tableInfo.pageCount); start += tableInfo.batchPageSize { + rangeChan <- pageRange{ + start: start, + end: start + tableInfo.batchPageSize, + } + } + + // wait for all table ranges to complete + close(rangeChan) + return errGroup.Wait() +} + +func (r *ctidReader) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan pageRange) error { + for pageRange := range pageRangeChan { + if err := r.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil { + return err + } + } + return nil +} + +var pageRangeQuery = "SELECT * FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'" + +func (r *ctidReader) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error { + return execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error { + r.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{ + "schema": table.schema, "table": table.name, "snapshotID": snapshotID, + }) + + query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(table.schema, table.name), pageRange.start, pageRange.end) + rows, err := tx.Query(ctx, query) + if err != nil { + return fmt.Errorf("querying table rows: %w", err) + } + defer rows.Close() + + // resolve the column metadata (names/types) and timestamp once per page + // range, since the field descriptions are identical for every row in the + // result set. + rowAdapter := r.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions()) + rowCount := uint(0) + for rows.Next() { + rowCount++ + select { + case <-ctx.Done(): + return ctx.Err() + default: + values, err := rows.Values() + if err != nil { + return fmt.Errorf("retrieving rows values: %w", err) + } + + event := rowAdapter.rowToWalEvent(values) + if event == nil { + continue + } + + if err := r.processor.ProcessWALEvent(ctx, event); err != nil { + return fmt.Errorf("processing snapshot row: %w", err) + } + } + } + + r.progress.advance(table.schema, int64(rowCount)*table.rowSize) + + r.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{ + "schema": table.schema, "table": table.name, "snapshotID": snapshotID, + }) + + return rows.Err() + }) +} + +const ( + // use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices. + // pg_relation_size will return only the size of the table itself, without toast tables. + tableInfoQuery = `SELECT + (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes, + CASE + WHEN c.reltuples > 0 THEN + ROUND(pg_table_size(c.oid) / c.reltuples) + ELSE + 0 + END AS avg_row_size +FROM + pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace +WHERE + c.relname = $1 + AND n.nspname = $2 + AND c.relkind = 'r';` + + // select the max page for the relation instead of using pg_class.relpages, it may not contain an accurate value if + // the table is small, the table has active inserts, or the database has not been vacuumed/analyzed recently. + maxPageQuery = `SELECT MAX(ctid) FROM ONLY %s;` +) + +func (r *ctidReader) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) { + tableInfo := &tableInfo{} + err := execInSnapshotTx(ctx, r.conn, snapshotID, func(tx pglib.Tx) error { + // make sure the schema and table names are unquoted since the system + // catalogs store unquoted names + err := tx.QueryRow(ctx, + []any{&tableInfo.avgPageBytes, &tableInfo.avgRowBytes}, + tableInfoQuery, + pglib.UnquoteIdentifier(tableName), + pglib.UnquoteIdentifier(schemaName)) + if err != nil { + return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err) + } + + var ctid pgtype.TID + if err := tx.QueryRow(ctx, []any{&ctid}, fmt.Sprintf(maxPageQuery, pglib.QuoteQualifiedIdentifier(schemaName, tableName))); err != nil { + return fmt.Errorf("getting max page for table %s.%s: %w", schemaName, tableName, err) + } + tableInfo.pageCount = int(ctid.BlockNumber) + + tableInfo.calculateBatchPageSize(r.batchBytes) + + r.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{ + "schema": schemaName, "table": tableName, "snapshotID": snapshotID, + }) + return nil + }) + if err != nil { + return nil, err + } + + return tableInfo, nil +} diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go new file mode 100644 index 00000000..dc12335b --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + + "github.com/xataio/pgstream/pkg/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +type instrumentedTableReader struct { + tracer trace.Tracer + reader tableReader +} + +func newInstrumentedTableReader(reader tableReader, i *otel.Instrumentation) *instrumentedTableReader { + return &instrumentedTableReader{ + tracer: i.Tracer, + reader: reader, + } +} + +func (i *instrumentedTableReader) beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error { + return i.reader.beginSchema(ctx, st, fn) +} + +func (i *instrumentedTableReader) readTable(ctx context.Context, session *readSession, table *table) (err error) { + ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.ReadTable", trace.WithAttributes([]attribute.KeyValue{ + {Key: "schema", Value: attribute.StringValue(table.schema)}, + {Key: "table", Value: attribute.StringValue(table.name)}, + }...)) + defer otel.CloseSpan(span, err) + return i.reader.readTable(ctx, session, table) +} diff --git a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go deleted file mode 100644 index 6ed5e39f..00000000 --- a/pkg/snapshot/generator/postgres/data/instrumented_table_snapshot_generator.go +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 - -package postgres - -import ( - "context" - - "github.com/xataio/pgstream/pkg/otel" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/trace" -) - -type instrumentedTableSnapshotGenerator struct { - tracer trace.Tracer - snapshotTableFn snapshotTableFn -} - -func newInstrumentedTableSnapshotGenerator(fn snapshotTableFn, i *otel.Instrumentation) *instrumentedTableSnapshotGenerator { - return &instrumentedTableSnapshotGenerator{ - tracer: i.Tracer, - snapshotTableFn: fn, - } -} - -func (i *instrumentedTableSnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) (err error) { - ctx, span := otel.StartSpan(ctx, i.tracer, "tableSnapshotGenerator.SnapshotTable", trace.WithAttributes([]attribute.KeyValue{ - {Key: "schema", Value: attribute.StringValue(table.schema)}, - {Key: "table", Value: attribute.StringValue(table.name)}, - }...)) - defer otel.CloseSpan(span, err) - return i.snapshotTableFn(ctx, snapshotID, table) -} diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go index 47e95abe..20b8e00c 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go @@ -8,11 +8,9 @@ import ( "fmt" "sync" - "github.com/jackc/pgx/v5/pgtype" pglib "github.com/xataio/pgstream/internal/postgres" pglibinstrumentation "github.com/xataio/pgstream/internal/postgres/instrumentation" "github.com/xataio/pgstream/internal/progress" - synclib "github.com/xataio/pgstream/internal/sync" loglib "github.com/xataio/pgstream/pkg/log" "github.com/xataio/pgstream/pkg/otel" "github.com/xataio/pgstream/pkg/snapshot" @@ -21,24 +19,22 @@ import ( ) type SnapshotGenerator struct { - logger loglib.Logger - conn pglib.Querier - adapter *adapter + logger loglib.Logger + conn pglib.Querier + processor processor.Processor + // reader encapsulates the strategy used to read a schema's tables (ctid + // range scan by default). + reader tableReader + // instrumentation is captured while applying options and used to wrap the + // reader once it has been built. + instrumentation *otel.Instrumentation // workers per snapshot, parallelise the snapshot creation for each schema snapshotWorkers uint // workers per schema, parallelise the snapshot creation for each table schemaWorkers uint - // workers per table, parallelise the snapshot creation for each page range - tableWorkers uint - batchBytes uint64 - // Function called for processing produced rows. - processor processor.Processor - tableSnapshotGenerator snapshotTableFn - - progressTracking bool - progressBars *synclib.Map[string, progress.Bar] + progress progressTracker progressBarBuilder func(totalBytes int64, description string) progress.Bar } @@ -69,8 +65,6 @@ type table struct { rowSize int64 } -type snapshotTableFn func(ctx context.Context, snapshotID string, table *table) error - type Option func(sg *SnapshotGenerator) func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor.Processor, opts ...Option) (*SnapshotGenerator, error) { @@ -87,19 +81,27 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, processor processor. logger: loglib.NewNoopLogger(), conn: conn, processor: processor, - batchBytes: cfg.batchBytes(), - tableWorkers: cfg.tableWorkers(), schemaWorkers: cfg.schemaWorkers(), snapshotWorkers: cfg.snapshotWorkers(), } - sg.tableSnapshotGenerator = sg.snapshotTable - for _, opt := range opts { opt(sg) } - sg.adapter = newAdapter(pglib.NewMapper(conn), sg.logger) + sg.reader = &ctidReader{ + conn: sg.conn, + logger: sg.logger, + adapter: newAdapter(pglib.NewMapper(conn), sg.logger), + processor: sg.processor, + tableWorkers: cfg.tableWorkers(), + batchBytes: cfg.batchBytes(), + progress: sg.progress, + } + + if sg.instrumentation != nil { + sg.reader = newInstrumentedTableReader(sg.reader, sg.instrumentation) + } return sg, nil } @@ -121,15 +123,13 @@ func WithInstrumentation(i *otel.Instrumentation) Option { panic(err) } - ig := newInstrumentedTableSnapshotGenerator(sg.tableSnapshotGenerator, i) - sg.tableSnapshotGenerator = ig.snapshotTable + sg.instrumentation = i } } func WithProgressTracking() Option { return func(sg *SnapshotGenerator) { - sg.progressTracking = true - sg.progressBars = synclib.NewMap[string, progress.Bar]() + sg.progress = newProgressTracker() sg.progressBarBuilder = progress.NewBytesBar } } @@ -188,23 +188,14 @@ func (sg *SnapshotGenerator) Close() error { } func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTables *schemaTables) error { - // use a transaction snapshot to ensure the table rows can be parallelised. - // The transaction snapshot is available for use only until the end of the - // transaction that exported it. - // https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION - return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) (err error) { - snapshotID, err := sg.exportSnapshot(ctx, tx) - if err != nil { - return snapshot.NewSchemaErrors(schemaTables.schema, err) - } - - if sg.progressTracking { - if err := sg.addProgressBar(ctx, snapshotID, schemaTables); err != nil { + return sg.reader.beginSchema(ctx, schemaTables, func(ctx context.Context, session *readSession) (err error) { + if sg.progress.enabled { + if err := sg.addProgressBar(ctx, session.snapshotID, schemaTables); err != nil { return err } defer func() { if err == nil { - sg.markProgressBarCompleted(schemaTables.schema) + sg.progress.complete(schemaTables.schema) } }() } @@ -217,7 +208,7 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab for i := uint(0); i < sg.schemaWorkers; i++ { wg.Add(1) workerTableErrs[i] = make(map[string]error, len(schemaTables.tables)) - go sg.createSnapshotWorker(ctx, wg, snapshotID, tableChan, workerTableErrs[i]) + go sg.createSnapshotWorker(ctx, wg, session, tableChan, workerTableErrs[i]) } for _, tableName := range schemaTables.tables { @@ -231,16 +222,16 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab wg.Wait() return sg.collectTableErrors(schemaTables.schema, workerTableErrs) - }, snapshotTxOptions()) + }) } -func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, snapshotID string, tableChan <-chan *table, tableErrMap map[string]error) { +func (sg *SnapshotGenerator) createSnapshotWorker(ctx context.Context, wg *sync.WaitGroup, session *readSession, tableChan <-chan *table, tableErrMap map[string]error) { defer wg.Done() for t := range tableChan { - logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": snapshotID} + logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": session.snapshotID} sg.logger.Debug("snapshotting table", logFields) - if err := sg.tableSnapshotGenerator(ctx, snapshotID, t); err != nil { + if err := sg.reader.readTable(ctx, session, t); err != nil { sg.logger.Error(err, "snapshotting table", logFields) // errors will get notified unless the table doesn't exist if !errors.Is(err, pglib.ErrNoRows) { @@ -288,108 +279,6 @@ func (sg *SnapshotGenerator) collectSchemaErrors(workerSchemaErrs map[string]err return nil } -func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, table *table) error { - tableInfo, err := sg.getTableInfo(ctx, table.schema, table.name, snapshotID) - if err != nil { - return err - } - if tableInfo.isEmpty() { - return nil - } - table.rowSize = tableInfo.avgRowBytes - - // If one page range fails, we abort the entire table snapshot. The - // snapshot relies on the transaction snapshot id to ensure all workers - // have the same table view, which allows us to use the ctid to - // parallelise the work. - rangeChan := make(chan pageRange, tableInfo.pageCount) - errGroup, ctx := errgroup.WithContext(ctx) - for i := uint(0); i < sg.tableWorkers; i++ { - errGroup.Go(func() error { - return sg.snapshotTableRangeWorker(ctx, snapshotID, table, rangeChan) - }) - } - - // page count returned by postgres starts at 0, so we need to include it - // when creating the page ranges. - for start := uint(0); start <= uint(tableInfo.pageCount); start += tableInfo.batchPageSize { - rangeChan <- pageRange{ - start: start, - end: start + tableInfo.batchPageSize, - } - } - - // wait for all table ranges to complete - close(rangeChan) - return errGroup.Wait() -} - -func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snapshotID string, table *table, pageRangeChan <-chan pageRange) error { - for pageRange := range pageRangeChan { - if err := sg.snapshotTableRange(ctx, snapshotID, table, pageRange); err != nil { - return err - } - } - return nil -} - -var pageRangeQuery = "SELECT * FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'" - -func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID string, table *table, pageRange pageRange) error { - return sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error { - sg.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{ - "schema": table.schema, "table": table.name, "snapshotID": snapshotID, - }) - - query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(table.schema, table.name), pageRange.start, pageRange.end) - rows, err := tx.Query(ctx, query) - if err != nil { - return fmt.Errorf("querying table rows: %w", err) - } - defer rows.Close() - - // resolve the column metadata (names/types) and timestamp once per page - // range, since the field descriptions are identical for every row in the - // result set. - rowAdapter := sg.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions()) - rowCount := uint(0) - for rows.Next() { - rowCount++ - select { - case <-ctx.Done(): - return ctx.Err() - default: - values, err := rows.Values() - if err != nil { - return fmt.Errorf("retrieving rows values: %w", err) - } - - event := rowAdapter.rowToWalEvent(values) - if event == nil { - continue - } - - if err := sg.processor.ProcessWALEvent(ctx, event); err != nil { - return fmt.Errorf("processing snapshot row: %w", err) - } - } - } - - if sg.progressTracking { - bar, found := sg.progressBars.Get(table.schema) - if found { - bar.Add64(int64(rowCount) * table.rowSize) - } - } - - sg.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{ - "schema": table.schema, "table": table.name, "snapshotID": snapshotID, - }) - - return rows.Err() - }) -} - func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID string, schemaTables *schemaTables) error { totalBytes, err := sg.getSnapshotSchemaTotalBytes(ctx, snapshotID, schemaTables.schema, schemaTables.tables) if err != nil { @@ -397,76 +286,10 @@ func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, snapshotID stri } bar := sg.progressBarBuilder(totalBytes, fmt.Sprintf("[cyan][%s][reset] Snapshotting data...", schemaTables.schema)) - sg.progressBars.Set(schemaTables.schema, bar) + sg.progress.set(schemaTables.schema, bar) return nil } -func (sg *SnapshotGenerator) markProgressBarCompleted(schema string) { - bar, found := sg.progressBars.Get(schema) - if found { - bar.Close() - } - sg.progressBars.Delete(schema) -} - -const ( - // use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices. - // pg_relation_size will return only the size of the table itself, without toast tables. - tableInfoQuery = `SELECT - (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes, - CASE - WHEN c.reltuples > 0 THEN - ROUND(pg_table_size(c.oid) / c.reltuples) - ELSE - 0 - END AS avg_row_size -FROM - pg_class c - JOIN pg_namespace n ON n.oid = c.relnamespace -WHERE - c.relname = $1 - AND n.nspname = $2 - AND c.relkind = 'r';` - - // select the max page for the relation instead of using pg_class.relpages, it may not contain an accurate value if - // the table is small, the table has active inserts, or the database has not been vacuumed/analyzed recently. - maxPageQuery = `SELECT MAX(ctid) FROM ONLY %s;` -) - -func (sg *SnapshotGenerator) getTableInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tableInfo, error) { - tableInfo := &tableInfo{} - err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error { - // make sure the schema and table names are unquoted since the system - // catalogs store unquoted names - err := tx.QueryRow(ctx, - []any{&tableInfo.avgPageBytes, &tableInfo.avgRowBytes}, - tableInfoQuery, - pglib.UnquoteIdentifier(tableName), - pglib.UnquoteIdentifier(schemaName)) - if err != nil { - return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err) - } - - var ctid pgtype.TID - if err := tx.QueryRow(ctx, []any{&ctid}, fmt.Sprintf(maxPageQuery, pglib.QuoteQualifiedIdentifier(schemaName, tableName))); err != nil { - return fmt.Errorf("getting max page for table %s.%s: %w", schemaName, tableName, err) - } - tableInfo.pageCount = int(ctid.BlockNumber) - - tableInfo.calculateBatchPageSize(sg.batchBytes) - - sg.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{ - "schema": schemaName, "table": tableName, "snapshotID": snapshotID, - }) - return nil - }) - if err != nil { - return nil, err - } - - return tableInfo, nil -} - const tablesBytesQuery = `SELECT SUM(pg_table_size(c.oid)) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = ANY($2) AND c.relkind = 'r';` func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, snapshotID, schema string, tables []string) (int64, error) { @@ -482,7 +305,7 @@ func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, sn unquotedTables[i] = pglib.UnquoteIdentifier(table) } - err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error { + err := execInSnapshotTx(ctx, sg.conn, snapshotID, func(tx pglib.Tx) error { err := tx.QueryRow(ctx, []any{&totalBytes}, tablesBytesQuery, pglib.UnquoteIdentifier(schema), unquotedTables) if err != nil { return fmt.Errorf("retrieving total bytes for schema: %w", err) @@ -493,41 +316,6 @@ func (sg *SnapshotGenerator) getSnapshotSchemaTotalBytes(ctx context.Context, sn return totalBytes, err } -const exportSnapshotQuery = `SELECT pg_export_snapshot()` - -func (sg *SnapshotGenerator) exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) { - var snapshotID string - if err := tx.QueryRow(ctx, []any{&snapshotID}, exportSnapshotQuery); err != nil { - return "", fmt.Errorf("exporting snapshot: %w", err) - } - return snapshotID, nil -} - -func (sg *SnapshotGenerator) setTransactionSnapshot(ctx context.Context, tx pglib.Tx, snapshotID string) error { - _, err := tx.Exec(ctx, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", snapshotID)) - if err != nil { - return fmt.Errorf("setting transaction snapshot: %w", err) - } - return nil -} - -func (sg *SnapshotGenerator) execInSnapshotTx(ctx context.Context, snapshotID string, fn func(tx pglib.Tx) error) error { - return sg.conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error { - if err := sg.setTransactionSnapshot(ctx, tx, snapshotID); err != nil { - return err - } - - return fn(tx) - }, snapshotTxOptions()) -} - -func snapshotTxOptions() pglib.TxOptions { - return pglib.TxOptions{ - IsolationLevel: pglib.RepeatableRead, - AccessMode: pglib.ReadOnly, - } -} - // calculateBatchPageSize will automatically determine the batch page size based // on the average page size and the configured batch bytes limit. func (t *tableInfo) calculateBatchPageSize(bytes uint64) { diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go index bb3aeda4..8c67c22e 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go @@ -1132,32 +1132,42 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) { t.Parallel() eventChan := make(chan *wal.Event, 10) - sg := SnapshotGenerator{ - logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{ - LogLevel: "debug", - })), - conn: tc.querier, - adapter: newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()), - processor: &processormocks.Processor{ - ProcessWALEventFn: func(ctx context.Context, e *wal.Event) error { - eventChan <- e - return nil - }, - CloseFn: func() error { - return tc.processorCloseErr - }, + logger := zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{ + LogLevel: "debug", + })) + processor := &processormocks.Processor{ + ProcessWALEventFn: func(ctx context.Context, e *wal.Event) error { + eventChan <- e + return nil + }, + CloseFn: func() error { + return tc.processorCloseErr }, - schemaWorkers: 1, - tableWorkers: 1, - batchBytes: 1024 * 1024, // 1MB - snapshotWorkers: 1, - progressTracking: tc.progressBar != nil, - progressBars: synclib.NewMap[string, progress.Bar](), + } + pt := progressTracker{ + enabled: tc.progressBar != nil, + bars: synclib.NewMap[string, progress.Bar](), + } + sg := SnapshotGenerator{ + logger: logger, + conn: tc.querier, + processor: processor, + schemaWorkers: 1, + snapshotWorkers: 1, + progress: pt, progressBarBuilder: func(totalBytes int64, description string) progress.Bar { return tc.progressBar }, + reader: &ctidReader{ + conn: tc.querier, + logger: logger, + adapter: newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()), + processor: processor, + tableWorkers: 1, + batchBytes: 1024 * 1024, // 1MB + progress: pt, + }, } - sg.tableSnapshotGenerator = sg.snapshotTable if tc.schemaWorkers != 0 { sg.schemaWorkers = tc.schemaWorkers @@ -1667,7 +1677,7 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { }, } - sg := SnapshotGenerator{ + reader := &ctidReader{ logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{ LogLevel: "debug", })), @@ -1684,15 +1694,17 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { return nil }, }, - progressTracking: tc.name == "ok - with progress tracking", - progressBars: synclib.NewMap[string, progress.Bar](), + progress: progressTracker{ + enabled: tc.name == "ok - with progress tracking", + bars: synclib.NewMap[string, progress.Bar](), + }, } - if sg.progressTracking { - sg.progressBars.Set(tc.table.schema, progressBar) + if reader.progress.enabled { + reader.progress.set(tc.table.schema, progressBar) } - err := sg.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange) + err := reader.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange) require.Equal(t, tc.wantErr, err) close(eventChan) diff --git a/pkg/snapshot/generator/postgres/data/progress_tracker.go b/pkg/snapshot/generator/postgres/data/progress_tracker.go new file mode 100644 index 00000000..3d4e52a1 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/progress_tracker.go @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "github.com/xataio/pgstream/internal/progress" + synclib "github.com/xataio/pgstream/internal/sync" +) + +// progressTracker owns the per-schema progress bars shared between the snapshot +// generator, which creates and completes the bars, and the table readers, which +// advance them as rows are processed. The enabled flag and the bars map always +// travel together, so they are grouped here instead of being threaded as a pair +// through every constructor. +// +// The bars map is a pointer, so copies of a progressTracker value share the same +// underlying bars. A zero-value tracker is disabled and every method is a no-op, +// which lets callers advance progress without guarding each call. +type progressTracker struct { + enabled bool + bars *synclib.Map[string, progress.Bar] +} + +func newProgressTracker() progressTracker { + return progressTracker{ + enabled: true, + bars: synclib.NewMap[string, progress.Bar](), + } +} + +// set registers the bar tracking the given schema's progress. +func (p progressTracker) set(schema string, bar progress.Bar) { + if !p.enabled { + return + } + p.bars.Set(schema, bar) +} + +// advance adds the given number of bytes to the given schema's bar, if one is +// being tracked. +func (p progressTracker) advance(schema string, bytes int64) { + if !p.enabled { + return + } + if bar, found := p.bars.Get(schema); found { + bar.Add64(bytes) + } +} + +// complete closes the given schema's bar and stops tracking it. +func (p progressTracker) complete(schema string) { + if !p.enabled { + return + } + if bar, found := p.bars.Get(schema); found { + bar.Close() + } + p.bars.Delete(schema) +} diff --git a/pkg/snapshot/generator/postgres/data/snapshot_tx.go b/pkg/snapshot/generator/postgres/data/snapshot_tx.go new file mode 100644 index 00000000..cec143b0 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/snapshot_tx.go @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + pglib "github.com/xataio/pgstream/internal/postgres" +) + +// snapshot_tx.go groups the helpers used to read from a shared transaction +// snapshot. A transaction snapshot is exported once per schema and imported by +// every reader transaction so that all workers observe the same stable view of +// the database, which is what allows the ctid based reader to parallelise the +// work. +// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-SNAPSHOT-SYNCHRONIZATION + +const exportSnapshotQuery = `SELECT pg_export_snapshot()` + +func exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) { + var snapshotID string + if err := tx.QueryRow(ctx, []any{&snapshotID}, exportSnapshotQuery); err != nil { + return "", fmt.Errorf("exporting snapshot: %w", err) + } + return snapshotID, nil +} + +func setTransactionSnapshot(ctx context.Context, tx pglib.Tx, snapshotID string) error { + _, err := tx.Exec(ctx, fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s'", snapshotID)) + if err != nil { + return fmt.Errorf("setting transaction snapshot: %w", err) + } + return nil +} + +// execInSnapshotTx runs fn in a read only repeatable read transaction that +// imports the given transaction snapshot, so it observes the same view of the +// database as the transaction that exported it. +func execInSnapshotTx(ctx context.Context, conn pglib.Querier, snapshotID string, fn func(tx pglib.Tx) error) error { + return conn.ExecInTxWithOptions(ctx, func(tx pglib.Tx) error { + if err := setTransactionSnapshot(ctx, tx, snapshotID); err != nil { + return err + } + + return fn(tx) + }, snapshotTxOptions()) +} + +func snapshotTxOptions() pglib.TxOptions { + return pglib.TxOptions{ + IsolationLevel: pglib.RepeatableRead, + AccessMode: pglib.ReadOnly, + } +} diff --git a/pkg/snapshot/generator/postgres/data/table_reader.go b/pkg/snapshot/generator/postgres/data/table_reader.go new file mode 100644 index 00000000..196289b9 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/table_reader.go @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import "context" + +// readSession carries the per-schema state that a tableReader needs to snapshot +// individual tables. It is only valid for the duration of the beginSchema +// callback that produced it. +type readSession struct { + // snapshotID is the exported transaction snapshot shared by all of a + // schema's table read transactions. It is empty when the schema has no + // tables to read, in which case readTable must not be called (an empty id + // would make SET TRANSACTION SNAPSHOT fail). + snapshotID string +} + +// tableReader abstracts the strategy used to read a schema's tables during a +// data snapshot. +type tableReader interface { + // beginSchema prepares the reader to snapshot the given schema tables and + // invokes fn with the readSession that must be used for every readTable call + // belonging to that schema. + // The session is only valid for the duration of fn. + beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, *readSession) error) error + // readTable snapshots a single table using the session provided by beginSchema. + readTable(ctx context.Context, session *readSession, table *table) error +}