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..bb01a627 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/ctid_table_reader.go @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "errors" + "fmt" + "strings" + + "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" + "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 + sink rowSink + tableWorkers uint + batchBytes uint64 +} + +// beginSchema opens the transaction that exports the shared transaction +// snapshot and keeps it open for the duration of fn, so that every read the +// session performs 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 := &ctidSession{ + reader: r, + schemaTables: st, + } + 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()) +} + +// ctidSession reads the tables of a single schema over the transaction snapshot +// exported by beginSchema. Every transaction it opens imports that snapshot, so +// all the workers observe the same view of the schema, which is what makes it +// safe to split a table into page ranges read in parallel. +type ctidSession struct { + reader *ctidReader + schemaTables *schemaTables + // snapshotID is empty when the schema has no tables to read, in which case + // no read is performed: an empty id would make SET TRANSACTION SNAPSHOT + // fail. + snapshotID string +} + +func (s *ctidSession) logFields() loglib.Fields { + return loglib.Fields{"snapshotID": s.snapshotID} +} + +// execInTx runs fn in a transaction that observes the same view of the database +// as the transaction that exported this session's snapshot. +func (s *ctidSession) execInTx(ctx context.Context, fn func(tx pglib.Tx) error) error { + return execInSnapshotTx(ctx, s.reader.conn, s.snapshotID, fn) +} + +func (s *ctidSession) readTable(ctx context.Context, table *table) error { + tableInfo, err := s.getTableInfo(ctx, table.schema, table.name) + if err != nil { + return err + } + if tableInfo.isEmpty() { + return nil + } + table.rowSize = tableInfo.avgRowBytes + + // an empty intersection must not fall back to SELECT * + columns, pinLost := readableColumns(table.columns, tableInfo.columns) + if pinLost { + return fmt.Errorf("%w: no captured column of %s.%s exists on the source", + ErrSchemaChangedDuringSnapshot, + pglib.UnquoteIdentifier(table.schema), pglib.UnquoteIdentifier(table.name)) + } + table.columns = columns + + // 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 < s.reader.tableWorkers; i++ { + errGroup.Go(func() error { + return s.snapshotTableRangeWorker(ctx, 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 (s *ctidSession) snapshotTableRangeWorker(ctx context.Context, table *table, pageRangeChan <-chan pageRange) error { + for pageRange := range pageRangeChan { + if err := s.snapshotTableRange(ctx, table, pageRange); err != nil { + return err + } + } + return nil +} + +const pageRangeQuery = "SELECT %s FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'" + +// buildPageRangeQuery spells columns out. +func buildPageRangeQuery(t *table, r pageRange) string { + quotedTable := pglib.QuoteQualifiedIdentifier(t.schema, t.name) + if len(t.columns) == 0 { + return fmt.Sprintf(pageRangeQuery, allColumns, quotedTable, r.start, r.end) + } + + quotedColumns := make([]string, len(t.columns)) + for i, column := range t.columns { + quotedColumns[i] = pglib.QuoteRawIdentifier(column) + } + return fmt.Sprintf(pageRangeQuery, strings.Join(quotedColumns, ", "), quotedTable, r.start, r.end) +} + +func (s *ctidSession) snapshotTableRange(ctx context.Context, table *table, pageRange pageRange) error { + return s.execInTx(ctx, func(tx pglib.Tx) error { + s.reader.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{ + "schema": table.schema, "table": table.name, "snapshotID": s.snapshotID, + }) + + query := buildPageRangeQuery(table, pageRange) + rows, err := tx.Query(ctx, query) + if err != nil { + // something this query names vanished + var relationErr *pglib.ErrRelationDoesNotExist + if errors.As(err, &relationErr) { + return fmt.Errorf("%w: querying table rows: %w", ErrSchemaChangedDuringSnapshot, err) + } + return fmt.Errorf("querying table rows: %w", err) + } + defer rows.Close() + + rowCount, err := s.reader.sink.emit(ctx, table, rows) + if err != nil { + return err + } + + s.reader.logger.Debug(fmt.Sprintf("%d rows processed", rowCount), loglib.Fields{ + "schema": table.schema, "table": table.name, "snapshotID": s.snapshotID, + }) + + return nil + }) +} + +// tableInfoQuery shares the capture rule. +var tableInfoQuery = fmt.Sprintf(tableInfoQueryFmt, pglib.SelectStarColumnPredicate) + +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. + tableInfoQueryFmt = `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, + ARRAY( + SELECT a.attname::text FROM pg_catalog.pg_attribute a + WHERE a.attrelid = c.oid AND %s + ORDER BY a.attnum + ) AS columns +FROM + pg_catalog.pg_class c + JOIN pg_catalog.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;` + + 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 (s *ctidSession) getTableInfo(ctx context.Context, schemaName, tableName string) (*tableInfo, error) { + tableInfo := &tableInfo{} + err := s.execInTx(ctx, 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, &tableInfo.columns}, + 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(s.reader.batchBytes) + + s.reader.logger.Debug(fmt.Sprintf("table page count: %d, batch page size: %d", tableInfo.pageCount, tableInfo.batchPageSize), loglib.Fields{ + "schema": schemaName, "table": tableName, "snapshotID": s.snapshotID, + }) + return nil + }) + if err != nil { + return nil, err + } + + return tableInfo, nil +} + +// totalBytes returns the on disk size of the session's schema tables, as seen +// by the session's transaction snapshot. +func (s *ctidSession) totalBytes(ctx context.Context) (int64, error) { + schema, tables := s.schemaTables.schema, s.schemaTables.tables + + totalBytes := int64(0) + s.reader.logger.Debug("querying total bytes for schema", loglib.Fields{ + "schema": schema, "tables": tables, "snapshotID": s.snapshotID, + }) + + // make sure the schema and table names are unquoted since the system + // catalogs store unquoted names + unquotedTables := make([]string, len(tables)) + for i, table := range tables { + unquotedTables[i] = pglib.UnquoteIdentifier(table) + } + + err := s.execInTx(ctx, 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) + } + return nil + }) + + return totalBytes, err +} 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..597bb168 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/instrumented_table_reader.go @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + + loglib "github.com/xataio/pgstream/pkg/log" + "github.com/xataio/pgstream/pkg/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +// instrumentedTableReader traces the reads of the reader it wraps. Since the +// reader hands out the sessions that do the reading, it also wraps every +// session it produces, so that the table reads are traced as children of the +// schema they belong to. +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) (err error) { + ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.BeginSchema", trace.WithAttributes([]attribute.KeyValue{ + {Key: "schema", Value: attribute.StringValue(st.schema)}, + }...)) + defer otel.CloseSpan(span, err) + + return i.reader.beginSchema(ctx, st, func(ctx context.Context, session readSession) error { + return fn(ctx, &instrumentedReadSession{tracer: i.tracer, session: session}) + }) +} + +type instrumentedReadSession struct { + tracer trace.Tracer + session readSession +} + +func (i *instrumentedReadSession) readTable(ctx context.Context, 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.session.readTable(ctx, table) +} + +func (i *instrumentedReadSession) totalBytes(ctx context.Context) (bytes int64, err error) { + ctx, span := otel.StartSpan(ctx, i.tracer, "tableReader.TotalBytes") + defer otel.CloseSpan(span, err) + return i.session.totalBytes(ctx) +} + +func (i *instrumentedReadSession) logFields() loglib.Fields { + return i.session.logFields() +} 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 36a6d23e..93d9e54a 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go @@ -6,14 +6,12 @@ import ( "context" "errors" "fmt" - "strings" + "maps" "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" @@ -27,24 +25,22 @@ const allColumns = "*" var ErrSchemaChangedDuringSnapshot = errors.New("source schema changed during the snapshot") 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 decorate + // 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 } @@ -80,8 +76,6 @@ type table struct { columns []string } -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) { @@ -98,19 +92,16 @@ 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) + sink := newRowSink(pglib.NewMapper(conn), sg.processor, sg.logger, sg.progress) + sg.reader = newTableReader(sg.conn, sg.logger, sink, cfg, sg.instrumentation) return sg, nil } @@ -132,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 } } @@ -200,23 +189,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, schemaTables); err != nil { return err } defer func() { if err == nil { - sg.markProgressBarCompleted(schemaTables.schema) + sg.progress.complete(schemaTables.schema) } }() } @@ -229,7 +209,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 { @@ -244,7 +224,7 @@ func (sg *SnapshotGenerator) createSchemaSnapshot(ctx context.Context, schemaTab wg.Wait() return sg.collectTableErrors(schemaTables.schema, workerTableErrs) - }, snapshotTxOptions()) + }) } // pinnedColumns warns on uncaptured tables. @@ -282,13 +262,15 @@ func readableColumns(pinned, live []string) ([]string, bool) { return readable, len(readable) == 0 } -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() + sessionFields := session.logFields() for t := range tableChan { - logFields := loglib.Fields{"schema": t.schema, "table": t.name, "snapshotID": snapshotID} + logFields := loglib.Fields{"schema": t.schema, "table": t.name} + maps.Copy(logFields, sessionFields) sg.logger.Debug("snapshotting table", logFields) - if err := sg.tableSnapshotGenerator(ctx, snapshotID, t); err != nil { + if err := session.readTable(ctx, 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) { @@ -336,287 +318,20 @@ 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 - - // an empty intersection must not fall back to SELECT * - columns, pinLost := readableColumns(table.columns, tableInfo.columns) - if pinLost { - return fmt.Errorf("%w: no captured column of %s.%s exists on the source", - ErrSchemaChangedDuringSnapshot, - pglib.UnquoteIdentifier(table.schema), pglib.UnquoteIdentifier(table.name)) - } - table.columns = columns - - // 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 -} - -const pageRangeQuery = "SELECT %s FROM ONLY %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'" - -// buildPageRangeQuery spells columns out. -func buildPageRangeQuery(t *table, r pageRange) string { - quotedTable := pglib.QuoteQualifiedIdentifier(t.schema, t.name) - if len(t.columns) == 0 { - return fmt.Sprintf(pageRangeQuery, allColumns, quotedTable, r.start, r.end) - } - - quotedColumns := make([]string, len(t.columns)) - for i, column := range t.columns { - quotedColumns[i] = pglib.QuoteRawIdentifier(column) - } - return fmt.Sprintf(pageRangeQuery, strings.Join(quotedColumns, ", "), quotedTable, r.start, r.end) -} - -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 := buildPageRangeQuery(table, pageRange) - rows, err := tx.Query(ctx, query) - if err != nil { - // something this query names vanished - var relationErr *pglib.ErrRelationDoesNotExist - if errors.As(err, &relationErr) { - return fmt.Errorf("%w: querying table rows: %w", ErrSchemaChangedDuringSnapshot, err) - } - 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) +// addProgressBar sizes the schema's progress bar with the total bytes the read +// session reports for it. How those bytes are measured is the reading +// strategy's business; the generator only owns the bar. +func (sg *SnapshotGenerator) addProgressBar(ctx context.Context, session readSession, schemaTables *schemaTables) error { + totalBytes, err := session.totalBytes(ctx) if err != nil { return err } bar := sg.progressBarBuilder(totalBytes, fmt.Sprintf("[cyan][%s][reset] Snapshotting data...", schemaTables.schema)) - sg.progressBars.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) -} - -// tableInfoQuery shares the capture rule. -var tableInfoQuery = fmt.Sprintf(tableInfoQueryFmt, pglib.SelectStarColumnPredicate) - -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. - tableInfoQueryFmt = `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, - ARRAY( - SELECT a.attname::text FROM pg_catalog.pg_attribute a - WHERE a.attrelid = c.oid AND %s - ORDER BY a.attnum - ) AS columns -FROM - pg_catalog.pg_class c - JOIN pg_catalog.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, &tableInfo.columns}, - 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) { - totalBytes := int64(0) - sg.logger.Debug("querying total bytes for schema", loglib.Fields{ - "schema": schema, "tables": tables, "snapshotID": snapshotID, - }) - - // make sure the schema and table names are unquoted since the system - // catalogs store unquoted names - unquotedTables := make([]string, len(tables)) - for i, table := range tables { - unquotedTables[i] = pglib.UnquoteIdentifier(table) - } - - err := sg.execInSnapshotTx(ctx, 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) - } - return nil - }) - - 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 { - // An instance-local exported snapshot that a worker connection can't see - // means the source spans multiple instances; join the actionable cause. - if pglib.IsExportedSnapshotMissing(err) { - return fmt.Errorf("setting transaction snapshot: %w: %w", err, pglib.ErrLoadBalancedSource) - } - return fmt.Errorf("setting transaction snapshot: %w", err) - } + sg.progress.set(schemaTables.schema, bar) 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 a4b62433..843b37e3 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go @@ -302,7 +302,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) { QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) { require.Equal(t, fmt.Sprintf( `SELECT "id" FROM ONLY %s WHERE ctid BETWEEN '(0,0)' AND '(1,0)'`, - quotedSchemaTable1), query) + quotedSchemaTable1, + ), query) return &pgmocks.Rows{ CloseFn: func() {}, NextFn: func(i uint) bool { return i == 1 }, @@ -1382,32 +1383,40 @@ 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 }, - schemaWorkers: 1, - tableWorkers: 1, - batchBytes: 1024 * 1024, // 1MB - snapshotWorkers: 1, - progressTracking: tc.progressBar != nil, - progressBars: synclib.NewMap[string, progress.Bar](), + CloseFn: func() error { + return tc.processorCloseErr + }, + } + 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, + sink: newRowSink(pglib.NewMapper(tc.querier), processor, loglib.NewNoopLogger(), pt), + tableWorkers: 1, + batchBytes: 1024 * 1024, // 1MB + }, } - sg.tableSnapshotGenerator = sg.snapshotTable if tc.schemaWorkers != 0 { sg.schemaWorkers = tc.schemaWorkers @@ -2018,13 +2027,20 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { }, } - sg := SnapshotGenerator{ + pt := progressTracker{ + enabled: tc.name == "ok - with progress tracking", + bars: synclib.NewMap[string, progress.Bar](), + } + if pt.enabled { + pt.set(tc.table.schema, progressBar) + } + + reader := &ctidReader{ logger: zerolog.NewStdLogger(zerolog.NewLogger(&zerolog.Config{ LogLevel: "debug", })), - conn: tc.querier, - adapter: newAdapter(pglib.NewMapper(tc.querier), loglib.NewNoopLogger()), - processor: &processormocks.Processor{ + conn: tc.querier, + sink: newRowSink(pglib.NewMapper(tc.querier), &processormocks.Processor{ ProcessWALEventFn: func(ctx context.Context, walEvent *wal.Event) error { if tc.processor != nil { if err := tc.processor.ProcessWALEvent(ctx, walEvent); err != nil { @@ -2034,16 +2050,15 @@ func TestSnapshotGenerator_snapshotTableRange(t *testing.T) { eventChan <- walEvent return nil }, - }, - progressTracking: tc.name == "ok - with progress tracking", - progressBars: synclib.NewMap[string, progress.Bar](), + }, loglib.NewNoopLogger(), pt), } - - if sg.progressTracking { - sg.progressBars.Set(tc.table.schema, progressBar) + session := &ctidSession{ + reader: reader, + schemaTables: &schemaTables{schema: tc.table.schema, tables: []string{tc.table.name}}, + snapshotID: testSnapshotID, } - err := sg.snapshotTableRange(context.Background(), testSnapshotID, tc.table, tc.pageRange) + err := session.snapshotTableRange(context.Background(), tc.table, tc.pageRange) require.Equal(t, tc.wantErr, err) close(eventChan) diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_loadbalanced_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_loadbalanced_test.go index 608e26bb..1d4a3685 100644 --- a/pkg/snapshot/generator/postgres/data/pg_snapshot_loadbalanced_test.go +++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_loadbalanced_test.go @@ -13,7 +13,7 @@ import ( pgmocks "github.com/xataio/pgstream/internal/postgres/mocks" ) -func TestSnapshotGenerator_setTransactionSnapshot(t *testing.T) { +func TestSetTransactionSnapshot(t *testing.T) { t.Parallel() snapshotMissing := &pglib.ErrRelationDoesNotExist{Details: `snapshot "abc" does not exist`} @@ -32,13 +32,12 @@ func TestSnapshotGenerator_setTransactionSnapshot(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - sg := &SnapshotGenerator{} tx := &pgmocks.Tx{ ExecFn: func(context.Context, uint, string, ...any) (pglib.CommandTag, error) { return pglib.CommandTag{}, tt.execErr }, } - err := sg.setTransactionSnapshot(context.Background(), tx, "abc") + err := setTransactionSnapshot(context.Background(), tx, "abc") if !tt.wantErr { require.NoError(t, err) return 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/row_sink.go b/pkg/snapshot/generator/postgres/data/row_sink.go new file mode 100644 index 00000000..0bd952de --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/row_sink.go @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + "fmt" + + pglib "github.com/xataio/pgstream/internal/postgres" + loglib "github.com/xataio/pgstream/pkg/log" + "github.com/xataio/pgstream/pkg/wal/processor" +) + +// rowSink turns the rows a reading strategy queries into the wal events the +// snapshot emits. It owns everything that happens to a row once it has been +// read: adapting it to a wal event, handing it to the processor and reporting +// the bytes read. A strategy is therefore only responsible for deciding which +// rows to read, and every strategy consumes them the same way. +type rowSink struct { + adapter *adapter + processor processor.Processor + progress progressTracker +} + +func newRowSink(mapper mapper, processor processor.Processor, logger loglib.Logger, progress progressTracker) rowSink { + return rowSink{ + adapter: newAdapter(mapper, logger), + processor: processor, + progress: progress, + } +} + +// emit converts every row of the result set into a wal event and hands it to +// the processor, reporting the bytes read against the table's schema once the +// result set has been fully consumed. It returns the number of rows read. +func (s rowSink) emit(ctx context.Context, table *table, rows pglib.Rows) (uint, error) { + // resolve the column metadata (names/types) and timestamp once per result + // set, since the field descriptions are identical for every row in it. + rowAdapter := s.adapter.newRowEventAdapter(ctx, table.schema, table.name, rows.FieldDescriptions()) + rowCount := uint(0) + for rows.Next() { + rowCount++ + select { + case <-ctx.Done(): + return rowCount, ctx.Err() + default: + values, err := rows.Values() + if err != nil { + return rowCount, fmt.Errorf("retrieving rows values: %w", err) + } + + event := rowAdapter.rowToWalEvent(values) + if event == nil { + continue + } + + if err := s.processor.ProcessWALEvent(ctx, event); err != nil { + return rowCount, fmt.Errorf("processing snapshot row: %w", err) + } + } + } + + s.progress.advance(table.schema, int64(rowCount)*table.rowSize) + + return rowCount, rows.Err() +} 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..fc3467bf --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/snapshot_tx.go @@ -0,0 +1,60 @@ +// 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 { + // An instance-local exported snapshot that a worker connection can't see + // means the source spans multiple instances; join the actionable cause. + if pglib.IsExportedSnapshotMissing(err) { + return fmt.Errorf("setting transaction snapshot: %w: %w", err, pglib.ErrLoadBalancedSource) + } + 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..7486c521 --- /dev/null +++ b/pkg/snapshot/generator/postgres/data/table_reader.go @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 + +package postgres + +import ( + "context" + + pglib "github.com/xataio/pgstream/internal/postgres" + loglib "github.com/xataio/pgstream/pkg/log" + "github.com/xataio/pgstream/pkg/otel" +) + +// tableReader abstracts the strategy used to read a schema's tables during a +// data snapshot. It does not read tables itself: it opens a read session per +// schema, and the session does the reading. This keeps the per-schema state a +// strategy needs (a shared transaction snapshot, in the ctid reader's case) +// inside the strategy that understands it, instead of travelling through the +// snapshot generator. +type tableReader interface { + // beginSchema opens a read session for the given schema tables and invokes + // fn with it. The session, and any resource backing it, is only valid for + // the duration of fn. + beginSchema(ctx context.Context, st *schemaTables, fn func(context.Context, readSession) error) error +} + +// readSession reads the tables of a single schema. It is created by a +// tableReader in beginSchema, already bound to the schema it reads and to +// whatever per-schema state its strategy needs, and it is only valid for the +// duration of the beginSchema callback that produced it. +type readSession interface { + // readTable snapshots a single table of the session's schema. + readTable(ctx context.Context, table *table) error + // totalBytes returns the on disk size of the schema tables this session + // reads, so the caller can size its progress reporting. + totalBytes(ctx context.Context) (int64, error) + // logFields returns the session specific fields to attach to the log + // entries of the schema being read. + logFields() loglib.Fields +} + +// newTableReader builds the strategy used to read the snapshot tables and wraps +// it with any active decorator. Picking between strategies belongs here, so +// that the snapshot generator only ever sees a tableReader. +func newTableReader(conn pglib.Querier, logger loglib.Logger, sink rowSink, cfg *Config, instrumentation *otel.Instrumentation) tableReader { + var reader tableReader = &ctidReader{ + conn: conn, + logger: logger, + sink: sink, + tableWorkers: cfg.tableWorkers(), + batchBytes: cfg.batchBytes(), + } + + if instrumentation != nil { + reader = newInstrumentedTableReader(reader, instrumentation) + } + + return reader +}