Forward query logs without control-plane batching (#959)

Author: bill-ph

controlplane/session_mgr.go

@@ -585,7 +585,11 @@ func (sm *SessionManager) createSessionOnWorker(ctx context.Context, username st
 			"pid", pid, "worker", worker.ID, "user", username, "warning", w)
 	}
 
-	executor := flightclient.NewFlightExecutorFromClient(worker.client, sessionToken)
+	executor := flightclient.NewFlightExecutorFromClientWithQueryLogLimiter(
+		worker.client,
+		sessionToken,
+		worker.workerQueryLogLimiter(),
+	)
 	executor.SetControlMetadata(worker.ID, worker.OwnerCPInstanceID(), worker.OwnerEpoch())
 
 	if pid == 0 {

controlplane/session_mgr_drain_test.go

@@ -6,7 +6,6 @@ import (
 	"context"
 	"errors"
 	"io"
-	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -718,7 +717,6 @@ func TestDestroyAllSessions_ReleasesAllLeasesAndClearsSessions(t *testing.T) {
 }
 
 func TestDestroyAllSessionsRejectsInFlightCreateBeforeRegistration(t *testing.T) {
-	queryLogForwardersBefore := countFlightExecutorQueryLogForwarders()
 	flightClient := &blockingCreateSessionFlightClient{
 		createStarted:   make(chan struct{}),
 		allowCreate:     make(chan struct{}),
@@ -792,7 +790,6 @@ func TestDestroyAllSessionsRejectsInFlightCreateBeforeRegistration(t *testing.T)
 	if got := flightClient.destroyCalls.Load(); got != 1 {
 		t.Fatalf("expected worker session to be destroyed once, got %d", got)
 	}
-	waitForFlightExecutorQueryLogForwardersAtMost(t, queryLogForwardersBefore)
 }
 
 func TestDestroyAllSessionsWaitsForCreateBlockedInLimiterAcquire(t *testing.T) {
@@ -979,24 +976,6 @@ func waitForSessionManagerDraining(t *testing.T, sm *SessionManager) {
 	t.Fatal("timed out waiting for session manager to start draining")
 }
 
-func waitForFlightExecutorQueryLogForwardersAtMost(t *testing.T, want int) {
-	t.Helper()
-	deadline := time.Now().Add(time.Second)
-	for time.Now().Before(deadline) {
-		if got := countFlightExecutorQueryLogForwarders(); got <= want {
-			return
-		}
-		time.Sleep(time.Millisecond)
-	}
-	t.Fatalf("flight executor query-log forwarder count = %d, want <= %d", countFlightExecutorQueryLogForwarders(), want)
-}
-
-func countFlightExecutorQueryLogForwarders() int {
-	buf := make([]byte, 1<<20)
-	n := runtime.Stack(buf, true)
-	return strings.Count(string(buf[:n]), "(*FlightExecutor).queryLogForwardLoop")
-}
-
 func countEvents(events []string, want string) int {
 	count := 0
 	for _, event := range events {

controlplane/worker_mgr.go

@@ -26,25 +26,27 @@ import (
 
 // ManagedWorker represents a duckdb-service worker process.
 type ManagedWorker struct {
-	ID             int
-	podName        string
-	nodeName       string        //nolint:unused // only set in kubernetes remote backend; drives cache-locality-aware scheduling
-	image          string        //nolint:unused // only set in kubernetes remote backend; carried through runtime store records
-	profile        WorkerProfile //nolint:unused // only set in kubernetes remote backend; pod-shape this worker was spawned with (zero = default exclusive)
-	cmd            *exec.Cmd
-	socketPath     string
-	bearerToken    string
-	client         *flightsql.Client
-	parentListener net.Listener    // CP-side listener; lifecycle managed by releaseSocket
-	prebound       *preboundSocket // non-nil if using a pre-bound socket slot
-	releaseOnce    sync.Once       // ensures releaseWorkerSocket body runs exactly once
-	done           chan struct{}   // closed when process exits
-	exitErr        error
-	activeSessions int       // Number of sessions currently assigned to this worker
-	lastUsed       time.Time // Last time a session was destroyed on this worker
-	sharedState    SharedWorkerState
-	reservedAt     time.Time //nolint:unused // only set in kubernetes remote backend reservation path
-	peakSessions   int       // High-water mark of concurrent sessions (for retirement metrics)
+	ID                  int
+	podName             string
+	nodeName            string        //nolint:unused // only set in kubernetes remote backend; drives cache-locality-aware scheduling
+	image               string        //nolint:unused // only set in kubernetes remote backend; carried through runtime store records
+	profile             WorkerProfile //nolint:unused // only set in kubernetes remote backend; pod-shape this worker was spawned with (zero = default exclusive)
+	cmd                 *exec.Cmd
+	socketPath          string
+	bearerToken         string
+	client              *flightsql.Client
+	queryLogLimiterOnce sync.Once
+	queryLogLimiter     *flightclient.QueryLogLimiter
+	parentListener      net.Listener    // CP-side listener; lifecycle managed by releaseSocket
+	prebound            *preboundSocket // non-nil if using a pre-bound socket slot
+	releaseOnce         sync.Once       // ensures releaseWorkerSocket body runs exactly once
+	done                chan struct{}   // closed when process exits
+	exitErr             error
+	activeSessions      int       // Number of sessions currently assigned to this worker
+	lastUsed            time.Time // Last time a session was destroyed on this worker
+	sharedState         SharedWorkerState
+	reservedAt          time.Time //nolint:unused // only set in kubernetes remote backend reservation path
+	peakSessions        int       // High-water mark of concurrent sessions (for retirement metrics)
 	// ownerEpoch is guarded by epochMu so cred-refresh's
 	// "RefreshLease then SetOwnerEpoch" sequence appears atomic to
 	// concurrent readers (notably ShutdownAll's lease minting).
@@ -58,6 +60,13 @@ type ManagedWorker struct {
 	cachedActivationPayload any  //nolint:unused // *TenantActivationPayload, cached in kubernetes activation path
 }
 
+func (w *ManagedWorker) workerQueryLogLimiter() *flightclient.QueryLogLimiter {
+	w.queryLogLimiterOnce.Do(func() {
+		w.queryLogLimiter = flightclient.NewQueryLogLimiter()
+	})
+	return w.queryLogLimiter
+}
+
 // Profile returns the pod-shape profile this worker was spawned with (zero =
 // default exclusive profile). Exposes the unexported field for the admin API.
 func (w *ManagedWorker) Profile() WorkerProfile {