Fix CloudFetch goroutine leak that retains Arrow buffers after Close (#357)

## Summary

Fixes #356.

Under high CloudFetch concurrency (≥6 simultaneous downloads), in-flight
`cloudFetchDownloadTask` goroutines could leak when the consumer closed
the iterator before draining all results. Each leaked goroutine pinned a
downloaded chunk in the Go heap, producing the multi-GiB heap plateau
described in the issue that only released on process restart.

## Root cause

`cloudFetchDownloadTask.Run` sends the download result on an
**unbuffered** channel without honoring context cancellation:

```go
cft.resultChan <- cloudFetchDownloadTaskResult{data: bytes.NewReader(buf), ...}
```

Sequence that triggers the leak:

1. `cloudIPCStreamIterator.Next` schedules `MaxDownloadThreads` (default
10) tasks concurrently.
2. The consumer dequeues task 1, gets its result, returns.
3. Tasks 2..N have completed their HTTP read in parallel and are now
**blocked** on the unbuffered send, holding their downloaded buffer.
4. The consumer abandons the iterator (timeout, error, early close,
etc.) and calls `iterator.Close()`.
5. `Close` calls `task.cancel()` on each remaining task. But context
cancellation does **not** unblock an in-flight channel send — the
goroutines stay blocked forever, retaining their buffers.

In v1.7.1 (the version the reporter is on) the goroutine had already
decoded the bytes into Arrow records *before* the send, so the leaked
memory was Arrow-allocator buffers — matching the stack trace in the
issue:

```
(*cloudFetchDownloadTask).Run.func1
  getArrowRecords → (*ipc.Reader).Next → newRecord → loadArray
    → loadBinary → buffer → (*ipcSource).buffer → NewResizableBuffer
      → (*Buffer).Resize → (*GoAllocator).Allocate
```

In the current code (v1.11.0) the decode happens later in
`batchIterator.Next`, so the leak is the raw decompressed `buf` instead
— same shape, smaller per-goroutine retention, same plateau pattern.

## Fix

Route every channel send through a helper that selects on `ctx.Done()`:

```go
func (cft *cloudFetchDownloadTask) sendResult(result cloudFetchDownloadTaskResult) {
    select {
    case cft.resultChan <- result:
    case <-cft.ctx.Done():
    }
}
```

`cloudIPCStreamIterator.Close` already calls `task.cancel()` for every
queued task, so cancellation now correctly drains stuck goroutines and
lets their buffers be GC'd.

## Test plan

- [x] New unit test
`TestCloudFetchIterator_CloseReleasesInFlightDownloads` reproduces the
leak: spawns `MaxDownloadThreads` concurrent downloads, releases them
after the iterator has consumed only the first, then calls `Close()` and
asserts that no `cloudFetchDownloadTask.Run` goroutines remain.
  - Fails on `main` (~9 leaked goroutines after `Close`).
  - Passes with this change.
- [x] Full `go test ./...` passes locally.
- [x] `go vet` and `gofmt` clean.

## Who is affected

Any user with CloudFetch enabled (default since v1.7.0) whose query
context can be cancelled or whose result set can be abandoned mid-stream
— i.e., basically everyone running large CloudFetch queries with
timeouts.

This pull request and its description were written by Isaac.

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>

Author: vikrantpuppala

internal/rows/arrowbased/batchloader.go

@@ -297,7 +297,7 @@ func (cft *cloudFetchDownloadTask) Run() {
 		downloadStart := time.Now()
 		data, err := fetchBatchBytes(cft.ctx, cft.link, cft.minTimeToExpiry, cft.speedThresholdMbps, cft.httpClient)
 		if err != nil {
-			cft.resultChan <- cloudFetchDownloadTaskResult{data: nil, err: err}
+			cft.sendResult(cloudFetchDownloadTaskResult{data: nil, err: err})
 			return
 		}
 
@@ -306,7 +306,7 @@ func (cft *cloudFetchDownloadTask) Run() {
 		data.Close() //nolint:errcheck,gosec // G104: close after reading data
 		downloadMs := time.Since(downloadStart).Milliseconds()
 		if err != nil {
-			cft.resultChan <- cloudFetchDownloadTaskResult{data: nil, err: err}
+			cft.sendResult(cloudFetchDownloadTaskResult{data: nil, err: err})
 			return
 		}
 
@@ -316,10 +316,21 @@ func (cft *cloudFetchDownloadTask) Run() {
 			cft.link.RowCount,
 		)
 
-		cft.resultChan <- cloudFetchDownloadTaskResult{data: bytes.NewReader(buf), err: nil, downloadMs: downloadMs}
+		cft.sendResult(cloudFetchDownloadTaskResult{data: bytes.NewReader(buf), err: nil, downloadMs: downloadMs})
 	}()
 }
 
+// sendResult delivers the download result to the consumer, but drops it if the
+// task's context has already been cancelled. Without this guard, a goroutine
+// that finishes its work after the iterator is closed blocks forever on the
+// unbuffered resultChan and pins the downloaded buffer in the heap (issue #356).
+func (cft *cloudFetchDownloadTask) sendResult(result cloudFetchDownloadTaskResult) {
+	select {
+	case cft.resultChan <- result:
+	case <-cft.ctx.Done():
+	}
+}
+
 // logCloudFetchSpeed calculates and logs download speed metrics
 func logCloudFetchSpeed(fullURL string, contentLength int64, duration time.Duration, speedThresholdMbps float64) {
 	if contentLength > 0 && duration.Seconds() > 0 {

internal/rows/arrowbased/batchloader_test.go

@@ -6,7 +6,10 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"runtime"
+	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -604,3 +607,97 @@ func generateMockArrowBytes(record arrow.Record) []byte {
 	}
 	return buf.Bytes()
 }
+
+// TestCloudFetchIterator_CloseReleasesInFlightDownloads reproduces issue #356:
+// when the consumer closes the iterator while downloads are still in flight,
+// goroutines that completed their HTTP fetch get permanently blocked sending
+// to the unbuffered resultChan. They retain the downloaded buffers (Arrow
+// allocations in earlier versions, raw bytes in current code) until process
+// exit, producing a heap plateau that only releases on restart.
+//
+// The test schedules many concurrent downloads, lets them complete, and then
+// closes the iterator without consuming the queued results. After Close
+// returns, no cloudFetchDownloadTask goroutines must remain.
+func TestCloudFetchIterator_CloseReleasesInFlightDownloads(t *testing.T) {
+	arrowBytes := generateMockArrowBytes(generateArrowRecord())
+
+	// Track in-flight downloads. The server signals each request as it starts
+	// and waits on a release channel so the test can hold downloads in the
+	// queued-but-not-yet-consumed state before closing the iterator.
+	var inFlight atomic.Int64
+	release := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		inFlight.Add(1)
+		<-release
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write(arrowBytes)
+	}))
+	defer server.Close()
+
+	const nLinks = 20
+	links := make([]*cli_service.TSparkArrowResultLink, nLinks)
+	for i := range links {
+		links[i] = &cli_service.TSparkArrowResultLink{
+			FileLink:       server.URL,
+			ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
+			StartRowOffset: int64(i),
+			RowCount:       1,
+		}
+	}
+
+	cfg := config.WithDefaults()
+	cfg.UseLz4Compression = false
+	cfg.MaxDownloadThreads = 10
+
+	bi, err := NewCloudBatchIterator(context.Background(), links, 0, nil, cfg, nil)
+	assert.Nil(t, err)
+
+	// Kick off the first batch download. The iterator schedules
+	// MaxDownloadThreads concurrent fetches behind the scenes.
+	go func() { _, _ = bi.Next() }()
+
+	// Wait for all MaxDownloadThreads goroutines to be blocked inside the
+	// server handler (they've issued the GET and are waiting for the body).
+	assert.Eventually(t, func() bool {
+		return inFlight.Load() == int64(cfg.MaxDownloadThreads)
+	}, 5*time.Second, 10*time.Millisecond, "expected %d in-flight downloads", cfg.MaxDownloadThreads)
+
+	// Release the downloads so each goroutine finishes its HTTP read and
+	// attempts to send its result on the unbuffered resultChan. Only the
+	// first task's result will be read (by the Next() call above); the rest
+	// will be queued, blocked on the send.
+	close(release)
+
+	// Give the goroutines time to finish their HTTP work and reach the
+	// channel send.
+	time.Sleep(200 * time.Millisecond)
+
+	// Close the iterator without consuming the remaining batches.
+	bi.Close()
+
+	// After Close, every cloudFetchDownloadTask goroutine must exit. We don't
+	// compare against the total goroutine count because httptest keeps
+	// persistent server/transport goroutines around — we look only for our
+	// own download goroutines.
+	assert.Eventually(t, func() bool {
+		return countDownloadTaskGoroutines() == 0
+	}, 5*time.Second, 50*time.Millisecond,
+		"cloudFetchDownloadTask goroutines leaked after Close: have %d",
+		countDownloadTaskGoroutines())
+}
+
+// countDownloadTaskGoroutines returns the number of live goroutines whose
+// stack includes cloudFetchDownloadTask.Run. Used to detect the leak in
+// issue #356.
+func countDownloadTaskGoroutines() int {
+	buf := make([]byte, 64*1024)
+	for {
+		n := runtime.Stack(buf, true)
+		if n < len(buf) {
+			buf = buf[:n]
+			break
+		}
+		buf = make([]byte, 2*len(buf))
+	}
+	return strings.Count(string(buf), "cloudFetchDownloadTask).Run")
+}