(arrow-go) branch main updated: fix(arrow/flight): avoid blocking after stream cancellation (#989)

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git


The following commit(s) were added to refs/heads/main by this push:
     new bab58105 fix(arrow/flight): avoid blocking after stream cancellation (#989)
bab58105 is described below

commit bab58105e336430c782bba4a04ecfc1a4915f58b
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Thu Jul 23 18:09:43 2026 +0200

    fix(arrow/flight): avoid blocking after stream cancellation (#989)
    
    ### Rationale for this change
    
    `StreamChunksFromReader` observes context cancellation while sending
    record batches, but its final reader-error send still wrote directly to
    the output channel. If the consumer stopped receiving after
    cancellation, that send could block forever, preventing reader release
    and channel closure. This is a follow-up to the cancellation handling
    added in #615.
    
    ### What changes are included in this PR?
    
    Race final error delivery against context cancellation, matching the
    existing record-batch send behavior. Capture the reader error once
    before delivery. Add a regression with an already-canceled context, an
    unbuffered channel, and a reader that fails immediately.
    
    ### Are these changes tested?
    
    Yes:
    
    - `go test ./arrow/flight`
    - `go test -race -count=1 ./arrow/flight`
    
    ### Are there any user-facing changes?
    
    `StreamChunksFromReader` now returns promptly and releases its reader
    when cancellation prevents final error delivery. Error delivery is
    unchanged while the context remains active.
---
 arrow/flight/flight_test.go         | 37 +++++++++++++++++++++++++++++++++++++
 arrow/flight/record_batch_reader.go |  7 +++++--
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/arrow/flight/flight_test.go b/arrow/flight/flight_test.go
index 8d75aac2..e6b8a25b 100644
--- a/arrow/flight/flight_test.go
+++ b/arrow/flight/flight_test.go
@@ -24,6 +24,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
@@ -635,6 +636,42 @@ func TestStreamChunksFromReader_OK(t *testing.T) {
 
 }
 
+type immediateErrorRecordReader struct {
+	err      error
+	released atomic.Bool
+}
+
+func (*immediateErrorRecordReader) Retain()                        {}
+func (r *immediateErrorRecordReader) Release()                     { r.released.Store(true) }
+func (*immediateErrorRecordReader) Schema() *arrow.Schema          { return nil }
+func (*immediateErrorRecordReader) Next() bool                     { return false }
+func (*immediateErrorRecordReader) RecordBatch() arrow.RecordBatch { return nil }
+func (*immediateErrorRecordReader) Record() arrow.RecordBatch      { return nil }
+func (r *immediateErrorRecordReader) Err() error                   { return r.err }
+
+func TestStreamChunksFromReader_CancellationWhileSendingError(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	rdr := &immediateErrorRecordReader{err: errors.New("read failed")}
+	ch := make(chan flight.StreamChunk)
+	done := make(chan struct{})
+	go func() {
+		flight.StreamChunksFromReader(ctx, rdr, ch)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("StreamChunksFromReader blocked sending an error after cancellation")
+	}
+
+	if !rdr.released.Load() {
+		t.Fatal("reader was not released")
+	}
+}
+
 // TestStreamChunksFromReader_HandlesCancellation verifies that context cancellation
 // causes StreamChunksFromReader to exit cleanly and release the reader.
 func TestStreamChunksFromReader_HandlesCancellation(t *testing.T) {
diff --git a/arrow/flight/record_batch_reader.go b/arrow/flight/record_batch_reader.go
index e6990a57..8c9d3876 100644
--- a/arrow/flight/record_batch_reader.go
+++ b/arrow/flight/record_batch_reader.go
@@ -248,8 +248,11 @@ func StreamChunksFromReader(ctx context.Context, rdr array.RecordReader, ch chan
 	}
 
 	if e, ok := rdr.(haserr); ok {
-		if e.Err() != nil {
-			ch <- StreamChunk{Err: e.Err()}
+		if err := e.Err(); err != nil {
+			select {
+			case ch <- StreamChunk{Err: err}:
+			case <-ctx.Done():
+			}
 		}
 	}
 }