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 <[email protected]>
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():
+ }
}
}
}
