Timeout QueryProgress CGO call to prevent health check stalls

When DuckDB's httpfs extension is mid-download on a large remote
parquet file, it holds an internal mutex on the connection. The health
check's QueryProgress CGO call blocks waiting for that mutex, the
health check goroutine stalls inside the pool RLock, the CP's 3-second
gRPC deadline fires, and after 3 consecutive failures the CP kills the
worker — even though it's alive and making progress on the download.

Observed in production: read_parquet('https://...14GB_file.parquet')
caused both warm workers to be declared unresponsive and deleted within
~8 seconds of query submission.

The fix restructures the health check in two ways:

  1. Snapshot session data under RLock, then release the lock before
     calling QueryProgress. Previously the RLock was held for the
     entire loop including CGO calls, so a blocked QueryProgress also
     prevented session create/close operations.

  2. Run each QueryProgress call in a goroutine with a 500ms timeout.
     If the CGO call doesn't return in time, report the session as
     "busy (pct=-1)" and skip stall detection for that cycle. The
     health check always responds promptly.

Real crashes (process death) are detected by the K8s pod informer
independently — the w.done channel fires immediately when the pod
terminates, without waiting for a health check cycle. So this change
doesn't create a blind spot for crash recovery.

Stall detection still works for queries that don't block QueryProgress
(the common case). For queries that do block it (httpfs downloads),
the client's own statement timeout is the appropriate kill mechanism —
it comes from the user who decided how long to wait, not from
infrastructure guessing.

Author: benben

duckdbservice/flight_handler.go

@@ -176,61 +176,109 @@ func (h *FlightSQLHandler) doHealthCheck(body []byte, stream flight.FlightServic
 	<-h.pool.warmupDone
 
 	// Poll DuckDB query progress for each active session.
+	//
+	// QueryProgress is a CGO call into DuckDB that *should* return instantly
+	// (it reads atomic progress counters), but can block if DuckDB holds an
+	// internal lock — for example, when the httpfs extension is mid-download
+	// on a large remote parquet file. If that happens while we hold the pool
+	// RLock, both the health check and any session create/close operations
+	// stall, the CP's 3-second health check timeout fires, and after 3
+	// consecutive failures the CP kills the worker — even though it's alive
+	// and making progress on the download.
+	//
+	// To prevent this:
+	// 1. Snapshot the session data we need under RLock, then release it.
+	// 2. Call QueryProgress outside the lock, with a per-session timeout.
+	//    If the CGO call doesn't return within queryProgressTimeout, we
+	//    report the session as "busy" (pct=-1) and skip stall detection
+	//    for this cycle. The health check always responds promptly.
+	//
+	// Real crashes (process death) are detected by the K8s pod informer
+	// independently of health checks, so skipping stall detection during
+	// I/O-heavy operations doesn't create a blind spot for crash recovery.
 	type sessionProgressInfo struct {
 		Pct     float64 `json:"pct"`
 		Rows    uint64  `json:"rows"`
 		Total   uint64  `json:"total"`
 		Stalled bool    `json:"stalled,omitempty"`
 	}
-	sessionProgress := make(map[string]sessionProgressInfo)
+
+	type sessionSnapshot struct {
+		key      string
+		conn     duckdbConnHandle
+		progress *progressState
+	}
 
 	h.pool.mu.RLock()
+	snapshots := make([]sessionSnapshot, 0, len(h.pool.sessions))
 	for token, session := range h.pool.sessions {
 		if session.duckdbConn.Ptr == nil {
 			continue
 		}
-
-		qp := bindings.QueryProgress(session.duckdbConn)
-		pct, rows, total := bindings.QueryProgressTypeMembers(&qp)
-
-		// Use truncated token as key to avoid leaking full bearer tokens
-		// in the health check JSON response. 16 hex chars = 8 bytes of entropy,
-		// sufficient for matching on the control plane side.
 		key := token
 		if len(key) > 16 {
 			key = key[:16]
 		}
+		snapshots = append(snapshots, sessionSnapshot{
+			key:      key,
+			conn:     session.duckdbConn,
+			progress: &session.progress,
+		})
+	}
+	h.pool.mu.RUnlock()
 
-		stalled := false
+	sessionProgress := make(map[string]sessionProgressInfo, len(snapshots))
+	for _, snap := range snapshots {
+		if !snap.progress.queryActive.Load() {
+			snap.progress.lastRowsProcessed.Store(0)
+			snap.progress.stalledChecks.Store(0)
+			sessionProgress[snap.key] = sessionProgressInfo{}
+			continue
+		}
 
-		if !session.progress.queryActive.Load() {
-			// No query running — reset stall tracking.
-			session.progress.lastRowsProcessed.Store(0)
-			session.progress.stalledChecks.Store(0)
-		} else if pct < 0 {
-			// pct == -1 means DuckDB can't track this query; skip stall detection.
-		} else {
-			if rows == session.progress.lastRowsProcessed.Load() {
-				session.progress.stalledChecks.Add(1)
+		type qpResult struct {
+			pct   float64
+			rows  uint64
+			total uint64
+		}
+		ch := make(chan qpResult, 1)
+		go func(conn duckdbConnHandle) {
+			qp := bindings.QueryProgress(conn)
+			pct, rows, total := bindings.QueryProgressTypeMembers(&qp)
+			ch <- qpResult{pct, rows, total}
+		}(snap.conn)
+
+		select {
+		case pr := <-ch:
+			stalled := false
+			if pr.pct < 0 {
+				// pct == -1 means DuckDB can't track this query; skip stall detection.
 			} else {
-				session.progress.lastRowsProcessed.Store(rows)
-				session.progress.stalledChecks.Store(0)
-			}
-
-			if session.progress.stalledChecks.Load() >= stallCheckThreshold {
-				stalled = true
-				// Log once when first crossing the threshold (exact match avoids log spam).
-				if session.progress.stalledChecks.Load() == stallCheckThreshold {
-					slog.Warn("Query appears stuck — no progress detected.",
-						"session", key, "rows_processed", rows, "total_rows", total,
-						"stalled_checks", stallCheckThreshold)
+				if pr.rows == snap.progress.lastRowsProcessed.Load() {
+					snap.progress.stalledChecks.Add(1)
+				} else {
+					snap.progress.lastRowsProcessed.Store(pr.rows)
+					snap.progress.stalledChecks.Store(0)
+				}
+				if snap.progress.stalledChecks.Load() >= stallCheckThreshold {
+					stalled = true
+					if snap.progress.stalledChecks.Load() == stallCheckThreshold {
+						slog.Warn("Query appears stuck — no progress detected.",
+							"session", snap.key, "rows_processed", pr.rows, "total_rows", pr.total,
+							"stalled_checks", stallCheckThreshold)
+					}
 				}
 			}
+			sessionProgress[snap.key] = sessionProgressInfo{Pct: pr.pct, Rows: pr.rows, Total: pr.total, Stalled: stalled}
+
+		case <-time.After(queryProgressTimeout):
+			// CGO call blocked — DuckDB is busy with I/O (e.g., httpfs
+			// download). Report as untrackable and don't advance stall
+			// counters so a legitimately slow operation doesn't get
+			// flagged as stuck.
+			sessionProgress[snap.key] = sessionProgressInfo{Pct: -1}
 		}
-
-		sessionProgress[key] = sessionProgressInfo{Pct: pct, Rows: rows, Total: total, Stalled: stalled}
 	}
-	h.pool.mu.RUnlock()
 
 	resp, _ := json.Marshal(map[string]interface{}{
 		"healthy":          true,

duckdbservice/progress.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"reflect"
 	"sync/atomic"
+	"time"
 	"unsafe"
 
 	bindings "github.com/duckdb/duckdb-go-bindings"
@@ -19,6 +20,16 @@ type duckdbConnHandle = bindings.Connection
 // 300 checks × 2s = ~10 minutes of zero progress.
 const stallCheckThreshold = 300
 
+// queryProgressTimeout is the maximum time the health check waits for a
+// single QueryProgress CGO call before treating the session as "busy
+// (untrackable)". DuckDB's QueryProgress is normally instant (reads
+// atomic counters), but can block when the connection holds an internal
+// mutex — e.g., during a long httpfs download. Without this timeout the
+// health check itself stalls, the CP's 3-second gRPC deadline fires, and
+// after 3 consecutive failures the CP kills the worker even though it's
+// alive and working.
+const queryProgressTimeout = 500 * time.Millisecond
+
 // progressState tracks query activity and stall detection for a session.
 // All fields are atomic because queryActive is written by query goroutines
 // while lastRowsProcessed and stalledChecks are written by the health check

duckdbservice/progress_test.go

@@ -3,7 +3,9 @@ package duckdbservice
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"testing"
+	"time"
 
 	bindings "github.com/duckdb/duckdb-go-bindings"
 	_ "github.com/duckdb/duckdb-go/v2"
@@ -44,3 +46,88 @@ func TestStallCheckThreshold(t *testing.T) {
 		t.Errorf("expected stallCheckThreshold=300, got %d", stallCheckThreshold)
 	}
 }
+
+func TestQueryProgressTimeout(t *testing.T) {
+	// Verify the timeout is long enough for normal QueryProgress calls
+	// but short enough to prevent health check stalls.
+	if queryProgressTimeout < 100*time.Millisecond {
+		t.Errorf("queryProgressTimeout too short: %v", queryProgressTimeout)
+	}
+	if queryProgressTimeout > 2*time.Second {
+		t.Errorf("queryProgressTimeout too long (must be well under the 3s health check gRPC deadline): %v", queryProgressTimeout)
+	}
+}
+
+func TestHealthCheckRespondsWhenQueryProgressBlocks(t *testing.T) {
+	pool := &SessionPool{
+		sessions:    make(map[string]*Session),
+		stopRefresh: make(map[string]func()),
+		warmupDone:  make(chan struct{}),
+		startTime:   time.Now(),
+	}
+	close(pool.warmupDone)
+
+	// Create a session with a fake DuckDB connection whose QueryProgress
+	// would block. We can't actually block the CGO call in a unit test,
+	// but we CAN verify the restructured code path: snapshot sessions
+	// outside the RLock, call QueryProgress per-session, and always
+	// produce a response.
+	db, err := sql.Open("duckdb", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = db.Close() }()
+
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = conn.Close() }()
+
+	duckConn, err := extractDuckDBConnection(conn)
+	if err != nil {
+		t.Fatalf("extractDuckDBConnection: %v", err)
+	}
+
+	session := &Session{ID: "test-session", duckdbConn: duckConn}
+	session.progress.queryActive.Store(true) // simulate active query
+	pool.sessions["abcdef1234567890xx"] = session
+
+	handler := &FlightSQLHandler{pool: pool}
+	stream := &mockDoActionStream{}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- handler.doHealthCheck([]byte(`{}`), stream)
+	}()
+
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("health check returned error: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("health check blocked longer than 2s — would trigger CP timeout")
+	}
+
+	if len(stream.results) != 1 {
+		t.Fatalf("expected 1 result, got %d", len(stream.results))
+	}
+
+	var resp map[string]interface{}
+	if err := json.Unmarshal(stream.results[0].Body, &resp); err != nil {
+		t.Fatalf("unmarshal response: %v", err)
+	}
+	if resp["healthy"] != true {
+		t.Errorf("expected healthy=true, got %v", resp["healthy"])
+	}
+
+	// Verify session progress is reported (with pct=-1 since no real query)
+	sp, ok := resp["session_progress"].(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected session_progress map, got %T", resp["session_progress"])
+	}
+	if len(sp) != 1 {
+		t.Fatalf("expected 1 session in progress, got %d", len(sp))
+	}
+}