Make target connection limits configurable

cmd/config/config_env.go

@@ -74,6 +74,7 @@ func init() {
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_DISABLE_PROGRESS_TRACKING")
 
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_TARGET_URL")
+	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_MAX_CONNECTIONS")
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT")
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_BATCH_BYTES")
 	viper.BindEnv("PGSTREAM_POSTGRES_WRITER_BATCH_SIZE")
@@ -582,7 +583,8 @@ func parsePostgresProcessorConfig() (*stream.PostgresProcessorConfig, error) {
 	bulkIngestEnabled := viper.GetBool("PGSTREAM_POSTGRES_WRITER_BULK_INGEST_ENABLED")
 	cfg := &stream.PostgresProcessorConfig{
 		BatchWriter: postgres.Config{
-			URL: targetPostgresURL,
+			URL:            targetPostgresURL,
+			MaxConnections: viper.GetUint("PGSTREAM_POSTGRES_WRITER_MAX_CONNECTIONS"),
 			BatchConfig: batch.Config{
 				BatchTimeout:     viper.GetDuration("PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT"),
 				MaxBatchBytes:    maxBatchBytes,

cmd/config/config_env_test.go

@@ -68,6 +68,7 @@ func Test_EnvVarsToStreamConfig(t *testing.T) {
 	os.Setenv("PGSTREAM_KAFKA_TLS_CLIENT_KEY_FILE", "/path/to/client.key")
 
 	os.Setenv("PGSTREAM_POSTGRES_WRITER_TARGET_URL", "postgresql://user:password@localhost:5432/mytargetdatabase")
+	os.Setenv("PGSTREAM_POSTGRES_WRITER_MAX_CONNECTIONS", "60")
 	os.Setenv("PGSTREAM_POSTGRES_WRITER_BATCH_SIZE", "100")
 	os.Setenv("PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT", "1s")
 	os.Setenv("PGSTREAM_POSTGRES_WRITER_MAX_QUEUE_BYTES", "204800")

cmd/config/config_yaml.go

@@ -191,6 +191,7 @@ type ConstantBackoffConfig struct {
 
 type PostgresTargetConfig struct {
 	URL              string            `mapstructure:"url" yaml:"url"`
+	MaxConnections   uint              `mapstructure:"max_connections" yaml:"max_connections"`
 	Batch            *BatchConfig      `mapstructure:"batch" yaml:"batch"`
 	BulkIngest       *BulkIngestConfig `mapstructure:"bulk_ingest" yaml:"bulk_ingest"`
 	DisableTriggers  bool              `mapstructure:"disable_triggers" yaml:"disable_triggers"`
@@ -702,6 +703,7 @@ func (c *YAMLConfig) parsePostgresProcessorConfig() *stream.PostgresProcessorCon
 	cfg := &stream.PostgresProcessorConfig{
 		BatchWriter: postgres.Config{
 			URL:              c.Target.Postgres.URL,
+			MaxConnections:   c.Target.Postgres.MaxConnections,
 			BatchConfig:      c.Target.Postgres.Batch.parseBatchConfig(),
 			DisableTriggers:  c.Target.Postgres.DisableTriggers,
 			OnConflictAction: c.Target.Postgres.OnConflictAction,

cmd/config/helper_test.go

@@ -132,7 +132,8 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) {
 		Processor: stream.ProcessorConfig{
 			Postgres: &stream.PostgresProcessorConfig{
 				BatchWriter: postgres.Config{
-					URL: "postgresql://user:password@localhost:5432/mytargetdatabase",
+					URL:            "postgresql://user:password@localhost:5432/mytargetdatabase",
+					MaxConnections: 60,
 					BatchConfig: batch.Config{
 						MaxBatchSize:     100,
 						BatchTimeout:     time.Second,

cmd/config/test/test_config.env

@@ -51,6 +51,7 @@ PGSTREAM_KAFKA_TLS_CLIENT_KEY_FILE="/path/to/client.key"
 
 # Postgres
 PGSTREAM_POSTGRES_WRITER_TARGET_URL="postgresql://user:password@localhost:5432/mytargetdatabase"
+PGSTREAM_POSTGRES_WRITER_MAX_CONNECTIONS=60
 PGSTREAM_POSTGRES_WRITER_BATCH_SIZE=100
 PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT="1s"
 PGSTREAM_POSTGRES_WRITER_MAX_QUEUE_BYTES=204800

cmd/config/test/test_config.yaml

@@ -66,6 +66,7 @@ source:
 target:
   postgres:
     url: "postgresql://user:password@localhost:5432/mytargetdatabase"
+    max_connections: 60 # maximum number of connections per pool to the target database
     batch:
       timeout: 1000 # batch timeout in milliseconds
       size: 100 # number of messages in a batch

config_template.yaml

@@ -79,6 +79,7 @@ source:
 target:
   postgres:
     url: "postgresql://user:password@localhost:5432/mytargetdatabase"
+    max_connections: 50 # maximum number of connections per pool to the target database. Defaults to 50; overrides pool_max_conns in the URL when set.
     batch:
       timeout: 1000 # batch timeout in milliseconds. Defaults to 30s
       size: 100 # number of messages in a batch. Defaults to 20000

docs/configuration.md

@@ -91,6 +91,7 @@ source:
 target:
   postgres:
     url: "postgresql://user:password@localhost:5432/mytargetdatabase"
+    max_connections: 50 # maximum number of connections per pool to the target database. Defaults to 50; overrides pool_max_conns in the URL when set.
     batch:
       timeout: 1000 # batch timeout in milliseconds. Defaults to 30s
       size: 100 # number of messages in a batch. Defaults to 20000
@@ -365,6 +366,7 @@ One of exponential/constant/disable retries backoff policies can be provided for
 | Environment Variable                                           | Default                         | Required | Description                                                                                                                                                                                                    |
 | -------------------------------------------------------------- | ------------------------------- | -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | PGSTREAM_POSTGRES_WRITER_TARGET_URL                            | N/A                             | Yes      | URL for the PostgreSQL store to connect to                                                                                                                                                                     |
+| PGSTREAM_POSTGRES_WRITER_MAX_CONNECTIONS                       | 50                              | No       | Maximum number of connections per pool to the target PostgreSQL database. Overrides `pool_max_conns` in the target URL when set.                                                                               |
 | PGSTREAM_POSTGRES_WRITER_BATCH_TIMEOUT                         | 30s                             | No       | Max time interval at which the batch sending to PostgreSQL is triggered.                                                                                                                                       |
 | PGSTREAM_POSTGRES_WRITER_BATCH_SIZE                            | 20000                           | No       | Max number of messages to be sent per batch. When this size is reached, the batch is sent to PostgreSQL.                                                                                                       |
 | PGSTREAM_POSTGRES_WRITER_MAX_QUEUE_BYTES                       | 104857600 (100MiB)              | No       | Max memory used by the postgres batch writer for inflight batches.                                                                                                                                             |

internal/postgres/pg_conn_pool.go

@@ -17,20 +17,41 @@ type Pool struct {
 type PoolOption func(*pgxpool.Config)
 
 // MaxConns is the default maximum number of connections in a Postgres
-// connection pool. It also bounds the global concurrent-COPY budget in the
-// bulk-ingest writer, so the two never drift.
+// connection pool.
 const MaxConns = 50
 
 func NewConnPool(ctx context.Context, url string, opts ...PoolOption) (*Pool, error) {
+	pgCfg, err := newConnPoolConfig(url, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	pool, err := pgxpool.NewWithConfig(ctx, pgCfg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create a postgres connection pool: %w", MapError(err))
+	}
+
+	return &Pool{Pool: pool}, nil
+}
+
+func newConnPoolConfig(url string, opts ...PoolOption) (*pgxpool.Config, error) {
 	escapedURL, err := escapeConnectionURL(url)
 	if err != nil {
 		return nil, fmt.Errorf("failed to escape connection URL: %w", err)
 	}
+	connCfg, err := pgx.ParseConfig(escapedURL)
+	if err != nil {
+		return nil, fmt.Errorf("failed parsing postgres connection string: %w", MapError(err))
+	}
+	_, maxConnsConfigured := connCfg.RuntimeParams["pool_max_conns"]
+
 	pgCfg, err := pgxpool.ParseConfig(escapedURL)
 	if err != nil {
 		return nil, fmt.Errorf("failed parsing postgres connection string: %w", MapError(err))
 	}
-	pgCfg.MaxConns = MaxConns
+	if !maxConnsConfigured {
+		pgCfg.MaxConns = MaxConns
+	}
 	pgCfg.AfterConnect = registerTypesToConnMap
 
 	configureTCPKeepalive(pgCfg.ConnConfig)
@@ -39,12 +60,17 @@ func NewConnPool(ctx context.Context, url string, opts ...PoolOption) (*Pool, er
 		opt(pgCfg)
 	}
 
-	pool, err := pgxpool.NewWithConfig(ctx, pgCfg)
+	return pgCfg, nil
+}
+
+// ConnPoolMaxConnections returns the resolved maximum number of connections
+// after applying the connection string, the pgstream default, and pool options.
+func ConnPoolMaxConnections(url string, opts ...PoolOption) (int32, error) {
+	pgCfg, err := newConnPoolConfig(url, opts...)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create a postgres connection pool: %w", MapError(err))
+		return 0, err
 	}
-
-	return &Pool{Pool: pool}, nil
+	return pgCfg.MaxConns, nil
 }
 
 func WithMaxConnections(maxConns int32) PoolOption {

internal/postgres/pg_conn_pool_test.go

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewConnPool_maxConnections(t *testing.T) {
+	tests := []struct {
+		name     string
+		url      string
+		opts     []PoolOption
+		expected int32
+	}{
+		{
+			name:     "default",
+			url:      "postgresql://user:password@localhost:5432/database",
+			expected: MaxConns,
+		},
+		{
+			name:     "connection URL",
+			url:      "postgresql://user:password@localhost:5432/database?pool_max_conns=12",
+			expected: 12,
+		},
+		{
+			name:     "pool option overrides connection URL",
+			url:      "postgresql://user:password@localhost:5432/database?pool_max_conns=12",
+			opts:     []PoolOption{WithMaxConnections(24)},
+			expected: 24,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pool, err := NewConnPool(t.Context(), tt.url, tt.opts...)
+			require.NoError(t, err)
+			t.Cleanup(func() { require.NoError(t, pool.Close(t.Context())) })
+
+			require.Equal(t, tt.expected, pool.Config().MaxConns)
+		})
+	}
+}

pkg/wal/processor/postgres/config.go

@@ -5,12 +5,14 @@ package postgres
 import (
 	"time"
 
+	pglib "github.com/xataio/pgstream/internal/postgres"
 	"github.com/xataio/pgstream/pkg/backoff"
 	"github.com/xataio/pgstream/pkg/wal/processor/batch"
 )
 
 type Config struct {
 	URL               string
+	MaxConnections    uint
 	BatchConfig       batch.Config
 	DisableTriggers   bool
 	OnConflictAction  string
@@ -36,3 +38,10 @@ func (c *Config) retryPolicy() backoff.Config {
 		},
 	}
 }
+
+func (c *Config) poolOptions() []pglib.PoolOption {
+	if c.MaxConnections == 0 {
+		return nil
+	}
+	return []pglib.PoolOption{pglib.WithMaxConnections(int32(c.MaxConnections))}
+}

pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go

@@ -26,7 +26,7 @@ type BulkIngestWriter struct {
 	batchSenderBuilder func(ctx context.Context, schema, table string) (queryBatchSender, error)
 	// copyBudget caps the total number of concurrent COPYs across all tables
 	// (and all their send drainers) so they never exhaust the target
-	// connection pool. It is sized from the same pool max-connections constant,
+	// connection pool. It is sized from the resolved pool max-connections value,
 	// minus a reserve for the writer's own metadata/type queries.
 	copyBudget synclib.WeightedSemaphore
 }
@@ -56,7 +56,7 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti
 	biw := &BulkIngestWriter{
 		Writer:         w,
 		batchSenderMap: synclib.NewMap[string, queryBatchSender](),
-		copyBudget:     synclib.NewWeightedSemaphore(int64(pglib.MaxConns - copyBudgetReserve)),
+		copyBudget:     synclib.NewWeightedSemaphore(copyBudgetSize(w.maxConnections)),
 	}
 
 	biw.batchSenderBuilder = func(ctx context.Context, schema, table string) (queryBatchSender, error) {
@@ -70,6 +70,10 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti
 	return biw, nil
 }
 
+func copyBudgetSize(maxConnections int32) int64 {
+	return max(1, int64(maxConnections)-copyBudgetReserve)
+}
+
 // ProcessWALEvent is called on every new message from the wal. It can be called
 // concurrently.
 func (w *BulkIngestWriter) ProcessWALEvent(ctx context.Context, walEvent *wal.Event) (err error) {

pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go

@@ -14,6 +14,7 @@ import (
 	pglib "github.com/xataio/pgstream/internal/postgres"
 	pgmocks "github.com/xataio/pgstream/internal/postgres/mocks"
 	synclib "github.com/xataio/pgstream/internal/sync"
+	"github.com/xataio/pgstream/pkg/backoff"
 	loglib "github.com/xataio/pgstream/pkg/log"
 	"github.com/xataio/pgstream/pkg/wal"
 	"github.com/xataio/pgstream/pkg/wal/processor"
@@ -488,3 +489,77 @@ func TestBulkIngestWriter_sendBatch_copyBudget(t *testing.T) {
 	require.NoError(t, eg.Wait())
 	require.Equal(t, int32(budget), maxActive.Load(), "concurrent COPYs must not exceed the budget")
 }
+
+func TestCopyBudgetSize(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name           string
+		maxConnections int32
+		expected       int64
+	}{
+		{name: "default pool", maxConnections: pglib.MaxConns, expected: 45},
+		{name: "configured pool", maxConnections: 12, expected: 7},
+		{name: "reserve matches pool", maxConnections: copyBudgetReserve, expected: 1},
+		{name: "pool smaller than reserve", maxConnections: 2, expected: 1},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			require.Equal(t, tt.expected, copyBudgetSize(tt.maxConnections))
+		})
+	}
+}
+
+func TestNewBulkIngestWriter_maxConnections(t *testing.T) {
+	tests := []struct {
+		name           string
+		url            string
+		maxConnections uint
+		expected       int32
+	}{
+		{
+			name:     "connection URL",
+			url:      "postgresql://user:password@localhost:5432/database?pool_max_conns=12",
+			expected: 12,
+		},
+		{
+			name:           "writer config overrides connection URL",
+			url:            "postgresql://user:password@localhost:5432/database?pool_max_conns=12",
+			maxConnections: 20,
+			expected:       20,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			writer, err := NewBulkIngestWriter(t.Context(), &Config{
+				URL:            tt.url,
+				MaxConnections: tt.maxConnections,
+				RetryPolicy:    backoff.Config{DisableRetries: true},
+			})
+			require.NoError(t, err)
+			t.Cleanup(func() { require.NoError(t, writer.Close()) })
+
+			writerPool, ok := writer.pgConn.(*pglib.Pool)
+			require.True(t, ok)
+			require.Equal(t, tt.expected, writerPool.Config().MaxConns)
+
+			adapter, ok := writer.adapter.(*adapter)
+			require.True(t, ok)
+			observer, ok := adapter.schemaObserver.(*pgSchemaObserver)
+			require.True(t, ok)
+			observerPool, ok := observer.pgConn.(*pglib.Pool)
+			require.True(t, ok)
+			require.Equal(t, tt.expected, observerPool.Config().MaxConns)
+
+			budget := copyBudgetSize(tt.expected)
+			for range budget {
+				require.True(t, writer.copyBudget.TryAcquire(1))
+			}
+			require.False(t, writer.copyBudget.TryAcquire(1))
+			writer.copyBudget.Release(budget)
+		})
+	}
+}

pkg/wal/processor/postgres/postgres_schema_observer.go

@@ -56,7 +56,7 @@ type pgSchemaObserver struct {
 // DDL event is received through the WAL.
 func newPGSchemaObserver(ctx context.Context, cfg *Config, logger loglib.Logger) (*pgSchemaObserver, error) {
 	newConnPool := func(ctx context.Context) (pglib.Querier, error) {
-		return pglib.NewConnPool(ctx, cfg.URL)
+		return pglib.NewConnPool(ctx, cfg.URL, cfg.poolOptions()...)
 	}
 
 	// the observer sits on the hot path: an unretried transient failure here

pkg/wal/processor/postgres/postgres_writer.go

@@ -26,6 +26,7 @@ type Writer struct {
 	writerType      string
 	disableTriggers bool
 	strictMode      bool
+	maxConnections  int32
 
 	droppedQueries       atomic.Uint64
 	instrumentation      *otel.Instrumentation
@@ -49,11 +50,18 @@ type walMessageBatchSender interface {
 type WriterOption func(*Writer)
 
 func newWriter(ctx context.Context, config *Config, writerType string, opts ...WriterOption) (*Writer, error) {
+	poolOpts := config.poolOptions()
+	maxConnections, err := pglib.ConnPoolMaxConnections(config.URL, poolOpts...)
+	if err != nil {
+		return nil, err
+	}
+
 	w := &Writer{
 		logger:          loglib.NewNoopLogger(),
 		writerType:      writerType,
 		disableTriggers: config.DisableTriggers,
 		strictMode:      config.StrictMode,
+		maxConnections:  maxConnections,
 		dropped:         batch.NewDroppedCounter(),
 	}
 
@@ -78,14 +86,13 @@ func newWriter(ctx context.Context, config *Config, writerType string, opts ...W
 		w.logger.Info("strict_mode is disabled: non-internal query failures will be dropped and counted rather than stopping the pipeline")
 	}
 
-	var err error
 	if config.RetryPolicy.DisableRetries {
-		w.pgConn, err = pglib.NewConnPool(ctx, config.URL)
+		w.pgConn, err = pglib.NewConnPool(ctx, config.URL, poolOpts...)
 	} else {
 		// unless retries are disabled, wrap the Postgres querier with a retrier
 		// and apply default retry policy if none is set
 		w.pgConn, err = pglibretrier.NewQuerier(ctx, config.retryPolicy(), func(ctx context.Context) (pglib.Querier, error) {
-			return pglib.NewConnPool(ctx, config.URL)
+			return pglib.NewConnPool(ctx, config.URL, poolOpts...)
 		}, w.logger)
 	}
 	if err != nil {