PG: clean up replication pings, unify wal_sender_timeout handling (#4363)

* Make ReplPing private as it's only called within the PG connector
* Unify wal_sender_timeout handling - started as deduping the constant
but we'd really want to reduce catalog accesses on every batch and
sticking to a canoncial value seems good. The struct thing still seems
messy but that's the least messy shape I could think of without losing
nuance

Author: ilidemi

docs/peerdb-architecture.md

@@ -228,7 +228,6 @@ classDiagram
         +EnsurePullability()
         +ExportTxSnapshot()
         +SetupReplication()
-        +ReplPing()
         +PullFlowCleanup()
     }
 

flow/connectors/bigquery/qrep_object_pull.go

@@ -499,10 +499,6 @@ func (c *BigQueryConnector) SetupReplConn(context.Context, map[string]string) er
 	return nil
 }
 
-func (c *BigQueryConnector) ReplPing(context.Context) error {
-	return nil
-}
-
 func (c *BigQueryConnector) UpdateReplStateLastOffset(context.Context, model.CdcCheckpoint) error {
 	return nil
 }

flow/connectors/core.go

@@ -115,9 +115,6 @@ type CDCPullConnectorCore interface {
 	// Methods related to retrieving and pushing records for this connector as a source and destination.
 	SetupReplConn(context.Context, map[string]string) error
 
-	// Ping source to keep connection alive. Can be called concurrently with pulling records; skips ping in that case.
-	ReplPing(context.Context) error
-
 	// Called when offset has been confirmed to destination
 	UpdateReplStateLastOffset(ctx context.Context, lastOffset model.CdcCheckpoint) error
 

flow/connectors/mongo/cdc.go

@@ -551,10 +551,6 @@ func (c *MongoConnector) SetupReplConn(context.Context, map[string]string) error
 	return nil
 }
 
-func (c *MongoConnector) ReplPing(context.Context) error {
-	return nil
-}
-
 func (c *MongoConnector) UpdateReplStateLastOffset(ctx context.Context, lastOffset model.CdcCheckpoint) error {
 	return nil
 }

flow/connectors/mysql/cdc.go

@@ -352,10 +352,6 @@ func (c *MySqlConnector) closeSyncerWithTimeout(syncer *replication.BinlogSyncer
 	}
 }
 
-func (c *MySqlConnector) ReplPing(context.Context) error {
-	return nil
-}
-
 func (c *MySqlConnector) UpdateReplStateLastOffset(ctx context.Context, lastOffset model.CdcCheckpoint) error {
 	flowName := ctx.Value(shared.FlowNameKey).(string)
 	return c.SetLastOffset(ctx, flowName, lastOffset)

flow/connectors/postgres/cdc.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"log/slog"
 	"maps"
-	"regexp"
 	"slices"
 	"strings"
 	"sync"
@@ -476,8 +475,6 @@ func clientSidePingPeriod(walSenderTimeout time.Duration) time.Duration {
 	return 3 * walSenderTimeout / 4
 }
 
-const NoWALSenderTimeoutSetting = "NONE"
-
 // PullCdcRecords pulls records from req's cdc stream
 func PullCdcRecords[Items model.Items](
 	ctx context.Context,
@@ -502,34 +499,16 @@ func PullCdcRecords[Items model.Items](
 	}
 
 	// determine message wait period in function of idle and wal_sender timeouts
-	var walSenderTimeout time.Duration
-	if walSenderTimeoutStr, err := internal.PeerDBPostgresWalSenderTimeout(ctx, req.Env); err != nil {
-		return fmt.Errorf("could't get wal_sender_timeout parameter: %w", err)
-	} else if walSenderTimeoutStr != NoWALSenderTimeoutSetting {
-		// If no units are provided, milliseconds are assumed
-		if regexp.MustCompile(`^\d+$`).MatchString(walSenderTimeoutStr) {
-			walSenderTimeoutStr += "ms"
-		}
-		if walSenderTimeout, err = time.ParseDuration(walSenderTimeoutStr); err != nil {
-			return fmt.Errorf("failed to parse wal_sender_timeout value: %w", err)
-		}
-		if walSenderTimeout < 0 {
-			return fmt.Errorf("invalid wal_sender_timeout value: %s", walSenderTimeout)
-		}
-	}
 	// this value controls for how long the main message loop is blocked waiting for new messages from Postgres.
-	var messageWaitPeriod time.Duration
-	if walSenderTimeout <= 0 {
-		// If no WAL sender timeout is set
-		messageWaitPeriod = req.IdleTimeout
-	} else {
+	messageWaitPeriod := req.IdleTimeout
+	if p.walSenderTimeout.isSet && p.walSenderTimeout.duration > 0 {
 		// If set, consider for ping interval
-		messageWaitPeriod = min(req.IdleTimeout, clientSidePingPeriod(walSenderTimeout))
+		messageWaitPeriod = min(req.IdleTimeout, clientSidePingPeriod(p.walSenderTimeout.duration))
 	}
 
 	logger.Debug("Message wait period determined",
 		slog.Duration("messageWaitPeriod", messageWaitPeriod),
-		slog.Duration("wal_sender_timeout", walSenderTimeout),
+		slog.Duration("wal_sender_timeout", p.walSenderTimeout.duration),
 		slog.Duration("req.IdleTimeout", req.IdleTimeout),
 	)
 

flow/connectors/postgres/client.go

@@ -546,7 +546,7 @@ func (c *PostgresConnector) createSlotAndPublication(
 
 	// create slot only after we succeeded in creating publication.
 	if !s.SlotExists {
-		conn, err := c.CreateReplConn(ctx, env)
+		conn, _, err := c.CreateReplConn(ctx, env)
 		if err != nil {
 			return model.SetupReplicationResult{}, fmt.Errorf("[slot] error acquiring connection: %w", err)
 		}

flow/connectors/postgres/postgres.go

@@ -47,6 +47,7 @@ type PostgresConnector struct {
 	rdsAuth                *utils.RDSAuth
 	connStr                string
 	metadataSchema         string
+	walSenderTimeout       walSenderTimeout
 	replLock               sync.Mutex
 	closeLock              sync.Mutex
 	pgVersion              shared.PGVersion

flow/connectors/postgres/postgres_source.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 	"maps"
+	"regexp"
 	"slices"
 	"strings"
 	"sync/atomic"
@@ -35,12 +36,50 @@ type SlotCheckResult struct {
 	PublicationExists bool
 }
 
-func (c *PostgresConnector) CreateReplConn(ctx context.Context, env map[string]string) (*pgx.Conn, error) {
+type walSenderTimeout struct {
+	str      string
+	duration time.Duration
+	isSet    bool
+}
+
+var walSenderTimeoutNumbersOnly = regexp.MustCompile(`^\d+$`)
+
+func walSenderTimeoutOverride(ctx context.Context, env map[string]string) (walSenderTimeout, error) {
+	str, err := internal.PeerDBPostgresWalSenderTimeout(ctx, env)
+	if err != nil {
+		return walSenderTimeout{}, fmt.Errorf("failed to get wal_sender_timeout value: %w", err)
+	}
+	if strings.EqualFold(str, "NONE") {
+		return walSenderTimeout{
+			str:      str,
+			duration: 0,
+			isSet:    false,
+		}, nil
+	}
+	parseStr := str
+	if walSenderTimeoutNumbersOnly.MatchString(parseStr) {
+		parseStr += "ms" // A bare integer is interpreted as milliseconds, per Postgres
+	}
+	duration, err := time.ParseDuration(parseStr)
+	if err != nil {
+		return walSenderTimeout{}, fmt.Errorf("failed to parse wal_sender_timeout value %q: %w", parseStr, err)
+	}
+	if duration < 0 {
+		return walSenderTimeout{}, fmt.Errorf("invalid wal_sender_timeout value: %s", duration)
+	}
+	return walSenderTimeout{
+		str:      str,
+		duration: duration,
+		isSet:    true,
+	}, nil
+}
+
+func (c *PostgresConnector) CreateReplConn(ctx context.Context, env map[string]string) (*pgx.Conn, walSenderTimeout, error) {
 	// create a separate connection for non-replication queries as replication connections cannot
 	// be used for extended query protocol, i.e. prepared statements
 	replConfig, err := ParseConfig(c.connStr, c.Config)
 	if err != nil {
-		return nil, fmt.Errorf("failed to parse connection string: %w", err)
+		return nil, walSenderTimeout{}, fmt.Errorf("failed to parse connection string: %w", err)
 	}
 
 	replConfig.Config.RuntimeParams["timezone"] = "UTC"
@@ -55,32 +94,33 @@ func (c *PostgresConnector) CreateReplConn(ctx context.Context, env map[string]s
 	replConfig.Config.RuntimeParams["standard_conforming_strings"] = "on"
 	replConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol
 
-	walSenderTimeout, err := internal.PeerDBPostgresWalSenderTimeout(ctx, env)
+	wst, err := walSenderTimeoutOverride(ctx, env)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get wal_sender_timeout value: %w", err)
+		return nil, walSenderTimeout{}, err
 	}
-	if !strings.EqualFold(walSenderTimeout, "NONE") {
-		c.logger.Info("set wal_sender_timeout", slog.String("wal_sender_timeout", walSenderTimeout))
-		replConfig.Config.RuntimeParams["wal_sender_timeout"] = walSenderTimeout
+	if wst.isSet {
+		c.logger.Info("set wal_sender_timeout", slog.String("wal_sender_timeout", wst.str))
+		replConfig.Config.RuntimeParams["wal_sender_timeout"] = wst.str
 	} else {
 		c.logger.Info("not setting wal_sender_timeout")
 	}
 
 	conn, err := NewPostgresConnFromConfig(ctx, replConfig, c.Config.TlsHost, c.rdsAuth, c.ssh)
 	if err != nil {
 		internal.LoggerFromCtx(ctx).Error("failed to create replication connection", slog.Any("error", err))
-		return nil, fmt.Errorf("failed to create replication connection: %w", err)
+		return nil, walSenderTimeout{}, fmt.Errorf("failed to create replication connection: %w", err)
 	}
-	return conn, nil
+	return conn, wst, nil
 }
 
 func (c *PostgresConnector) SetupReplConn(ctx context.Context, env map[string]string) error {
-	conn, err := c.CreateReplConn(ctx, env)
+	conn, wst, err := c.CreateReplConn(ctx, env)
 	if err != nil {
 		return err
 	}
 	c.replConn = conn
-	go c.runReplConnKeepalive(ctx, 15*time.Second, c.ReplPing)
+	c.walSenderTimeout = wst
+	go c.runReplConnKeepalive(ctx, 15*time.Second, c.replPing)
 	return nil
 }
 
@@ -89,7 +129,7 @@ func (c *PostgresConnector) SetupReplConn(ctx context.Context, env map[string]st
 // from the PullRecords stop: (1) between PullRecords calls (2) within
 // PullRecords when AddRecord blocks on a full record stream.
 //
-// While PullRecords holds replLock inside ReceiveMessage, ReplPing's TryLock
+// While PullRecords holds replLock inside ReceiveMessage, replPing's TryLock
 // makes it a no-op. This is safe because ReceiveMessage will always run its
 // own keepalive on timeout.
 func (c *PostgresConnector) runReplConnKeepalive(
@@ -111,8 +151,8 @@ func (c *PostgresConnector) runReplConnKeepalive(
 	}
 }
 
-// ReplPing sends a standby status update so Postgres doesn't terminate the walsender for inactivity.
-func (c *PostgresConnector) ReplPing(ctx context.Context) error {
+// replPing sends a standby status update so Postgres doesn't terminate the walsender for inactivity.
+func (c *PostgresConnector) replPing(ctx context.Context) error {
 	if c.replLock.TryLock() {
 		defer c.replLock.Unlock()
 		if c.replState != nil {

flow/connectors/postgres/validate.go

@@ -204,7 +204,7 @@ func (c *PostgresConnector) CheckReplicationPermissions(ctx context.Context, use
 
 func (c *PostgresConnector) CheckReplicationConnectivity(ctx context.Context, env map[string]string) error {
 	// Check if we can create a replication connection
-	conn, err := c.CreateReplConn(ctx, env)
+	conn, _, err := c.CreateReplConn(ctx, env)
 	if err != nil {
 		return fmt.Errorf("failed to create replication connection: %v", err)
 	}