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 603dfb0a fix(arrow/ipc): avoid deadlock on compression errors (#1142)
603dfb0a is described below
commit 603dfb0aeacf8e678ddb73914b9e4b1c452a687a
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 14 18:07:27 2026 +0200
fix(arrow/ipc): avoid deadlock on compression errors (#1142)
### Rationale for this change
Parallel body compression sent worker errors through an unbuffered
channel before the parent waited for the workers. A compressor failure
could therefore leave the writer blocked forever, and the producer could
continue sending work after the workers had stopped.
### What changes are included in this PR?
Buffer the first worker error, make producer sends cancellation-aware,
and add a regression test that uses a failing compressor with parallel
compression.
### Are these changes tested?
- `go test ./arrow/ipc`
- `go test -race ./arrow/ipc -run
TestRecordEncoderCompressionErrorDoesNotDeadlock`
### Are there any user-facing changes?
Parallel IPC compression now returns compressor failures instead of
hanging.
---
arrow/ipc/writer.go | 29 ++++++++++++++++++-----
arrow/ipc/writer_test.go | 61 ++++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 84 insertions(+), 6 deletions(-)
diff --git a/arrow/ipc/writer.go b/arrow/ipc/writer.go
index aebdb768..32edaf25 100644
--- a/arrow/ipc/writer.go
+++ b/arrow/ipc/writer.go
@@ -436,9 +436,11 @@ func (w *recordEncoder) compressBodyBuffers(p *Payload)
error {
n, err := codec.Write(p.body[idx].Bytes())
if err != nil {
+ buf.Release()
return err
}
if err := codec.Close(); err != nil {
+ buf.Release()
return err
}
@@ -471,7 +473,7 @@ func (w *recordEncoder) compressBodyBuffers(p *Payload)
error {
var (
wg sync.WaitGroup
ch = make(chan int)
- errch = make(chan error)
+ errch = make(chan error, 1)
ctx, cancel = context.WithCancel(context.Background())
)
defer cancel()
@@ -490,7 +492,10 @@ func (w *recordEncoder) compressBodyBuffers(p *Payload)
error {
}
if err := compress(idx, codec); err !=
nil {
- errch <- err
+ select {
+ case errch <- err:
+ default:
+ }
cancel()
return
}
@@ -502,15 +507,24 @@ func (w *recordEncoder) compressBodyBuffers(p *Payload)
error {
}(workerID)
}
+send:
for idx := range p.body {
- ch <- idx
+ select {
+ case ch <- idx:
+ case <-ctx.Done():
+ break send
+ }
}
close(ch)
wg.Wait()
- close(errch)
- return <-errch
+ select {
+ case err := <-errch:
+ return err
+ default:
+ return nil
+ }
}
func (w *recordEncoder) encode(p *Payload, rec arrow.RecordBatch) error {
@@ -528,7 +542,9 @@ func (w *recordEncoder) encode(p *Payload, rec
arrow.RecordBatch) error {
return fmt.Errorf("%w: minSpaceSavings not in range
[0,1]. Provided %.05f",
arrow.ErrInvalid, w.minSpaceSavings)
}
- w.compressBodyBuffers(p)
+ if err := w.compressBodyBuffers(p); err != nil {
+ return err
+ }
}
// position for the start of a buffer relative to the passed frame of
reference.
@@ -1153,6 +1169,7 @@ func GetRecordBatchPayload(batch arrow.RecordBatch, opts
...Option) (Payload, er
err := enc.Encode(&data, batch)
if err != nil {
+ data.Release()
return Payload{}, err
}
diff --git a/arrow/ipc/writer_test.go b/arrow/ipc/writer_test.go
index 32bc5d90..dc13032b 100644
--- a/arrow/ipc/writer_test.go
+++ b/arrow/ipc/writer_test.go
@@ -25,6 +25,7 @@ import (
"math"
"strings"
"testing"
+ "time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -46,6 +47,18 @@ type failingPayloadWriter struct {
type shortWriteWriter struct{}
+type failingCompressor struct {
+ err error
+}
+
+func (failingCompressor) MaxCompressedLen(n int) int { return n }
+func (failingCompressor) Reset(io.Writer) {}
+func (f failingCompressor) Write([]byte) (int, error) { return 0, f.err }
+func (failingCompressor) Close() error { return nil }
+func (failingCompressor) Type() flatbuf.CompressionType {
+ return flatbuf.CompressionTypeZSTD
+}
+
func (shortWriteWriter) Write(p []byte) (int, error) {
return len(p) - 1, io.ErrShortWrite
}
@@ -333,6 +346,54 @@ func TestNewWriterWithMinSpaceSavings(t *testing.T) {
assert.Equal(t, minSpaceSavings, writer.minSpaceSavings)
}
+func TestRecordEncoderCompressionErrorDoesNotDeadlock(t *testing.T) {
+ want := errors.New("compression failed")
+ body := make([]*memory.Buffer, 64)
+ for i := range body {
+ body[i] = memory.NewBufferBytes([]byte("payload"))
+ }
+ payload := Payload{body: body}
+ defer payload.Release()
+
+ encoder := newRecordEncoder(memory.DefaultAllocator, 0,
kMaxNestingDepth, true,
+ flatbuf.CompressionTypeZSTD, 2, 0, []compressor{
+ failingCompressor{err: want},
+ failingCompressor{err: want},
+ })
+
+ result := make(chan error, 1)
+ go func() {
+ result <- encoder.compressBodyBuffers(&payload)
+ }()
+
+ select {
+ case err := <-result:
+ require.ErrorIs(t, err, want)
+ case <-time.After(time.Second):
+ t.Fatal("compression did not return after timeout")
+ }
+}
+
+func TestRecordEncoderReturnsCompressionError(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type:
arrow.PrimitiveTypes.Int8}}, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+ builder.Field(0).(*array.Int8Builder).Append(1)
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ want := errors.New("compression failed")
+ encoder := newRecordEncoder(mem, 0, kMaxNestingDepth, true,
+ flatbuf.CompressionTypeZSTD, 1, 0,
[]compressor{failingCompressor{err: want}})
+ var payload Payload
+ defer payload.Release()
+
+ require.ErrorIs(t, encoder.Encode(&payload, record), want)
+}
+
func TestWriteWithCompressionAndMinSavings(t *testing.T) {
mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
defer mem.AssertSize(t, 0)
