feat: implement global concurrent connection limit and set default to CPUs * 2 (#222)

* feat: implement global concurrent connection limit and set default to CPUs * 2

* fix: increase MaxConnections in integration test harness

* fix: apply rate limit defaults per-field to preserve partial config

The all-or-nothing defaulting logic was silently overwriting the test
harness's MaxConnections: 100 because MaxFailedAttempts was zero,
causing CI to reject connections on low-CPU runners.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: fuziontech

config_resolution.go

@@ -31,6 +31,7 @@ type configCLIInputs struct {
 	ACMEDomain                string
 	ACMEEmail                 string
 	ACMECacheDir              string
+	MaxConnections            int
 }
 
 type resolvedConfig struct {
@@ -126,6 +127,9 @@ func resolveEffectiveConfig(fileCfg *FileConfig, cli configCLIInputs, getenv fun
 		if fileCfg.RateLimit.MaxConnectionsPerIP > 0 {
 			cfg.RateLimit.MaxConnectionsPerIP = fileCfg.RateLimit.MaxConnectionsPerIP
 		}
+		if fileCfg.RateLimit.MaxConnections > 0 {
+			cfg.RateLimit.MaxConnections = fileCfg.RateLimit.MaxConnections
+		}
 		if fileCfg.RateLimit.FailedAttemptWindow != "" {
 			if d, err := time.ParseDuration(fileCfg.RateLimit.FailedAttemptWindow); err == nil {
 				cfg.RateLimit.FailedAttemptWindow = d
@@ -370,6 +374,13 @@ func resolveEffectiveConfig(fileCfg *FileConfig, cli configCLIInputs, getenv fun
 	if v := getenv("DUCKGRES_ACME_CACHE_DIR"); v != "" {
 		cfg.ACMECacheDir = v
 	}
+	if v := getenv("DUCKGRES_MAX_CONNECTIONS"); v != "" {
+		if n, err := strconv.Atoi(v); err == nil {
+			cfg.RateLimit.MaxConnections = n
+		} else {
+			warn("Invalid DUCKGRES_MAX_CONNECTIONS: " + err.Error())
+		}
+	}
 
 	if cli.Set["host"] {
 		cfg.Host = cli.Host
@@ -454,6 +465,9 @@ func resolveEffectiveConfig(fileCfg *FileConfig, cli configCLIInputs, getenv fun
 	if cli.Set["acme-cache-dir"] {
 		cfg.ACMECacheDir = cli.ACMECacheDir
 	}
+	if cli.Set["max-connections"] {
+		cfg.RateLimit.MaxConnections = cli.MaxConnections
+	}
 
 	// Validate memory_limit format if explicitly set
 	if cfg.MemoryLimit != "" && !server.ValidateMemoryLimit(cfg.MemoryLimit) {

controlplane/memory_rebalancer.go

@@ -207,11 +207,10 @@ func memoryLimit(budget uint64) uint64 {
 	return budget
 }
 
-// DefaultMaxWorkers returns a reasonable default for max_workers based on
-// the memory budget. This is a concurrency cap (budget / 256MB), not a
-// memory-safety guarantee — each session gets the full budget.
+// DefaultMaxWorkers returns a reasonable default for max_workers.
+// Defaults to number of CPUs * 2.
 func (r *MemoryRebalancer) DefaultMaxWorkers() int {
-	return int(r.memoryBudget / minMemoryPerSession)
+	return runtime.NumCPU() * 2
 }
 
 // SetInitialLimits sets memory_limit and threads on a single session synchronously.

duckgres.example.yaml

@@ -116,3 +116,5 @@ rate_limit:
   ban_duration: "15m"
   # Max concurrent connections from a single IP (0 = unlimited)
   max_connections_per_ip: 100
+  # Max total concurrent connections (0 = unlimited, default: CPUs * 2)
+  # max_connections: 16

main.go

@@ -62,6 +62,7 @@ type RateLimitFileConfig struct {
 	FailedAttemptWindow string `yaml:"failed_attempt_window"` // e.g., "5m"
 	BanDuration         string `yaml:"ban_duration"`          // e.g., "15m"
 	MaxConnectionsPerIP int    `yaml:"max_connections_per_ip"`
+	MaxConnections      int    `yaml:"max_connections"`
 }
 
 type DuckLakeFileConfig struct {
@@ -161,6 +162,9 @@ func main() {
 	showVersion := flag.Bool("version", false, "Show version and exit")
 	showHelp := flag.Bool("help", false, "Show help message")
 
+	// Rate limiting flags
+	maxConnections := flag.Int("max-connections", 0, "Max concurrent connections, 0=unlimited (env: DUCKGRES_MAX_CONNECTIONS)")
+
 	// Control plane flags
 	mode := flag.String("mode", "standalone", "Run mode: standalone, control-plane, or duckdb-service")
 	minWorkers := flag.Int("min-workers", 0, "Pre-warm worker count at startup (control-plane mode) (env: DUCKGRES_MIN_WORKERS)")
@@ -206,6 +210,7 @@ func main() {
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_ACME_DOMAIN        Domain for ACME/Let's Encrypt certificate\n")
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_ACME_EMAIL         Contact email for Let's Encrypt notifications\n")
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_ACME_CACHE_DIR     Directory for ACME certificate cache\n")
+		fmt.Fprintf(os.Stderr, "  DUCKGRES_MAX_CONNECTIONS    Max concurrent connections (default: CPUs * 2)\n")
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_DUCKDB_LISTEN      DuckDB service listen address (duckdb-service mode)\n")
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_DUCKDB_TOKEN       DuckDB service bearer token (duckdb-service mode)\n")
 		fmt.Fprintf(os.Stderr, "  DUCKGRES_DUCKDB_MAX_SESSIONS  DuckDB service max sessions (duckdb-service mode)\n")
@@ -289,6 +294,7 @@ func main() {
 		ACMEDomain:                *acmeDomain,
 		ACMEEmail:                 *acmeEmail,
 		ACMECacheDir:              *acmeCacheDir,
+		MaxConnections:            *maxConnections,
 	}, os.Getenv, func(msg string) {
 		slog.Warn(msg)
 	})

server/ratelimit.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"net"
+	"runtime"
 	"sync"
 	"time"
 )
@@ -16,6 +17,8 @@ type RateLimitConfig struct {
 	BanDuration time.Duration
 	// MaxConnectionsPerIP is the max concurrent connections from a single IP (0 = unlimited)
 	MaxConnectionsPerIP int
+	// MaxConnections is the total max concurrent connections (0 = unlimited)
+	MaxConnections int
 }
 
 // DefaultRateLimitConfig returns sensible defaults for rate limiting
@@ -25,6 +28,7 @@ func DefaultRateLimitConfig() RateLimitConfig {
 		FailedAttemptWindow: 5 * time.Minute,
 		BanDuration:         15 * time.Minute,
 		MaxConnectionsPerIP: 100,
+		MaxConnections:      runtime.NumCPU() * 2,
 	}
 }
 
@@ -37,9 +41,10 @@ type ipRecord struct {
 
 // RateLimiter tracks and limits connections per IP
 type RateLimiter struct {
-	mu      sync.Mutex
-	config  RateLimitConfig
-	records map[string]*ipRecord
+	mu               sync.Mutex
+	config           RateLimitConfig
+	records          map[string]*ipRecord
+	totalActiveConns int
 }
 
 // NewRateLimiter creates a new rate limiter with the given config
@@ -76,6 +81,11 @@ func (rl *RateLimiter) CheckConnection(addr net.Addr) string {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
+	// Check global connection limit
+	if rl.config.MaxConnections > 0 && rl.totalActiveConns >= rl.config.MaxConnections {
+		return "too many concurrent connections"
+	}
+
 	record := rl.getOrCreateRecord(ip)
 
 	// Check if IP is banned
@@ -103,6 +113,11 @@ func (rl *RateLimiter) RegisterConnection(addr net.Addr) bool {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
+	// Check global connection limit
+	if rl.config.MaxConnections > 0 && rl.totalActiveConns >= rl.config.MaxConnections {
+		return false
+	}
+
 	record := rl.getOrCreateRecord(ip)
 
 	// Check if banned
@@ -116,6 +131,7 @@ func (rl *RateLimiter) RegisterConnection(addr net.Addr) bool {
 	}
 
 	record.activeConns++
+	rl.totalActiveConns++
 	return true
 }
 
@@ -129,6 +145,11 @@ func (rl *RateLimiter) UnregisterConnection(addr net.Addr) {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
+	rl.totalActiveConns--
+	if rl.totalActiveConns < 0 {
+		rl.totalActiveConns = 0
+	}
+
 	if record, ok := rl.records[ip]; ok {
 		record.activeConns--
 		if record.activeConns < 0 {

server/ratelimit_test.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"net"
+	"runtime"
 	"testing"
 	"time"
 )
@@ -316,6 +317,47 @@ func TestRateLimiter_UnlimitedConnections(t *testing.T) {
 	}
 }
 
+func TestRateLimiter_GlobalConnectionLimit(t *testing.T) {
+	cfg := RateLimitConfig{
+		MaxFailedAttempts:   5,
+		FailedAttemptWindow: 1 * time.Minute,
+		BanDuration:         5 * time.Minute,
+		MaxConnectionsPerIP: 100,
+		MaxConnections:      2,
+	}
+	rl := NewRateLimiter(cfg)
+
+	addr1 := mockAddr{"tcp", "192.168.1.1:5432"}
+	addr2 := mockAddr{"tcp", "192.168.1.2:5432"}
+	addr3 := mockAddr{"tcp", "192.168.1.3:5432"}
+
+	// Register 2 connections (at limit)
+	if !rl.RegisterConnection(addr1) {
+		t.Fatal("RegisterConnection(addr1) should succeed")
+	}
+	if !rl.RegisterConnection(addr2) {
+		t.Fatal("RegisterConnection(addr2) should succeed")
+	}
+
+	t.Run("exceeds global limit", func(t *testing.T) {
+		if rl.RegisterConnection(addr3) {
+			t.Error("RegisterConnection(addr3) should fail when at global limit")
+		}
+
+		msg := rl.CheckConnection(addr3)
+		if msg != "too many concurrent connections" {
+			t.Errorf("unexpected error message: %s", msg)
+		}
+	})
+
+	t.Run("allows connection after unregister", func(t *testing.T) {
+		rl.UnregisterConnection(addr1)
+		if !rl.RegisterConnection(addr3) {
+			t.Error("RegisterConnection(addr3) should succeed after unregistering addr1")
+		}
+	})
+}
+
 func TestDefaultRateLimitConfig(t *testing.T) {
 	cfg := DefaultRateLimitConfig()
 
@@ -331,4 +373,7 @@ func TestDefaultRateLimitConfig(t *testing.T) {
 	if cfg.MaxConnectionsPerIP != 100 {
 		t.Errorf("MaxConnectionsPerIP = %d, want 100", cfg.MaxConnectionsPerIP)
 	}
+	if cfg.MaxConnections != runtime.NumCPU()*2 {
+		t.Errorf("MaxConnections = %d, want %d", cfg.MaxConnections, runtime.NumCPU()*2)
+	}
 }

server/server.go

@@ -239,9 +239,22 @@ type Server struct {
 }
 
 func New(cfg Config) (*Server, error) {
-	// Use default rate limit config if not specified
+	// Apply default rate limit config for any unset fields
+	defaults := DefaultRateLimitConfig()
 	if cfg.RateLimit.MaxFailedAttempts == 0 {
-		cfg.RateLimit = DefaultRateLimitConfig()
+		cfg.RateLimit.MaxFailedAttempts = defaults.MaxFailedAttempts
+	}
+	if cfg.RateLimit.FailedAttemptWindow == 0 {
+		cfg.RateLimit.FailedAttemptWindow = defaults.FailedAttemptWindow
+	}
+	if cfg.RateLimit.BanDuration == 0 {
+		cfg.RateLimit.BanDuration = defaults.BanDuration
+	}
+	if cfg.RateLimit.MaxConnectionsPerIP == 0 {
+		cfg.RateLimit.MaxConnectionsPerIP = defaults.MaxConnectionsPerIP
+	}
+	if cfg.RateLimit.MaxConnections == 0 {
+		cfg.RateLimit.MaxConnections = defaults.MaxConnections
 	}
 
 	// Use default shutdown timeout if not specified

tests/integration/harness.go

@@ -133,6 +133,9 @@ func (h *TestHarness) startDuckgres(harnessCfg HarnessConfig) error {
 			"testuser": "testpass",
 		},
 		Extensions: []string{"ducklake"},
+		RateLimit: server.RateLimitConfig{
+			MaxConnections: 100,
+		},
 	}
 
 	// Configure DuckLake if enabled