one-sync

flow/activities/flowable.go

@@ -24,6 +24,7 @@ import (
 	"github.com/PeerDB-io/peer-flow/connectors/utils"
 	"github.com/PeerDB-io/peer-flow/connectors/utils/monitoring"
 	"github.com/PeerDB-io/peer-flow/generated/protos"
+	"github.com/PeerDB-io/peer-flow/logger"
 	"github.com/PeerDB-io/peer-flow/model"
 	"github.com/PeerDB-io/peer-flow/peerdbenv"
 	"github.com/PeerDB-io/peer-flow/shared"
@@ -44,6 +45,8 @@ type SlotSnapshotSignal struct {
 type FlowableActivity struct {
 	CatalogPool *pgxpool.Pool
 	Alerter     *alerting.Alerter
+	CdcCacheRw  sync.RWMutex
+	CdcCache    map[string]connectors.CDCPullConnector
 }
 
 func (a *FlowableActivity) CheckConnection(
@@ -206,7 +209,93 @@ func (a *FlowableActivity) CreateNormalizedTable(
 	}, nil
 }
 
-func (a *FlowableActivity) StartFlow(ctx context.Context,
+func (a *FlowableActivity) MaintainPull(
+	ctx context.Context,
+	config *protos.FlowConnectionConfigs,
+	sessionID string,
+) error {
+	srcConn, err := connectors.GetCDCPullConnector(ctx, config.Source)
+	if err != nil {
+		return err
+	}
+	defer connectors.CloseConnector(ctx, srcConn)
+
+	if err := srcConn.SetupReplConn(ctx); err != nil {
+		return err
+	}
+
+	a.CdcCacheRw.Lock()
+	a.CdcCache[sessionID] = srcConn
+	a.CdcCacheRw.Unlock()
+
+	ticker := time.NewTicker(15 * time.Second)
+	defer ticker.Stop()
+
+	slotNameForMetrics := fmt.Sprintf("peerflow_slot_%s", config.FlowJobName)
+	if config.ReplicationSlotName != "" {
+		slotNameForMetrics = config.ReplicationSlotName
+	}
+
+	go func() {
+		err := srcConn.HandleSlotInfo(ctx, a.Alerter, a.CatalogPool, slotNameForMetrics, config.Source.Name)
+		if err != nil {
+			logger.LoggerFromCtx(ctx).Error("failed to record slot size", "error", err)
+		}
+
+		// separate goroutine so recordSlotSize doesn't block heartbeats
+		// meanwhile recordSlotSize synchronously for ticker backpressure
+		slotSizeTicker := time.NewTicker(5 * time.Minute)
+		defer slotSizeTicker.Stop()
+		for {
+			select {
+			case <-slotSizeTicker.C:
+				err := srcConn.HandleSlotInfo(ctx, a.Alerter, a.CatalogPool, slotNameForMetrics, config.Source.Name)
+				if err != nil {
+					logger.LoggerFromCtx(ctx).Error("failed to record slot size", slog.Any("error", err))
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	for {
+		select {
+		case <-ticker.C:
+			activity.RecordHeartbeat(ctx, "keep session alive")
+			if err := srcConn.ReplPing(ctx); err != nil {
+				activity.GetLogger(ctx).Error("Failed to send keep alive ping to replication connection", slog.Any("error", err))
+			}
+		case <-ctx.Done():
+			a.CdcCacheRw.Lock()
+			delete(a.CdcCache, sessionID)
+			a.CdcCacheRw.Unlock()
+			return nil
+		}
+	}
+}
+
+func (a *FlowableActivity) WaitForSourceConnector(
+	ctx context.Context,
+	sessionID string,
+) error {
+	for {
+		a.CdcCacheRw.RLock()
+		_, ok := a.CdcCache[sessionID]
+		a.CdcCacheRw.RUnlock()
+		if ok {
+			return nil
+		}
+		activity.RecordHeartbeat(ctx, "wait another second for source connector")
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		time.Sleep(time.Second)
+	}
+}
+
+func (a *FlowableActivity) StartFlow(
+	ctx context.Context,
 	input *protos.StartFlowInput,
 ) (*model.SyncResponse, error) {
 	ctx = context.WithValue(ctx, shared.FlowNameKey, input.FlowConnectionConfigs.FlowJobName)
@@ -225,15 +314,14 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 		tblNameMapping[v.SourceTableIdentifier] = model.NewNameAndExclude(v.DestinationTableIdentifier, v.Exclude)
 	}
 
-	srcConn, err := connectors.GetCDCPullConnector(ctx, config.Source)
-	if err != nil {
-		return nil, fmt.Errorf("failed to get source connector: %w", err)
+	a.CdcCacheRw.RLock()
+	srcConn, ok := a.CdcCache[input.SessionId]
+	a.CdcCacheRw.RUnlock()
+	if !ok {
+		return nil, errors.New("source connector missing from CdcCache")
 	}
-	defer connectors.CloseConnector(ctx, srcConn)
-
-	slotNameForMetrics := fmt.Sprintf("peerflow_slot_%s", input.FlowConnectionConfigs.FlowJobName)
-	if input.FlowConnectionConfigs.ReplicationSlotName != "" {
-		slotNameForMetrics = input.FlowConnectionConfigs.ReplicationSlotName
+	if err := srcConn.ConnectionActive(ctx); err != nil {
+		return nil, err
 	}
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
@@ -242,23 +330,22 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 	})
 	defer shutdown()
 
-	errGroup, errCtx := errgroup.WithContext(ctx)
-	go a.recordSlotSizePeriodically(errCtx, srcConn, slotNameForMetrics, input.FlowConnectionConfigs.Source.Name)
-
 	batchSize := input.SyncFlowOptions.BatchSize
 	if batchSize <= 0 {
 		batchSize = 1_000_000
 	}
 
-	lastOffset, err := dstConn.GetLastOffset(ctx, input.FlowConnectionConfigs.FlowJobName)
-	if err != nil {
-		return nil, err
-	}
-
 	// start a goroutine to pull records from the source
 	recordBatch := model.NewCDCRecordStream()
 	startTime := time.Now()
 	flowName := input.FlowConnectionConfigs.FlowJobName
+
+	lastOffset, err := dstConn.GetLastOffset(ctx, flowName)
+	if err != nil {
+		return nil, err
+	}
+
+	errGroup, errCtx := errgroup.WithContext(ctx)
 	errGroup.Go(func() error {
 		return srcConn.PullRecords(errCtx, a.CatalogPool, &model.PullRecordsRequest{
 			FlowJobName:           flowName,

flow/cmd/worker.go

@@ -16,6 +16,7 @@ import (
 	"go.temporal.io/sdk/worker"
 
 	"github.com/PeerDB-io/peer-flow/activities"
+	"github.com/PeerDB-io/peer-flow/connectors"
 	utils "github.com/PeerDB-io/peer-flow/connectors/utils/catalog"
 	"github.com/PeerDB-io/peer-flow/logger"
 	"github.com/PeerDB-io/peer-flow/shared"
@@ -127,9 +128,10 @@ func WorkerMain(opts *WorkerOptions) error {
 		return queueErr
 	}
 
-	w := worker.New(c, taskQueue, worker.Options{})
+	w := worker.New(c, taskQueue, worker.Options{
+		EnableSessionWorker: true,
+	})
 	w.RegisterWorkflow(peerflow.CDCFlowWorkflowWithConfig)
-	w.RegisterWorkflow(peerflow.SyncFlowWorkflow)
 	w.RegisterWorkflow(peerflow.SetupFlowWorkflow)
 	w.RegisterWorkflow(peerflow.NormalizeFlowWorkflow)
 	w.RegisterWorkflow(peerflow.QRepFlowWorkflow)
@@ -146,6 +148,7 @@ func WorkerMain(opts *WorkerOptions) error {
 	w.RegisterActivity(&activities.FlowableActivity{
 		CatalogPool: conn,
 		Alerter:     alerter,
+		CdcCache:    make(map[string]connectors.CDCPullConnector),
 	})
 
 	err = w.Run(worker.InterruptCh())

flow/connectors/core.go

@@ -39,6 +39,10 @@ type CDCPullConnector interface {
 		*protos.EnsurePullabilityBatchOutput, error)
 
 	// Methods related to retrieving and pushing records for this connector as a source and destination.
+	SetupReplConn(context.Context) error
+
+	// Ping source to keep connection alive. Can be called concurrently with PullRecords; skips ping in that case.
+	ReplPing(context.Context) error
 
 	// PullRecords pulls records from the source, and returns a RecordBatch.
 	// This method should be idempotent, and should be able to be called multiple times with the same request.

flow/connectors/postgres/cdc.go

@@ -3,7 +3,6 @@ package connpostgres
 import (
 	"context"
 	"crypto/sha256"
-	"errors"
 	"fmt"
 	"log/slog"
 	"time"
@@ -28,7 +27,6 @@ import (
 
 type PostgresCDCSource struct {
 	*PostgresConnector
-	replConn               *pgx.Conn
 	SrcTableIDNameMapping  map[uint32]string
 	TableNameMapping       map[string]model.NameAndExclude
 	slot                   string
@@ -64,37 +62,35 @@ type startReplicationOpts struct {
 }
 
 // Create a new PostgresCDCSource
-func (c *PostgresConnector) NewPostgresCDCSource(ctx context.Context, cdcConfig *PostgresCDCConfig) *PostgresCDCSource {
+func (c *PostgresConnector) NewPostgresCDCSource(cdcConfig *PostgresCDCConfig) *PostgresCDCSource {
 	return &PostgresCDCSource{
 		PostgresConnector:         c,
-		replConn:                  cdcConfig.Connection,
 		SrcTableIDNameMapping:     cdcConfig.SrcTableIDNameMapping,
 		TableNameMapping:          cdcConfig.TableNameMapping,
 		slot:                      cdcConfig.Slot,
 		publication:               cdcConfig.Publication,
 		relationMessageMapping:    cdcConfig.RelationMessageMapping,
-		typeMap:                   pgtype.NewMap(),
 		childToParentRelIDMapping: cdcConfig.ChildToParentRelIDMap,
+		typeMap:                   pgtype.NewMap(),
 		commitLock:                false,
 		catalogPool:               cdcConfig.CatalogPool,
 		flowJobName:               cdcConfig.FlowJobName,
 	}
 }
 
-func getChildToParentRelIDMap(ctx context.Context, conn *pgx.Conn) (map[uint32]uint32, error) {
+func GetChildToParentRelIDMap(ctx context.Context, conn *pgx.Conn) (map[uint32]uint32, error) {
 	query := `
 		SELECT parent.oid AS parentrelid, child.oid AS childrelid
 		FROM pg_inherits
 		JOIN pg_class parent ON pg_inherits.inhparent = parent.oid
-		JOIN pg_class child  ON pg_inherits.inhrelid = child.oid
+		JOIN pg_class child ON pg_inherits.inhrelid = child.oid
 		WHERE parent.relkind='p';
 	`
 
 	rows, err := conn.Query(ctx, query, pgx.QueryExecModeSimpleProtocol)
 	if err != nil {
 		return nil, fmt.Errorf("error querying for child to parent relid map: %w", err)
 	}
-
 	defer rows.Close()
 
 	childToParentRelIDMap := make(map[uint32]uint32)
@@ -113,83 +109,14 @@ func getChildToParentRelIDMap(ctx context.Context, conn *pgx.Conn) (map[uint32]u
 
 // PullRecords pulls records from the cdc stream
 func (p *PostgresCDCSource) PullRecords(ctx context.Context, req *model.PullRecordsRequest) error {
-	replicationOpts, err := p.replicationOptions()
-	if err != nil {
-		return fmt.Errorf("error getting replication options: %w", err)
-	}
-
-	pgConn := p.replConn.PgConn()
-
-	// start replication
-	var clientXLogPos, startLSN pglogrepl.LSN
-	if req.LastOffset > 0 {
-		p.logger.Info("starting replication from last sync state", slog.Int64("last checkpoint", req.LastOffset))
-		clientXLogPos = pglogrepl.LSN(req.LastOffset)
-		startLSN = clientXLogPos + 1
-	}
-
-	opts := startReplicationOpts{
-		conn:            pgConn,
-		startLSN:        startLSN,
-		replicationOpts: *replicationOpts,
-	}
-
-	err = p.startReplication(ctx, opts)
-	if err != nil {
-		return fmt.Errorf("error starting replication: %w", err)
-	}
-
-	p.logger.Info(fmt.Sprintf("started replication on slot %s at startLSN: %d", p.slot, startLSN))
-
-	return p.consumeStream(ctx, pgConn, req, clientXLogPos, req.RecordStream)
-}
-
-func (p *PostgresCDCSource) startReplication(ctx context.Context, opts startReplicationOpts) error {
-	err := pglogrepl.StartReplication(ctx, opts.conn, p.slot, opts.startLSN, opts.replicationOpts)
-	if err != nil {
-		p.logger.Error("error starting replication", slog.Any("error", err))
-		return fmt.Errorf("error starting replication at startLsn - %d: %w", opts.startLSN, err)
-	}
-
-	p.logger.Info(fmt.Sprintf("started replication on slot %s at startLSN: %d", p.slot, opts.startLSN))
-	return nil
-}
-
-func (p *PostgresCDCSource) replicationOptions() (*pglogrepl.StartReplicationOptions, error) {
-	pluginArguments := []string{
-		"proto_version '1'",
-	}
-
-	if p.publication != "" {
-		pubOpt := fmt.Sprintf("publication_names '%s'", p.publication)
-		pluginArguments = append(pluginArguments, pubOpt)
-	} else {
-		return nil, errors.New("publication name is not set")
-	}
-
-	return &pglogrepl.StartReplicationOptions{PluginArgs: pluginArguments}, nil
-}
-
-// start consuming the cdc stream
-func (p *PostgresCDCSource) consumeStream(
-	ctx context.Context,
-	conn *pgconn.PgConn,
-	req *model.PullRecordsRequest,
-	clientXLogPos pglogrepl.LSN,
-	records *model.CDCRecordStream,
-) error {
-	defer func() {
-		err := conn.Close(ctx)
-		if err != nil {
-			p.logger.Error("error closing replication connection", slog.Any("error", err))
-		}
-	}()
-
+	conn := p.replConn.PgConn()
+	records := req.RecordStream
 	// clientXLogPos is the last checkpoint id, we need to ack that we have processed
 	// until clientXLogPos each time we send a standby status update.
 	// consumedXLogPos is the lsn that has been committed on the destination.
-	consumedXLogPos := pglogrepl.LSN(0)
-	if clientXLogPos > 0 {
+	var clientXLogPos, consumedXLogPos pglogrepl.LSN
+	if req.LastOffset > 0 {
+		clientXLogPos = pglogrepl.LSN(req.LastOffset)
 		consumedXLogPos = clientXLogPos
 
 		err := pglogrepl.SendStandbyStatusUpdate(ctx, conn,
@@ -302,7 +229,6 @@ func (p *PostgresCDCSource) consumeStream(
 
 		var receiveCtx context.Context
 		var cancel context.CancelFunc
-
 		if cdcRecordsStorage.IsEmpty() {
 			receiveCtx, cancel = context.WithCancel(ctx)
 		} else {