zeroshade commented on code in PR #1331:
URL: https://github.com/apache/arrow-go/pull/1331#discussion_r4065754530
##########
arrow/ipc/compression.go:
##########
@@ -81,29 +84,63 @@ func getCompressor(codec flatbuf.CompressionType)
compressor {
}
type decompressor interface {
- io.Reader
- Reset(io.Reader)
+ Decompress(dst, src []byte) error
Close()
}
Review Comment:
`Close()` has quietly become an asymmetric contract: for zstd it now means
"return to the pool", while for lz4 it still means "reset". A caller that calls
`Close()` twice, or keeps using the value afterwards, hands one decoder to two
owners.
Both current call sites are a single `defer codec.Close()`
(`file_reader.go:535`, `:923`), so this is correct today — and even a
double-`Put` wouldn't be a data race, since `DecodeAll` serializes on the
`decoders` channel; it would just contend. But there's no guardrail, and the
interface reads like an `io.Closer`.
Could you add a doc comment on the interface along the lines of: `Close
releases the decompressor; it must not be used afterwards.`
##########
arrow/ipc/compression.go:
##########
@@ -81,29 +84,63 @@ func getCompressor(codec flatbuf.CompressionType)
compressor {
}
type decompressor interface {
- io.Reader
- Reset(io.Reader)
+ Decompress(dst, src []byte) error
Close()
}
+var zstdDecompressorPool = sync.Pool{
+ New: func() any {
+ // WithDecoderConcurrency(1): Each pooled decoder is used by
one goroutine at a time, so a single
+ // block decoder is enough. The default would create up to four
that
+ // could never run in parallel;
+ //
+ // WithDecodeAllCapLimit(true): The cap limit bounds DecodeAll
to cap(dst), so a frame claiming a
+ // larger content size can't allocate beyond the buffer the
caller
+ // sized from the uncompressed length prefix.
+ dec, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1),
zstd.WithDecodeAllCapLimit(true))
+ if err != nil {
+ panic(err)
+ }
+ return &zstdDecompressor{Decoder: dec}
+ },
+}
+
type zstdDecompressor struct {
*zstd.Decoder
}
-func (z *zstdDecompressor) Reset(r io.Reader) {
- if err := z.Decoder.Reset(r); err != nil {
- panic(err)
+func (z *zstdDecompressor) Decompress(dst, src []byte) error {
+ if len(dst) == 0 {
+ return nil
+ }
+
+ // The decoder was created with WithDecodeAllCapLimit, so it decodes
into
+ // dst's own capacity and fails rather than allocating a larger slice.
+ out, err := z.DecodeAll(src, dst[:0])
+ if err != nil {
+ return err
}
+ // Catch cases where the prefix says fewer bytes than the content, but
the content fits in the dst's spare capacity
+ if len(out) != len(dst) {
+ return fmt.Errorf("arrow/ipc: zstd decompressed to %d bytes,
expected %d", len(out), len(dst))
+ }
+ return nil
Review Comment:
Good catch making this check unconditional — it's genuinely load-bearing,
not belt-and-braces. `memory.Buffer.Bytes()` returns `b.buf[:b.length]` where
`b.buf` is `roundUpToMultipleOf64`, so `cap(dst)` routinely exceeds `len(dst)`
by up to 63 bytes; without this a frame declaring `len+k` (k ≤ 63) would decode
"successfully" into the padding. `TestZstdDecompressorSpareCapacity` covers
exactly that.
Question on the resulting strictness: a frame carrying *more* data than the
uncompressed-length prefix claims previously truncated silently (`io.ReadFull`
stops at `len(dst)`) and now errors. That looks like a fix rather than a
regression — no Arrow producer should be emitting that — but it is an
observable behavior change and isn't mentioned in the PR description. Was it
deliberate? If so, could you note it under "Are there any user-facing changes?"
##########
arrow/ipc/compression.go:
##########
@@ -81,29 +84,63 @@ func getCompressor(codec flatbuf.CompressionType)
compressor {
}
type decompressor interface {
- io.Reader
- Reset(io.Reader)
+ Decompress(dst, src []byte) error
Close()
}
+var zstdDecompressorPool = sync.Pool{
+ New: func() any {
+ // WithDecoderConcurrency(1): Each pooled decoder is used by
one goroutine at a time, so a single
+ // block decoder is enough. The default would create up to four
that
+ // could never run in parallel;
+ //
+ // WithDecodeAllCapLimit(true): The cap limit bounds DecodeAll
to cap(dst), so a frame claiming a
+ // larger content size can't allocate beyond the buffer the
caller
+ // sized from the uncompressed length prefix.
+ dec, err := zstd.NewReader(nil, zstd.WithDecoderConcurrency(1),
zstd.WithDecodeAllCapLimit(true))
+ if err != nil {
+ panic(err)
+ }
+ return &zstdDecompressor{Decoder: dec}
+ },
+}
+
type zstdDecompressor struct {
*zstd.Decoder
}
-func (z *zstdDecompressor) Reset(r io.Reader) {
- if err := z.Decoder.Reset(r); err != nil {
- panic(err)
+func (z *zstdDecompressor) Decompress(dst, src []byte) error {
+ if len(dst) == 0 {
+ return nil
+ }
+
+ // The decoder was created with WithDecodeAllCapLimit, so it decodes
into
+ // dst's own capacity and fails rather than allocating a larger slice.
+ out, err := z.DecodeAll(src, dst[:0])
+ if err != nil {
+ return err
}
+ // Catch cases where the prefix says fewer bytes than the content, but
the content fits in the dst's spare capacity
+ if len(out) != len(dst) {
+ return fmt.Errorf("arrow/ipc: zstd decompressed to %d bytes,
expected %d", len(out), len(dst))
+ }
+ return nil
}
func (z *zstdDecompressor) Close() {
- z.Decoder.Close()
+ zstdDecompressorPool.Put(z)
}
Review Comment:
Related to the interface comment above — worth a line here recording *why*
`z.Decoder.Close()` is deliberately never called, i.e. that
`zstd.NewReader(nil)` starts no goroutines (it returns before `d.Reset(r)`), so
a pooled decoder holds only GC-reclaimable state. Without that note this looks
like a missing `Close` and someone will eventually "fix" it back into a
per-batch teardown.
##########
arrow/ipc/compression_test.go:
##########
@@ -0,0 +1,285 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ipc
+
+import (
+ "bytes"
+ "encoding/binary"
+ "errors"
+ "fmt"
+ "math/rand"
+ "runtime"
+ "sync"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/internal/flatbuf"
+ "github.com/klauspost/compress/zstd"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+var testCodecs = []struct {
+ name string
+ codec flatbuf.CompressionType
+ opt Option
+}{
+ {"zstd", flatbuf.CompressionTypeZSTD, WithZstd()},
+ {"lz4", flatbuf.CompressionTypeLZ4_FRAME, WithLZ4()},
+}
+
+func compressBuffer(t *testing.T, codec flatbuf.CompressionType, src []byte)
[]byte {
+ t.Helper()
+
+ var out bytes.Buffer
+ c := getCompressor(codec)
+ c.Reset(&out)
+ _, err := c.Write(src)
+ require.NoError(t, err)
+ require.NoError(t, c.Close())
+ return out.Bytes()
+}
+
+// compressibleBytes returns n deterministic bytes drawn from a small alphabet,
+// so they compress but aren't trivially uniform.
+func compressibleBytes(n int, seed int64) []byte {
+ rng := rand.New(rand.NewSource(seed))
+ out := make([]byte, n)
+ for i := range out {
+ out[i] = byte(rng.Intn(16))
+ }
+ return out
+}
+
+func TestDecompressorRoundTrip(t *testing.T) {
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+
+ for i, n := range []int{1, 100, 4096, 300_000, 17,
300_000} {
+ want := compressibleBytes(n, int64(i))
+ got := make([]byte, n)
+ require.NoError(t, dec.Decompress(got,
compressBuffer(t, tc.codec, want)), "size %d", n)
+ assert.Equal(t, want, got, "size %d", n)
+ }
+ })
+ }
+}
+
+func TestDecompressorEmptyDestination(t *testing.T) {
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+ assert.NoError(t, dec.Decompress(nil, nil))
+ assert.NoError(t, dec.Decompress([]byte{}, []byte{1, 2,
3}))
+ })
+ }
+}
+
+func TestDecompressorSizeMismatch(t *testing.T) {
+ want := compressibleBytes(1000, 1)
+
+ for _, tc := range testCodecs {
+ t.Run(tc.name, func(t *testing.T) {
+ dec := getDecompressor(tc.codec)
+ defer dec.Close()
+ compressed := compressBuffer(t, tc.codec, want)
+
+ // the uncompressed length prefix claims more than the
frame holds
+ assert.Error(t, dec.Decompress(make([]byte,
len(want)+1), compressed))
+
+ // and the decoder is still usable afterwards
+ got := make([]byte, len(want))
+ require.NoError(t, dec.Decompress(got, compressed))
+ assert.Equal(t, want, got)
+ })
+ }
+
+ // zstd frames declare their content size, so a prefix that is too small
+ // is caught instead of silently truncating.
+ t.Run("zstd/short-destination", func(t *testing.T) {
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+ assert.Error(t, dec.Decompress(make([]byte, len(want)-1),
compressBuffer(t, flatbuf.CompressionTypeZSTD, want)))
+ })
+}
+
+func TestZstdDecompressorBoundedByDestination(t *testing.T) {
+ // 64 MiB of zeros compresses to a few bytes
+ compressed := compressBuffer(t, flatbuf.CompressionTypeZSTD,
make([]byte, 64<<20))
+ require.Less(t, len(compressed), 64<<10)
+
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+
+ dst := make([]byte, 16)
+ err := dec.Decompress(dst, compressed)
+ require.Error(t, err)
+ assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected
error: %v", err)
+}
+
+// A frame header can declare the content size up front. A hostile frame can
+// declare far more than it holds, hoping the decoder allocates that much
+// before it ever reads the (tiny) body. The cap limit has to reject it up
+// front, without allocating: check both the error and the allocation, since a
+// decoder that allocated first and failed afterwards would return an error
too.
+func TestZstdDecompressorRejectsHugeDeclaredSizeWithoutAllocating(t
*testing.T) {
+ const declared = 64 << 20
+
+ // Hand-built 17-byte zstd frame that is well formed but lies about its
+ // size: the header declares 64 MiB of content, the frame holds one
byte.
+ frame := []byte{0x28, 0xb5, 0x2f, 0xfd, 0xe0}
+ frame = binary.LittleEndian.AppendUint64(frame, declared)
+ frame = append(frame, 0x09, 0x00, 0x00, 0x41)
+
+ var hdr zstd.Header
+ require.NoError(t, hdr.Decode(frame))
+ require.True(t, hdr.HasFCS)
+ require.Equal(t, uint64(declared), hdr.FrameContentSize)
+
+ dec := getDecompressor(flatbuf.CompressionTypeZSTD)
+ defer dec.Close()
+
+ // warm up, so decoder state that is created lazily isn't counted below
+ dst := make([]byte, 100)
+ require.NoError(t, dec.Decompress(dst, compressBuffer(t,
flatbuf.CompressionTypeZSTD, compressibleBytes(100, 3))))
+
+ var before, after runtime.MemStats
+ runtime.ReadMemStats(&before)
+ err := dec.Decompress(dst, frame)
+ runtime.ReadMemStats(&after)
+
+ require.Error(t, err)
+ assert.True(t, errors.Is(err, zstd.ErrDecoderSizeExceeded), "unexpected
error: %v", err)
+ assert.Less(t, after.TotalAlloc-before.TotalAlloc, uint64(1<<20),
+ "decoding a frame declaring %d bytes allocated too much",
declared)
Review Comment:
`TotalAlloc` is process-wide and cumulative, so anything else allocating
inside this window inflates the delta. The margin here is wide (1 MiB budget
against a 64 MiB declared size) so it will almost always pass, but this is the
shape of test that fails once a quarter on a loaded CI runner.
The intent — distinguishing "rejected up front" from "allocated, then
failed" — is worth keeping; it's just the mechanism I'd like to be less
ambient. A `runtime.GC()` before the first `ReadMemStats` would at least keep
GC bookkeeping out of the window.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
