Free Arrow IPC resources from one defer, not two call sites (#733)

RELEASE_NOTES_2026.09.2.md

@@ -222,6 +222,35 @@ is published. Responsibly reported by **[@rexpository](https://github.com/rexpos
 
 ## Bug fixes
 
+### Arrow IPC cleanup no longer runs twice on the panic path ([#733](https://github.com/Basekick-Labs/arc/issues/733))
+
+No behaviour changes for anyone running Arc. This is recorded because the code
+was correct for a reason nothing in the repo could check.
+
+`POST /api/v1/query/arrow` freed its DuckDB reader, pooled connection and query
+timeout from two places: an ordinary statement at the end of the stream writer,
+and again from the panic handler. A panic raised after the first one, in the
+trailer or logging statements that follow it, ran the whole cleanup a second
+time.
+
+That was harmless, and it is worth being precise about why, because the first
+version of this report claimed otherwise. The DuckDB reader's `Release` has an
+explicit "already at zero" guard, `*sql.Conn.Close` returns
+`sql.ErrConnDone` rather than failing, and cancelling a context twice is
+allowed. Nothing was over-released and nothing leaked.
+
+The problem was that none of those three properties is promised by the
+interfaces Arc codes against, and none was pinned by a test. Arrow's own readers
+guard over-release behind an assertion that is compiled out unless the build
+sets `-tags assert`, which no Arc build does, so losing the invariant would not
+have failed anything anywhere. It would have gone unnoticed rather than
+corrupted anything, which is the point: there was no signal to rely on.
+
+Cleanup is now a single deferred call, which runs exactly once whether the
+writer returns normally or unwinds, and matches what every other stream writer
+in Arc already did. It also covers a panic raised inside the panic handler
+itself, ahead of where the release used to sit, which would have stranded a
+pooled connection.
 ### The measurement endpoint now appears in query management ([#731](https://github.com/Basekick-Labs/arc/issues/731))
 
 `GET /api/v1/query/:measurement` never registered with the query registry, so
@@ -448,7 +477,9 @@ service.
 Recovering also exposed what the unwind had been skipping. Those writers
 released their result set, pooled connection and query timeout in ordinary
 statements after the streaming call, so the fix moves each into a deferred block
-that keeps the original release order. Six of the writers also disposed of their
+that keeps the original release order. The Arrow IPC writer was the exception
+and kept a split between an inline release and the panic handler until the #733
+entry above, later in this same release, brought it into line. Six of the writers also disposed of their
 query-registry entry only on the normal path, which meant a panicking query would
 have sat in `GET /api/v1/queries/active` as `running` forever, holding its SQL
 text and inflating the active-query gauge, with no reaper to clear it; those
@@ -463,8 +494,11 @@ reader, returned the pooled connection and stopped the query timeout timer ran
 as ordinary statements after the streaming call, and a panic unwound straight
 past them. The connection was the costly one: it was never returned to the
 pool, so each occurrence permanently shrank the pool and repeated occurrences
-would starve the endpoint until a restart. Cleanup now also runs on the
-recovery path, and the per-batch records the stream owns (row-cap slices,
+would starve the endpoint until a restart. Cleanup now runs on the recovery
+path too, from a single deferred call that covers both paths (the original fix
+ran it from the panic handler as well as inline, which the #733 entry above
+replaced later in this same release). The per-batch records the stream owns
+(row-cap slices,
 decimal casts, dictionary-encoded batches) are released through a defer so a
 panic mid-batch cannot strand their buffers either, which for DuckDB-backed
 records are C-allocated and not reclaimed by the garbage collector.

internal/api/query_arrow.go

@@ -27,6 +27,10 @@ import (
 // response writer's panic path (#716). Production always uses the real one.
 var streamArrowIPCFunc = streamArrowIPC
 
+// releaseArrowStreamResourcesFunc indirects the cleanup so tests can count how
+// many times it runs (#733). Production always uses the real one.
+var releaseArrowStreamResourcesFunc = releaseArrowStreamResources
+
 // arrowTrailerWarnOnce gates the AddTrailer failure log so a fasthttp
 // upgrade that ever rejects the trailer name does not produce a Warn per
 // request.
@@ -119,11 +123,13 @@ func poisonArrowStream(w *bufio.Writer, logger zerolog.Logger) {
 // and statement that live on this connection, so returning the connection to
 // the pool first would hand another query a connection with an open result.
 //
-// It recovers on its own behalf (#716). It is called from the panic path,
-// where it runs on a bare fasthttp goroutine while a panic is already in
-// flight: a second panic raised here would either kill the process or, if
-// caught by the caller's recover, replace the root-cause panic value, since
-// recover() only ever yields the most recent one.
+// It recovers on its own behalf (#716). Its single caller is a defer inside the
+// Arrow IPC stream writer, so it runs on a bare fasthttp goroutine and, on the
+// panic path, during the unwind with a panic already in flight: a second panic
+// raised here would either kill the process or, if caught by the caller's
+// recover, replace the root-cause panic value, since recover() only ever yields
+// the most recent one. Do not remove the recover on the grounds that the happy
+// path cannot panic; the unwind is the case it exists for.
 func releaseArrowStreamResources(
 	reader array.RecordReader,
 	conn interface{ Close() error },
@@ -631,44 +637,53 @@ func (h *QueryHandler) executeQueryArrow(c *fiber.Ctx) error {
 	trailers := newResponseTrailers()
 
 	streamCtx := ctx
-	// Cleanup runs from safeStream's panic path rather than an ordinary defer
-	// inside the writer, so the resources are freed on both paths (#716).
 	// streamW lets the panic path reach the same writer the body used, so it
 	// can mark the stream truncated. safeStream's onPanic takes no arguments,
 	// and widening it would churn five other call sites that have no writer to
 	// mark.
 	//
-	// The ordering constraint this comment used to carry, that the release had
-	// to stay ahead of the trailer set because both goroutines touched the
-	// response header, is gone: trailer values now go into `trailers` and are
-	// published by the connection goroutine (#729). What remains is that a
-	// panic after the straight-line release below runs it a second time, which
-	// over-releases the reader and logs a spurious cleanup panic (#733). A
-	// plain defer would be strictly better and is left to that issue.
+	// Cleanup is an ordinary defer inside the writer, which is what every other
+	// stream writer in this package already does (#733). It used to run from
+	// onPanic as well as straight-line at the end of the writer, so a panic
+	// after the straight-line call ran it twice. That was harmless only because
+	// all three resources tolerate a second call, which no test pinned and the
+	// array.RecordReader interface does not promise.
+	//
+	// It also covers a panic inside onPanic ahead of where the release used to
+	// sit, which skipped it and stranded the pooled connection. That one is
+	// reasoned, not tested: nothing in onPanic can be made to panic from a test
+	// without adding a seam for it.
+	//
+	// One thing the old shape did that this does not: releaseArrowStreamResources
+	// installs its recover first, so a panic in reader.Release skipped the
+	// conn.Close and cancel below it, and a later panic then retried them
+	// through onPanic. That retry is gone. It is not worth restoring, but the
+	// change is not strictly better on every panic path.
 	var streamW *bufio.Writer
 	h.setBodyStreamWithTrailers(fctx, "query_arrow_ipc", trailers, func() {
 		// A recovered panic leaves a stream the client would otherwise read as
 		// complete: fasthttp still writes the terminating chunk, and a cut on
 		// a batch boundary decodes without error (#721).
 		//
 		// The body marker is the signal a client acts on, and it is always
-		// written: streamW is assigned as the writer closure's first statement,
-		// so it is never nil by the time a recovered panic reaches here. The
-		// trailer adds the reason, which this path carried not at all until
+		// written: streamW is assigned before anything in the writer that can
+		// panic, so it is never nil by the time a recovered panic reaches here.
+		// The trailer adds the reason, which this path carried not at all until
 		// #729 made setting one from this goroutine safe, and which the error
 		// path has always carried.
 		//
 		// setIfAbsent, so a stream that failed with a real error and then
 		// panicked on its way out still reports the error.
 		poisonArrowStream(streamW, h.logger)
 		trailers.setIfAbsent(arrowStreamTruncatedTrailer, "stream writer panicked")
-		releaseArrowStreamResources(reader, conn, cancel, h.logger)
 	}, func(w *bufio.Writer) {
+		// Registered before anything that can panic, so it runs exactly once
+		// whether the writer returns normally or unwinds (#733).
+		defer releaseArrowStreamResourcesFunc(reader, conn, cancel, h.logger)
 		streamW = w
 		totalRows, streamErr := streamArrowIPCFunc(
 			streamCtx, w, reader, schema, castInfo, dictEnabled, ipcCompression, governanceMaxRows, h.logger,
 		)
-		releaseArrowStreamResources(reader, conn, cancel, h.logger)
 
 		// Publish authoritative server-side timing as a chunked-transfer
 		// trailer. Set even on the error path so partial results carry

internal/api/query_arrow_cleanup_test.go

@@ -83,8 +83,10 @@ func TestReleaseArrowStreamResources(t *testing.T) {
 
 // TestExecuteQueryArrowReturnsConnectionOnPanic is the #716 regression proper.
 // It runs the real handler against a real DuckDB, forces the stream writer to
-// panic, and asserts the pooled connection goes back to the pool. Reverting the
-// cleanup to straight-line code leaves InUse at 1 and fails this test.
+// panic, and asserts the pooled connection goes back to the pool. Deleting the
+// deferred releaseArrowStreamResourcesFunc call at the top of the writer leaves
+// InUse at 1 and fails this test. It says nothing about how many times cleanup
+// runs; TestExecuteQueryArrowReleasesExactlyOnce covers that (#733).
 func TestExecuteQueryArrowReturnsConnectionOnPanic(t *testing.T) {
 	metrics.Init(zerolog.Nop())
 

internal/api/query_arrow_release_once_test.go

@@ -0,0 +1,132 @@
+//go:build duckdb_arrow
+
+package api
+
+import (
+	"bufio"
+	"context"
+	"io"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/basekick-labs/arc/internal/database"
+	"github.com/basekick-labs/arc/internal/metrics"
+	"github.com/basekick-labs/arc/internal/storage"
+	"github.com/gofiber/fiber/v2"
+	"github.com/rs/zerolog"
+)
+
+// Tests for #733. The Arrow IPC writer used to free its resources from two
+// places: straight-line at the end of the writer, and again from safeStream's
+// onPanic. A panic after the straight-line call therefore ran cleanup twice.
+//
+// That was harmless in practice, but only because all three resources tolerate
+// it: the DuckDB reader's Release has an explicit refCount <= 0 guard,
+// *sql.Conn.Close returns ErrConnDone, and cancel is idempotent by contract.
+// The array.RecordReader interface promises none of that, and arrow-go's own
+// readers guard over-release behind debug.Assert, which is compiled out unless
+// built with -tags assert. So no test and no CI job could ever have detected
+// the invariant being lost, which is why it is now structural: one defer.
+
+// panicOnMessageHook panics the first time zerolog emits the given message,
+// which is the only injection point that lands AFTER the old straight-line
+// release and before the writer returns.
+//
+// Keyed on the message rather than the level, because releaseArrowStreamResources
+// logs at Error on its own recovery path and would re-enter the hook.
+type panicOnMessageHook struct {
+	msg   string
+	fired atomic.Bool
+}
+
+func (h *panicOnMessageHook) Run(_ *zerolog.Event, _ zerolog.Level, msg string) {
+	if msg == h.msg && h.fired.CompareAndSwap(false, true) {
+		panic("injected panic after the Arrow IPC stream writer finished")
+	}
+}
+
+func TestExecuteQueryArrowReleasesExactlyOnce(t *testing.T) {
+	run := func(t *testing.T, hook zerolog.Hook) int32 {
+		t.Helper()
+		metrics.Init(zerolog.Nop())
+
+		tmpDir := t.TempDir()
+		logger := zerolog.Nop()
+		if hook != nil {
+			logger = zerolog.New(io.Discard).Hook(hook)
+		}
+
+		backend, err := storage.NewLocalBackend(tmpDir, zerolog.Nop())
+		if err != nil {
+			t.Fatal(err)
+		}
+		duckdb, err := database.New(&database.Config{
+			MemoryLimit: "256MB", ThreadCount: 2, MaxConnections: 2, LocalStorageRoot: tmpDir,
+		}, zerolog.Nop())
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() { duckdb.Close() })
+
+		h := newGovernanceTestHandler(newGovernanceTestManager(t, nil), 0)
+		h.db = duckdb
+		h.storage = backend
+		h.logger = logger
+
+		origStream := streamArrowIPCFunc
+		t.Cleanup(func() { streamArrowIPCFunc = origStream })
+		streamArrowIPCFunc = func(_ context.Context, w *bufio.Writer, _ array.RecordReader, _ *arrow.Schema,
+			_ *decimalCastInfo, _ bool, _ string, _ int, _ zerolog.Logger) (int64, error) {
+			_, _ = w.Write([]byte{0})
+			_ = w.Flush()
+			return 1, nil
+		}
+
+		var releases int32
+		origRelease := releaseArrowStreamResourcesFunc
+		t.Cleanup(func() { releaseArrowStreamResourcesFunc = origRelease })
+		releaseArrowStreamResourcesFunc = func(r array.RecordReader, c interface{ Close() error },
+			cancel context.CancelFunc, lg zerolog.Logger) {
+			atomic.AddInt32(&releases, 1)
+			origRelease(r, c, cancel, lg)
+		}
+
+		app := fiber.New(fiber.Config{DisableStartupMessage: true})
+		app.Post("/api/v1/query/arrow", h.executeQueryArrow)
+		req := httptest.NewRequest("POST", "/api/v1/query/arrow", strings.NewReader(`{"sql":"SELECT 1 AS id"}`))
+		req.Header.Set("Content-Type", "application/json")
+		resp, err := app.Test(req, 10000)
+		if err != nil {
+			t.Fatalf("app.Test: %v", err)
+		}
+		if _, err := io.ReadAll(resp.Body); err != nil {
+			t.Fatalf("draining body: %v", err)
+		}
+		resp.Body.Close()
+		return atomic.LoadInt32(&releases)
+	}
+
+	t.Run("normal completion releases once", func(t *testing.T) {
+		if got := run(t, nil); got != 1 {
+			t.Errorf("cleanup ran %d times, want exactly 1", got)
+		}
+	})
+
+	// The regression proper: on 17aaaef this reports 2, because the panic is
+	// raised after the straight-line release and onPanic then runs it again.
+	t.Run("a panic after the stream finishes still releases once", func(t *testing.T) {
+		hook := &panicOnMessageHook{msg: "Arrow streaming query completed"}
+		got := run(t, hook)
+		if !hook.fired.Load() {
+			t.Fatal("the injected panic never fired, so this asserts nothing")
+		}
+		if got != 1 {
+			t.Errorf("cleanup ran %d times, want exactly 1", got)
+		}
+	})
+
+}

internal/api/stream_safety.go

@@ -16,16 +16,21 @@ import (
 // is a process crash, not a failed request.
 //
 // onPanic runs only on the panic path, after the root cause has been logged,
-// and is for work the unwind skipped that a plain defer inside sw cannot do:
-// releasing resources whose release must not move relative to the rest of the
-// happy path, and disposing of a query-registry entry that would otherwise sit
-// in "running" forever. It runs inside its own recover, because it executes
-// while a panic is already in flight: a second panic here would either kill the
-// process or, being the most recent value, replace the root cause in the log.
+// and is for work that has no meaning on the happy path: disposing of a
+// query-registry entry that would otherwise sit in "running" forever, and
+// marking a response the client would otherwise read as complete.
 //
-// Resources that can simply be freed at the end of sw belong in an ordinary
-// defer inside sw instead, not here.
+// Freeing resources is NOT its job. That belongs in an ordinary defer inside
+// sw, which every caller now uses, and which covers strictly more than onPanic
+// does: a runtime.Goexit runs defers but leaves recover returning nil, so
+// onPanic never fires, and a panic inside onPanic ahead of a release would
+// strand whatever it had not freed yet (#733).
 //
+// It runs inside its own recover, because it executes while a panic is already
+// in flight: a second panic here would either kill the process or, being the
+// most recent value, replace the root cause in the log.
+//
+
 // Every body stream in this package goes through this wrapper, whether it is
 // installed with SetBodyStreamWriter directly or through
 // setBodyStreamWithTrailers (#729). CI greps for both, plus the underlying