From 070e68a2a88ae98042b69ad61497e147ab67022b Mon Sep 17 00:00:00 2001 From: Mason Sharp Date: Sun, 5 Apr 2026 19:54:46 -0700 Subject: [PATCH 1/3] Cap database connection pool size for table-diff and repset-diff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, pgxpool defaulted MaxConns to max(4, NumCPU), meaning the connection pool was unbounded relative to the concurrency factor. On high-core machines this could exhaust database connections. Pool size is now derived from the concurrency factor (NumCPU × ConcurrencyFactor, minimum 4). A new --max-connections / -m flag (and table_diff.max_connections in ace.yaml) lets users set a hard cap that overrides the derived value. Adds integration test that sets MaxConnections=2 with 16 workers and asserts the cap holds strictly. Co-Authored-By: Claude Opus 4.6 (1M context) --- ace.yaml | 1 + docs/api.md | 3 + docs/commands/diff/repset-diff.md | 1 + docs/commands/diff/table-diff.md | 1 + docs/configuration.md | 1 + docs/design/table_diff.md | 3 +- docs/http-api.md | 3 + docs/openapi.yaml | 9 + docs/performance.md | 3 + internal/cli/cli.go | 8 + internal/consistency/diff/repset_diff.go | 3 + internal/consistency/diff/table_diff.go | 15 +- internal/jobs/config.go | 12 + pkg/config/config.go | 1 + .../integration/repset_diff_conntrack_test.go | 240 ++++++++++++++++++ 15 files changed, 302 insertions(+), 2 deletions(-) create mode 100644 tests/integration/repset_diff_conntrack_test.go diff --git a/ace.yaml b/ace.yaml index 27b1643..3875e39 100644 --- a/ace.yaml +++ b/ace.yaml @@ -28,6 +28,7 @@ table_diff: diff_batch_size: 1 max_diff_batch_size: 1000 compare_unit_size: 10000 + max_connections: 0 # max DB connections per node (0 = derive from concurrency factor) mtree: cdc: diff --git a/docs/api.md b/docs/api.md index 30608c5..c5e3967 100644 --- a/docs/api.md +++ b/docs/api.md @@ -33,6 +33,7 @@ Compares a table between nodes and generates a diff report. | `--output` | `-o` | Output format (`json` or `html`) | json | | `--nodes` | `-n` | Nodes to include in the diff (comma-separated, or "all") | all | | `--table-filter` | `-F` | `WHERE` clause expression to use while diffing tables | | +| `--max-connections` | `-m` | Maximum database connections per node | derived | | `--override-block-size`| `-B` | Allow block sizes outside `ace.yaml` guardrails | false | | `--quiet` | `-q` | Suppress progress output | false | | `--debug` | `-v` | Enable debug logging | false | @@ -147,6 +148,7 @@ Compares schemas across nodes in a pgEdge cluster. By default, `schema-diff` per | `--output` | `-o` | Output format (`json` or `html`) | json | | `--nodes` | `-n` | Nodes to include in the diff (comma-separated, or "all") | all | | `--table-filter` | `-F` | `WHERE` clause expression to use while diffing tables | | +| `--max-connections` | `-m` | Maximum database connections per node | derived | | `--override-block-size`| `-B` | Allow block sizes outside `ace.yaml` guardrails | false | | `--quiet` | `-q` | Suppress progress output | false | | `--debug` | `-v` | Enable debug logging | false | @@ -189,6 +191,7 @@ Performs a `table-diff` on every table in a replication set and reports differen | `--concurrency-factor` | `-c` | CPU ratio for concurrency (0.0–4.0) | 0.5 | | `--compare-unit-size` | `-u` | Recursive split size for mismatched blocks | 10000 | | `--output` | `-o` | Per-table diff output format (`json` or `html`) | json | +| `--max-connections` | `-m` | Maximum database connections per node | derived | | `--override-block-size` | `-B` | Allow block sizes outside `ace.yaml` guardrails | false | | `--quiet` | `-q` | Suppress output | false | | `--debug` | `-v` | Enable debug logging | false | diff --git a/docs/commands/diff/repset-diff.md b/docs/commands/diff/repset-diff.md index 824663c..932bd4f 100644 --- a/docs/commands/diff/repset-diff.md +++ b/docs/commands/diff/repset-diff.md @@ -25,6 +25,7 @@ Runs `table-diff` on every table in a replication set and reports differences. | `--concurrency-factor ` | `-c` | CPU ratio for concurrency (0.0–4.0). Default `0.5`. | | `--compare-unit-size ` | `-u` | Recursive split size for mismatched blocks. Default `10000`. | | `--output ` | `-o` | Per-table diff report format. Default `json`. | +| `--max-connections ` | `-m` | Maximum database connections per node. Caps the pool regardless of concurrency factor. | derived | | `--override-block-size` | `-B` | Allow block sizes outside `ace.yaml` guardrails. | | `--quiet` | `-q` | Suppress output | `false` | | `--debug` | `-v` | Debug logging | `false` | diff --git a/docs/commands/diff/table-diff.md b/docs/commands/diff/table-diff.md index d1930a4..2720750 100644 --- a/docs/commands/diff/table-diff.md +++ b/docs/commands/diff/table-diff.md @@ -24,6 +24,7 @@ This command compares the data in the specified table across nodes in a cluster | `--table-filter ` | `-F` | Optional SQL `WHERE` clause applied on every node before hashing. | | `--against-origin ` | | Limit the diff to rows whose `node_origin` matches this Spock node id or name (useful for failed-node recovery). | | `--until ` | | Optional commit timestamp fence (RFC3339) applied with `--against-origin` and `--table-filter`; excludes newer rows. | +| `--max-connections ` | `-m` | Maximum database connections per node. When set, caps the connection pool regardless of concurrency factor. Default: derived from `--concurrency-factor`. | | `--override-block-size` | `-B` | Skip block-size safety checks defined in `ace.yaml`. | | `--quiet` | `-q` | Suppress progress output. Results still write to the diff file. | | `--debug` | `-v` | Enable verbose logging. | diff --git a/docs/configuration.md b/docs/configuration.md index dc968ea..a7ae7ac 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -32,6 +32,7 @@ The [`ace.yaml` file](https://github.com/pgEdge/ace/blob/main/ace.yaml) defines | table_diff --> min_diff_block_size | Minimum diff block (row chunk) size. **Default: 1** | | table_diff --> max_diff_block_size | Maximum diff block size. **Default: 1000000** | | table_diff --> compare_unit_size | Unit size for smallest comparison chunk. **Default: 10000** | +| table_diff --> max_connections | Maximum database connections per node during diff operations. When set, caps the connection pool regardless of concurrency factor. **Default: 0** (derive from concurrency factor) | | mtree → cdc --> slot_name | Logical decoding slot name for mtree CDC. **Default: "ace_mtree_slot"** | | mtree → cdc --> publication_name | Publication used for mtree CDC. **Default: "ace_mtree_pub"** | | mtree → cdc --> cdc_processing_timeout | CDC processing timeout (s). **Default: 30** | diff --git a/docs/design/table_diff.md b/docs/design/table_diff.md index d1759a1..d7f6623 100644 --- a/docs/design/table_diff.md +++ b/docs/design/table_diff.md @@ -162,7 +162,8 @@ flowchart LR ### Resource Utilisation and Tuning - **block_size**: Larger blocks reduce hash tasks and recursion but increase memory/IO per hash and slow mismatch localisation; smaller blocks do the opposite (more queries, finer locality). -- **concurrency_factor**: CPU ratio (0.0–4.0) that scales workers relative to `NumCPU` (e.g. 0.5 on a 16-CPU host spawns 8 workers). Higher = faster hashing but more load on DB backends, network, and local CPU; can contend with other workloads and connection limits. +- **concurrency_factor**: CPU ratio (0.0–4.0) that scales workers relative to `NumCPU` (e.g. 0.5 on a 16-CPU host spawns 8 workers). Higher = faster hashing but more load on DB backends, network, and local CPU; can contend with other workloads and connection limits. The connection pool per node is sized to match the worker count (minimum 4). +- **max_connections**: Hard cap on the connection pool size per node. When set, overrides the concurrency-derived pool size. Useful for environments with limited `max_connections` on the database server. Workers that exceed the pool size will queue for a connection rather than fail. - **compare_unit_size**: Lower values push recursion deeper (more queries, smaller fetches); higher values stop earlier (fewer queries, larger fetches on mismatched ranges). - **max_diff_rows**: Early-exit guardrail. Lower caps keep runs short and reports small on divergent tables; raising/removing can grow memory and report size when drift is large. - **table_filter**: Narrows scope and cost; enables accurate `COUNT(*)` on the filtered view. Must be identical across nodes to avoid false positives. diff --git a/docs/http-api.md b/docs/http-api.md index c3bc3f1..2b06b6c 100644 --- a/docs/http-api.md +++ b/docs/http-api.md @@ -93,6 +93,7 @@ Request body: | `concurrency_factor` | float | no | CPU ratio (0.0–4.0). Defaults to `table_diff.concurrency_factor` or `0.5`. | | `compare_unit_size` | int | no | Defaults to `table_diff.compare_unit_size` or `10000`. | | `max_diff_rows` | int64 | no | Defaults to `table_diff.max_diff_rows` or `0` (no limit). | +| `max_connections` | int | no | Max DB connections per node. Defaults to `table_diff.max_connections` or derived from concurrency factor. | | `table_filter` | string | no | SQL `WHERE` predicate (without `WHERE`). | | `override_block_size` | bool | no | Bypass block-size guardrails. | | `quiet` | bool | no | Suppress progress output. | @@ -192,6 +193,7 @@ Request body: | `block_size` | int | no | Defaults to `table_diff.diff_block_size` or `100000`. | | `concurrency_factor` | float | no | CPU ratio (0.0–4.0). Defaults to `table_diff.concurrency_factor` or `0.5`. | | `compare_unit_size` | int | no | Defaults to `table_diff.compare_unit_size` or `10000`. | +| `max_connections` | int | no | Max DB connections per node. Defaults to `table_diff.max_connections` or derived from concurrency factor. | | `output` | string | no | `json` (default) or `html`. | | `override_block_size` | bool | no | Bypass block-size guardrails. | | `quiet` | bool | no | Suppress output. | @@ -213,6 +215,7 @@ Request body: | `block_size` | int | no | Defaults to `table_diff.diff_block_size` or `100000`. | | `concurrency_factor` | float | no | CPU ratio (0.0–4.0). Defaults to `table_diff.concurrency_factor` or `0.5`. | | `compare_unit_size` | int | no | Defaults to `table_diff.compare_unit_size` or `10000`. | +| `max_connections` | int | no | Max DB connections per node. Defaults to `table_diff.max_connections` or derived from concurrency factor. | | `output` | string | no | `json` (default) or `html`. | | `override_block_size` | bool | no | Bypass block-size guardrails. | | `quiet` | bool | no | Suppress output. | diff --git a/docs/openapi.yaml b/docs/openapi.yaml index b565cdf..dc47879 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -465,6 +465,9 @@ components: max_diff_rows: type: integer description: Defaults to table_diff.max_diff_rows or 0 (no limit). + max_connections: + type: integer + description: Maximum database connections per node. Defaults to table_diff.max_connections or derived from concurrency factor. table_filter: type: string description: SQL WHERE predicate (without WHERE). @@ -567,6 +570,9 @@ components: compare_unit_size: type: integer description: Defaults to table_diff.compare_unit_size or 10000. + max_connections: + type: integer + description: Maximum database connections per node. Defaults to table_diff.max_connections or derived from concurrency factor. output: type: string enum: [json, html] @@ -601,6 +607,9 @@ components: compare_unit_size: type: integer description: Defaults to table_diff.compare_unit_size or 10000. + max_connections: + type: integer + description: Maximum database connections per node. Defaults to table_diff.max_connections or derived from concurrency factor. output: type: string enum: [json, html] diff --git a/docs/performance.md b/docs/performance.md index 12679f7..596f7e0 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -34,6 +34,9 @@ When invoking [ACE commands](commands/index.md), review the available command op - **`--compare-unit-size`** Sets the minimum block size used when ACE recursively drills into mismatched blocks. Smaller values provide more granular comparisons at the cost of additional round-trips. Default is `10000`. +- **`--max-connections`** + Caps the number of database connections ACE opens per node. By default, the pool size is derived from `--concurrency-factor` and the number of CPUs. On machines with many cores, this can result in a large number of connections. Use `--max-connections` to set a hard upper limit, or set `table_diff.max_connections` in `ace.yaml` to apply it globally. This is especially useful when running ACE against databases with limited `max_connections` or when sharing the database with other applications. + - **`--override-block-size` (`-B`)** Override the safety limits defined in `ace.yaml`. Use cautiously: extremely large blocks can lead to `array_agg` memory pressure. diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 2ed0e87..305744f 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -102,6 +102,12 @@ func SetupCLI() *cli.Command { Usage: "Override block size", Value: false, }, + &cli.IntFlag{ + Name: "max-connections", + Aliases: []string{"m"}, + Usage: "Maximum number of database connections per node (default: derived from concurrency factor)", + Value: 0, + }, } rerunOnlyFlags := []cli.Flag{ @@ -879,6 +885,7 @@ func TableDiffCLI(cmd *cli.Command) error { task.DBName = cmd.String("dbname") task.BlockSize = int(blockSizeInt) task.ConcurrencyFactor = cmd.Float64("concurrency-factor") + task.MaxConnections = cmd.Int("max-connections") task.CompareUnitSize = cmd.Int("compare-unit-size") task.Output = strings.ToLower(cmd.String("output")) task.Nodes = cmd.String("nodes") @@ -1350,6 +1357,7 @@ func RepsetDiffCLI(cmd *cli.Command) error { task.BlockSize = int(blockSizeInt) task.ConcurrencyFactor = cmd.Float64("concurrency-factor") + task.MaxConnections = cmd.Int("max-connections") task.CompareUnitSize = cmd.Int("compare-unit-size") task.Output = cmd.String("output") task.OverrideBlockSize = cmd.Bool("override-block-size") diff --git a/internal/consistency/diff/repset_diff.go b/internal/consistency/diff/repset_diff.go index f30ca2a..80335b5 100644 --- a/internal/consistency/diff/repset_diff.go +++ b/internal/consistency/diff/repset_diff.go @@ -50,6 +50,7 @@ type RepsetDiffCmd struct { database types.Database ConnectionPool *pgxpool.Pool ConcurrencyFactor float64 + MaxConnections int BlockSize int CompareUnitSize int Output string @@ -367,6 +368,7 @@ func RepsetDiff(task *RepsetDiffCmd) (err error) { tdTask.Nodes = task.Nodes tdTask.QualifiedTableName = tableName tdTask.ConcurrencyFactor = task.ConcurrencyFactor + tdTask.MaxConnections = task.MaxConnections tdTask.BlockSize = task.BlockSize tdTask.CompareUnitSize = task.CompareUnitSize tdTask.Output = task.Output @@ -422,6 +424,7 @@ func (task *RepsetDiffCmd) CloneForSchedule(ctx context.Context) *RepsetDiffCmd clone.Quiet = task.Quiet clone.BlockSize = task.BlockSize clone.ConcurrencyFactor = task.ConcurrencyFactor + clone.MaxConnections = task.MaxConnections clone.CompareUnitSize = task.CompareUnitSize clone.Output = task.Output clone.TableFilter = task.TableFilter diff --git a/internal/consistency/diff/table_diff.go b/internal/consistency/diff/table_diff.go index 4ca5bb1..2055c22 100644 --- a/internal/consistency/diff/table_diff.go +++ b/internal/consistency/diff/table_diff.go @@ -70,6 +70,7 @@ type TableDiffTask struct { BlockSize int ConcurrencyFactor float64 + MaxConnections int Output string TableFilter string QuietMode bool @@ -1027,6 +1028,7 @@ func (t *TableDiffTask) CloneForSchedule(ctx context.Context) *TableDiffTask { cloned.Nodes = t.Nodes cloned.BlockSize = t.BlockSize cloned.ConcurrencyFactor = t.ConcurrencyFactor + cloned.MaxConnections = t.MaxConnections cloned.Output = t.Output cloned.TableFilter = t.TableFilter cloned.QuietMode = t.QuietMode @@ -1046,8 +1048,19 @@ func (t *TableDiffTask) CloneForSchedule(ctx context.Context) *TableDiffTask { return cloned } +func (t *TableDiffTask) maxPoolSize() int { + poolSize := int(math.Round(float64(runtime.NumCPU()) * t.ConcurrencyFactor)) + if poolSize < 4 { + poolSize = 4 + } + if t.MaxConnections > 0 && t.MaxConnections < poolSize { + poolSize = t.MaxConnections + } + return poolSize +} + func (t *TableDiffTask) connOpts() auth.ConnectionOptions { - return auth.ConnectionOptions{} + return auth.ConnectionOptions{PoolSize: t.maxPoolSize()} } func (t *TableDiffTask) CheckColumnSize() error { diff --git a/internal/jobs/config.go b/internal/jobs/config.go index e547cc5..3e11093 100644 --- a/internal/jobs/config.go +++ b/internal/jobs/config.go @@ -109,6 +109,9 @@ func buildTableDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec) if v := intArg(def.Args, "max_diff_rows", 0); v > 0 { base.MaxDiffRows = int64(v) } + if v := intArg(def.Args, "max_connections", 0); v > 0 { + base.MaxConnections = v + } if out := stringArg(def.Args, "output"); out != "" { base.Output = out } @@ -133,6 +136,9 @@ func buildTableDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec) if base.CompareUnitSize == 0 && cfg.TableDiff.CompareUnitSize > 0 { base.CompareUnitSize = cfg.TableDiff.CompareUnitSize } + if base.MaxConnections == 0 && cfg.TableDiff.MaxConnections > 0 { + base.MaxConnections = cfg.TableDiff.MaxConnections + } return Job{ Name: jobName(def, "table-diff", base.QualifiedTableName), @@ -252,6 +258,9 @@ func buildRepsetDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if out := stringArg(def.Args, "output"); out != "" { base.Output = out } + if v := intArg(def.Args, "max_connections", 0); v > 0 { + base.MaxConnections = v + } base.TableFilter = stringArg(def.Args, "table_filter") base.OverrideBlockSize = boolArg(def.Args, "override_block_size", base.OverrideBlockSize) base.SkipDBUpdate = boolArg(def.Args, "skip_db_update", base.SkipDBUpdate) @@ -272,6 +281,9 @@ func buildRepsetDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if base.CompareUnitSize == 0 && cfg.TableDiff.CompareUnitSize > 0 { base.CompareUnitSize = cfg.TableDiff.CompareUnitSize } + if base.MaxConnections == 0 && cfg.TableDiff.MaxConnections > 0 { + base.MaxConnections = cfg.TableDiff.MaxConnections + } return Job{ Name: jobName(def, "repset-diff", base.RepsetName), diff --git a/pkg/config/config.go b/pkg/config/config.go index 1bd1e31..b6611d0 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -54,6 +54,7 @@ type DiffConfig struct { BatchSize int `yaml:"diff_batch_size"` MaxBatchSize int `yaml:"max_diff_batch_size"` CompareUnitSize int `yaml:"compare_unit_size"` + MaxConnections int `yaml:"max_connections"` } type MTreeConfig struct { diff --git a/tests/integration/repset_diff_conntrack_test.go b/tests/integration/repset_diff_conntrack_test.go new file mode 100644 index 0000000..3575ea2 --- /dev/null +++ b/tests/integration/repset_diff_conntrack_test.go @@ -0,0 +1,240 @@ +// /////////////////////////////////////////////////////////////////////////// +// +// # ACE - Active Consistency Engine +// +// Copyright (C) 2023 - 2026, pgEdge (https://www.pgedge.com/) +// +// This software is released under the PostgreSQL License: +// https://opensource.org/license/postgresql +// +// /////////////////////////////////////////////////////////////////////////// + +package integration + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "github.com/pgedge/ace/internal/consistency/diff" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// connSnapshot holds a point-in-time reading of ACE connections on a node. +type connSnapshot struct { + Time time.Time + Node string + ACEConns int + TotalConn int +} + +// monitorConnections polls pg_stat_activity on the given pools at the specified +// interval, recording the number of connections where application_name = 'ACE'. +// It stops when the returned cancel function is called, and returns all +// collected snapshots. +func monitorConnections( + ctx context.Context, + pools map[string]*pgxpool.Pool, + interval time.Duration, +) (snapshots *[]connSnapshot, stop func()) { + + var mu sync.Mutex + var snaps []connSnapshot + snapshots = &snaps + + monCtx, cancel := context.WithCancel(ctx) + var wg sync.WaitGroup + wg.Add(1) + + go func() { + defer wg.Done() + ticker := time.NewTicker(interval) + defer ticker.Stop() + + sample := func() { + for name, pool := range pools { + var aceCount, totalCount int + err := pool.QueryRow(monCtx, ` + SELECT + count(*) FILTER (WHERE application_name = 'ACE') AS ace_conns, + count(*) AS total_conns + FROM pg_stat_activity + WHERE pid <> pg_backend_pid() + `).Scan(&aceCount, &totalCount) + if err != nil { + continue // context cancelled or transient error + } + mu.Lock() + snaps = append(snaps, connSnapshot{ + Time: time.Now(), + Node: name, + ACEConns: aceCount, + TotalConn: totalCount, + }) + mu.Unlock() + } + } + + // Take an initial sample immediately. + sample() + for { + select { + case <-monCtx.Done(): + // Take a final sample after the diff has finished. + sample() + return + case <-ticker.C: + sample() + } + } + }() + + stop = func() { + cancel() + wg.Wait() + // Update the caller's pointer since append may have reallocated. + *snapshots = snaps + } + return snapshots, stop +} + +// TestRepsetDiff_MaxConnectionsCap verifies that setting MaxConnections limits +// the number of database connections per node, even when the concurrency factor +// would allow more. On a 4-vCPU GitHub Actions runner with concurrency factor +// 4.0, this creates 16 workers but caps the pool at 2 connections per node. +// Workers queue for connections rather than opening new ones. +func TestRepsetDiff_MaxConnectionsCap(t *testing.T) { + const ( + numTables = 5 + rowsPerTable = 50000 + repsetName = "default" + concurrencyFactor = 4.0 + maxConnections = 2 + pollInterval = 50 * time.Millisecond + ) + + ctx := context.Background() + + tableNames := make([]string, numTables) + for i := 0; i < numTables; i++ { + tableName := fmt.Sprintf("maxconn_%03d", i) + qualifiedName := fmt.Sprintf("%s.%s", testSchema, tableName) + tableNames[i] = qualifiedName + + createSQL := fmt.Sprintf(` + CREATE TABLE IF NOT EXISTS %s ( + id INT PRIMARY KEY, + val TEXT + )`, qualifiedName) + + pools := []*pgxpool.Pool{pgCluster.Node1Pool, pgCluster.Node2Pool} + nodeNames := []string{serviceN1, serviceN2} + + for j, pool := range pools { + _, err := pool.Exec(ctx, createSQL) + require.NoError(t, err, "create table %s on %s", tableName, nodeNames[j]) + + _, err = pool.Exec(ctx, fmt.Sprintf( + `INSERT INTO %s (id, val) + SELECT g, 'row_' || g FROM generate_series(1, %d) g + ON CONFLICT DO NOTHING`, qualifiedName, rowsPerTable)) + require.NoError(t, err, "insert rows into %s on %s", tableName, nodeNames[j]) + } + + // Diverge every other table before adding to repset. + if i%2 == 0 { + _, err := pgCluster.Node1Pool.Exec(ctx, fmt.Sprintf( + `UPDATE %s SET val = 'diverged_n1' WHERE id <= 50`, qualifiedName)) + require.NoError(t, err, "diverge %s on n1", tableName) + } + + for j, pool := range pools { + _, err := pool.Exec(ctx, fmt.Sprintf( + `SELECT spock.repset_add_table('%s', '%s')`, repsetName, qualifiedName)) + require.NoError(t, err, "add %s to repset on %s", tableName, nodeNames[j]) + } + } + + t.Cleanup(func() { + for _, qn := range tableNames { + for _, pool := range []*pgxpool.Pool{pgCluster.Node1Pool, pgCluster.Node2Pool} { + pool.Exec(ctx, fmt.Sprintf( + `SELECT spock.repset_remove_table('%s', '%s')`, repsetName, qn)) + pool.Exec(ctx, fmt.Sprintf(`DROP TABLE IF EXISTS %s CASCADE`, qn)) + } + } + }) + + monitorPools := map[string]*pgxpool.Pool{ + serviceN1: pgCluster.Node1Pool, + serviceN2: pgCluster.Node2Pool, + } + + baseline := make(map[string]int) + for name, pool := range monitorPools { + var aceCount int + err := pool.QueryRow(ctx, ` + SELECT count(*) FROM pg_stat_activity + WHERE application_name = 'ACE' AND pid <> pg_backend_pid() + `).Scan(&aceCount) + require.NoError(t, err) + baseline[name] = aceCount + } + + snapshots, stopMonitor := monitorConnections(ctx, monitorPools, pollInterval) + + task := newTestRepsetDiffTask(repsetName) + task.ConcurrencyFactor = concurrencyFactor + task.MaxConnections = maxConnections + task.BlockSize = 1000 + task.Quiet = true + + err := diff.RepsetDiff(task) + require.NoError(t, err, "repset-diff should succeed") + + stopMonitor() + time.Sleep(500 * time.Millisecond) + + t.Logf("Collected %d connection snapshots", len(*snapshots)) + + peakPerNode := make(map[string]int) + for _, snap := range *snapshots { + above := snap.ACEConns - baseline[snap.Node] + if above > peakPerNode[snap.Node] { + peakPerNode[snap.Node] = above + } + } + + for node, peak := range peakPerNode { + t.Logf("Node %s: peak ACE connections above baseline = %d", node, peak) + } + + // With MaxConnections=2, peak should never exceed 2 per node, regardless + // of concurrency factor (16 workers on 4 CPUs). The pool queues excess + // workers. Allow no margin — the cap must hold strictly. + for node, peak := range peakPerNode { + assert.LessOrEqual(t, peak, maxConnections, + "node %s peak ACE connections (%d) exceeds max_connections (%d) — "+ + "MaxConnections cap is not being applied", + node, peak, maxConnections) + } + + // Verify no leaks. + for name, pool := range monitorPools { + var aceCount int + err := pool.QueryRow(ctx, ` + SELECT count(*) FROM pg_stat_activity + WHERE application_name = 'ACE' AND pid <> pg_backend_pid() + `).Scan(&aceCount) + require.NoError(t, err) + leaked := aceCount - baseline[name] + t.Logf("Node %s: post-diff ACE connections = %d (baseline=%d, leaked=%d)", + name, aceCount, baseline[name], leaked) + assert.LessOrEqual(t, leaked, 1, + "node %s should not leak connections (leaked %d)", name, leaked) + } +} From 2c33ea3ae94c7140ee9e686437fc1e494686294b Mon Sep 17 00:00:00 2001 From: Mason Sharp Date: Sun, 5 Apr 2026 21:30:32 -0700 Subject: [PATCH 2/3] Wire max_connections through HTTP API, schema-diff, and CLI fallback Addresses CodeRabbit review findings: - Add MaxConnections field to HTTP request structs for table-diff, schema-diff, and repset-diff endpoints with resolveMaxConnections helper, so API clients can set the connection cap. - Add MaxConnections to SchemaDiffCmd (struct, per-table wiring, CloneForSchedule) and schema-diff CLI/jobs builders. - Apply ace.yaml table_diff.max_connections fallback in TableDiffTask.Validate() so the config cap is honored on all paths (CLI, HTTP, scheduler), not just scheduled jobs. - Validate max_connections >= 1 when set (reject negative values). - Change alias from -m to -M to avoid collision with mtree's --max-cpu-ratio alias. - Add TestRepsetDiff_MaxConnectionsCap to CI workflow. Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/workflows/test.yml | 3 +++ docs/api.md | 2 +- docs/commands/diff/repset-diff.md | 2 +- docs/commands/diff/table-diff.md | 2 +- internal/api/http/handler.go | 16 ++++++++++++++++ internal/cli/cli.go | 3 ++- internal/consistency/diff/schema_diff.go | 3 +++ internal/consistency/diff/table_diff.go | 6 ++++++ internal/jobs/config.go | 6 ++++++ 9 files changed, 39 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a40ce25..a7ee2c2 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -42,3 +42,6 @@ jobs: - name: Run crash recovery origin-filter test run: go test -count=1 -v ./tests/integration -run 'TestAdvancedRepairPlan' + + - name: Run connection pool cap test + run: go test -count=1 -v ./tests/integration -run 'TestRepsetDiff_MaxConnectionsCap' diff --git a/docs/api.md b/docs/api.md index c5e3967..4b8ed1e 100644 --- a/docs/api.md +++ b/docs/api.md @@ -33,7 +33,7 @@ Compares a table between nodes and generates a diff report. | `--output` | `-o` | Output format (`json` or `html`) | json | | `--nodes` | `-n` | Nodes to include in the diff (comma-separated, or "all") | all | | `--table-filter` | `-F` | `WHERE` clause expression to use while diffing tables | | -| `--max-connections` | `-m` | Maximum database connections per node | derived | +| `--max-connections` | `-M` | Maximum database connections per node | derived | | `--override-block-size`| `-B` | Allow block sizes outside `ace.yaml` guardrails | false | | `--quiet` | `-q` | Suppress progress output | false | | `--debug` | `-v` | Enable debug logging | false | diff --git a/docs/commands/diff/repset-diff.md b/docs/commands/diff/repset-diff.md index 932bd4f..396e02d 100644 --- a/docs/commands/diff/repset-diff.md +++ b/docs/commands/diff/repset-diff.md @@ -25,7 +25,7 @@ Runs `table-diff` on every table in a replication set and reports differences. | `--concurrency-factor ` | `-c` | CPU ratio for concurrency (0.0–4.0). Default `0.5`. | | `--compare-unit-size ` | `-u` | Recursive split size for mismatched blocks. Default `10000`. | | `--output ` | `-o` | Per-table diff report format. Default `json`. | -| `--max-connections ` | `-m` | Maximum database connections per node. Caps the pool regardless of concurrency factor. | derived | +| `--max-connections ` | `-M` | Maximum database connections per node. Caps the pool regardless of concurrency factor. | derived | | `--override-block-size` | `-B` | Allow block sizes outside `ace.yaml` guardrails. | | `--quiet` | `-q` | Suppress output | `false` | | `--debug` | `-v` | Debug logging | `false` | diff --git a/docs/commands/diff/table-diff.md b/docs/commands/diff/table-diff.md index 2720750..a3c3e1b 100644 --- a/docs/commands/diff/table-diff.md +++ b/docs/commands/diff/table-diff.md @@ -24,7 +24,7 @@ This command compares the data in the specified table across nodes in a cluster | `--table-filter ` | `-F` | Optional SQL `WHERE` clause applied on every node before hashing. | | `--against-origin ` | | Limit the diff to rows whose `node_origin` matches this Spock node id or name (useful for failed-node recovery). | | `--until ` | | Optional commit timestamp fence (RFC3339) applied with `--against-origin` and `--table-filter`; excludes newer rows. | -| `--max-connections ` | `-m` | Maximum database connections per node. When set, caps the connection pool regardless of concurrency factor. Default: derived from `--concurrency-factor`. | +| `--max-connections ` | `-M` | Maximum database connections per node. When set, caps the connection pool regardless of concurrency factor. Default: derived from `--concurrency-factor`. | | `--override-block-size` | `-B` | Skip block-size safety checks defined in `ace.yaml`. | | `--quiet` | `-q` | Suppress progress output. Results still write to the diff file. | | `--debug` | `-v` | Enable verbose logging. | diff --git a/internal/api/http/handler.go b/internal/api/http/handler.go index d978534..3c8409a 100644 --- a/internal/api/http/handler.go +++ b/internal/api/http/handler.go @@ -26,6 +26,7 @@ type tableDiffRequest struct { Concurrency float64 `json:"concurrency_factor"` CompareUnitSize int `json:"compare_unit_size"` MaxDiffRows int64 `json:"max_diff_rows"` + MaxConnections int `json:"max_connections"` TableFilter string `json:"table_filter"` OverrideBlockSize bool `json:"override_block_size"` Quiet bool `json:"quiet"` @@ -76,6 +77,7 @@ type schemaDiffRequest struct { BlockSize int `json:"block_size"` Concurrency float64 `json:"concurrency_factor"` CompareUnitSize int `json:"compare_unit_size"` + MaxConnections int `json:"max_connections"` Output string `json:"output"` OverrideBlockSize bool `json:"override_block_size"` Quiet bool `json:"quiet"` @@ -91,6 +93,7 @@ type repsetDiffRequest struct { BlockSize int `json:"block_size"` Concurrency float64 `json:"concurrency_factor"` CompareUnitSize int `json:"compare_unit_size"` + MaxConnections int `json:"max_connections"` Output string `json:"output"` OverrideBlockSize bool `json:"override_block_size"` Quiet bool `json:"quiet"` @@ -212,6 +215,7 @@ func (s *APIServer) handleTableDiff(w http.ResponseWriter, r *http.Request) { task.ConcurrencyFactor = s.resolveConcurrency(cfg, req.Concurrency) task.CompareUnitSize = s.resolveCompareUnitSize(cfg, req.CompareUnitSize) task.MaxDiffRows = s.resolveMaxDiffRows(cfg, req.MaxDiffRows) + task.MaxConnections = s.resolveMaxConnections(cfg, req.MaxConnections) task.Output = "json" task.Nodes = s.resolveNodes(req.Nodes) task.TableFilter = strings.TrimSpace(req.TableFilter) @@ -291,6 +295,16 @@ func (s *APIServer) resolveMaxDiffRows(cfg *config.Config, requested int64) int6 return 0 } +func (s *APIServer) resolveMaxConnections(cfg *config.Config, requested int) int { + if requested > 0 { + return requested + } + if cfg != nil && cfg.TableDiff.MaxConnections > 0 { + return cfg.TableDiff.MaxConnections + } + return 0 +} + func (s *APIServer) resolveNodes(nodes []string) string { if len(nodes) == 0 { return "all" @@ -577,6 +591,7 @@ func (s *APIServer) handleSchemaDiff(w http.ResponseWriter, r *http.Request) { task.BlockSize = s.resolveBlockSize(cfg, req.BlockSize) task.ConcurrencyFactor = s.resolveConcurrency(cfg, req.Concurrency) task.CompareUnitSize = s.resolveCompareUnitSize(cfg, req.CompareUnitSize) + task.MaxConnections = s.resolveMaxConnections(cfg, req.MaxConnections) task.Output = strings.TrimSpace(req.Output) if task.Output == "" { task.Output = "json" @@ -659,6 +674,7 @@ func (s *APIServer) handleRepsetDiff(w http.ResponseWriter, r *http.Request) { task.BlockSize = s.resolveBlockSize(cfg, req.BlockSize) task.ConcurrencyFactor = s.resolveConcurrency(cfg, req.Concurrency) task.CompareUnitSize = s.resolveCompareUnitSize(cfg, req.CompareUnitSize) + task.MaxConnections = s.resolveMaxConnections(cfg, req.MaxConnections) task.Output = strings.TrimSpace(req.Output) if task.Output == "" { task.Output = "json" diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 305744f..9a69236 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -104,7 +104,7 @@ func SetupCLI() *cli.Command { }, &cli.IntFlag{ Name: "max-connections", - Aliases: []string{"m"}, + Aliases: []string{"M"}, Usage: "Maximum number of database connections per node (default: derived from concurrency factor)", Value: 0, }, @@ -1282,6 +1282,7 @@ func SchemaDiffCLI(cmd *cli.Command) error { task.BlockSize = int(blockSizeInt) task.ConcurrencyFactor = cmd.Float64("concurrency-factor") + task.MaxConnections = cmd.Int("max-connections") task.CompareUnitSize = cmd.Int("compare-unit-size") task.Output = cmd.String("output") task.OverrideBlockSize = cmd.Bool("override-block-size") diff --git a/internal/consistency/diff/schema_diff.go b/internal/consistency/diff/schema_diff.go index 93ed2b9..7b1689b 100644 --- a/internal/consistency/diff/schema_diff.go +++ b/internal/consistency/diff/schema_diff.go @@ -52,6 +52,7 @@ type SchemaDiffCmd struct { database types.Database ConnectionPool *pgxpool.Pool ConcurrencyFactor float64 + MaxConnections int BlockSize int CompareUnitSize int Output string @@ -523,6 +524,7 @@ func (task *SchemaDiffCmd) SchemaTableDiff() (err error) { tdTask.Nodes = task.Nodes tdTask.QualifiedTableName = qualifiedTableName tdTask.ConcurrencyFactor = task.ConcurrencyFactor + tdTask.MaxConnections = task.MaxConnections tdTask.BlockSize = task.BlockSize tdTask.CompareUnitSize = task.CompareUnitSize tdTask.Output = task.Output @@ -577,6 +579,7 @@ func (task *SchemaDiffCmd) CloneForSchedule(ctx context.Context) *SchemaDiffCmd clone.DDLOnly = task.DDLOnly clone.BlockSize = task.BlockSize clone.ConcurrencyFactor = task.ConcurrencyFactor + clone.MaxConnections = task.MaxConnections clone.CompareUnitSize = task.CompareUnitSize clone.Output = task.Output clone.TableFilter = task.TableFilter diff --git a/internal/consistency/diff/table_diff.go b/internal/consistency/diff/table_diff.go index 2055c22..ce4499d 100644 --- a/internal/consistency/diff/table_diff.go +++ b/internal/consistency/diff/table_diff.go @@ -729,6 +729,12 @@ func (t *TableDiffTask) Validate() error { } cfg := config.Get() + if t.MaxConnections == 0 && cfg.TableDiff.MaxConnections > 0 { + t.MaxConnections = cfg.TableDiff.MaxConnections + } + if t.MaxConnections < 0 { + return fmt.Errorf("max_connections must be >= 1 (or 0 to derive from concurrency factor)") + } if t.BlockSize > cfg.TableDiff.MaxBlockSize && !t.OverrideBlockSize { return fmt.Errorf("block row size should be <= %d", cfg.TableDiff.MaxBlockSize) } diff --git a/internal/jobs/config.go b/internal/jobs/config.go index 3e11093..422d5a2 100644 --- a/internal/jobs/config.go +++ b/internal/jobs/config.go @@ -187,6 +187,9 @@ func buildSchemaDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if v := intArg(def.Args, "compare_unit_size", 0); v > 0 { base.CompareUnitSize = v } + if v := intArg(def.Args, "max_connections", 0); v > 0 { + base.MaxConnections = v + } if out := stringArg(def.Args, "output"); out != "" { base.Output = out } @@ -210,6 +213,9 @@ func buildSchemaDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if base.CompareUnitSize == 0 && cfg.TableDiff.CompareUnitSize > 0 { base.CompareUnitSize = cfg.TableDiff.CompareUnitSize } + if base.MaxConnections == 0 && cfg.TableDiff.MaxConnections > 0 { + base.MaxConnections = cfg.TableDiff.MaxConnections + } return Job{ Name: jobName(def, "schema-diff", base.SchemaName), From 092e5b859b952445a418935bd5cd2a7e0b2d013e Mon Sep 17 00:00:00 2001 From: Mason Sharp Date: Sun, 5 Apr 2026 21:49:47 -0700 Subject: [PATCH 3/3] Address CodeRabbit review: validate negative max_connections on all paths - Let negative values pass through resolvers and config fallbacks instead of silently treating them as unset, so Validate() rejects them consistently. - Add max_connections validation to RepsetDiffCmd.Validate() and SchemaDiffCmd.Validate() to fail fast before per-table fan-out. - Simplify jobs config assignments to unconditional intArg calls. Co-Authored-By: Claude Opus 4.6 (1M context) --- internal/api/http/handler.go | 3 +++ internal/consistency/diff/repset_diff.go | 3 +++ internal/consistency/diff/schema_diff.go | 4 ++++ internal/consistency/diff/table_diff.go | 2 +- internal/jobs/config.go | 12 +++--------- 5 files changed, 14 insertions(+), 10 deletions(-) diff --git a/internal/api/http/handler.go b/internal/api/http/handler.go index 3c8409a..088ed55 100644 --- a/internal/api/http/handler.go +++ b/internal/api/http/handler.go @@ -299,6 +299,9 @@ func (s *APIServer) resolveMaxConnections(cfg *config.Config, requested int) int if requested > 0 { return requested } + if requested < 0 { + return requested // rejected by Validate() + } if cfg != nil && cfg.TableDiff.MaxConnections > 0 { return cfg.TableDiff.MaxConnections } diff --git a/internal/consistency/diff/repset_diff.go b/internal/consistency/diff/repset_diff.go index 80335b5..644abb0 100644 --- a/internal/consistency/diff/repset_diff.go +++ b/internal/consistency/diff/repset_diff.go @@ -231,6 +231,9 @@ func (c *RepsetDiffCmd) Validate() error { if c.RepsetName == "" { return fmt.Errorf("repset name is required") } + if c.MaxConnections < 0 { + return fmt.Errorf("max_connections must be >= 1 (or 0 to derive from concurrency factor)") + } return nil } diff --git a/internal/consistency/diff/schema_diff.go b/internal/consistency/diff/schema_diff.go index 7b1689b..8402dfd 100644 --- a/internal/consistency/diff/schema_diff.go +++ b/internal/consistency/diff/schema_diff.go @@ -178,6 +178,10 @@ func (c *SchemaDiffCmd) Validate() error { return fmt.Errorf("schema-diff needs at least two nodes to compare") } + if c.MaxConnections < 0 { + return fmt.Errorf("max_connections must be >= 1 (or 0 to derive from concurrency factor)") + } + return nil } diff --git a/internal/consistency/diff/table_diff.go b/internal/consistency/diff/table_diff.go index ce4499d..2336665 100644 --- a/internal/consistency/diff/table_diff.go +++ b/internal/consistency/diff/table_diff.go @@ -729,7 +729,7 @@ func (t *TableDiffTask) Validate() error { } cfg := config.Get() - if t.MaxConnections == 0 && cfg.TableDiff.MaxConnections > 0 { + if t.MaxConnections == 0 { t.MaxConnections = cfg.TableDiff.MaxConnections } if t.MaxConnections < 0 { diff --git a/internal/jobs/config.go b/internal/jobs/config.go index 422d5a2..a89256b 100644 --- a/internal/jobs/config.go +++ b/internal/jobs/config.go @@ -109,9 +109,7 @@ func buildTableDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec) if v := intArg(def.Args, "max_diff_rows", 0); v > 0 { base.MaxDiffRows = int64(v) } - if v := intArg(def.Args, "max_connections", 0); v > 0 { - base.MaxConnections = v - } + base.MaxConnections = intArg(def.Args, "max_connections", 0) if out := stringArg(def.Args, "output"); out != "" { base.Output = out } @@ -187,9 +185,7 @@ func buildSchemaDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if v := intArg(def.Args, "compare_unit_size", 0); v > 0 { base.CompareUnitSize = v } - if v := intArg(def.Args, "max_connections", 0); v > 0 { - base.MaxConnections = v - } + base.MaxConnections = intArg(def.Args, "max_connections", 0) if out := stringArg(def.Args, "output"); out != "" { base.Output = out } @@ -264,9 +260,7 @@ func buildRepsetDiffJob(cfg *config.Config, def config.JobDef, spec scheduleSpec if out := stringArg(def.Args, "output"); out != "" { base.Output = out } - if v := intArg(def.Args, "max_connections", 0); v > 0 { - base.MaxConnections = v - } + base.MaxConnections = intArg(def.Args, "max_connections", 0) base.TableFilter = stringArg(def.Args, "table_filter") base.OverrideBlockSize = boolArg(def.Args, "override_block_size", base.OverrideBlockSize) base.SkipDBUpdate = boolArg(def.Args, "skip_db_update", base.SkipDBUpdate)