From 7d40862fbcc593d2bdb814f3aca8733f3cf46774 Mon Sep 17 00:00:00 2001 From: subotac <73706465+subotac@users.noreply.github.com> Date: Tue, 28 Jul 2026 22:21:06 +0300 Subject: [PATCH 1/2] Make target connection limits configurable --- cmd/config/config_env.go | 3 +- cmd/config/config_env_test.go | 1 + cmd/config/config_yaml.go | 2 + cmd/config/helper_test.go | 3 +- cmd/config/test/test_config.env | 1 + cmd/config/test/test_config.yaml | 1 + config_template.yaml | 1 + docs/configuration.md | 2 + internal/postgres/pg_conn_pool.go | 40 ++++++++-- internal/postgres/pg_conn_pool_test.go | 45 +++++++++++ pkg/wal/processor/postgres/config.go | 9 +++ .../postgres/postgres_bulk_ingest_writer.go | 8 +- .../postgres_bulk_ingest_writer_test.go | 75 +++++++++++++++++++ .../postgres/postgres_schema_observer.go | 2 +- pkg/wal/processor/postgres/postgres_writer.go | 13 +++- 15 files changed, 191 insertions(+), 15 deletions(-) create mode 100644 internal/postgres/pg_conn_pool_test.go diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go index 579c5c5f..bf6c7d5d 100644 --- a/cmd/config/config_env.go +++ b/cmd/config/config_env.go @@ -582,7 +582,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, diff --git a/cmd/config/config_env_test.go b/cmd/config/config_env_test.go index e1962adc..2bfa53c6 100644 --- a/cmd/config/config_env_test.go +++ b/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") diff --git a/cmd/config/config_yaml.go b/cmd/config/config_yaml.go index 63865635..5f6dd93a 100644 --- a/cmd/config/config_yaml.go +++ b/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, diff --git a/cmd/config/helper_test.go b/cmd/config/helper_test.go index 97113a02..8eab0fdc 100644 --- a/cmd/config/helper_test.go +++ b/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, diff --git a/cmd/config/test/test_config.env b/cmd/config/test/test_config.env index 88e5b6c8..6f7ec7ed 100644 --- a/cmd/config/test/test_config.env +++ b/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 diff --git a/cmd/config/test/test_config.yaml b/cmd/config/test/test_config.yaml index 2d874c8b..cffa785c 100644 --- a/cmd/config/test/test_config.yaml +++ b/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 diff --git a/config_template.yaml b/config_template.yaml index 8759ed26..86f4f3ee 100644 --- a/config_template.yaml +++ b/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 diff --git a/docs/configuration.md b/docs/configuration.md index a95e2fe8..825995b3 100644 --- a/docs/configuration.md +++ b/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. | diff --git a/internal/postgres/pg_conn_pool.go b/internal/postgres/pg_conn_pool.go index 899e94c3..db2b2455 100644 --- a/internal/postgres/pg_conn_pool.go +++ b/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 { diff --git a/internal/postgres/pg_conn_pool_test.go b/internal/postgres/pg_conn_pool_test.go new file mode 100644 index 00000000..12c30d35 --- /dev/null +++ b/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) + }) + } +} diff --git a/pkg/wal/processor/postgres/config.go b/pkg/wal/processor/postgres/config.go index 73bc1f3d..c26e1b92 100644 --- a/pkg/wal/processor/postgres/config.go +++ b/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))} +} diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go index 7785119f..d6623893 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go +++ b/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) { diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go index ffbba270..563cc276 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go +++ b/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) + }) + } +} diff --git a/pkg/wal/processor/postgres/postgres_schema_observer.go b/pkg/wal/processor/postgres/postgres_schema_observer.go index 5f50859a..48d215fd 100644 --- a/pkg/wal/processor/postgres/postgres_schema_observer.go +++ b/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 diff --git a/pkg/wal/processor/postgres/postgres_writer.go b/pkg/wal/processor/postgres/postgres_writer.go index 74654102..583c7213 100644 --- a/pkg/wal/processor/postgres/postgres_writer.go +++ b/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 { From 1cfefe38606d02fdfd5656508a1be4d649c8d431 Mon Sep 17 00:00:00 2001 From: subotac <73706465+subotac@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:42:33 +0300 Subject: [PATCH 2/2] fix(config): bind target max connections env var --- cmd/config/config_env.go | 1 + 1 file changed, 1 insertion(+) diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go index bf6c7d5d..3578ce8b 100644 --- a/cmd/config/config_env.go +++ b/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")