From 211d8aa49a0a156a8db9846511bc44a4b844ddd3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?No=C3=A9mi=20V=C3=A1nyi?= Date: Thu, 10 Sep 2026 15:02:15 +0200 Subject: [PATCH] Share the retry rule and copy budget helpers --- .../postgres/retrier/pg_querier_retrier.go | 7 +++++++ internal/sync/semaphore.go | 7 +++++++ internal/sync/semaphore_test.go | 18 ++++++++++++++++++ pkg/wal/processor/postgres/config.go | 3 +++ .../postgres/postgres_bulk_ingest_writer.go | 11 ++--------- .../postgres_bulk_ingest_writer_test.go | 8 ++++---- 6 files changed, 41 insertions(+), 13 deletions(-) create mode 100644 internal/sync/semaphore_test.go diff --git a/internal/postgres/retrier/pg_querier_retrier.go b/internal/postgres/retrier/pg_querier_retrier.go index c89f7b7e8..0ed8c22e4 100644 --- a/internal/postgres/retrier/pg_querier_retrier.go +++ b/internal/postgres/retrier/pg_querier_retrier.go @@ -152,6 +152,13 @@ func (q *Querier) resetConn(ctx context.Context) error { } func (q *Querier) isRetriableError(err error) bool { + return IsRetriableError(err) +} + +// IsRetriableError reports whether retrying an operation that failed with err +// could succeed. Callers retrying at a coarser granularity than a single +// query use it so their rule cannot drift from this one. +func IsRetriableError(err error) bool { mappedErr := postgres.MapError(err) permissionDenied := &postgres.ErrPermissionDenied{} diff --git a/internal/sync/semaphore.go b/internal/sync/semaphore.go index 7aac104f6..3d343347c 100644 --- a/internal/sync/semaphore.go +++ b/internal/sync/semaphore.go @@ -17,3 +17,10 @@ type WeightedSemaphore interface { func NewWeightedSemaphore(size int64) *semaphore.Weighted { return semaphore.NewWeighted(size) } + +// CopyBudgetReserve leaves room for non-copy connections +const CopyBudgetReserve = 5 + +func CopyBudgetSize(maxConnections int32) int64 { + return max(1, int64(maxConnections)-CopyBudgetReserve) +} diff --git a/internal/sync/semaphore_test.go b/internal/sync/semaphore_test.go new file mode 100644 index 000000000..de7758341 --- /dev/null +++ b/internal/sync/semaphore_test.go @@ -0,0 +1,18 @@ +// SPDX-License-Identifier: Apache-2.0 + +package sync + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCopyBudgetSize(t *testing.T) { + t.Parallel() + + // zero would block every copy forever + require.Equal(t, int64(1), CopyBudgetSize(1)) + require.Equal(t, int64(1), CopyBudgetSize(CopyBudgetReserve)) + require.Equal(t, int64(45), CopyBudgetSize(50)) +} diff --git a/pkg/wal/processor/postgres/config.go b/pkg/wal/processor/postgres/config.go index bf37b92d6..dd4ab0375 100644 --- a/pkg/wal/processor/postgres/config.go +++ b/pkg/wal/processor/postgres/config.go @@ -50,6 +50,9 @@ func (c *Config) retryPolicy() backoff.Config { } } +// EffectiveRetryPolicy returns the retry policy once the default is applied. +func (c *Config) EffectiveRetryPolicy() backoff.Config { return c.retryPolicy() } + func (c *Config) poolOptions() []pglib.PoolOption { if c.MaxConnections == 0 { return nil diff --git a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go index a214ba81f..886a9a607 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go +++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer.go @@ -27,15 +27,12 @@ type BulkIngestWriter struct { // 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 resolved pool max-connections - // value, minus copyBudgetReserve. + // value, minus synclib.CopyBudgetReserve. copyBudget synclib.WeightedSemaphore } const bulkIngestWriter = "postgres_bulk_ingest_writer" -// batch writer and retrier reset share this pool -const copyBudgetReserve = 5 - var errUnexpectedCopiedRows = errors.New("number of rows copied doesn't match the source rows") // NewBulkIngestWriter returns a postgres processor that batches and writes data @@ -54,7 +51,7 @@ func NewBulkIngestWriter(ctx context.Context, config *Config, opts ...WriterOpti biw := &BulkIngestWriter{ Writer: w, batchSenderMap: synclib.NewMap[string, queryBatchSender](), - copyBudget: synclib.NewWeightedSemaphore(copyBudgetSize(w.maxConnections)), + copyBudget: synclib.NewWeightedSemaphore(synclib.CopyBudgetSize(w.maxConnections)), } biw.batchSenderBuilder = func(ctx context.Context, schema, table string) (queryBatchSender, error) { @@ -68,10 +65,6 @@ 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 704b11541..3d1bebe54 100644 --- a/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go +++ b/pkg/wal/processor/postgres/postgres_bulk_ingest_writer_test.go @@ -424,7 +424,7 @@ func TestBulkIngestWriter_sendBatch(t *testing.T) { pgConn: tc.pgConn, disableTriggers: tc.disableTriggers, }, - copyBudget: synclib.NewWeightedSemaphore(pglib.MaxConns - copyBudgetReserve), + copyBudget: synclib.NewWeightedSemaphore(pglib.MaxConns - synclib.CopyBudgetReserve), } err := writer.sendBatch(context.Background(), tc.batch) @@ -508,14 +508,14 @@ func TestCopyBudgetSize(t *testing.T) { }{ {name: "default pool", maxConnections: pglib.MaxConns, expected: 45}, {name: "configured pool", maxConnections: 12, expected: 7}, - {name: "reserve matches pool", maxConnections: copyBudgetReserve, expected: 1}, + {name: "reserve matches pool", maxConnections: synclib.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)) + require.Equal(t, tt.expected, synclib.CopyBudgetSize(tt.maxConnections)) }) } } @@ -579,7 +579,7 @@ func TestNewBulkIngestWriter_maxConnections(t *testing.T) { require.True(t, ok) require.Equal(t, tt.expectedObserver, observerPool.Config().MaxConns) - budget := copyBudgetSize(tt.expected) + budget := synclib.CopyBudgetSize(tt.expected) for range budget { require.True(t, writer.copyBudget.TryAcquire(1)) }