From 9abd6b6e061fbc8fc4be5d8a9bfbd1be644b9f04 Mon Sep 17 00:00:00 2001
From: Kevin Biju <52661649+heavycrystal@users.noreply.github.com>
Date: Tue, 19 Mar 2024 01:58:46 +0530
Subject: [PATCH] [snapshot] fallback to full partitions in <PG12, cleanup
 (#1499)

separating slot and txn snapshots better, split state and make things
clearer
TID scans not supported in PG12, falling back to full table partitions
for now, might add support for custom partition keys later
---
 flow/activities/flowable.go             |  6 --
 flow/activities/snapshot_activity.go    | 70 ++++++++++++--------
 flow/cmd/snapshot_worker.go             |  6 +-
 flow/connectors/core.go                 |  2 +-
 flow/connectors/postgres/client.go      | 23 ++++---
 flow/connectors/postgres/postgres.go    | 39 ++++++++---
 flow/connectors/postgres/slot_signal.go |  7 +-
 flow/workflows/snapshot_flow.go         | 86 ++++++++++++++++++-------
 protos/flow.proto                       |  6 ++
 9 files changed, 165 insertions(+), 80 deletions(-)

diff --git a/flow/activities/flowable.go b/flow/activities/flowable.go
index 4924627fd9..68f764a890 100644
--- a/flow/activities/flowable.go
+++ b/flow/activities/flowable.go
@@ -35,12 +35,6 @@ type CheckConnectionResult struct {
 	NeedsSetupMetadataTables bool
 }
 
-type SlotSnapshotSignal struct {
-	signal       connpostgres.SlotSignal
-	snapshotName string
-	connector    connectors.CDCPullConnector
-}
-
 type FlowableActivity struct {
 	CatalogPool *pgxpool.Pool
 	Alerter     *alerting.Alerter
diff --git a/flow/activities/snapshot_activity.go b/flow/activities/snapshot_activity.go
index 1cd4dac3ef..d5aeb50ed7 100644
--- a/flow/activities/snapshot_activity.go
+++ b/flow/activities/snapshot_activity.go
@@ -16,21 +16,33 @@ import (
 	"github.com/PeerDB-io/peer-flow/shared"
 )
 
+type SlotSnapshotState struct {
+	snapshotName string
+	signal       connpostgres.SlotSignal
+	connector    connectors.CDCPullConnector
+}
+
+type TxSnapshotState struct {
+	SnapshotName     string
+	SupportsTIDScans bool
+}
+
 type SnapshotActivity struct {
-	SnapshotConnectionsMutex sync.Mutex
-	SnapshotConnections      map[string]SlotSnapshotSignal
-	Alerter                  *alerting.Alerter
+	SnapshotStatesMutex sync.Mutex
+	SlotSnapshotStates  map[string]SlotSnapshotState
+	TxSnapshotStates    map[string]TxSnapshotState
+	Alerter             *alerting.Alerter
 }
 
 // closes the slot signal
 func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName string) error {
-	a.SnapshotConnectionsMutex.Lock()
-	defer a.SnapshotConnectionsMutex.Unlock()
+	a.SnapshotStatesMutex.Lock()
+	defer a.SnapshotStatesMutex.Unlock()
 
-	if s, ok := a.SnapshotConnections[flowJobName]; ok {
+	if s, ok := a.SlotSnapshotStates[flowJobName]; ok {
 		close(s.signal.CloneComplete)
 		connectors.CloseConnector(ctx, s.connector)
-		delete(a.SnapshotConnections, flowJobName)
+		delete(a.SlotSnapshotStates, flowJobName)
 	}
 	a.Alerter.LogFlowEvent(ctx, flowJobName, "Ended Snapshot Flow Job")
 
@@ -94,18 +106,19 @@ func (a *SnapshotActivity) SetupReplication(
 		return nil, fmt.Errorf("slot error: %w", slotInfo.Err)
 	}
 
-	a.SnapshotConnectionsMutex.Lock()
-	defer a.SnapshotConnectionsMutex.Unlock()
+	a.SnapshotStatesMutex.Lock()
+	defer a.SnapshotStatesMutex.Unlock()
 
-	a.SnapshotConnections[config.FlowJobName] = SlotSnapshotSignal{
+	a.SlotSnapshotStates[config.FlowJobName] = SlotSnapshotState{
 		signal:       slotSignal,
 		snapshotName: slotInfo.SnapshotName,
 		connector:    conn,
 	}
 
 	return &protos.SetupReplicationOutput{
-		SlotName:     slotInfo.SlotName,
-		SnapshotName: slotInfo.SnapshotName,
+		SlotName:         slotInfo.SlotName,
+		SnapshotName:     slotInfo.SnapshotName,
+		SupportsTidScans: slotInfo.SupportsTIDScans,
 	}, nil
 }
 
@@ -116,41 +129,46 @@ func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, pee
 	}
 	defer connectors.CloseConnector(ctx, conn)
 
-	snapshotName, tx, err := conn.ExportSnapshot(ctx)
+	exportSnapshotOutput, tx, err := conn.ExportTxSnapshot(ctx)
 	if err != nil {
 		return err
 	}
 
-	sss := SlotSnapshotSignal{snapshotName: snapshotName}
-	a.SnapshotConnectionsMutex.Lock()
-	a.SnapshotConnections[sessionID] = sss
-	a.SnapshotConnectionsMutex.Unlock()
+	a.SnapshotStatesMutex.Lock()
+	a.TxSnapshotStates[sessionID] = TxSnapshotState{
+		SnapshotName:     exportSnapshotOutput.SnapshotName,
+		SupportsTIDScans: exportSnapshotOutput.SupportsTidScans,
+	}
+	a.SnapshotStatesMutex.Unlock()
 
 	logger := activity.GetLogger(ctx)
 	start := time.Now()
 	for {
 		msg := fmt.Sprintf("maintaining export snapshot transaction %s", time.Since(start).Round(time.Second))
 		logger.Info(msg)
+		// this function relies on context cancellation to exit
+		// context is not explicitly cancelled, but workflow exit triggers an implicit cancel
+		// from activity.RecordBeat
 		activity.RecordHeartbeat(ctx, msg)
 		if ctx.Err() != nil {
-			a.SnapshotConnectionsMutex.Lock()
-			delete(a.SnapshotConnections, sessionID)
-			a.SnapshotConnectionsMutex.Unlock()
+			a.SnapshotStatesMutex.Lock()
+			delete(a.TxSnapshotStates, sessionID)
+			a.SnapshotStatesMutex.Unlock()
 			return conn.FinishExport(tx)
 		}
 		time.Sleep(time.Minute)
 	}
 }
 
-func (a *SnapshotActivity) WaitForExportSnapshot(ctx context.Context, sessionID string) (string, error) {
+func (a *SnapshotActivity) WaitForExportSnapshot(ctx context.Context, sessionID string) (*TxSnapshotState, error) {
 	logger := activity.GetLogger(ctx)
 	attempt := 0
 	for {
-		a.SnapshotConnectionsMutex.Lock()
-		sss, ok := a.SnapshotConnections[sessionID]
-		a.SnapshotConnectionsMutex.Unlock()
+		a.SnapshotStatesMutex.Lock()
+		tsc, ok := a.TxSnapshotStates[sessionID]
+		a.SnapshotStatesMutex.Unlock()
 		if ok {
-			return sss.snapshotName, nil
+			return &tsc, nil
 		}
 		activity.RecordHeartbeat(ctx, "wait another second for snapshot export")
 		attempt += 1
@@ -158,7 +176,7 @@ func (a *SnapshotActivity) WaitForExportSnapshot(ctx context.Context, sessionID
 			logger.Info("waiting on snapshot export", slog.Int("attempt", attempt))
 		}
 		if err := ctx.Err(); err != nil {
-			return "", err
+			return nil, err
 		}
 		time.Sleep(time.Second)
 	}
diff --git a/flow/cmd/snapshot_worker.go b/flow/cmd/snapshot_worker.go
index d5b9d4b51f..b7418e9ceb 100644
--- a/flow/cmd/snapshot_worker.go
+++ b/flow/cmd/snapshot_worker.go
@@ -66,9 +66,11 @@ func SnapshotWorkerMain(opts *SnapshotWorkerOptions) (client.Client, worker.Work
 	})
 
 	w.RegisterWorkflow(peerflow.SnapshotFlowWorkflow)
+	// explicitly not initializing mutex, in line with design
 	w.RegisterActivity(&activities.SnapshotActivity{
-		SnapshotConnections: make(map[string]activities.SlotSnapshotSignal),
-		Alerter:             alerting.NewAlerter(context.Background(), conn),
+		SlotSnapshotStates: make(map[string]activities.SlotSnapshotState),
+		TxSnapshotStates:   make(map[string]activities.TxSnapshotState),
+		Alerter:            alerting.NewAlerter(context.Background(), conn),
 	})
 
 	return c, w, nil
diff --git a/flow/connectors/core.go b/flow/connectors/core.go
index 6adb3a2429..988b8e3f28 100644
--- a/flow/connectors/core.go
+++ b/flow/connectors/core.go
@@ -44,7 +44,7 @@ type CDCPullConnector interface {
 
 	// For InitialSnapshotOnly correctness without replication slot
 	// `any` is for returning transaction if necessary
-	ExportSnapshot(context.Context) (string, any, error)
+	ExportTxSnapshot(context.Context) (*protos.ExportTxSnapshotOutput, any, error)
 
 	// `any` from ExportSnapshot passed here when done, allowing transaction to commit
 	FinishExport(any) error
diff --git a/flow/connectors/postgres/client.go b/flow/connectors/postgres/client.go
index 898f962aec..f667107bb2 100644
--- a/flow/connectors/postgres/client.go
+++ b/flow/connectors/postgres/client.go
@@ -369,12 +369,12 @@ func (c *PostgresConnector) createSlotAndPublication(
 
 		c.logger.Warn(fmt.Sprintf("Creating replication slot '%s'", slot))
 
-		_, err = conn.Exec(ctx, "SET idle_in_transaction_session_timeout = 0")
+		_, err = conn.Exec(ctx, "SET LOCAL idle_in_transaction_session_timeout=0")
 		if err != nil {
 			return fmt.Errorf("[slot] error setting idle_in_transaction_session_timeout: %w", err)
 		}
 
-		_, err = conn.Exec(ctx, "SET lock_timeout = 0")
+		_, err = conn.Exec(ctx, "SET LOCAL lock_timeout=0")
 		if err != nil {
 			return fmt.Errorf("[slot] error setting lock_timeout: %w", err)
 		}
@@ -388,11 +388,17 @@ func (c *PostgresConnector) createSlotAndPublication(
 			return fmt.Errorf("[slot] error creating replication slot: %w", err)
 		}
 
+		ok, _, err := c.MajorVersionCheck(ctx, POSTGRES_13)
+		if err != nil {
+			return fmt.Errorf("[slot] error getting PG version: %w", err)
+		}
+
 		c.logger.Info(fmt.Sprintf("Created replication slot '%s'", slot))
 		slotDetails := SlotCreationResult{
-			SlotName:     res.SlotName,
-			SnapshotName: res.SnapshotName,
-			Err:          nil,
+			SlotName:         res.SlotName,
+			SnapshotName:     res.SnapshotName,
+			Err:              nil,
+			SupportsTIDScans: ok,
 		}
 		signal.SlotCreated <- slotDetails
 		c.logger.Info("Waiting for clone to complete")
@@ -405,9 +411,10 @@ func (c *PostgresConnector) createSlotAndPublication(
 			e = errors.New("slot already exists")
 		}
 		slotDetails := SlotCreationResult{
-			SlotName:     slot,
-			SnapshotName: "",
-			Err:          e,
+			SlotName:         slot,
+			SnapshotName:     "",
+			Err:              e,
+			SupportsTIDScans: false,
 		}
 		signal.SlotCreated <- slotDetails
 	}
diff --git a/flow/connectors/postgres/postgres.go b/flow/connectors/postgres/postgres.go
index 1bba4e8eca..bd6721c076 100644
--- a/flow/connectors/postgres/postgres.go
+++ b/flow/connectors/postgres/postgres.go
@@ -941,30 +941,49 @@ func (c *PostgresConnector) EnsurePullability(
 	return &protos.EnsurePullabilityBatchOutput{TableIdentifierMapping: tableIdentifierMapping}, nil
 }
 
-func (c *PostgresConnector) ExportSnapshot(ctx context.Context) (string, any, error) {
+func (c *PostgresConnector) ExportTxSnapshot(ctx context.Context) (*protos.ExportTxSnapshotOutput, any, error) {
 	var snapshotName string
 	tx, err := c.conn.Begin(ctx)
 	if err != nil {
-		return "", nil, err
+		return nil, nil, err
+	}
+	txNeedsRollback := true
+	defer func() {
+		if txNeedsRollback {
+			rollbackCtx, cancelFunc := context.WithTimeout(context.Background(), 5*time.Minute)
+			defer cancelFunc()
+			err := tx.Rollback(rollbackCtx)
+			if err != pgx.ErrTxClosed {
+				c.logger.Error("error while rolling back transaction for snapshot export")
+			}
+		}
+	}()
+
+	_, err = tx.Exec(ctx, "SET LOCAL idle_in_transaction_session_timeout=0")
+	if err != nil {
+		return nil, nil, fmt.Errorf("[export-snapshot] error setting idle_in_transaction_session_timeout: %w", err)
 	}
 
-	_, err = tx.Exec(ctx, "SET idle_in_transaction_session_timeout = 0")
+	_, err = tx.Exec(ctx, "SET LOCAL lock_timeout=0")
 	if err != nil {
-		return "", nil, fmt.Errorf("[export-snapshot] error setting idle_in_transaction_session_timeout: %w", err)
+		return nil, nil, fmt.Errorf("[export-snapshot] error setting lock_timeout: %w", err)
 	}
 
-	_, err = tx.Exec(ctx, "SET lock_timeout = 0")
+	ok, _, err := c.MajorVersionCheck(ctx, POSTGRES_13)
 	if err != nil {
-		return "", nil, fmt.Errorf("[export-snapshot] error setting lock_timeout: %w", err)
+		return nil, nil, fmt.Errorf("[export-snapshot] error getting PG version: %w", err)
 	}
 
-	err = tx.QueryRow(ctx, "select pg_export_snapshot()").Scan(&snapshotName)
+	err = tx.QueryRow(ctx, "SELECT pg_export_snapshot()").Scan(&snapshotName)
 	if err != nil {
-		_ = tx.Rollback(ctx)
-		return "", nil, err
+		return nil, nil, err
 	}
+	txNeedsRollback = false
 
-	return snapshotName, tx, err
+	return &protos.ExportTxSnapshotOutput{
+		SnapshotName:     snapshotName,
+		SupportsTidScans: ok,
+	}, tx, err
 }
 
 func (c *PostgresConnector) FinishExport(tx any) error {
diff --git a/flow/connectors/postgres/slot_signal.go b/flow/connectors/postgres/slot_signal.go
index 386b3af36b..be6395917b 100644
--- a/flow/connectors/postgres/slot_signal.go
+++ b/flow/connectors/postgres/slot_signal.go
@@ -1,9 +1,10 @@
 package connpostgres
 
 type SlotCreationResult struct {
-	SlotName     string
-	SnapshotName string
-	Err          error
+	SlotName         string
+	SnapshotName     string
+	SupportsTIDScans bool
+	Err              error
 }
 
 // This struct contains two signals.
diff --git a/flow/workflows/snapshot_flow.go b/flow/workflows/snapshot_flow.go
index ce9ab27d78..8ec844cc4b 100644
--- a/flow/workflows/snapshot_flow.go
+++ b/flow/workflows/snapshot_flow.go
@@ -12,6 +12,7 @@ import (
 	"go.temporal.io/sdk/temporal"
 	"go.temporal.io/sdk/workflow"
 
+	"github.com/PeerDB-io/peer-flow/activities"
 	"github.com/PeerDB-io/peer-flow/concurrency"
 	connpostgres "github.com/PeerDB-io/peer-flow/connectors/postgres"
 	"github.com/PeerDB-io/peer-flow/connectors/utils"
@@ -19,12 +20,28 @@ import (
 	"github.com/PeerDB-io/peer-flow/shared"
 )
 
+type snapshotType int
+
+const (
+	SNAPSHOT_TYPE_UNKNOWN snapshotType = iota
+	SNAPSHOT_TYPE_SLOT
+	SNAPSHOT_TYPE_TX
+)
+
 type SnapshotFlowExecution struct {
 	config                 *protos.FlowConnectionConfigs
 	logger                 log.Logger
 	tableNameSchemaMapping map[string]*protos.TableSchema
 }
 
+type cloneTablesInput struct {
+	snapshotType      snapshotType
+	slotName          string
+	snapshotName      string
+	supportsTIDScans  bool
+	maxParallelClones int
+}
+
 // ensurePullability ensures that the source peer is pullable.
 func (s *SnapshotFlowExecution) setupReplication(
 	ctx workflow.Context,
@@ -117,11 +134,6 @@ func (s *SnapshotFlowExecution) cloneTable(
 	sourcePostgres := s.config.Source
 	sourcePostgres.GetPostgresConfig().TransactionSnapshot = snapshotName
 
-	partitionCol := "ctid"
-	if mapping.PartitionKey != "" {
-		partitionCol = mapping.PartitionKey
-	}
-
 	parsedSrcTable, err := utils.ParseSchemaTable(srcName)
 	if err != nil {
 		s.logger.Error("unable to parse source table", slog.Any("error", err), cloneLog)
@@ -142,9 +154,13 @@ func (s *SnapshotFlowExecution) cloneTable(
 			}
 		}
 	}
-
-	query := fmt.Sprintf("SELECT %s FROM %s WHERE %s BETWEEN {{.start}} AND {{.end}}",
-		from, parsedSrcTable.String(), partitionCol)
+	var query string
+	if mapping.PartitionKey == "" {
+		query = fmt.Sprintf("SELECT %s FROM %s", from, parsedSrcTable.String())
+	} else {
+		query = fmt.Sprintf("SELECT %s FROM %s WHERE %s BETWEEN {{.start}} AND {{.end}}",
+			from, parsedSrcTable.String(), mapping.PartitionKey)
+	}
 
 	numWorkers := uint32(8)
 	if s.config.SnapshotMaxParallelWorkers > 0 {
@@ -161,7 +177,7 @@ func (s *SnapshotFlowExecution) cloneTable(
 		SourcePeer:                 sourcePostgres,
 		DestinationPeer:            s.config.Destination,
 		Query:                      query,
-		WatermarkColumn:            partitionCol,
+		WatermarkColumn:            mapping.PartitionKey,
 		WatermarkTable:             srcName,
 		InitialCopyOnly:            true,
 		DestinationTableIdentifier: dstName,
@@ -182,23 +198,36 @@ func (s *SnapshotFlowExecution) cloneTable(
 
 func (s *SnapshotFlowExecution) cloneTables(
 	ctx workflow.Context,
-	slotInfo *protos.SetupReplicationOutput,
-	maxParallelClones int,
+	cloneTablesInput *cloneTablesInput,
 ) error {
-	s.logger.Info(fmt.Sprintf("cloning tables for slot name %s and snapshotName %s",
-		slotInfo.SlotName, slotInfo.SnapshotName))
+	if cloneTablesInput.snapshotType == SNAPSHOT_TYPE_SLOT {
+		s.logger.Info(fmt.Sprintf("cloning tables for slot name %s and snapshotName %s",
+			cloneTablesInput.slotName, cloneTablesInput.snapshotName))
+	} else if cloneTablesInput.snapshotType == SNAPSHOT_TYPE_TX {
+		s.logger.Info("cloning tables in txn snapshot mode with snapshotName " +
+			cloneTablesInput.snapshotName)
+	}
 
-	boundSelector := concurrency.NewBoundSelector(maxParallelClones)
+	boundSelector := concurrency.NewBoundSelector(cloneTablesInput.maxParallelClones)
 
+	defaultPartitionCol := "ctid"
+	if !cloneTablesInput.supportsTIDScans {
+		s.logger.Info("Postgres version too old for TID scans, might use full table partitions!")
+		defaultPartitionCol = ""
+	}
+
+	snapshotName := cloneTablesInput.snapshotName
 	for _, v := range s.config.TableMappings {
 		source := v.SourceTableIdentifier
 		destination := v.DestinationTableIdentifier
-		snapshotName := slotInfo.SnapshotName
 		s.logger.Info(fmt.Sprintf(
 			"Cloning table with source table %s and destination table name %s",
 			source, destination),
 			slog.String("snapshotName", snapshotName),
 		)
+		if v.PartitionKey == "" {
+			v.PartitionKey = defaultPartitionCol
+		}
 		err := s.cloneTable(ctx, boundSelector, snapshotName, v)
 		if err != nil {
 			s.logger.Error("failed to start clone child workflow: ", err)
@@ -227,7 +256,13 @@ func (s *SnapshotFlowExecution) cloneTablesWithSlot(
 	}
 
 	logger.Info(fmt.Sprintf("cloning %d tables in parallel", numTablesInParallel))
-	if err := s.cloneTables(ctx, slotInfo, numTablesInParallel); err != nil {
+	if err := s.cloneTables(ctx, &cloneTablesInput{
+		snapshotType:      SNAPSHOT_TYPE_SLOT,
+		slotName:          slotInfo.SlotName,
+		snapshotName:      slotInfo.SnapshotName,
+		supportsTIDScans:  slotInfo.SupportsTidScans,
+		maxParallelClones: numTablesInParallel,
+	}); err != nil {
 		return fmt.Errorf("failed to clone tables: %w", err)
 	}
 
@@ -246,7 +281,8 @@ func SnapshotFlowWorkflow(
 	se := &SnapshotFlowExecution{
 		config:                 config,
 		tableNameSchemaMapping: tableNameSchemaMapping,
-		logger:                 log.With(workflow.GetLogger(ctx), slog.String(string(shared.FlowNameKey), config.FlowJobName)),
+		logger: log.With(workflow.GetLogger(ctx),
+			slog.String(string(shared.FlowNameKey), config.FlowJobName)),
 	}
 
 	numTablesInParallel := int(max(config.SnapshotNumTablesInParallel, 1))
@@ -299,7 +335,7 @@ func SnapshotFlowWorkflow(
 		)
 
 		var sessionError error
-		var snapshotName string
+		var txnSnapshotState *activities.TxSnapshotState
 		sessionSelector := workflow.NewNamedSelector(ctx, "ExportSnapshotSetup")
 		sessionSelector.AddFuture(fMaintain, func(f workflow.Future) {
 			// MaintainTx should never exit without an error before this point
@@ -307,7 +343,7 @@ func SnapshotFlowWorkflow(
 		})
 		sessionSelector.AddFuture(fExportSnapshot, func(f workflow.Future) {
 			// Happy path is waiting for this to return without error
-			sessionError = f.Get(exportCtx, &snapshotName)
+			sessionError = f.Get(exportCtx, &txnSnapshotState)
 		})
 		sessionSelector.AddReceive(ctx.Done(), func(_ workflow.ReceiveChannel, _ bool) {
 			sessionError = ctx.Err()
@@ -317,11 +353,13 @@ func SnapshotFlowWorkflow(
 			return sessionError
 		}
 
-		slotInfo := &protos.SetupReplicationOutput{
-			SlotName:     "peerdb_initial_copy_only",
-			SnapshotName: snapshotName,
-		}
-		if err := se.cloneTables(ctx, slotInfo, int(config.SnapshotNumTablesInParallel)); err != nil {
+		if err := se.cloneTables(ctx, &cloneTablesInput{
+			snapshotType:      SNAPSHOT_TYPE_TX,
+			slotName:          "",
+			snapshotName:      txnSnapshotState.SnapshotName,
+			supportsTIDScans:  txnSnapshotState.SupportsTIDScans,
+			maxParallelClones: numTablesInParallel,
+		}); err != nil {
 			return fmt.Errorf("failed to clone tables: %w", err)
 		}
 	} else if err := se.cloneTablesWithSlot(ctx, sessionCtx, numTablesInParallel); err != nil {
diff --git a/protos/flow.proto b/protos/flow.proto
index e872b0fa17..f33b0696db 100644
--- a/protos/flow.proto
+++ b/protos/flow.proto
@@ -144,6 +144,7 @@ message SetupReplicationInput {
 message SetupReplicationOutput {
   string slot_name = 1;
   string snapshot_name = 2;
+  bool supports_tid_scans = 3;
 }
 
 message CreateRawTableInput {
@@ -394,3 +395,8 @@ message IsQRepPartitionSyncedInput {
   string partition_id = 2;
 }
 
+message ExportTxSnapshotOutput {
+  string snapshot_name = 1;
+  bool supports_tid_scans = 2;
+}
+