Refactor PostgreSQL handshake for JDBC compatibility (#275)

* Refactor PostgreSQL handshake for JDBC compatibility and optimize session creation

* Fix GSS fallback handling and enforce default session limits

* fix: address PR 275 review comments

- Fix CancelRequest race condition by registering the query before CreateSession
- Fix worker_mgr reaper to respect minWorkers as an idle pool size rather than total pool size
- Fix uncancelled context leak for tmpCC in controlplane

* Fix warm worker floor and queued session cancellation

Author: fuziontech

controlplane/control.go

@@ -188,7 +188,7 @@ func RunControlPlane(cfg ControlPlaneConfig) {
 		minWorkers = maxWorkers
 	}
 
-	pool := NewFlightWorkerPool(cfg.SocketDir, cfg.ConfigPath, maxWorkers)
+	pool := NewFlightWorkerPool(cfg.SocketDir, cfg.ConfigPath, minWorkers, maxWorkers)
 	pool.idleTimeout = cfg.WorkerIdleTimeout
 
 	// Import pre-bound sockets from handover, or pre-bind new ones.
@@ -392,6 +392,21 @@ func (cp *ControlPlane) acceptLoop() {
 	}
 }
 
+func createSessionWithRegisteredCancel(
+	srv *server.Server,
+	timeout time.Duration,
+	key server.BackendKey,
+	createFn func(context.Context) (int32, *server.FlightExecutor, error),
+) (int32, *server.FlightExecutor, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	srv.RegisterQuery(key, cancel)
+	defer srv.UnregisterQuery(key)
+
+	return createFn(ctx)
+}
+
 func (cp *ControlPlane) handleConnection(conn net.Conn) {
 	remoteAddr := conn.RemoteAddr()
 	slog.Info("Connection accepted.", "remote_addr", remoteAddr)
@@ -412,7 +427,8 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) {
 		return
 	}
 
-	// Read startup message to determine SSL vs cancel
+	// Read startup message to determine SSL vs cancel.
+	// readStartupFromRaw handles GSSENC probes by replying 'N' and continuing.
 	params, err := readStartupFromRaw(conn)
 	if err != nil {
 		if err == io.EOF || errors.Is(err, io.EOF) {
@@ -543,14 +559,47 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) {
 	server.RecordSuccessfulAuthAttempt(cp.rateLimiter, remoteAddr)
 	slog.Info("User authenticated.", "user", username, "remote_addr", remoteAddr)
 
+	// Feed initial parameters and backend key data to the client IMMEDIATELY.
+	// This keeps JDBC drivers happy while we perform the slow worker acquisition.
+	pid := cp.sessions.ReservePID()
+	secretKey := server.GenerateSecretKey()
+
+	// Use a temporary clientConn just to send initial params
+	tmpCC := server.NewClientConn(cp.srv, nil, nil, writer, username, database, applicationName, nil, pid, secretKey, -1)
+	defer server.CancelClientConn(tmpCC)
+	server.SendInitialParams(tmpCC)
+	if err := writer.Flush(); err != nil {
+		slog.Error("Failed to flush initial params.", "remote_addr", remoteAddr, "error", err)
+		return
+	}
+
 	// Create session on a worker. The timeout controls how long we wait in the
 	// worker queue when all slots are occupied.
-	ctx, cancel := context.WithTimeout(context.Background(), cp.cfg.WorkerQueueTimeout)
-	pid, executor, err := cp.sessions.CreateSession(ctx, username)
-	cancel()
+	// Pass resource limits to be applied immediately by the worker (one RPC).
+	var (
+		memLimit string
+		threads  int
+	)
+	if cp.rebalancer != nil {
+		memLimit = cp.rebalancer.MemoryLimit()
+		threads = cp.rebalancer.PerSessionThreads()
+	}
+
+	_, executor, err := createSessionWithRegisteredCancel(
+		cp.srv,
+		cp.cfg.WorkerQueueTimeout,
+		server.BackendKey{Pid: pid, SecretKey: secretKey},
+		func(ctx context.Context) (int32, *server.FlightExecutor, error) {
+			return cp.sessions.CreateSession(ctx, username, pid, memLimit, threads)
+		},
+	)
 	if err != nil {
 		slog.Error("Failed to create session.", "user", username, "remote_addr", remoteAddr, "error", err)
-		_ = server.WriteErrorResponse(writer, "FATAL", "53300", "too many connections")
+		if errors.Is(err, context.Canceled) {
+			_ = server.WriteErrorResponse(writer, "FATAL", "57014", "canceling authentication due to user request")
+		} else {
+			_ = server.WriteErrorResponse(writer, "FATAL", "53300", "too many connections")
+		}
 		_ = writer.Flush()
 		return
 	}
@@ -560,14 +609,11 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) {
 	// the message loop if the backing worker dies.
 	cp.sessions.SetConnCloser(pid, tlsConn)
 
-	secretKey := server.GenerateSecretKey()
-
-	// Create clientConn with FlightExecutor
+	// Create real clientConn with FlightExecutor and worker assignment
 	workerID := cp.sessions.WorkerIDForPID(pid)
 	cc := server.NewClientConn(cp.srv, tlsConn, reader, writer, username, database, applicationName, executor, pid, secretKey, workerID)
 
-	// Send initial parameters and ReadyForQuery
-	server.SendInitialParams(cc)
+	// Send ReadyForQuery to signal that the handshake is complete
 	if err := server.WriteReadyForQuery(writer, 'I'); err != nil {
 		slog.Error("Failed to send ReadyForQuery.", "remote_addr", remoteAddr, "error", err)
 		return
@@ -589,6 +635,7 @@ func (cp *ControlPlane) handleConnection(conn net.Conn) {
 // startupResult holds the parsed initial startup message.
 type startupResult struct {
 	sslRequest      bool
+	gssRequest      bool
 	cancelRequest   bool
 	cancelPid       int32
 	cancelSecretKey int32
@@ -649,6 +696,28 @@ func readStartupFromRaw(conn net.Conn) (startupResult, error) {
 	return startupResult{}, fmt.Errorf("too many negotiation rounds")
 }
 
+// readStartupWithGSSFallback accepts a GSSAPI probe, rejects it with 'N',
+// and keeps reading startup packets on the same connection so clients can
+// continue with SSLRequest/startup without reconnecting.
+func readStartupWithGSSFallback(conn net.Conn) (startupResult, error) {
+	for i := 0; i < 4; i++ {
+		params, err := readStartupFromRaw(conn)
+		if err != nil {
+			return startupResult{}, err
+		}
+
+		if !params.gssRequest {
+			return params, nil
+		}
+
+		if _, err := conn.Write([]byte{'N'}); err != nil {
+			return startupResult{}, fmt.Errorf("write GSSAPI rejection: %w", err)
+		}
+	}
+
+	return startupResult{}, fmt.Errorf("too many GSSAPI startup requests")
+}
+
 func fullRead(conn net.Conn, buf []byte) (int, error) {
 	total := 0
 	for total < len(buf) {

controlplane/control_cancel_test.go

@@ -0,0 +1,58 @@
+package controlplane
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/posthog/duckgres/server"
+)
+
+func TestCreateSessionWithRegisteredCancel_CancelQueryCancelsWait(t *testing.T) {
+	srv := &server.Server{}
+	server.InitMinimalServer(srv, server.Config{}, nil)
+
+	key := server.BackendKey{Pid: 1234, SecretKey: 5678}
+
+	started := make(chan struct{})
+	errCh := make(chan error, 1)
+	go func() {
+		_, _, err := createSessionWithRegisteredCancel(
+			srv,
+			200*time.Millisecond,
+			key,
+			func(ctx context.Context) (int32, *server.FlightExecutor, error) {
+				close(started)
+				<-ctx.Done()
+				return 0, nil, ctx.Err()
+			},
+		)
+		errCh <- err
+	}()
+
+	select {
+	case <-started:
+	case <-time.After(2 * time.Second):
+		t.Fatal("create function did not start")
+	}
+
+	if !srv.CancelQuery(key) {
+		t.Fatal("expected CancelQuery to find registered query")
+	}
+
+	select {
+	case err := <-errCh:
+		if err == nil {
+			t.Fatal("expected context cancellation error")
+		}
+		if err != context.Canceled {
+			t.Fatalf("expected context.Canceled, got %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("createSessionWithRegisteredCancel did not return after cancel")
+	}
+
+	if srv.CancelQuery(key) {
+		t.Fatal("expected query to be unregistered after return")
+	}
+}

controlplane/control_startup_test.go

@@ -0,0 +1,62 @@
+package controlplane
+
+import (
+	"encoding/binary"
+	"net"
+	"testing"
+)
+
+func startupPacket(protocolVersion uint32) []byte {
+	pkt := make([]byte, 8)
+	binary.BigEndian.PutUint32(pkt[0:4], uint32(len(pkt)))
+	binary.BigEndian.PutUint32(pkt[4:8], protocolVersion)
+	return pkt
+}
+
+func TestReadStartupWithGSSFallback(t *testing.T) {
+	serverConn, clientConn := net.Pipe()
+	defer func() { _ = serverConn.Close() }()
+	defer func() { _ = clientConn.Close() }()
+
+	errCh := make(chan error, 1)
+	go func() {
+		// Send GSSENCRequest first.
+		if _, err := clientConn.Write(startupPacket(80877104)); err != nil {
+			errCh <- err
+			return
+		}
+
+		// Server should reject GSS with a single 'N' byte.
+		resp := make([]byte, 1)
+		if _, err := clientConn.Read(resp); err != nil {
+			errCh <- err
+			return
+		}
+		if resp[0] != 'N' {
+			errCh <- net.InvalidAddrError("expected GSS rejection byte 'N'")
+			return
+		}
+
+		// Continue negotiation on the same connection with SSLRequest.
+		if _, err := clientConn.Write(startupPacket(80877103)); err != nil {
+			errCh <- err
+			return
+		}
+		errCh <- nil
+	}()
+
+	params, err := readStartupWithGSSFallback(serverConn)
+	if err != nil {
+		t.Fatalf("readStartupWithGSSFallback returned error: %v", err)
+	}
+	if !params.sslRequest {
+		t.Fatalf("expected sslRequest=true after GSS fallback, got %+v", params)
+	}
+	if params.gssRequest {
+		t.Fatalf("expected gssRequest=false final result, got %+v", params)
+	}
+
+	if err := <-errCh; err != nil {
+		t.Fatalf("client side negotiation failed: %v", err)
+	}
+}

controlplane/session_mgr.go

@@ -44,17 +44,25 @@ func NewSessionManager(pool *FlightWorkerPool, rebalancer *MemoryRebalancer) *Se
 	return sm
 }
 
+// ReservePID generates a new unique PID for a session.
+func (sm *SessionManager) ReservePID() int32 {
+	return sm.nextPID.Add(1)
+}
+
 // CreateSession acquires a worker (reusing an idle one or spawning a new one),
 // creates a session on it, and rebalances memory/thread limits across all active sessions.
-func (sm *SessionManager) CreateSession(ctx context.Context, username string) (int32, *server.FlightExecutor, error) {
+// If pid is 0, a new one is generated.
+func (sm *SessionManager) CreateSession(ctx context.Context, username string, pid int32, memoryLimit string, threads int) (int32, *server.FlightExecutor, error) {
+	memoryLimit, threads = sm.resolveSessionLimits(memoryLimit, threads)
+
 	// Acquire a worker: reuses idle pre-warmed workers or spawns a new one.
 	// When max-workers is set, this blocks until a slot is available.
 	worker, err := sm.pool.AcquireWorker(ctx)
 	if err != nil {
 		return 0, nil, fmt.Errorf("acquire worker: %w", err)
 	}
 
-	sessionToken, err := worker.CreateSession(ctx, username)
+	sessionToken, err := worker.CreateSession(ctx, username, memoryLimit, threads)
 	if err != nil {
 		// Clean up the worker we just spawned (but not if it was a pre-warmed idle worker
 		// that has sessions from other concurrent requests).
@@ -65,7 +73,9 @@ func (sm *SessionManager) CreateSession(ctx context.Context, username string) (i
 	// Create FlightExecutor sharing the worker's existing gRPC connection
 	executor := server.NewFlightExecutorFromClient(worker.client, sessionToken)
 
-	pid := sm.nextPID.Add(1)
+	if pid == 0 {
+		pid = sm.nextPID.Add(1)
+	}
 
 	session := &ManagedSession{
 		PID:          pid,
@@ -81,16 +91,27 @@ func (sm *SessionManager) CreateSession(ctx context.Context, username string) (i
 
 	slog.Debug("Session created.", "pid", pid, "worker", worker.ID, "user", username)
 
-	// Set memory/thread limits on this session synchronously so it never
-	// runs with unlimited resources.
+	// Update other sessions if rebalancing is enabled.
 	if sm.rebalancer != nil {
-		sm.rebalancer.SetInitialLimits(ctx, session)
 		sm.rebalancer.RequestRebalance()
 	}
 
 	return pid, executor, nil
 }
 
+func (sm *SessionManager) resolveSessionLimits(memoryLimit string, threads int) (string, int) {
+	if sm.rebalancer == nil {
+		return memoryLimit, threads
+	}
+	if memoryLimit == "" {
+		memoryLimit = sm.rebalancer.MemoryLimit()
+	}
+	if threads <= 0 {
+		threads = sm.rebalancer.PerSessionThreads()
+	}
+	return memoryLimit, threads
+}
+
 // DestroySession destroys a session, retires its dedicated worker, and rebalances
 // memory/thread limits across remaining sessions.
 func (sm *SessionManager) DestroySession(pid int32) {
@@ -239,4 +260,3 @@ func (sm *SessionManager) AllSessions() []*ManagedSession {
 	}
 	return result
 }
-

controlplane/session_mgr_test.go

@@ -251,3 +251,29 @@ func TestDestroySessionAfterOnWorkerCrash(t *testing.T) {
 		t.Fatal("expected 0 sessions after DestroySession")
 	}
 }
+
+func TestResolveSessionLimits_UsesRebalancerDefaultsWhenUnset(t *testing.T) {
+	r := NewMemoryRebalancer(24*1024*1024*1024, 8, nil, false)
+	sm := NewSessionManager(&FlightWorkerPool{workers: make(map[int]*ManagedWorker)}, r)
+
+	mem, threads := sm.resolveSessionLimits("", 0)
+	if mem != "24576MB" {
+		t.Fatalf("expected memory limit 24576MB, got %q", mem)
+	}
+	if threads != 8 {
+		t.Fatalf("expected threads 8, got %d", threads)
+	}
+}
+
+func TestResolveSessionLimits_PreservesExplicitValues(t *testing.T) {
+	r := NewMemoryRebalancer(24*1024*1024*1024, 8, nil, false)
+	sm := NewSessionManager(&FlightWorkerPool{workers: make(map[int]*ManagedWorker)}, r)
+
+	mem, threads := sm.resolveSessionLimits("1024MB", 2)
+	if mem != "1024MB" {
+		t.Fatalf("expected memory limit 1024MB, got %q", mem)
+	}
+	if threads != 2 {
+		t.Fatalf("expected threads 2, got %d", threads)
+	}
+}