diff --git a/purego/NEXT_CHANGELOG.md b/purego/NEXT_CHANGELOG.md index c01a2f97..e501a211 100644 --- a/purego/NEXT_CHANGELOG.md +++ b/purego/NEXT_CHANGELOG.md @@ -14,6 +14,10 @@ ### Bug Fixes +- Accept an explicit `null` for the optional `close_stream_duration_ms` field in + Arrow Flight acknowledgment metadata. It was rejected as malformed, which + turned an otherwise valid acknowledgment into a stream failure. + ### Documentation - Flush recovery no longer treats every flush error as terminal. The JSON single @@ -32,6 +36,28 @@ and encoder, ack-model, and opener seams let a protocol instantiate the core over its own payload type. Proto and JSON behavior is unchanged, and nothing is exposed through a public API yet. +- Add `internal/arrowproto`, the Arrow IPC payload and Flight frame encoder for + the upcoming Arrow ingestion path. It fills the stream core's encoder seam with + the two hooks that are trivial for proto and JSON but not for Arrow: row counts + as durability units, and real row-range slicing so a partially acknowledged + batch replays only its unacknowledged suffix. Frames are chunked to 2 MiB based + on measured protobuf size. Nothing is exposed through a public API yet. +- Charge Arrow payloads against the buffered-bytes limit before decoding them. + Compressed Arrow IPC input is inspected for its declared uncompressed buffer + sizes, so a highly compressible payload cannot pass admission and then expand + past the limit while Arrow materializes it. Because canonicalizing re-encodes + the input, the encoding seam also reports what a payload retains once it + exists, and the reservation is reconciled against that instead of the input + length: a payload whose encoded form does not track its input size can no + longer be admitted for a fraction of the memory it holds. Proto and JSON keep + charging their existing estimate. +- Size a RecordBatch for admission from the rows it covers rather than from the + whole buffers it points at. A slice shares its parent's buffers, so the old + measurement charged a small slice of a large batch for the entire parent and + rejected it as too large. +- Add `github.com/apache/arrow-go/v18` and `github.com/google/flatbuffers` + dependencies. arrow-go requires `google.golang.org/grpc` v1.82.0, which raises + this module's grpc minimum from v1.81.1. ### Breaking Changes diff --git a/purego/go.mod b/purego/go.mod index 17015b55..0cc952de 100644 --- a/purego/go.mod +++ b/purego/go.mod @@ -3,13 +3,21 @@ module github.com/databricks/zerobus-sdk/purego go 1.25.0 require ( - google.golang.org/grpc v1.81.1 + github.com/apache/arrow-go/v18 v18.7.0 + github.com/google/flatbuffers v25.12.19+incompatible + google.golang.org/grpc v1.82.0 google.golang.org/protobuf v1.36.11 ) require ( - golang.org/x/net v0.51.0 // indirect - golang.org/x/sys v0.42.0 // indirect - golang.org/x/text v0.34.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 // indirect + github.com/goccy/go-json v0.10.6 // indirect + github.com/klauspost/compress v1.19.0 // indirect + github.com/klauspost/cpuid/v2 v2.4.0 // indirect + github.com/pierrec/lz4/v4 v4.1.27 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect + golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.47.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect ) diff --git a/purego/go.sum b/purego/go.sum index 44c671d1..6b517da9 100644 --- a/purego/go.sum +++ b/purego/go.sum @@ -1,15 +1,41 @@ +github.com/andybalholm/brotli v1.2.2 h1:HzTuoo2ErYQqf5qvcJInB8uvqSVxRttzkFexPWtnceM= +github.com/andybalholm/brotli v1.2.2/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/apache/arrow-go/v18 v18.7.0 h1:Vw/i+cJyebUofT7JlqFpe65LrmwxULn166jjwStM4HY= +github.com/apache/arrow-go/v18 v18.7.0/go.mod h1:PM6IigLJkdMwIpeHXnymo+xZ52f42a9EYiLtRel4p/A= +github.com/apache/thrift v0.24.0 h1:zy31L1a49QTNB2bG1BBfMXol3yJrTH975G3pPubQVLQ= +github.com/apache/thrift v0.24.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= +github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= +github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/compress v1.19.0 h1:sXLILfc9jV2QYWkzFOPWStmcUVH2RHEB1JCdY2oVvCQ= +github.com/klauspost/compress v1.19.0/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/cpuid/v2 v2.4.0 h1:S6Hrbc7+ywsr0r+RLapfGBHfyefhCTwEh3A0tV913Dw= +github.com/klauspost/cpuid/v2 v2.4.0/go.mod h1:19jmZ9mjzoF//ddRSUsv0zfBTJWh3QJh9FNxZTMrGxU= +github.com/pierrec/lz4/v4 v4.1.27 h1:+PhzhWDrjRj89TH2sw43nE3+4+W8lSxIuQadEHZyjUk= +github.com/pierrec/lz4/v4 v4.1.27/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= @@ -22,17 +48,21 @@ go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfC go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= -golang.org/x/net v0.51.0 h1:94R/GTO7mt3/4wIKpcR5gkGmRLOuE/2hNGeWq/GBIFo= -golang.org/x/net v0.51.0/go.mod h1:aamm+2QF5ogm02fjy5Bb7CQ0WMt1/WVM7FtyaTLlA9Y= -golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= -golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= -golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 h1:ggcbiqK8WWh6l1dnltU4BgWGIGo+EVYxCaAPih/zQXQ= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= -google.golang.org/grpc v1.81.1/go.mod h1:xGH9GfzOyMTGIOXBJmXt+BX/V0kcdQbdcuwQ/zNw42I= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.82.0 h1:vguDnZUPjE26w09A63VoxZPnvPjB5Riyc0mkXPFmAIU= +google.golang.org/grpc v1.82.0/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/purego/internal/arrowproto/encoder.go b/purego/internal/arrowproto/encoder.go new file mode 100644 index 00000000..d3b36f95 --- /dev/null +++ b/purego/internal/arrowproto/encoder.go @@ -0,0 +1,1214 @@ +// Package arrowproto implements the Arrow IPC payload and the encoder that turns +// it into Arrow Flight frames. Payloads hold only self-contained IPC bytes, so +// caller-owned Arrow arrays are never retained by the stream core. +package arrowproto + +import ( + "bytes" + "encoding/json" + "fmt" + "math" + "sync/atomic" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/flight" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" + "google.golang.org/protobuf/proto" + + "github.com/databricks/zerobus-sdk/purego/internal/stream" + "github.com/databricks/zerobus-sdk/purego/internal/transport" +) + +// TargetFlightDataBytes caps the protobuf size of a non-schema FlightData frame. +const TargetFlightDataBytes = 2 * 1024 * 1024 + +const payloadOverheadBytes = int64(64) +const admissionSlopBytes = int64(64 * 1024) + +// Compression selects the compression used when serializing Arrow IPC batches. +type Compression uint8 + +const ( + CompressionNone Compression = iota + CompressionLZ4 + CompressionZstd +) + +// Options configures Arrow IPC materialization. +type Options struct { + Compression Compression + + // Allocator is mainly a test seam. Nil uses Arrow's default. + Allocator memory.Allocator +} + +// Payload is one eagerly materialized, non-empty Arrow RecordBatch held as a +// canonical self-contained IPC stream: schema, dictionaries, batch, end marker. +// chunkRows is the frame plan the sender replays after decoding that stream once. +type Payload struct { + ipcBytes []byte + rows uint64 + chunkRows []int64 +} + +// UnitCount returns the number of row durability units in the payload. +func (p *Payload) UnitCount() uint64 { + if p == nil { + return 0 + } + return p.rows +} + +// RetainedSize returns a conservative heap charge for the payload. +func (p *Payload) RetainedSize() int64 { + if p == nil { + return 0 + } + return payloadOverheadBytes + int64(cap(p.ipcBytes)) + + int64(cap(p.chunkRows))*8 +} + +// IPCBytes returns a caller-owned copy of the self-contained IPC stream. +func (p *Payload) IPCBytes() []byte { + if p == nil { + return nil + } + return bytes.Clone(p.ipcBytes) +} + +// Protocol owns one exact Arrow schema and the IPC/Flight encoding policy. +type Protocol struct { + schema *arrow.Schema + compression Compression + allocator memory.Allocator + schemaFrame *flight.FlightData + admissionBaseBytes int64 + chunkProbe func(rowStart, rowCount int64) +} + +// EncodeSchemaIPC serializes schema as a schema-only Arrow IPC stream. +func EncodeSchemaIPC(schema *arrow.Schema) ([]byte, error) { + if schema == nil { + return nil, fmt.Errorf("arrow protocol: schema is required") + } + var output bytes.Buffer + writer := ipc.NewWriter(&output, ipc.WithSchema(schema)) + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("arrow protocol: serialize schema: %w", err) + } + return output.Bytes(), nil +} + +// DecodeSchemaIPC parses a schema-only Arrow IPC stream and returns an +// independent schema. Record batches in the input are rejected. +func DecodeSchemaIPC(data []byte) (*arrow.Schema, error) { + if len(data) == 0 { + return nil, fmt.Errorf("arrow protocol: schema IPC input is empty") + } + source := bytes.NewReader(data) + reader, err := ipc.NewReader(source) + if err != nil { + return nil, fmt.Errorf("arrow protocol: invalid schema IPC stream: %w", err) + } + defer reader.Release() + if reader.Next() { + return nil, fmt.Errorf("arrow protocol: schema IPC stream must not contain RecordBatches") + } + if err := reader.Err(); err != nil { + return nil, fmt.Errorf("arrow protocol: read schema IPC stream: %w", err) + } + if source.Len() != 0 { + return nil, fmt.Errorf( + "arrow protocol: schema IPC stream contains %d trailing bytes", + source.Len(), + ) + } + return cloneSchemaThroughIPC(reader.Schema(), memory.DefaultAllocator) +} + +// New constructs an Arrow protocol encoder, copying schema so the caller's +// Schema object is not retained. +func New(schema *arrow.Schema, options Options) (*Protocol, error) { + if schema == nil { + return nil, fmt.Errorf("arrow protocol: schema is required") + } + if _, err := compressionOption(options.Compression); err != nil { + return nil, err + } + allocator := options.Allocator + if allocator == nil { + allocator = memory.DefaultAllocator + } + ownedSchema, err := cloneSchemaThroughIPC(schema, allocator) + if err != nil { + return nil, err + } + p := &Protocol{ + schema: ownedSchema, + compression: options.Compression, + allocator: allocator, + } + frame, err := p.makeSchemaFrame() + if err != nil { + return nil, err + } + p.schemaFrame = frame + p.admissionBaseBytes = admissionSlopBytes + 2*int64(proto.Size(frame)) + return p, nil +} + +// SchemaFlightData returns a caller-owned schema frame for the first DoPut, with +// no body and no app metadata. +func (p *Protocol) SchemaFlightData() *flight.FlightData { + if p == nil || p.schemaFrame == nil { + return nil + } + return proto.Clone(p.schemaFrame).(*flight.FlightData) +} + +// EncodeRecordBatch serializes batch immediately. The returned payload owns no +// references to batch, its columns, or their buffers. +func (p *Protocol) EncodeRecordBatch(batch arrow.RecordBatch) (*Payload, error) { + if batch == nil { + return nil, fmt.Errorf("arrow protocol: RecordBatch is required") + } + if err := p.validateBatch(batch); err != nil { + return nil, err + } + serialized, err := p.serialize(batch) + if err != nil { + return nil, fmt.Errorf("arrow protocol: serialize RecordBatch: %w", err) + } + chunkRows, err := p.planChunks(batch) + if err != nil { + return nil, err + } + return p.payloadFromCanonicalIPC( + serialized, + uint64(batch.NumRows()), + chunkRows, + ) +} + +// EncodeIPC validates one non-empty RecordBatch against the exact schema and +// reserializes it, hydrating dictionary state into a standalone stream rather +// than depending on the caller's frames. +func (p *Protocol) EncodeIPC(data []byte) (*Payload, error) { + batch, err := p.decodeOne(data) + if err != nil { + return nil, err + } + defer batch.Release() + + serialized, err := p.serialize(batch) + if err != nil { + return nil, fmt.Errorf("arrow protocol: canonicalize IPC RecordBatch: %w", err) + } + chunkRows, err := p.planChunks(batch) + if err != nil { + return nil, err + } + return p.payloadFromCanonicalIPC( + serialized, + uint64(batch.NumRows()), + chunkRows, + ) +} + +// EstimateRecordBatchRetainedSize returns a conservative pre-materialization +// reservation. Arrow's buffer total already covers nested children and +// dictionaries; doubling it covers framing, compression, and buffer growth. +func (p *Protocol) EstimateRecordBatchRetainedSize( + batch arrow.RecordBatch, +) (int64, error) { + if batch == nil { + return 0, fmt.Errorf("arrow protocol: RecordBatch is required") + } + if err := p.validateBatch(batch); err != nil { + return 0, err + } + rowPlanBytes := int64(batch.NumRows()) + if rowPlanBytes > math.MaxInt64/8 { + return math.MaxInt64, nil + } + metadataBytes := recordBatchMetadataSize(batch) + inputBytes, err := addInt64Saturating( + totalRecordBufferSize(batch), + metadataBytes, + ) + if err != nil { + return math.MaxInt64, nil + } + return p.admissionEstimate(inputBytes, 0, rowPlanBytes*8), nil +} + +// EstimateIPCRetainedSize reserves for canonicalizing one caller-owned IPC +// stream. Compressed metadata is inspected before Arrow sees the input, so +// declared uncompressed sizes are part of the reservation. +func (p *Protocol) EstimateIPCRetainedSize(data []byte) (int64, error) { + expandedBytes, err := preflightIPCExpansion(data) + if err != nil { + return 0, err + } + return p.admissionEstimate(int64(len(data)), expandedBytes, 0), nil +} + +func (p *Protocol) admissionEstimate( + inputBytes, expandedBytes, extraBytes int64, +) int64 { + if inputBytes < 0 || + expandedBytes < 0 || + extraBytes < 0 || + p.admissionBaseBytes > math.MaxInt64-payloadOverheadBytes { + return math.MaxInt64 + } + total := payloadOverheadBytes + p.admissionBaseBytes + for _, component := range []struct { + value int64 + scale int64 + }{ + {value: inputBytes, scale: 2}, + {value: expandedBytes, scale: 2}, + {value: extraBytes, scale: 1}, + } { + if component.value > math.MaxInt64/component.scale { + return math.MaxInt64 + } + scaled := component.value * component.scale + if scaled > math.MaxInt64-total { + return math.MaxInt64 + } + total += scaled + } + return total +} + +// DecodeIPCRecordBatch parses exactly one non-empty RecordBatch with the +// protocol schema. The caller owns the returned reference and must Release it. +func (p *Protocol) DecodeIPCRecordBatch(data []byte) (arrow.RecordBatch, error) { + if p == nil { + return nil, fmt.Errorf("arrow protocol: protocol is nil") + } + return p.decodeOne(data) +} + +// Slice drops an acknowledged row prefix and reserializes the remaining suffix +// as a standalone IPC payload. +func (p *Protocol) Slice(payload *Payload, acknowledgedPrefix uint64) (*Payload, error) { + if payload == nil { + return nil, fmt.Errorf("arrow protocol: payload is nil") + } + // The core only slices a partially acknowledged payload, so a prefix outside + // (0,rows) means its accounting is wrong. Say so rather than invent a payload. + if acknowledgedPrefix == 0 || acknowledgedPrefix >= payload.rows { + return nil, fmt.Errorf( + "arrow protocol: acknowledged prefix %d is invalid for %d rows", + acknowledgedPrefix, + payload.rows, + ) + } + + batch, err := p.decodeOne(payload.ipcBytes) + if err != nil { + return nil, fmt.Errorf("arrow protocol: decode payload for slicing: %w", err) + } + defer batch.Release() + // The prefix was checked against the header count, so a decoded batch that + // disagrees would make NewSlice panic on an out-of-range bound. + if uint64(batch.NumRows()) != payload.rows { + return nil, fmt.Errorf( + "arrow protocol: payload row count changed: header=%d decoded=%d", + payload.rows, + batch.NumRows(), + ) + } + + suffix := batch.NewSlice(int64(acknowledgedPrefix), batch.NumRows()) + defer suffix.Release() + serialized, err := p.serialize(suffix) + if err != nil { + return nil, fmt.Errorf("arrow protocol: serialize sliced suffix: %w", err) + } + chunkRows, err := p.planChunks(suffix) + if err != nil { + return nil, err + } + return p.payloadFromCanonicalIPC( + serialized, + uint64(suffix.NumRows()), + chunkRows, + ) +} + +// Decode returns the payload as caller-owned IPC bytes shaped for +// stream.EncoderHooks. +func (p *Protocol) Decode(payload *Payload) [][]byte { + if payload == nil { + return nil + } + return [][]byte{payload.IPCBytes()} +} + +// EncoderHooks returns the stream-core hooks for Arrow payloads. The []byte entry +// points take one self-contained IPC RecordBatch; typed callers should prefer +// EncodeRecordBatch with CoreStream.EnqueuePayload. +func (p *Protocol) EncoderHooks() stream.EncoderHooks[*Payload] { + return stream.EncoderHooks[*Payload]{ + EncodeRecord: p.EncodeIPC, + EncodeBatch: func(records [][]byte) (*Payload, error) { + if len(records) != 1 { + return nil, fmt.Errorf( + "arrow protocol: IPC batch input must contain exactly one payload, got %d", + len(records), + ) + } + return p.EncodeIPC(records[0]) + }, + StampOffset: func(*Payload, int64) {}, + UnitCount: func(payload *Payload) uint64 { + return payload.UnitCount() + }, + Slice: p.Slice, + Decode: p.Decode, + MaxWireSize: func(payload *Payload) int { + if payload == nil { + return 0 + } + return TargetFlightDataBytes + }, + RetainedSize: func(rawBytes, recordCount int) int64 { + return payloadOverheadBytes + int64(rawBytes) + int64(recordCount)*8 + }, + // Canonicalizing re-serializes with this protocol's compression, so the + // input length says nothing about what the payload retains: compressed + // input expands, uncompressed input may shrink. Charge the real figure. + ActualRetainedSize: func(payload *Payload) int64 { + return payload.RetainedSize() + }, + } +} + +// EncodeFlightData materializes payload into caller-owned non-schema frames. The +// wire path drives the same emitter with a live sink, so it never buffers a whole +// payload before the first Send. +func (p *Protocol) EncodeFlightData( + payload *Payload, + startOffset int64, +) ([]*flight.FlightData, int64, error) { + if startOffset < 0 { + return nil, startOffset, fmt.Errorf( + "arrow protocol: Flight frame offset must be non-negative, got %d", + startOffset, + ) + } + var frames []*flight.FlightData + emitter := &flightFrameEmitter{ + send: func(frame *flight.FlightData) error { + frames = append(frames, proto.Clone(frame).(*flight.FlightData)) + return nil + }, + nextOffset: startOffset, + skipSchema: true, + } + options, err := p.flightWriterOptions() + if err != nil { + return nil, startOffset, err + } + writer := flight.NewRecordWriter(emitter, options...) + if err := p.emitPayload(payload, writer, emitter); err != nil { + _ = writer.Close() + return nil, startOffset, err + } + if err := writer.Close(); err != nil { + return nil, startOffset, fmt.Errorf("arrow protocol: close Flight encoder: %w", err) + } + return frames, emitter.nextOffset, nil +} + +// flightFrameEmitter stamps sequential metadata and hands each frame straight to +// its sink. One connection-scoped ipc.Writer feeds it, so Arrow Go caches an +// unchanged dictionary instead of re-emitting it per chunk. +type flightFrameEmitter struct { + send func(*flight.FlightData) error + nextOffset int64 + exhausted bool + skipSchema bool + highest *atomic.Int64 + + currentRowEnd uint64 + receipt *stream.SubmissionReceipt +} + +func (e *flightFrameEmitter) Send(frame *flight.FlightData) error { + if frame == nil { + return fmt.Errorf("arrow protocol: Flight encoder emitted a nil frame") + } + messageType, err := flightMessageType(frame) + if err != nil { + return err + } + if messageType == ipc.MessageSchema { + if !e.skipSchema { + return fmt.Errorf("arrow protocol: Flight encoder emitted a duplicate schema frame") + } + e.skipSchema = false + if len(frame.GetDataBody()) != 0 { + return fmt.Errorf("arrow protocol: encoded Flight schema has a body") + } + return nil + } + if e.skipSchema { + return fmt.Errorf("arrow protocol: Flight encoder omitted its schema frame") + } + if e.exhausted { + return fmt.Errorf("arrow protocol: Flight frame offset space exhausted") + } + metadata, err := json.Marshal(transport.FlightBatchMetadata{ + OffsetID: e.nextOffset, + }) + if err != nil { + return fmt.Errorf("arrow protocol: marshal Flight frame metadata: %w", err) + } + frame.AppMetadata = metadata + if size := proto.Size(frame); size > TargetFlightDataBytes { + return fmt.Errorf( + "arrow protocol: encoded FlightData is %d bytes, exceeds %d-byte target", + size, + TargetFlightDataBytes, + ) + } + if e.highest != nil { + // Publish before Send: the server may respond as soon as gRPC accepts the + // frame, concurrently with Send returning. + e.highest.Store(e.nextOffset) + } + if err := e.send(frame); err != nil { + return err + } + if messageType == ipc.MessageRecordBatch && e.receipt != nil { + e.receipt.SubmittedUnits = e.currentRowEnd + } + if e.nextOffset == math.MaxInt64 { + e.exhausted = true + } else { + e.nextOffset++ + } + return nil +} + +func flightMessageType(frame *flight.FlightData) (ipc.MessageType, error) { + if len(frame.GetDataHeader()) == 0 { + return ipc.MessageNone, fmt.Errorf("arrow protocol: Flight frame has no IPC header") + } + meta := memory.NewBufferBytes(frame.GetDataHeader()) + body := memory.NewBufferBytes(frame.GetDataBody()) + message := ipc.NewMessage(meta, body) + meta.Release() + body.Release() + defer message.Release() + return message.Type(), nil +} + +func (p *Protocol) payloadFromCanonicalIPC( + data []byte, + rows uint64, + chunkRows []int64, +) (*Payload, error) { + if len(data) == 0 || rows == 0 || len(chunkRows) == 0 { + return nil, fmt.Errorf("arrow protocol: canonical IPC payload is empty") + } + // A plan that does not cover exactly the rows the payload reports would + // otherwise surface as a short submission mid-send, long after encoding. + var planned int64 + for _, rowCount := range chunkRows { + if rowCount <= 0 || planned > math.MaxInt64-rowCount { + return nil, fmt.Errorf("arrow protocol: invalid Flight chunk plan") + } + planned += rowCount + } + if uint64(planned) != rows { + return nil, fmt.Errorf( + "arrow protocol: Flight chunk plan covers %d of %d rows", + planned, + rows, + ) + } + // data is the serializer's own buffer, which nothing else aliases, so the + // payload can adopt it without a second full copy. + return &Payload{ipcBytes: data, rows: rows, chunkRows: chunkRows}, nil +} + +func (p *Protocol) validateBatch(batch arrow.RecordBatch) error { + if !exactSchemaEqual(batch.Schema(), p.schema) { + return fmt.Errorf( + "arrow protocol: RecordBatch schema does not exactly match stream schema", + ) + } + if batch.NumRows() <= 0 { + return fmt.Errorf("arrow protocol: RecordBatch must contain at least one row") + } + return nil +} + +func exactSchemaEqual(left, right *arrow.Schema) bool { + return left != nil && right != nil && + left.Equal(right) && + left.Metadata().Equal(right.Metadata()) +} + +// totalRecordBufferSize sums the bytes a batch's columns hold. A slice shares +// its parent's buffers, so a buffer's own length reports the parent's whole +// extent: charging that would reject a ten-row slice of a million-row batch. +// Layouts with a derivable per-row extent are charged for their own rows only; +// the rest fall back to whole buffers, an over-estimate that reconciliation +// corrects once the payload exists. +func totalRecordBufferSize(batch arrow.RecordBatch) int64 { + seen := make(map[*memory.Buffer]struct{}) + var total int64 + for _, column := range batch.Columns() { + size, err := addInt64Saturating(total, arrayDataSize(column.Data(), seen)) + if err != nil { + return math.MaxInt64 + } + total = size + } + return total +} + +func arrayDataSize(data arrow.ArrayData, seen map[*memory.Buffer]struct{}) int64 { + if data == nil { + return 0 + } + if concrete, ok := data.(*array.Data); ok && concrete == nil { + return 0 + } + total := ownedBufferSize(data, seen) + for _, child := range data.Children() { + size, err := addInt64Saturating(total, arrayDataSize(child, seen)) + if err != nil { + return math.MaxInt64 + } + total = size + } + // A dictionary is shared whole rather than sliced per row, so it is charged + // in full through the same recursion. + size, err := addInt64Saturating(total, arrayDataSize(data.Dictionary(), seen)) + if err != nil { + return math.MaxInt64 + } + return size +} + +// ownedBufferSize charges one array node for the rows it actually covers. +func ownedBufferSize(data arrow.ArrayData, seen map[*memory.Buffer]struct{}) int64 { + rows := int64(data.Len()) + buffers := data.Buffers() + if rows < 0 { + return wholeBufferSize(buffers, seen) + } + var validityBytes int64 + if len(buffers) > 0 && buffers[0] != nil { + validityBytes = bitmapBytes(rows) + } + switch dataType := data.DataType().(type) { + case *arrow.StringType, *arrow.BinaryType: + if valueBytes, ok := variableWidthValueBytes(data, rows, false); ok { + return validityBytes + (rows+1)*4 + valueBytes + } + case *arrow.LargeStringType, *arrow.LargeBinaryType: + if valueBytes, ok := variableWidthValueBytes(data, rows, true); ok { + return validityBytes + (rows+1)*8 + valueBytes + } + case arrow.FixedWidthDataType: + // A wider layout would put row data in a buffer this arithmetic does not + // know about, so only the canonical validity+values shape is derived. + if len(buffers) <= 2 { + return validityBytes + fixedWidthValueBytes(rows, dataType.BitWidth()) + } + } + return wholeBufferSize(buffers, seen) +} + +// variableWidthValueBytes reads the offsets buffer to size exactly the values +// this array's rows reference. +func variableWidthValueBytes( + data arrow.ArrayData, + rows int64, + large bool, +) (int64, bool) { + buffers := data.Buffers() + if len(buffers) != 3 || buffers[1] == nil { + return 0, false + } + start := int64(data.Offset()) + end := start + rows + if start < 0 || end < start { + return 0, false + } + var first, last int64 + if large { + offsets := arrow.Int64Traits.CastFromBytes(buffers[1].Bytes()) + if int64(len(offsets)) <= end { + return 0, false + } + first, last = offsets[start], offsets[end] + } else { + offsets := arrow.Int32Traits.CastFromBytes(buffers[1].Bytes()) + if int64(len(offsets)) <= end { + return 0, false + } + first, last = int64(offsets[start]), int64(offsets[end]) + } + if first < 0 || last < first { + return 0, false + } + return last - first, true +} + +func fixedWidthValueBytes(rows int64, bitWidth int) int64 { + if bitWidth <= 0 { + return 0 + } + if bitWidth < 8 { + return bitmapBytes(rows) + } + width := int64(bitWidth / 8) + if rows > math.MaxInt64/width { + return math.MaxInt64 + } + return rows * width +} + +func bitmapBytes(rows int64) int64 { + return (rows + 7) / 8 +} + +func wholeBufferSize( + buffers []*memory.Buffer, + seen map[*memory.Buffer]struct{}, +) int64 { + var total int64 + for _, buffer := range buffers { + if buffer == nil { + continue + } + if _, exists := seen[buffer]; exists { + continue + } + seen[buffer] = struct{}{} + size, err := addInt64Saturating(total, int64(buffer.Len())) + if err != nil { + return math.MaxInt64 + } + total = size + } + return total +} + +func recordBatchMetadataSize(batch arrow.RecordBatch) int64 { + withMetadata, ok := batch.(arrow.RecordBatchWithMetadata) + if !ok { + return 0 + } + metadata := withMetadata.Metadata() + keys := metadata.Keys() + values := metadata.Values() + var total int64 + for index, key := range keys { + value := values[index] + // Count FlatBuffer vector/table/offset overhead on top of UTF-8 contents. + // The full estimate doubles this for materialization. + entryBytes, err := addInt64Saturating(int64(len(key)), int64(len(value))) + if err == nil { + entryBytes, err = addInt64Saturating(entryBytes, 32) + } + if err == nil { + total, err = addInt64Saturating(total, entryBytes) + } + if err != nil { + return math.MaxInt64 + } + } + return total +} + +// cloneSchemaThroughIPC rebuilds nested and dictionary DataTypes independently. A +// shallow field copy would still alias mutable pointer-backed types such as +// DictionaryType. +func cloneSchemaThroughIPC( + schema *arrow.Schema, + allocator memory.Allocator, +) (*arrow.Schema, error) { + var output bytes.Buffer + writer := ipc.NewWriter( + &output, + ipc.WithSchema(schema), + ipc.WithAllocator(allocator), + ) + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("arrow protocol: serialize schema: %w", err) + } + reader, err := ipc.NewReader( + bytes.NewReader(output.Bytes()), + ipc.WithAllocator(allocator), + ) + if err != nil { + return nil, fmt.Errorf("arrow protocol: deserialize schema: %w", err) + } + defer reader.Release() + decoded := reader.Schema() + metadata := decoded.Metadata() + return arrow.NewSchemaWithEndian( + decoded.Fields(), + &metadata, + decoded.Endianness(), + ), nil +} + +func (p *Protocol) decodeOne(data []byte) (arrow.RecordBatch, error) { + if len(data) == 0 { + return nil, fmt.Errorf("arrow protocol: IPC input is empty") + } + // Cross-check declared message and buffer sizes against the input before + // Arrow allocates against them: every decode entry point runs through here. + if _, err := preflightIPCExpansion(data); err != nil { + return nil, err + } + source := bytes.NewReader(data) + reader, err := ipc.NewReader( + source, + ipc.WithAllocator(p.allocator), + ) + if err != nil { + return nil, fmt.Errorf("arrow protocol: invalid IPC stream: %w", err) + } + defer reader.Release() + if !exactSchemaEqual(reader.Schema(), p.schema) { + return nil, fmt.Errorf("arrow protocol: IPC schema does not exactly match stream schema") + } + if !reader.Next() { + if err := reader.Err(); err != nil { + return nil, fmt.Errorf("arrow protocol: read IPC RecordBatch: %w", err) + } + return nil, fmt.Errorf("arrow protocol: IPC stream contains no RecordBatch") + } + batch := reader.RecordBatch() + batch.Retain() + if batch.NumRows() <= 0 { + batch.Release() + return nil, fmt.Errorf("arrow protocol: IPC RecordBatch must contain at least one row") + } + if reader.Next() { + batch.Release() + return nil, fmt.Errorf( + "arrow protocol: IPC stream must contain exactly one RecordBatch", + ) + } + if err := reader.Err(); err != nil { + batch.Release() + return nil, fmt.Errorf("arrow protocol: read trailing IPC data: %w", err) + } + if source.Len() != 0 { + batch.Release() + return nil, fmt.Errorf( + "arrow protocol: IPC stream contains %d trailing bytes", + source.Len(), + ) + } + return batch, nil +} + +func (p *Protocol) serialize(batch arrow.RecordBatch) ([]byte, error) { + var output bytes.Buffer + options, err := p.ipcOptions() + if err != nil { + return nil, err + } + options = append(options, ipc.WithSchema(p.schema)) + writer := ipc.NewWriter(&output, options...) + if err := writer.Write(batch); err != nil { + _ = writer.Close() + return nil, err + } + if err := writer.Close(); err != nil { + return nil, err + } + return output.Bytes(), nil +} + +func (p *Protocol) makeSchemaFrame() (*flight.FlightData, error) { + collector := new(flightDataCollector) + options, err := p.ipcOptions() + if err != nil { + return nil, err + } + options = append(options, ipc.WithSchema(p.schema)) + writer := flight.NewRecordWriter(collector, options...) + if err := writer.Close(); err != nil { + return nil, fmt.Errorf("arrow protocol: encode Flight schema: %w", err) + } + if len(collector.frames) != 1 { + return nil, fmt.Errorf( + "arrow protocol: Flight schema encoding produced %d frames, want 1", + len(collector.frames), + ) + } + frame := collector.frames[0] + if len(frame.GetDataHeader()) == 0 || + len(frame.GetDataBody()) != 0 || + len(frame.GetAppMetadata()) != 0 { + return nil, fmt.Errorf("arrow protocol: malformed encoded Flight schema frame") + } + return frame, nil +} + +func (p *Protocol) emitPayload( + payload *Payload, + writer *flight.Writer, + emitter *flightFrameEmitter, +) error { + if payload == nil { + return fmt.Errorf("arrow protocol: payload is nil") + } + batch, err := p.decodeOne(payload.ipcBytes) + if err != nil { + return fmt.Errorf("arrow protocol: decode Flight payload: %w", err) + } + defer batch.Release() + if uint64(batch.NumRows()) != payload.rows { + return fmt.Errorf( + "arrow protocol: payload row count changed: header=%d decoded=%d", + payload.rows, + batch.NumRows(), + ) + } + + rowStart := int64(0) + for _, rowCount := range payload.chunkRows { + if rowCount <= 0 || rowStart > batch.NumRows()-rowCount { + return fmt.Errorf("arrow protocol: invalid Flight chunk plan") + } + rowEnd := rowStart + rowCount + chunk := batch.NewSlice(rowStart, rowEnd) + emitter.currentRowEnd = uint64(rowEnd) + err = writer.Write(chunk) + chunk.Release() + if err != nil { + return fmt.Errorf( + "arrow protocol: encode/send Flight rows [%d,%d): %w", + rowStart, + rowEnd, + err, + ) + } + rowStart = rowEnd + } + if rowStart != batch.NumRows() { + return fmt.Errorf( + "arrow protocol: Flight chunk plan covers %d of %d rows", + rowStart, + batch.NumRows(), + ) + } + return nil +} + +type chunkSearch struct { + bytesPerRow float64 + priorRows int64 +} + +type chunkMeasurement struct { + recordBytes int +} + +func (p *Protocol) planChunks(batch arrow.RecordBatch) ([]int64, error) { + average := float64(totalRecordBufferSize(batch)) / float64(batch.NumRows()) + if average < 1 { + average = 1 + } + search := chunkSearch{bytesPerRow: average} + chunks := make([]int64, 0, 1) + for rowStart := int64(0); rowStart < batch.NumRows(); { + rowCount, err := p.findChunkRows(batch, rowStart, &search) + if err != nil { + return nil, err + } + chunks = append(chunks, rowCount) + rowStart += rowCount + } + return chunks, nil +} + +func (p *Protocol) findChunkRows( + batch arrow.RecordBatch, + rowStart int64, + search *chunkSearch, +) (int64, error) { + remaining := batch.NumRows() - rowStart + if remaining <= 0 { + return 0, fmt.Errorf("arrow protocol: no rows remain for Flight chunk") + } + estimated := search.priorRows + if estimated <= 0 { + estimated = int64(float64(TargetFlightDataBytes-1024) / search.bytesPerRow) + } + if estimated < 1 { + estimated = 1 + } + if estimated > remaining { + estimated = remaining + } + + cache := make(map[int64]chunkMeasurement) + measure := func(rows int64) (chunkMeasurement, error) { + if measured, ok := cache[rows]; ok { + return measured, nil + } + measured, err := p.measureChunk(batch, rowStart, rows) + if err == nil { + cache[rows] = measured + } + return measured, err + } + fits := func(measured chunkMeasurement) bool { + return measured.recordBytes <= TargetFlightDataBytes + } + + measured, err := measure(estimated) + if err != nil { + return 0, err + } + var best int64 + var low, high int64 + if fits(measured) { + best = estimated + if best == remaining { + search.priorRows = best + return best, nil + } + // Probe outward from the local prior with exponentially growing steps: this + // brackets the failure without rescanning the whole batch per chunk. + step := int64(1) + for { + candidate := estimated + step + if candidate < estimated || candidate > remaining { + candidate = remaining + } + candidateMeasurement, err := measure(candidate) + if err != nil { + return 0, err + } + if fits(candidateMeasurement) { + best = candidate + if candidate == remaining { + low, high = 1, 0 + break + } + if step > math.MaxInt64/2 { + step = remaining - estimated + } else { + step *= 2 + } + continue + } + low, high = best+1, candidate-1 + break + } + } else { + if estimated == 1 { + return 0, fmt.Errorf( + "arrow protocol: one-row FlightData exceeds %d-byte target", + TargetFlightDataBytes, + ) + } + failed := estimated + step := int64(1) + for { + candidate := estimated - step + if candidate < 1 || candidate > estimated { + candidate = 1 + } + candidateMeasurement, err := measure(candidate) + if err != nil { + return 0, err + } + if fits(candidateMeasurement) { + best = candidate + low, high = candidate+1, failed-1 + break + } + if candidate == 1 { + return 0, fmt.Errorf( + "arrow protocol: one-row FlightData exceeds %d-byte target", + TargetFlightDataBytes, + ) + } + failed = candidate + if step > math.MaxInt64/2 { + step = estimated - 1 + } else { + step *= 2 + } + } + } + + // Actual compressed protobuf size is authoritative; search only the bracket. + for low <= high { + mid := low + (high-low)/2 + candidate, err := measure(mid) + if err != nil { + return 0, err + } + if fits(candidate) { + best = mid + low = mid + 1 + } else { + high = mid - 1 + } + } + accepted, err := measure(best) + if err != nil { + return 0, err + } + if accepted.recordBytes > 0 { + search.bytesPerRow = float64(accepted.recordBytes) / float64(best) + if search.bytesPerRow < 1 { + search.bytesPerRow = 1 + } + } + search.priorRows = best + return best, nil +} + +func (p *Protocol) measureChunk( + batch arrow.RecordBatch, + rowStart, rowCount int64, +) (chunkMeasurement, error) { + if rowCount <= 0 || rowStart < 0 || rowStart > batch.NumRows()-rowCount { + return chunkMeasurement{}, fmt.Errorf("arrow protocol: invalid Flight chunk range") + } + if p.chunkProbe != nil { + p.chunkProbe(rowStart, rowCount) + } + sizer := &flightDataSizer{skipSchema: true} + options, err := p.flightWriterOptions() + if err != nil { + return chunkMeasurement{}, err + } + writer := flight.NewRecordWriter(sizer, options...) + chunk := batch.NewSlice(rowStart, rowStart+rowCount) + err = writer.Write(chunk) + chunk.Release() + if err != nil { + _ = writer.Close() + return chunkMeasurement{}, fmt.Errorf("arrow protocol: size Flight row chunk: %w", err) + } + if err := writer.Close(); err != nil { + return chunkMeasurement{}, fmt.Errorf("arrow protocol: close Flight chunk sizer: %w", err) + } + if sizer.recordFrames != 1 { + return chunkMeasurement{}, fmt.Errorf( + "arrow protocol: Flight row chunk produced %d record frames, want 1", + sizer.recordFrames, + ) + } + // Slicing rows never shrinks a dictionary, so an oversized one is a property + // of the batch. Report it here rather than letting the row search bottom out + // at one row and blame the row. + if sizer.dictionaryBytes > TargetFlightDataBytes { + return chunkMeasurement{}, fmt.Errorf( + "arrow protocol: dictionary FlightData is %d bytes, exceeds %d-byte target", + sizer.dictionaryBytes, + TargetFlightDataBytes, + ) + } + return chunkMeasurement{recordBytes: sizer.recordBytes}, nil +} + +type flightDataSizer struct { + skipSchema bool + recordBytes int + recordFrames int + dictionaryBytes int +} + +func (s *flightDataSizer) Send(frame *flight.FlightData) error { + messageType, err := flightMessageType(frame) + if err != nil { + return err + } + if messageType == ipc.MessageSchema { + if !s.skipSchema { + return fmt.Errorf("arrow protocol: chunk sizer received duplicate schema") + } + s.skipSchema = false + return nil + } + frame.AppMetadata = []byte(`{"offset_id":9223372036854775807}`) + size := proto.Size(frame) + // The record frame alone decides the row count. A dictionary rides in its own + // frame, and the live writer emits it once per connection rather than once + // per chunk, so folding it in here would shrink every chunk for nothing. + if messageType == ipc.MessageRecordBatch { + s.recordFrames++ + s.recordBytes = size + } else if size > s.dictionaryBytes { + s.dictionaryBytes = size + } + return nil +} + +func (p *Protocol) flightWriterOptions() ([]ipc.Option, error) { + options, err := p.ipcOptions() + if err != nil { + return nil, err + } + return append( + options, + ipc.WithSchema(p.schema), + ipc.WithDictionaryDeltas(true), + ), nil +} + +func (p *Protocol) ipcOptions() ([]ipc.Option, error) { + compression, err := compressionOption(p.compression) + if err != nil { + return nil, err + } + options := []ipc.Option{ipc.WithAllocator(p.allocator)} + if compression != nil { + options = append(options, compression) + } + return options, nil +} + +func compressionOption(compression Compression) (ipc.Option, error) { + switch compression { + case CompressionNone: + return nil, nil + case CompressionLZ4: + return ipc.WithLZ4(), nil + case CompressionZstd: + return ipc.WithZstd(), nil + default: + return nil, fmt.Errorf( + "arrow protocol: unsupported IPC compression %d", + compression, + ) + } +} + +type flightDataCollector struct { + frames []*flight.FlightData +} + +func (c *flightDataCollector) Send(frame *flight.FlightData) error { + if frame == nil { + return fmt.Errorf("arrow protocol: Flight encoder emitted a nil frame") + } + c.frames = append(c.frames, proto.Clone(frame).(*flight.FlightData)) + return nil +} diff --git a/purego/internal/arrowproto/encoder_test.go b/purego/internal/arrowproto/encoder_test.go new file mode 100644 index 00000000..56de6f12 --- /dev/null +++ b/purego/internal/arrowproto/encoder_test.go @@ -0,0 +1,1392 @@ +package arrowproto + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "io" + "math" + "strconv" + "strings" + "testing" + "time" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/flight" + "github.com/apache/arrow-go/v18/arrow/ipc" + "github.com/apache/arrow-go/v18/arrow/memory" + flatbuffers "github.com/google/flatbuffers/go" + "google.golang.org/protobuf/proto" + + "github.com/databricks/zerobus-sdk/purego/internal/stream" + "github.com/databricks/zerobus-sdk/purego/internal/transport" +) + +func idSchema(metadata *arrow.Metadata) *arrow.Schema { + return arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false, + }}, metadata) +} + +func idBatch( + t *testing.T, + allocator memory.Allocator, + schema *arrow.Schema, + values []int32, +) arrow.RecordBatch { + t.Helper() + builder := array.NewInt32Builder(allocator) + builder.AppendValues(values, nil) + column := builder.NewArray() + builder.Release() + record := array.NewRecordBatch(schema, []arrow.Array{column}, int64(len(values))) + column.Release() + return record +} + +func binaryBatch( + t *testing.T, + allocator memory.Allocator, + rows int, + valueBytes int, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "value", Type: arrow.BinaryTypes.String, Nullable: false}, + }, nil) + idBuilder := array.NewInt32Builder(allocator) + valueBuilder := array.NewStringBuilder(allocator) + value := strings.Repeat("x", valueBytes) + for i := range rows { + idBuilder.Append(int32(i)) + valueBuilder.Append(value) + } + ids := idBuilder.NewArray() + values := valueBuilder.NewArray() + idBuilder.Release() + valueBuilder.Release() + record := array.NewRecordBatch( + schema, + []arrow.Array{ids, values}, + int64(rows), + ) + ids.Release() + values.Release() + return schema, record +} + +func dictionaryBatch( + t *testing.T, + allocator memory.Allocator, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + dictionaryType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.BinaryTypes.String, + } + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "category", Type: dictionaryType, Nullable: false}, + }, nil) + + idBuilder := array.NewInt32Builder(allocator) + idBuilder.AppendValues([]int32{1, 2, 3, 4}, nil) + ids := idBuilder.NewArray() + idBuilder.Release() + + dictionaryBuilder := array.NewDictionaryBuilder(allocator, dictionaryType) + stringBuilder := dictionaryBuilder.(*array.BinaryDictionaryBuilder) + for _, value := range []string{"alpha", "beta", "alpha", "gamma"} { + if err := stringBuilder.AppendString(value); err != nil { + t.Fatalf("append dictionary value: %v", err) + } + } + categories := dictionaryBuilder.NewDictionaryArray() + dictionaryBuilder.Release() + + record := array.NewRecordBatch( + schema, + []arrow.Array{ids, categories}, + 4, + ) + ids.Release() + categories.Release() + return schema, record +} + +func chunkedDictionaryBatch( + t *testing.T, + allocator memory.Allocator, + rows int, + valueBytes int, +) (*arrow.Schema, arrow.RecordBatch) { + t.Helper() + dictionaryType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.BinaryTypes.String, + } + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + {Name: "category", Type: dictionaryType, Nullable: false}, + {Name: "value", Type: arrow.BinaryTypes.String, Nullable: false}, + }, nil) + idBuilder := array.NewInt32Builder(allocator) + dictionaryBuilder := array.NewDictionaryBuilder(allocator, dictionaryType) + categoryBuilder := dictionaryBuilder.(*array.BinaryDictionaryBuilder) + valueBuilder := array.NewStringBuilder(allocator) + value := strings.Repeat("x", valueBytes) + categories := []string{"alpha", "beta", "gamma"} + for row := range rows { + idBuilder.Append(int32(row)) + if err := categoryBuilder.AppendString(categories[row%len(categories)]); err != nil { + t.Fatalf("append dictionary value: %v", err) + } + valueBuilder.Append(value) + } + ids := idBuilder.NewArray() + categoryArray := dictionaryBuilder.NewDictionaryArray() + values := valueBuilder.NewArray() + idBuilder.Release() + dictionaryBuilder.Release() + valueBuilder.Release() + record := array.NewRecordBatch( + schema, + []arrow.Array{ids, categoryArray, values}, + int64(rows), + ) + ids.Release() + categoryArray.Release() + values.Release() + return schema, record +} + +func serializeRecords( + t *testing.T, + schema *arrow.Schema, + records ...arrow.RecordBatch, +) []byte { + return serializeRecordsWithOptions(t, schema, nil, records...) +} + +func serializeRecordsWithOptions( + t *testing.T, + schema *arrow.Schema, + options []ipc.Option, + records ...arrow.RecordBatch, +) []byte { + t.Helper() + var output bytes.Buffer + writerOptions := append([]ipc.Option{ipc.WithSchema(schema)}, options...) + writer := ipc.NewWriter(&output, writerOptions...) + for _, record := range records { + if err := writer.Write(record); err != nil { + t.Fatalf("write IPC record: %v", err) + } + } + if err := writer.Close(); err != nil { + t.Fatalf("close IPC writer: %v", err) + } + return bytes.Clone(output.Bytes()) +} + +func readIDs(t *testing.T, data []byte) []int32 { + t.Helper() + reader, err := ipc.NewReader(bytes.NewReader(data)) + if err != nil { + t.Fatalf("new IPC reader: %v", err) + } + defer reader.Release() + if !reader.Next() { + t.Fatalf("read IPC record: %v", reader.Err()) + } + column, ok := reader.RecordBatch().Column(0).(*array.Int32) + if !ok { + t.Fatalf("first column type = %T, want *array.Int32", reader.RecordBatch().Column(0)) + } + values := append([]int32(nil), column.Int32Values()...) + if reader.Next() { + t.Fatal("decoded IPC contains more than one record") + } + if err := reader.Err(); err != nil { + t.Fatalf("finish IPC reader: %v", err) + } + return values +} + +type flightSliceReader struct { + frames []*flight.FlightData + next int +} + +func (r *flightSliceReader) Recv() (*flight.FlightData, error) { + if r.next == len(r.frames) { + return nil, io.EOF + } + frame := proto.Clone(r.frames[r.next]).(*flight.FlightData) + r.next++ + return frame, nil +} + +func readFlightIDs( + t *testing.T, + schema *flight.FlightData, + frames []*flight.FlightData, +) []int32 { + t.Helper() + all := make([]*flight.FlightData, 0, len(frames)+1) + all = append(all, schema) + all = append(all, frames...) + reader, err := flight.NewRecordReader(&flightSliceReader{frames: all}) + if err != nil { + t.Fatalf("new Flight record reader: %v", err) + } + defer reader.Release() + + var values []int32 + for reader.Next() { + column := reader.RecordBatch().Column(0).(*array.Int32) + values = append(values, column.Int32Values()...) + } + if err := reader.Err(); err != nil { + t.Fatalf("read Flight records: %v", err) + } + return values +} + +func TestTypedRecordBatchRoundTripOwnsIPC(t *testing.T) { + allocator := memory.NewCheckedAllocator(memory.DefaultAllocator) + schema := idSchema(nil) + protocol, err := New(schema, Options{Allocator: allocator}) + if err != nil { + t.Fatalf("New: %v", err) + } + if allocator.CurrentAlloc() != 0 { + t.Fatalf("schema construction retained %d allocator bytes", allocator.CurrentAlloc()) + } + + record := idBatch(t, allocator, schema, []int32{10, 20, 30}) + estimate, err := protocol.EstimateRecordBatchRetainedSize(record) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + record.Release() + + if got := payload.UnitCount(); got != 3 { + t.Fatalf("UnitCount = %d, want 3", got) + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{10, 20, 30}) { + t.Fatalf("round-trip ids = %v", got) + } + if payload.RetainedSize() < int64(len(payload.ipcBytes)) { + t.Fatalf( + "RetainedSize = %d, smaller than IPC bytes %d", + payload.RetainedSize(), + len(payload.ipcBytes), + ) + } + if payload.RetainedSize() > estimate { + t.Fatalf( + "actual retained size %d exceeds admission estimate %d", + payload.RetainedSize(), + estimate, + ) + } + + first := protocol.Decode(payload)[0] + first[0] ^= 0xff + if bytes.Equal(first, protocol.Decode(payload)[0]) { + t.Fatal("Decode returned bytes aliasing the retained payload") + } + allocator.AssertSize(t, 0) +} + +func TestCanonicalPayloadTakesSerializerOwnership(t *testing.T) { + protocol, err := New(idSchema(nil), Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + canonical := []byte{1, 2, 3, 4} + payload, err := protocol.payloadFromCanonicalIPC(canonical, 1, []int64{1}) + if err != nil { + t.Fatalf("payloadFromCanonicalIPC: %v", err) + } + if &payload.ipcBytes[0] != &canonical[0] { + t.Fatal("canonical IPC was copied instead of ownership being transferred") + } +} + +func TestIPCValidationAndCanonicalOwnership(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1, 2}) + defer record.Release() + + input := serializeRecords(t, schema, record) + payload, err := protocol.EncodeIPC(input) + if err != nil { + t.Fatalf("EncodeIPC: %v", err) + } + for i := range input { + input[i] = 0 + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{1, 2}) { + t.Fatalf("canonical ids = %v", got) + } + + if _, err := protocol.EncodeIPC([]byte("not IPC")); err == nil { + t.Fatal("invalid IPC accepted") + } +} + +func TestIPCRejectsNoEmptyAndMultipleBatches(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + empty := idBatch(t, memory.DefaultAllocator, schema, nil) + defer empty.Release() + one := idBatch(t, memory.DefaultAllocator, schema, []int32{1}) + defer one.Release() + + tests := []struct { + name string + data []byte + }{ + {name: "no batch", data: serializeRecords(t, schema)}, + {name: "empty batch", data: serializeRecords(t, schema, empty)}, + {name: "multiple batches", data: serializeRecords(t, schema, one, one)}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if _, err := protocol.EncodeIPC(test.data); err == nil { + t.Fatalf("EncodeIPC accepted %s", test.name) + } + }) + } +} + +func TestIPCDecodersRejectTrailingBytesAndConcatenatedStreams(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1}) + defer record.Release() + + batchStream := serializeRecords(t, schema, record) + schemaStream, err := EncodeSchemaIPC(schema) + if err != nil { + t.Fatalf("EncodeSchemaIPC: %v", err) + } + for _, test := range []struct { + name string + input []byte + decode func([]byte) error + }{ + { + name: "batch trailing bytes", + input: append(bytes.Clone(batchStream), []byte("trailing")...), + decode: func(data []byte) error { + _, err := protocol.EncodeIPC(data) + return err + }, + }, + { + name: "batch concatenated stream", + input: append(bytes.Clone(batchStream), batchStream...), + decode: func(data []byte) error { + _, err := protocol.EncodeIPC(data) + return err + }, + }, + { + name: "schema trailing bytes", + input: append(bytes.Clone(schemaStream), []byte("trailing")...), + decode: func(data []byte) error { + _, err := DecodeSchemaIPC(data) + return err + }, + }, + { + name: "schema concatenated stream", + input: append(bytes.Clone(schemaStream), schemaStream...), + decode: func(data []byte) error { + _, err := DecodeSchemaIPC(data) + return err + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + if err := test.decode(test.input); err == nil { + t.Fatal("decoder accepted trailing IPC data") + } + }) + } +} + +func TestExactSchemaMatchIncludesMetadata(t *testing.T) { + expectedMetadata := arrow.NewMetadata([]string{"owner"}, []string{"expected"}) + actualMetadata := arrow.NewMetadata([]string{"owner"}, []string{"different"}) + expected := idSchema(&expectedMetadata) + actual := idSchema(&actualMetadata) + protocol, err := New(expected, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, actual, []int32{1}) + defer record.Release() + + if _, err := protocol.EncodeRecordBatch(record); err == nil { + t.Fatal("typed RecordBatch with different schema metadata accepted") + } + if _, err := protocol.EncodeIPC(serializeRecords(t, actual, record)); err == nil { + t.Fatal("IPC RecordBatch with different schema metadata accepted") + } +} + +func TestProtocolOwnsPointerBackedSchemaTypes(t *testing.T) { + dictionaryType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.BinaryTypes.String, + } + schema := arrow.NewSchema([]arrow.Field{{ + Name: "category", Type: dictionaryType, Nullable: false, + }}, nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + + dictionaryType.ValueType = arrow.BinaryTypes.Binary + owned := protocol.schema.Field(0).Type.(*arrow.DictionaryType) + if owned.ValueType.ID() != arrow.STRING { + t.Fatalf("owned dictionary value type = %v, want string", owned.ValueType) + } +} + +func TestSliceReserializesUnacknowledgedSuffix(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1, 2, 3, 4, 5}) + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + suffix, err := protocol.Slice(payload, 2) + if err != nil { + t.Fatalf("Slice: %v", err) + } + if suffix.UnitCount() != 3 { + t.Fatalf("suffix units = %d, want 3", suffix.UnitCount()) + } + if got := readIDs(t, suffix.IPCBytes()); !equalInt32(got, []int32{3, 4, 5}) { + t.Fatalf("suffix ids = %v", got) + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{1, 2, 3, 4, 5}) { + t.Fatalf("original payload changed after Slice: %v", got) + } + if _, err := protocol.Slice(payload, 5); err == nil { + t.Fatal("Slice accepted a fully acknowledged prefix") + } + // The core never slices an unacknowledged payload, so a zero prefix is a bug + // to surface, not a payload to pass through. + if _, err := protocol.Slice(payload, 0); err == nil { + t.Fatal("Slice accepted an unacknowledged prefix") + } +} + +func TestIPCCompressionOptionsRoundTrip(t *testing.T) { + schema, record := binaryBatch(t, memory.DefaultAllocator, 2_000, 200) + defer record.Release() + + sizes := make(map[Compression]int) + for _, compression := range []Compression{ + CompressionNone, + CompressionLZ4, + CompressionZstd, + } { + protocol, err := New(schema, Options{Compression: compression}) + if err != nil { + t.Fatalf("New compression %d: %v", compression, err) + } + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch compression %d: %v", compression, err) + } + sizes[compression] = len(payload.ipcBytes) + reader, err := ipc.NewReader(bytes.NewReader(payload.IPCBytes())) + if err != nil { + t.Fatalf("read compression %d: %v", compression, err) + } + if !reader.Next() || reader.RecordBatch().NumRows() != 2_000 { + t.Fatalf("compression %d did not round-trip: %v", compression, reader.Err()) + } + reader.Release() + } + if sizes[CompressionLZ4] >= sizes[CompressionNone] { + t.Errorf("LZ4 size %d >= uncompressed %d", sizes[CompressionLZ4], sizes[CompressionNone]) + } + if sizes[CompressionZstd] >= sizes[CompressionNone] { + t.Errorf("Zstd size %d >= uncompressed %d", sizes[CompressionZstd], sizes[CompressionNone]) + } + if _, err := New(schema, Options{Compression: Compression(99)}); err == nil { + t.Fatal("unsupported compression accepted") + } +} + +func TestCompressedIPCAdmissionUsesDeclaredUncompressedSizes(t *testing.T) { + const ( + rows = 4_096 + valueBytes = 2_048 + memoryCap = 4 * 1024 * 1024 + ) + schema, record := binaryBatch( + t, + memory.DefaultAllocator, + rows, + valueBytes, + ) + defer record.Release() + + for _, test := range []struct { + name string + option ipc.Option + }{ + {name: "LZ4", option: ipc.WithLZ4()}, + {name: "Zstd", option: ipc.WithZstd()}, + } { + t.Run(test.name, func(t *testing.T) { + input := serializeRecordsWithOptions( + t, + schema, + []ipc.Option{test.option}, + record, + ) + if len(input) >= rows*valueBytes/4 { + t.Fatalf( + "compressed IPC size = %d, input is not highly compressible", + len(input), + ) + } + + allocator := memory.NewCheckedAllocator(memory.DefaultAllocator) + protocol, err := New(schema, Options{Allocator: allocator}) + if err != nil { + t.Fatalf("New: %v", err) + } + estimate, err := protocol.EstimateIPCRetainedSize(input) + if err != nil { + t.Fatalf("EstimateIPCRetainedSize: %v", err) + } + if estimate <= memoryCap { + t.Fatalf( + "compressed IPC estimate = %d, want above %d-byte limit", + estimate, + memoryCap, + ) + } + + cfg := stream.DefaultConfig() + cfg.MaxBufferedPayloadBytes = memoryCap + cfg.Recovery = stream.RecoveryDisabled + cfg.RecoveryTimeout = time.Hour + open := stream.OpenFunc[*Payload, *flight.PutResult](func( + ctx context.Context, + _ stream.StreamParams, + ) (stream.WireStream[*Payload, *flight.PutResult], error) { + <-ctx.Done() + return nil, ctx.Err() + }) + core, err := stream.NewCoreStreamWithHooks( + context.Background(), + stream.StreamParams{}, + cfg, + open, + protocol.EncoderHooks(), + unreachableAckModel(t), + nil, + ) + if err != nil { + t.Fatalf("NewCoreStreamWithHooks: %v", err) + } + t.Cleanup(func() { _ = core.Terminate() }) + + builderCalled := false + offset, err := core.EnqueuePayloadBuilder( + context.Background(), + estimate, + func() (*Payload, uint64, int64, error) { + builderCalled = true + payload, encodeErr := protocol.EncodeIPC(input) + if encodeErr != nil { + return nil, 0, 0, encodeErr + } + return payload, payload.UnitCount(), payload.RetainedSize(), nil + }, + ) + if offset != -1 || !errors.Is(err, stream.ErrPayloadTooLarge) { + t.Fatalf( + "compressed admission = (%d,%v), want (-1, ErrPayloadTooLarge)", + offset, + err, + ) + } + if builderCalled { + t.Fatal("Arrow decoder ran before compressed IPC admission") + } + allocator.AssertSize(t, 0) + }) + } +} + +func TestCompressedIPCDeclaredSizeOverflowIsRejected(t *testing.T) { + builder := flatbuffers.NewBuilder(128) + + builder.StartObject(2) + compression := builder.EndObject() + + builder.StartVector(16, 2, 8) + for index := 1; index >= 0; index-- { + builder.Prep(8, 16) + builder.PrependInt64(8) + builder.PrependInt64(int64(index * 8)) + } + buffers := builder.EndVector(2) + + builder.StartObject(5) + builder.PrependUOffsetTSlot(2, buffers, 0) + builder.PrependUOffsetTSlot(3, compression, 0) + recordBatch := builder.EndObject() + builder.Finish(recordBatch) + + body := make([]byte, 16) + binary.LittleEndian.PutUint64(body[:8], math.MaxInt64) + binary.LittleEndian.PutUint64(body[8:], 1) + if _, err := ipcCompressedExpansion( + ipcRootTable(builder.FinishedBytes()), + body, + ); err == nil || !strings.Contains(err.Error(), "overflow") { + t.Fatalf("ipcCompressedExpansion overflow error = %v", err) + } +} + +func TestTypedAdmissionIncludesRecordBatchMetadata(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + plain := idBatch(t, memory.DefaultAllocator, schema, []int32{1}) + defer plain.Release() + + metadataValue := strings.Repeat("m", 512*1024) + metadata := arrow.NewMetadata([]string{"large"}, []string{metadataValue}) + withMetadata := array.NewRecordBatchWithMetadata( + schema, + plain.Columns(), + plain.NumRows(), + metadata, + ) + defer withMetadata.Release() + + plainEstimate, err := protocol.EstimateRecordBatchRetainedSize(plain) + if err != nil { + t.Fatalf("plain estimate: %v", err) + } + metadataEstimate, err := protocol.EstimateRecordBatchRetainedSize(withMetadata) + if err != nil { + t.Fatalf("metadata estimate: %v", err) + } + if increase := metadataEstimate - plainEstimate; increase < int64(len(metadataValue))*2 { + t.Fatalf( + "metadata increased estimate by %d bytes, want at least %d", + increase, + len(metadataValue)*2, + ) + } + payload, err := protocol.EncodeRecordBatch(withMetadata) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if payload.RetainedSize() > metadataEstimate { + t.Fatalf( + "metadata payload retained size %d exceeds estimate %d", + payload.RetainedSize(), + metadataEstimate, + ) + } + decoded, err := protocol.DecodeIPCRecordBatch(payload.IPCBytes()) + if err != nil { + t.Fatalf("DecodeIPCRecordBatch: %v", err) + } + defer decoded.Release() + decodedMetadata := decoded.(arrow.RecordBatchWithMetadata).Metadata() + if !decodedMetadata.Equal(metadata) { + t.Fatal("RecordBatch custom metadata did not round-trip") + } +} + +func TestFlightChunkingUsesActualProtoSizeAndSequentialMetadata(t *testing.T) { + const rows = 10_500 + schema, record := binaryBatch(t, memory.DefaultAllocator, rows, 256) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + frames, next, err := protocol.EncodeFlightData(payload, 0) + if err != nil { + t.Fatalf("EncodeFlightData: %v", err) + } + if len(frames) < 2 { + t.Fatalf("Flight frames = %d, want at least 2", len(frames)) + } + if next != int64(len(frames)) { + t.Fatalf("next offset = %d, want %d", next, len(frames)) + } + for offset, frame := range frames { + if size := proto.Size(frame); size > TargetFlightDataBytes { + t.Errorf("frame %d size = %d, exceeds %d", offset, size, TargetFlightDataBytes) + } + metadata, err := transport.ParseFlightBatchMetadata(frame.GetAppMetadata()) + if err != nil { + t.Fatalf("frame %d metadata: %v", offset, err) + } + if metadata.OffsetID != int64(offset) { + t.Errorf("frame %d offset = %d", offset, metadata.OffsetID) + } + } + + ids := readFlightIDs(t, protocol.SchemaFlightData(), frames) + if len(ids) != rows { + t.Fatalf("decoded Flight rows = %d, want %d", len(ids), rows) + } + for i, id := range ids { + if id != int32(i) { + t.Fatalf("decoded Flight id[%d] = %d", i, id) + } + } +} + +func TestFlightChunkPlanningBoundsProbeWorkAcrossManyChunks(t *testing.T) { + const rows = 40_000 + schema, record := binaryBatch(t, memory.DefaultAllocator, rows, 512) + protocol, err := New(schema, Options{}) + if err != nil { + record.Release() + t.Fatalf("New: %v", err) + } + probes := 0 + protocol.chunkProbe = func(int64, int64) { probes++ } + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if chunks := len(payload.chunkRows); chunks < 8 { + t.Fatalf("planned chunks = %d, want at least 8", chunks) + } else if limit := chunks*6 + 20; probes > limit { + t.Fatalf( + "chunk probes = %d for %d chunks, exceeds local-search bound %d", + probes, + chunks, + limit, + ) + } +} + +func TestFlightRejectsOneRowOverTarget(t *testing.T) { + schema, record := binaryBatch( + t, + memory.DefaultAllocator, + 1, + TargetFlightDataBytes+1024, + ) + protocol, err := New(schema, Options{}) + if err != nil { + record.Release() + t.Fatalf("New: %v", err) + } + _, err = protocol.EncodeRecordBatch(record) + record.Release() + if err == nil || + !strings.Contains(err.Error(), "one-row FlightData exceeds") { + t.Fatalf("EncodeRecordBatch oversize error = %v", err) + } +} + +func TestDictionaryPayloadsAreSelfContainedOnIPCAndFlight(t *testing.T) { + allocator := memory.NewCheckedAllocator(memory.DefaultAllocator) + schema, record := dictionaryBatch(t, allocator) + protocol, err := New(schema, Options{Allocator: allocator}) + if err != nil { + t.Fatalf("New: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + // A dictionary payload encodes to a dictionary frame plus a record frame, each + // taking the next sequential offset. + typedFrames, nextOffset, err := protocol.EncodeFlightData(payload, 0) + if err != nil { + t.Fatalf("EncodeFlightData dictionary: %v", err) + } + if len(typedFrames) < 2 { + t.Fatalf( + "dictionary payload frames = %d, want dictionary plus record", + len(typedFrames), + ) + } + if nextOffset != int64(len(typedFrames)) { + t.Fatalf("next offset = %d, want %d", nextOffset, len(typedFrames)) + } + ids := readFlightIDs(t, protocol.SchemaFlightData(), typedFrames) + if !equalInt32(ids, []int32{1, 2, 3, 4}) { + t.Fatalf("dictionary Flight ids = %v", ids) + } + if got := readIDs(t, payload.IPCBytes()); !equalInt32(got, []int32{1, 2, 3, 4}) { + t.Fatalf("dictionary IPC ids = %v", got) + } + + // The IPC input path reserializes the dictionary too, so its Flight output + // survives the caller's IPC bytes and RecordBatch going away. + sourceSchema, sourceRecord := dictionaryBatch(t, allocator) + input := serializeRecords(t, sourceSchema, sourceRecord) + sourceRecord.Release() + fromIPC, err := protocol.EncodeIPC(input) + if err != nil { + t.Fatalf("EncodeIPC dictionary: %v", err) + } + clear(input) + ipcFrames, _, err := protocol.EncodeFlightData(fromIPC, 0) + if err != nil { + t.Fatalf("EncodeFlightData dictionary IPC: %v", err) + } + if got := readFlightIDs(t, protocol.SchemaFlightData(), ipcFrames); !equalInt32(got, []int32{1, 2, 3, 4}) { + t.Fatalf("dictionary IPC Flight ids = %v", got) + } + + suffix, err := protocol.Slice(fromIPC, 2) + if err != nil { + t.Fatalf("Slice dictionary payload: %v", err) + } + suffixFrames, _, err := protocol.EncodeFlightData(suffix, 0) + if err != nil { + t.Fatalf("EncodeFlightData dictionary suffix: %v", err) + } + if got := readFlightIDs(t, protocol.SchemaFlightData(), suffixFrames); !equalInt32(got, []int32{3, 4}) { + t.Fatalf("dictionary suffix ids = %v", got) + } + allocator.AssertSize(t, 0) +} + +func TestChunkedFlightEmitsUnchangedDictionaryOnce(t *testing.T) { + const rows = 6_000 + schema, record := chunkedDictionaryBatch( + t, + memory.DefaultAllocator, + rows, + 512, + ) + protocol, err := New(schema, Options{}) + if err != nil { + record.Release() + t.Fatalf("New: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + allFrames, nextOffset, err := protocol.EncodeFlightData(payload, 0) + if err != nil { + t.Fatalf("EncodeFlightData: %v", err) + } + if nextOffset != int64(len(allFrames)) { + t.Fatalf("next offset = %d, want %d", nextOffset, len(allFrames)) + } + + var dictionaryFrames, recordFrames int + for _, frame := range allFrames { + messageType, err := flightMessageType(frame) + if err != nil { + t.Fatalf("flightMessageType: %v", err) + } + switch messageType { + case ipc.MessageDictionaryBatch: + dictionaryFrames++ + case ipc.MessageRecordBatch: + recordFrames++ + } + } + if dictionaryFrames != 1 || recordFrames < 2 { + t.Fatalf( + "chunked frames dictionaries=%d records=%d, want 1 dictionary and multiple records", + dictionaryFrames, + recordFrames, + ) + } + if ids := readFlightIDs(t, protocol.SchemaFlightData(), allFrames); len(ids) != rows { + t.Fatalf("decoded chunked dictionary rows = %d, want %d", len(ids), rows) + } +} + +// A chunked payload spans several record frames, so a send failing partway has +// still submitted a row prefix. The emitter reports that prefix so the core can +// bound an acknowledgment to it rather than to the whole payload. +func TestFlightEmitterReportsSubmittedRowPrefixOnPartialSend(t *testing.T) { + const rows = 6_000 + schema, record := chunkedDictionaryBatch(t, memory.DefaultAllocator, rows, 512) + protocol, err := New(schema, Options{}) + if err != nil { + record.Release() + t.Fatalf("New: %v", err) + } + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + emit := func(failAtRecordFrame int) (stream.SubmissionReceipt, error) { + var receipt stream.SubmissionReceipt + recordFrames := 0 + sendErr := errors.New("send failed") + emitter := &flightFrameEmitter{ + skipSchema: true, + receipt: &receipt, + } + emitter.send = func(frame *flight.FlightData) error { + messageType, err := flightMessageType(frame) + if err != nil { + return err + } + if messageType != ipc.MessageRecordBatch { + return nil + } + recordFrames++ + if failAtRecordFrame > 0 && recordFrames == failAtRecordFrame { + return sendErr + } + return nil + } + options, err := protocol.flightWriterOptions() + if err != nil { + return receipt, err + } + writer := flight.NewRecordWriter(emitter, options...) + emitErr := protocol.emitPayload(payload, writer, emitter) + _ = writer.Close() + return receipt, emitErr + } + + complete, err := emit(0) + if err != nil { + t.Fatalf("complete emit: %v", err) + } + if complete.SubmittedUnits != rows { + t.Fatalf("complete submitted rows = %d, want %d", complete.SubmittedUnits, rows) + } + + partial, err := emit(2) + if err == nil { + t.Fatal("emit with a failing second record frame returned no error") + } + if partial.SubmittedUnits == 0 || partial.SubmittedUnits >= rows { + t.Fatalf( + "partial submitted rows = %d, want a prefix in (0,%d)", + partial.SubmittedUnits, + rows, + ) + } +} + +func TestSchemaSeparateAndOffsetsCanRestartPerConnection(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + schemaFrame := protocol.SchemaFlightData() + if len(schemaFrame.GetDataHeader()) == 0 || + len(schemaFrame.GetDataBody()) != 0 || + len(schemaFrame.GetAppMetadata()) != 0 { + t.Fatalf("schema frame = %+v", schemaFrame) + } + + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1}) + payload, err := protocol.EncodeRecordBatch(record) + record.Release() + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + for connection := 0; connection < 2; connection++ { + frames, _, err := protocol.EncodeFlightData(payload, 0) + if err != nil { + t.Fatalf("connection %d EncodeFlightData: %v", connection, err) + } + metadata, err := transport.ParseFlightBatchMetadata(frames[0].GetAppMetadata()) + if err != nil { + t.Fatalf("connection %d metadata: %v", connection, err) + } + if metadata.OffsetID != 0 { + t.Fatalf("connection %d first offset = %d", connection, metadata.OffsetID) + } + } +} + +// TestIngestChargesMaterializedSizeNotCompressedInput covers the []byte hook +// path, whose admission estimate can only be derived from the input length. +// Canonicalizing decompresses, so charging the input would admit a payload that +// retains orders of magnitude more than the byte limit allows. +func TestIngestChargesMaterializedSizeNotCompressedInput(t *testing.T) { + const ( + rows = 4_096 + valueBytes = 2_048 + memoryCap = 4 * 1024 * 1024 + ) + schema, record := binaryBatch(t, memory.DefaultAllocator, rows, valueBytes) + defer record.Release() + input := serializeRecordsWithOptions( + t, + schema, + []ipc.Option{ipc.WithZstd()}, + record, + ) + + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + hooks := protocol.EncoderHooks() + if estimate := hooks.RetainedSize(len(input), 1); estimate > memoryCap { + t.Fatalf( + "compressed input estimate = %d, want under the %d-byte cap", + estimate, + memoryCap, + ) + } + payload, err := protocol.EncodeIPC(input) + if err != nil { + t.Fatalf("EncodeIPC: %v", err) + } + if payload.RetainedSize() <= memoryCap { + t.Fatalf( + "materialized payload = %d bytes, want above the %d-byte cap", + payload.RetainedSize(), + memoryCap, + ) + } + if hooks.ActualRetainedSize == nil { + t.Fatal("Arrow hooks report no actual retained size, so the estimate stands") + } + if got := hooks.ActualRetainedSize(payload); got != payload.RetainedSize() { + t.Fatalf("ActualRetainedSize = %d, want %d", got, payload.RetainedSize()) + } + + cfg := stream.DefaultConfig() + cfg.MaxBufferedPayloadBytes = memoryCap + cfg.Recovery = stream.RecoveryDisabled + cfg.RecoveryTimeout = time.Hour + open := stream.OpenFunc[*Payload, *flight.PutResult](func( + ctx context.Context, + _ stream.StreamParams, + ) (stream.WireStream[*Payload, *flight.PutResult], error) { + <-ctx.Done() + return nil, ctx.Err() + }) + core, err := stream.NewCoreStreamWithHooks( + context.Background(), + stream.StreamParams{}, + cfg, + open, + hooks, + unreachableAckModel(t), + nil, + ) + if err != nil { + t.Fatalf("NewCoreStreamWithHooks: %v", err) + } + t.Cleanup(func() { _ = core.Terminate() }) + + offset, err := core.Ingest(context.Background(), input) + if offset != -1 || !errors.Is(err, stream.ErrPayloadTooLarge) { + t.Fatalf( + "Ingest of compressed payload = (%d,%v), want (-1, ErrPayloadTooLarge)", + offset, + err, + ) + } +} + +// TestTypedAdmissionChargesSliceNotParentBuffers covers a small slice of a large +// batch. A slice shares the parent's buffers, so charging a buffer's own length +// would reject a payload that is a few kilobytes on the wire. +func TestTypedAdmissionChargesSliceNotParentBuffers(t *testing.T) { + const ( + rows = 20_000 + valueBytes = 512 + sliceRows = 10 + ) + schema, record := binaryBatch(t, memory.DefaultAllocator, rows, valueBytes) + defer record.Release() + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + + slice := record.NewSlice(rows-sliceRows, rows) + defer slice.Release() + estimate, err := protocol.EstimateRecordBatchRetainedSize(slice) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(slice) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if estimate < payload.RetainedSize() { + t.Fatalf( + "slice estimate %d under-reserves its %d-byte payload", + estimate, + payload.RetainedSize(), + ) + } + if estimate > 64*payload.RetainedSize() { + t.Fatalf( + "slice estimate %d is disproportionate to its %d-byte payload", + estimate, + payload.RetainedSize(), + ) + } + parentEstimate, err := protocol.EstimateRecordBatchRetainedSize(record) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize(parent): %v", err) + } + if estimate > parentEstimate/100 { + t.Fatalf( + "slice estimate %d tracks the %d-byte parent rather than its own rows", + estimate, + parentEstimate, + ) + } +} + +// TestAdmissionCoversMaterializedPayloadAcrossLayouts pins the direction that +// matters for the byte limit: a reservation may be generous, never short. Each +// layout reaches a different branch of the per-row sizing, and a slice exercises +// the branch that cannot read its extent off a buffer length. +func TestAdmissionCoversMaterializedPayloadAcrossLayouts(t *testing.T) { + allocator := memory.DefaultAllocator + fixedSchema := idSchema(nil) + fixedRecord := idBatch(t, allocator, fixedSchema, []int32{1, 2, 3, 4, 5, 6}) + defer fixedRecord.Release() + variableSchema, variableRecord := binaryBatch(t, allocator, 1_000, 64) + defer variableRecord.Release() + dictionarySchema, dictionaryRecord := dictionaryBatch(t, allocator) + defer dictionaryRecord.Release() + + for _, test := range []struct { + name string + schema *arrow.Schema + record arrow.RecordBatch + }{ + {name: "fixed width", schema: fixedSchema, record: fixedRecord}, + {name: "variable width", schema: variableSchema, record: variableRecord}, + {name: "dictionary", schema: dictionarySchema, record: dictionaryRecord}, + } { + t.Run(test.name, func(t *testing.T) { + protocol, err := New(test.schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + rows := test.record.NumRows() + for _, span := range []struct { + name string + start, end int64 + }{ + {name: "whole", start: 0, end: rows}, + {name: "suffix", start: rows / 2, end: rows}, + } { + t.Run(span.name, func(t *testing.T) { + batch := test.record.NewSlice(span.start, span.end) + defer batch.Release() + estimate, err := protocol.EstimateRecordBatchRetainedSize(batch) + if err != nil { + t.Fatalf("EstimateRecordBatchRetainedSize: %v", err) + } + payload, err := protocol.EncodeRecordBatch(batch) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + if estimate < payload.RetainedSize() { + t.Fatalf( + "estimate %d under-reserves its %d-byte payload", + estimate, + payload.RetainedSize(), + ) + } + }) + } + }) + } +} + +// TestSliceRejectsPayloadWhoseRowCountDrifted keeps a corrupted invariant on the +// error path. Slicing against a bound past the decoded batch panics, and that +// panic would escape through recovery into the caller's goroutine. +func TestSliceRejectsPayloadWhoseRowCountDrifted(t *testing.T) { + schema := idSchema(nil) + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + record := idBatch(t, memory.DefaultAllocator, schema, []int32{1, 2, 3, 4}) + defer record.Release() + payload, err := protocol.EncodeRecordBatch(record) + if err != nil { + t.Fatalf("EncodeRecordBatch: %v", err) + } + + drifted := &Payload{ + ipcBytes: payload.ipcBytes, + rows: 64, + chunkRows: []int64{64}, + } + _, err = protocol.Slice(drifted, 32) + if err == nil || !strings.Contains(err.Error(), "row count changed") { + t.Fatalf("Slice of a drifted payload = %v, want a row-count error", err) + } +} + +// TestCanonicalPayloadRejectsIncompleteChunkPlan fails a plan at encode time. +// Left alone it surfaces as a short submission mid-send, which recovery has to +// reconcile against a payload that can never cover its rows. +func TestCanonicalPayloadRejectsIncompleteChunkPlan(t *testing.T) { + protocol, err := New(idSchema(nil), Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + _, err = protocol.payloadFromCanonicalIPC([]byte("ipc"), 10, []int64{4, 5}) + if err == nil || !strings.Contains(err.Error(), "covers 9 of 10 rows") { + t.Fatalf("payloadFromCanonicalIPC with a short plan = %v", err) + } + _, err = protocol.payloadFromCanonicalIPC([]byte("ipc"), 10, []int64{4, -5}) + if err == nil || !strings.Contains(err.Error(), "invalid Flight chunk plan") { + t.Fatalf("payloadFromCanonicalIPC with a negative chunk = %v", err) + } +} + +// TestDecodeRejectsDeclaredSizeBeyondInput covers the decode entry points that +// do not go through EncodeIPC. Arrow allocates against a declared length before +// reading it, so the cross-check has to run for every decode. +func TestDecodeRejectsDeclaredSizeBeyondInput(t *testing.T) { + protocol, err := New(idSchema(nil), Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + data := make([]byte, 12) + binary.LittleEndian.PutUint32(data[0:4], 0xffffffff) + binary.LittleEndian.PutUint32(data[4:8], 1<<26) + _, err = protocol.DecodeIPCRecordBatch(data) + if err == nil || !strings.Contains(err.Error(), "extends beyond") { + t.Fatalf("DecodeIPCRecordBatch of an overlong declared size = %v", err) + } +} + +// TestOversizedDictionaryIsReportedAsDictionary covers a dictionary too large to +// frame. Slicing rows never shrinks it, so the row search bottoms out at one row +// and would otherwise blame the row for a property of the batch. +func TestOversizedDictionaryIsReportedAsDictionary(t *testing.T) { + const ( + distinct = 24_000 + valueBytes = 128 + ) + dictionaryType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int32, + ValueType: arrow.BinaryTypes.String, + } + schema := arrow.NewSchema([]arrow.Field{{ + Name: "category", Type: dictionaryType, Nullable: false, + }}, nil) + + builder := array.NewDictionaryBuilder(memory.DefaultAllocator, dictionaryType) + values := builder.(*array.BinaryDictionaryBuilder) + filler := strings.Repeat("x", valueBytes) + for index := range distinct { + if err := values.AppendString(strconv.Itoa(index) + filler); err != nil { + t.Fatalf("append dictionary value: %v", err) + } + } + categories := builder.NewDictionaryArray() + builder.Release() + record := array.NewRecordBatch(schema, []arrow.Array{categories}, distinct) + categories.Release() + defer record.Release() + + protocol, err := New(schema, Options{}) + if err != nil { + t.Fatalf("New: %v", err) + } + _, err = protocol.EncodeRecordBatch(record) + if err == nil || !strings.Contains(err.Error(), "dictionary FlightData") { + t.Fatalf("EncodeRecordBatch with an oversized dictionary = %v", err) + } +} + +// unreachableAckModel satisfies NewCoreStreamWithHooks for tests whose payload is +// rejected at admission, before any connection opens. Reaching either hook means +// it was admitted, which is what the caller is asserting against. +func unreachableAckModel(t *testing.T) stream.AckModelHooks[*flight.PutResult] { + t.Helper() + return stream.AckModelHooks[*flight.PutResult]{ + Classify: func(*flight.PutResult) stream.ResponseClassification { + t.Error("ack model classified a response for a rejected payload") + return stream.ResponseClassification{Status: stream.ResponseMalformed} + }, + Resolve: func( + *flight.PutResult, + stream.AckState, + ) (stream.AckResolution, error) { + t.Error("ack model resolved a response for a rejected payload") + return stream.AckResolution{}, errors.New("unreachable") + }, + } +} + +func equalInt32(left, right []int32) bool { + if len(left) != len(right) { + return false + } + for i := range left { + if left[i] != right[i] { + return false + } + } + return true +} diff --git a/purego/internal/arrowproto/ipc_preflight.go b/purego/internal/arrowproto/ipc_preflight.go new file mode 100644 index 00000000..94fdfbc0 --- /dev/null +++ b/purego/internal/arrowproto/ipc_preflight.go @@ -0,0 +1,288 @@ +package arrowproto + +import ( + "encoding/binary" + "fmt" + "math" + + flatbuffers "github.com/google/flatbuffers/go" +) + +const ( + ipcContinuationToken = uint32(0xffffffff) + + ipcHeaderSchema = byte(1) + ipcHeaderDictionaryBatch = byte(2) + ipcHeaderRecordBatch = byte(3) +) + +// preflightIPCExpansion walks IPC message metadata without building Arrow arrays. +// For compressed batches it sums the per-buffer uncompressed-size prefixes that +// Arrow would otherwise allocate against unchecked. +func preflightIPCExpansion(data []byte) (expandedBytes int64, err error) { + defer func() { + if recovered := recover(); recovered != nil { + expandedBytes = 0 + err = fmt.Errorf("arrow protocol: invalid IPC metadata: %v", recovered) + } + }() + if len(data) == 0 { + return 0, fmt.Errorf("arrow protocol: IPC input is empty") + } + + position := int64(0) + sawSchema := false + sawRecordBatch := false + for position < int64(len(data)) { + first, next, takeErr := takeIPCBytes(data, position, 4) + if takeErr != nil { + return 0, takeErr + } + position = next + indicator := binary.LittleEndian.Uint32(first) + + var metadataLength uint32 + switch indicator { + case 0: + if position != int64(len(data)) { + return 0, fmt.Errorf( + "arrow protocol: IPC stream contains %d trailing bytes", + int64(len(data))-position, + ) + } + position = int64(len(data)) + continue + case ipcContinuationToken: + lengthBytes, next, lengthErr := takeIPCBytes(data, position, 4) + if lengthErr != nil { + return 0, lengthErr + } + position = next + metadataLength = binary.LittleEndian.Uint32(lengthBytes) + if metadataLength == 0 { + if position != int64(len(data)) { + return 0, fmt.Errorf( + "arrow protocol: IPC stream contains %d trailing bytes", + int64(len(data))-position, + ) + } + position = int64(len(data)) + continue + } + default: + metadataLength = indicator + } + if metadataLength < 4 { + return 0, fmt.Errorf( + "arrow protocol: invalid IPC message metadata length %d", + metadataLength, + ) + } + metadata, next, metadataErr := takeIPCBytes( + data, + position, + int64(metadataLength), + ) + if metadataErr != nil { + return 0, metadataErr + } + position = next + + message := ipcRootTable(metadata) + headerType := message.GetByteSlot(6, 0) + bodyLength := message.GetInt64Slot(10, 0) + if bodyLength < 0 { + return 0, fmt.Errorf( + "arrow protocol: invalid IPC message body length %d", + bodyLength, + ) + } + body, next, bodyErr := takeIPCBytes(data, position, bodyLength) + if bodyErr != nil { + return 0, bodyErr + } + position = next + + switch headerType { + case ipcHeaderSchema: + if sawSchema || sawRecordBatch { + return 0, fmt.Errorf("arrow protocol: IPC stream contains an unexpected schema message") + } + sawSchema = true + case ipcHeaderDictionaryBatch: + if !sawSchema || sawRecordBatch { + return 0, fmt.Errorf("arrow protocol: IPC dictionary message is out of order") + } + header, headerErr := ipcMessageHeader(message) + if headerErr != nil { + return 0, headerErr + } + recordBatch, recordErr := ipcDictionaryRecordBatch(header) + if recordErr != nil { + return 0, recordErr + } + expanded, expansionErr := ipcCompressedExpansion(recordBatch, body) + if expansionErr != nil { + return 0, expansionErr + } + expandedBytes, err = addInt64Saturating(expandedBytes, expanded) + if err != nil { + return math.MaxInt64, err + } + case ipcHeaderRecordBatch: + if !sawSchema || sawRecordBatch { + return 0, fmt.Errorf( + "arrow protocol: IPC stream must contain exactly one RecordBatch", + ) + } + sawRecordBatch = true + recordBatch, headerErr := ipcMessageHeader(message) + if headerErr != nil { + return 0, headerErr + } + expanded, expansionErr := ipcCompressedExpansion(recordBatch, body) + if expansionErr != nil { + return 0, expansionErr + } + expandedBytes, err = addInt64Saturating(expandedBytes, expanded) + if err != nil { + return math.MaxInt64, err + } + default: + return 0, fmt.Errorf( + "arrow protocol: unsupported IPC message header type %d", + headerType, + ) + } + } + if !sawSchema { + return 0, fmt.Errorf("arrow protocol: IPC stream contains no schema") + } + if !sawRecordBatch { + return 0, fmt.Errorf("arrow protocol: IPC stream contains no RecordBatch") + } + return expandedBytes, nil +} + +func takeIPCBytes(data []byte, position, length int64) ([]byte, int64, error) { + if position < 0 || length < 0 || + position > int64(len(data)) || + length > int64(len(data))-position { + return nil, position, fmt.Errorf( + "arrow protocol: IPC message extends beyond %d-byte input", + len(data), + ) + } + end := position + length + return data[int(position):int(end)], end, nil +} + +func ipcRootTable(metadata []byte) flatbuffers.Table { + root := flatbuffers.GetUOffsetT(metadata) + return flatbuffers.Table{Bytes: metadata, Pos: root} +} + +func ipcMessageHeader(message flatbuffers.Table) (flatbuffers.Table, error) { + offset := flatbuffers.UOffsetT(message.Offset(8)) + if offset == 0 { + return flatbuffers.Table{}, fmt.Errorf("arrow protocol: IPC message has no header") + } + var header flatbuffers.Table + message.Union(&header, offset) + return header, nil +} + +func ipcDictionaryRecordBatch( + dictionary flatbuffers.Table, +) (flatbuffers.Table, error) { + offset := flatbuffers.UOffsetT(dictionary.Offset(6)) + if offset == 0 { + return flatbuffers.Table{}, fmt.Errorf( + "arrow protocol: IPC dictionary message has no RecordBatch header", + ) + } + position := dictionary.Indirect(offset + dictionary.Pos) + return flatbuffers.Table{Bytes: dictionary.Bytes, Pos: position}, nil +} + +func ipcCompressedExpansion(recordBatch flatbuffers.Table, body []byte) (int64, error) { + compressionOffset := flatbuffers.UOffsetT(recordBatch.Offset(10)) + if compressionOffset == 0 { + return 0, nil + } + compressionPosition := recordBatch.Indirect(compressionOffset + recordBatch.Pos) + compression := flatbuffers.Table{ + Bytes: recordBatch.Bytes, + Pos: compressionPosition, + } + codec := compression.GetInt8Slot(4, 0) + method := compression.GetInt8Slot(6, 0) + if codec != 0 && codec != 1 { + return 0, fmt.Errorf("arrow protocol: unsupported IPC compression codec %d", codec) + } + if method != 0 { + return 0, fmt.Errorf("arrow protocol: unsupported IPC compression method %d", method) + } + + buffersOffset := flatbuffers.UOffsetT(recordBatch.Offset(8)) + if buffersOffset == 0 { + return 0, nil + } + buffersLength := recordBatch.VectorLen(buffersOffset) + buffersStart := recordBatch.Vector(buffersOffset) + var total int64 + for index := range buffersLength { + bufferPosition := buffersStart + flatbuffers.UOffsetT(index*16) + offset := recordBatch.GetInt64(bufferPosition) + length := recordBatch.GetInt64(bufferPosition + 8) + if offset < 0 || length < 0 || offset > int64(len(body)) || + length > int64(len(body))-offset { + return 0, fmt.Errorf( + "arrow protocol: compressed IPC buffer %d range [%d,%d) exceeds %d-byte body", + index, + offset, + offset+length, + len(body), + ) + } + if length == 0 { + continue + } + if length < 8 { + return 0, fmt.Errorf( + "arrow protocol: compressed IPC buffer %d is %d bytes, smaller than its size prefix", + index, + length, + ) + } + prefix := int64(binary.LittleEndian.Uint64( + body[int(offset) : int(offset)+8], + )) + if prefix == -1 { + continue + } + if prefix < 0 { + return 0, fmt.Errorf( + "arrow protocol: compressed IPC buffer %d declares invalid uncompressed size %d", + index, + prefix, + ) + } + var err error + total, err = addInt64Saturating(total, prefix) + if err != nil { + return math.MaxInt64, fmt.Errorf( + "arrow protocol: declared uncompressed IPC buffer sizes overflow: %w", + err, + ) + } + } + return total, nil +} + +func addInt64Saturating(left, right int64) (int64, error) { + if left < 0 || right < 0 || right > math.MaxInt64-left { + return math.MaxInt64, fmt.Errorf("int64 size overflow") + } + return left + right, nil +} diff --git a/purego/internal/stream/core.go b/purego/internal/stream/core.go index 633b13c9..1536c1f3 100644 --- a/purego/internal/stream/core.go +++ b/purego/internal/stream/core.go @@ -563,7 +563,12 @@ func (cs *CoreStream[Req, Resp]) enqueuePayload( ) (int64, error) { return cs.enqueuePayloadReserved(ctx, weight, func() (Req, uint64, int64, error) { payload, err := build() - return payload, explicitUnits, weight, err + if err != nil { + return payload, explicitUnits, weight, err + } + // weight was derived from the input bytes. An encoding that expands or + // compresses its input must be charged for what it actually holds. + return payload, explicitUnits, cs.enc.actualRetainedSize(payload, weight), nil }) } diff --git a/purego/internal/stream/encoder.go b/purego/internal/stream/encoder.go index f9be7158..74bf97a8 100644 --- a/purego/internal/stream/encoder.go +++ b/purego/internal/stream/encoder.go @@ -54,6 +54,10 @@ type encoder[Req any] interface { // request-object overhead so aggregate backpressure cannot be bypassed by // batches of tiny or empty records. retainedSize(rawBytes, recordCount int) int64 + // actualRetainedSize reports what msg retains now that it exists, replacing + // the pre-encode estimate. An encoding whose output size is not a function + // of its input size returns the true figure here; the rest return estimate. + actualRetainedSize(msg Req, estimate int64) int64 } const encodedRequestOverhead = int64(512) @@ -135,6 +139,10 @@ func (protoEncoder) retainedSize(rawBytes, recordCount int) int64 { return ephemeralRetainedSize(rawBytes, recordCount) } +func (protoEncoder) actualRetainedSize(_ encodedMsg, estimate int64) int64 { + return estimate +} + // jsonEncoder builds EphemeralStream payloads for JSON-encoded records, single // and batched. type jsonEncoder struct{} @@ -201,6 +209,10 @@ func (jsonEncoder) retainedSize(rawBytes, recordCount int) int64 { return ephemeralRetainedSize(rawBytes, recordCount) } +func (jsonEncoder) actualRetainedSize(_ encodedMsg, estimate int64) int64 { + return estimate +} + func stampEphemeralOffset(msg encodedMsg, offset int64) { if msg == nil { return @@ -246,6 +258,10 @@ type EncoderHooks[Req any] struct { Decode func(msg Req) [][]byte MaxWireSize func(msg Req) int RetainedSize func(rawBytes, recordCount int) int64 + // ActualRetainedSize is optional. Set it when RetainedSize cannot predict the + // encoded size from the input size, so the reservation is reconciled against + // what the payload really holds instead of the estimate. + ActualRetainedSize func(msg Req) int64 } type hookEncoder[Req any] struct { @@ -284,6 +300,13 @@ func (e hookEncoder[Req]) retainedSize(rawBytes, recordCount int) int64 { return e.hooks.RetainedSize(rawBytes, recordCount) } +func (e hookEncoder[Req]) actualRetainedSize(msg Req, estimate int64) int64 { + if e.hooks.ActualRetainedSize == nil { + return estimate + } + return e.hooks.ActualRetainedSize(msg) +} + // extractEphemeralRecords recovers the raw record bytes from an EphemeralStream // wire message. A single-record message yields one entry; a batch yields all of // its records. Shared by the proto and JSON encoders' decode. diff --git a/purego/internal/transport/flight_metadata.go b/purego/internal/transport/flight_metadata.go new file mode 100644 index 00000000..bc283eba --- /dev/null +++ b/purego/internal/transport/flight_metadata.go @@ -0,0 +1,204 @@ +package transport + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "math" + "time" +) + +// FlightStreamReadyOffset is the sentinel in the first DoPut PutResult: auth, +// table access, and schema validation all succeeded. +const FlightStreamReadyOffset int64 = -1 + +// FlightBatchMetadata rides in FlightData.AppMetadata. OffsetID is +// connection-local and must be sequential from zero. +type FlightBatchMetadata struct { + OffsetID int64 `json:"offset_id"` +} + +// FlightAckMetadata rides in PutResult.AppMetadata. +type FlightAckMetadata struct { + // Highest durable frame offset. FlightStreamReadyOffset is reserved for setup. + AckUpToOffset int64 `json:"ack_up_to_offset"` + // Cumulative durable row count. + AckUpToRecords uint64 `json:"ack_up_to_records"` + // Requests connection rotation after this grace period. + CloseStreamDurationMS *uint64 `json:"close_stream_duration_ms,omitempty"` +} + +// IsStreamReady reports whether this is the setup-ready response. +func (m FlightAckMetadata) IsStreamReady() bool { + return m.AckUpToOffset == FlightStreamReadyOffset +} + +// ParseFlightBatchMetadata decodes one strict FlightData metadata object. +func ParseFlightBatchMetadata(data []byte) (FlightBatchMetadata, error) { + if len(bytes.TrimSpace(data)) == 0 { + return FlightBatchMetadata{}, fmt.Errorf("parse Flight batch metadata: metadata is empty") + } + var metadata FlightBatchMetadata + if err := json.Unmarshal(data, &metadata); err != nil { + return FlightBatchMetadata{}, fmt.Errorf("parse Flight batch metadata: %w", err) + } + return metadata, nil +} + +// ParseFlightAckMetadata decodes one strict PutResult metadata object. +func ParseFlightAckMetadata(data []byte) (FlightAckMetadata, error) { + if len(bytes.TrimSpace(data)) == 0 { + return FlightAckMetadata{}, fmt.Errorf("parse Flight ack metadata: metadata is empty") + } + var metadata FlightAckMetadata + if err := json.Unmarshal(data, &metadata); err != nil { + return FlightAckMetadata{}, fmt.Errorf("parse Flight ack metadata: %w", err) + } + return metadata, nil +} + +// UnmarshalJSON is strict because a permissive zero-value decode would turn +// malformed metadata such as {} into a real offset-zero frame. +func (m *FlightBatchMetadata) UnmarshalJSON(data []byte) error { + var parsed FlightBatchMetadata + seen, err := decodeStrictJSONObject(data, map[string]func(json.RawMessage) error{ + "offset_id": func(raw json.RawMessage) error { + return decodeRequiredJSONNumber(raw, &parsed.OffsetID) + }, + }) + if err != nil { + return err + } + if !seen["offset_id"] { + return fmt.Errorf("missing required field %q", "offset_id") + } + if parsed.OffsetID < 0 { + return fmt.Errorf("field %q must be non-negative", "offset_id") + } + *m = parsed + return nil +} + +// UnmarshalJSON is as strict as FlightBatchMetadata.UnmarshalJSON, except that +// the close signal is optional. +func (m *FlightAckMetadata) UnmarshalJSON(data []byte) error { + var parsed FlightAckMetadata + seen, err := decodeStrictJSONObject(data, map[string]func(json.RawMessage) error{ + "ack_up_to_offset": func(raw json.RawMessage) error { + return decodeRequiredJSONNumber(raw, &parsed.AckUpToOffset) + }, + "ack_up_to_records": func(raw json.RawMessage) error { + return decodeRequiredJSONNumber(raw, &parsed.AckUpToRecords) + }, + "close_stream_duration_ms": func(raw json.RawMessage) error { + // An explicit null is how many encoders write an absent optional, so + // it must not fail an otherwise valid acknowledgment. + if isJSONNull(raw) { + return nil + } + var duration uint64 + if err := decodeRequiredJSONNumber(raw, &duration); err != nil { + return err + } + parsed.CloseStreamDurationMS = &duration + return nil + }, + }) + if err != nil { + return err + } + for _, required := range []string{"ack_up_to_offset", "ack_up_to_records"} { + if !seen[required] { + return fmt.Errorf("missing required field %q", required) + } + } + if parsed.AckUpToOffset < FlightStreamReadyOffset { + return fmt.Errorf( + "field %q must be at least %d", + "ack_up_to_offset", + FlightStreamReadyOffset, + ) + } + if parsed.CloseStreamDurationMS != nil { + maxMillis := uint64(math.MaxInt64 / int64(time.Millisecond)) + if *parsed.CloseStreamDurationMS > maxMillis { + return fmt.Errorf( + "field %q exceeds maximum representable duration", + "close_stream_duration_ms", + ) + } + } + *m = parsed + return nil +} + +// decodeStrictJSONObject decodes exactly one JSON object. Values stay raw so +// each field decoder can check its own numeric type. Unknown fields are ignored +// for additive compatibility; duplicate names are rejected even then. +func decodeStrictJSONObject( + data []byte, + decoders map[string]func(json.RawMessage) error, +) (map[string]bool, error) { + decoder := json.NewDecoder(bytes.NewReader(data)) + start, err := decoder.Token() + if err != nil { + if err == io.EOF { + return nil, fmt.Errorf("metadata is empty") + } + return nil, err + } + if delim, ok := start.(json.Delim); !ok || delim != '{' { + return nil, fmt.Errorf("metadata must be a JSON object") + } + + seen := make(map[string]bool, len(decoders)) + for decoder.More() { + token, err := decoder.Token() + if err != nil { + return nil, err + } + key, ok := token.(string) + if !ok { + return nil, fmt.Errorf("metadata field name is not a string") + } + if seen[key] { + return nil, fmt.Errorf("duplicate field %q", key) + } + seen[key] = true + + var raw json.RawMessage + if err := decoder.Decode(&raw); err != nil { + return nil, fmt.Errorf("field %q: %w", key, err) + } + if decode, ok := decoders[key]; ok { + if err := decode(raw); err != nil { + return nil, fmt.Errorf("field %q: %w", key, err) + } + } + } + if _, err := decoder.Token(); err != nil { + return nil, err + } + if _, err := decoder.Token(); err != io.EOF { + if err == nil { + return nil, fmt.Errorf("metadata contains trailing JSON") + } + return nil, fmt.Errorf("metadata contains trailing data: %w", err) + } + return seen, nil +} + +func isJSONNull(raw json.RawMessage) bool { + return bytes.Equal(bytes.TrimSpace(raw), []byte("null")) +} + +func decodeRequiredJSONNumber[T int64 | uint64](raw json.RawMessage, dst *T) error { + if isJSONNull(raw) { + return fmt.Errorf("must not be null") + } + if err := json.Unmarshal(raw, dst); err != nil { + return err + } + return nil +} diff --git a/purego/internal/transport/flight_metadata_test.go b/purego/internal/transport/flight_metadata_test.go new file mode 100644 index 00000000..97936480 --- /dev/null +++ b/purego/internal/transport/flight_metadata_test.go @@ -0,0 +1,180 @@ +package transport + +import ( + "strings" + "testing" +) + +func TestParseFlightBatchMetadata(t *testing.T) { + tests := []struct { + name string + input string + wantOffset int64 + wantErr string + }{ + {name: "offset", input: `{"offset_id":7}`, wantOffset: 7}, + {name: "zero offset", input: `{"offset_id":0}`}, + { + name: "unknown fields stay additive", + input: `{"offset_id":3,"future_field":"x"}`, + wantOffset: 3, + }, + {name: "empty object", input: `{}`, wantErr: "missing required field"}, + {name: "empty input", input: ``, wantErr: "metadata is empty"}, + {name: "null offset", input: `{"offset_id":null}`, wantErr: "must not be null"}, + { + name: "negative offset", + input: `{"offset_id":-1}`, + wantErr: "must be non-negative", + }, + { + name: "duplicate field", + input: `{"offset_id":1,"offset_id":2}`, + wantErr: "duplicate field", + }, + {name: "not an object", input: `[1]`, wantErr: "must be a JSON object"}, + { + name: "trailing json", + input: `{"offset_id":1} {"offset_id":2}`, + wantErr: "invalid character", + }, + { + name: "string offset", + input: `{"offset_id":"1"}`, + wantErr: "offset_id", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + metadata, err := ParseFlightBatchMetadata([]byte(test.input)) + if test.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), test.wantErr) { + t.Fatalf("error = %v, want it to mention %q", err, test.wantErr) + } + return + } + if err != nil { + t.Fatalf("ParseFlightBatchMetadata: %v", err) + } + if metadata.OffsetID != test.wantOffset { + t.Fatalf("OffsetID = %d, want %d", metadata.OffsetID, test.wantOffset) + } + }) + } +} + +func TestParseFlightAckMetadata(t *testing.T) { + const closeMillis = uint64(2_500) + tests := []struct { + name string + input string + wantOffset int64 + wantRecords uint64 + wantClose *uint64 + wantReady bool + wantErr string + }{ + { + name: "ack", + input: `{"ack_up_to_offset":4,"ack_up_to_records":900}`, + wantOffset: 4, + wantRecords: 900, + }, + { + name: "ack with rotation request", + input: `{"ack_up_to_offset":4,"ack_up_to_records":900,"close_stream_duration_ms":2500}`, + wantOffset: 4, + wantRecords: 900, + wantClose: &[]uint64{closeMillis}[0], + }, + { + // An absent optional is written as an explicit null by many encoders, + // and dropping the acknowledgment over it would fail the stream. + name: "explicit null rotation request is absent", + input: `{"ack_up_to_offset":4,"ack_up_to_records":900,"close_stream_duration_ms":null}`, + wantOffset: 4, + wantRecords: 900, + }, + { + name: "stream ready sentinel", + input: `{"ack_up_to_offset":-1,"ack_up_to_records":0}`, + wantOffset: FlightStreamReadyOffset, + wantReady: true, + }, + { + name: "offset below the sentinel", + input: `{"ack_up_to_offset":-2,"ack_up_to_records":0}`, + wantErr: "must be at least", + }, + { + name: "missing record count", + input: `{"ack_up_to_offset":4}`, + wantErr: "missing required field", + }, + { + name: "null record count", + input: `{"ack_up_to_offset":4,"ack_up_to_records":null}`, + wantErr: "must not be null", + }, + { + name: "negative record count", + input: `{"ack_up_to_offset":4,"ack_up_to_records":-1}`, + wantErr: "ack_up_to_records", + }, + { + name: "rotation request overflows a duration", + input: `{"ack_up_to_offset":4,"ack_up_to_records":1,"close_stream_duration_ms":9223372036854775807}`, + wantErr: "maximum representable duration", + }, + {name: "empty input", input: ``, wantErr: "metadata is empty"}, + {name: "empty object", input: `{}`, wantErr: "missing required field"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + metadata, err := ParseFlightAckMetadata([]byte(test.input)) + if test.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), test.wantErr) { + t.Fatalf("error = %v, want it to mention %q", err, test.wantErr) + } + return + } + if err != nil { + t.Fatalf("ParseFlightAckMetadata: %v", err) + } + if metadata.AckUpToOffset != test.wantOffset { + t.Errorf( + "AckUpToOffset = %d, want %d", + metadata.AckUpToOffset, + test.wantOffset, + ) + } + if metadata.AckUpToRecords != test.wantRecords { + t.Errorf( + "AckUpToRecords = %d, want %d", + metadata.AckUpToRecords, + test.wantRecords, + ) + } + if metadata.IsStreamReady() != test.wantReady { + t.Errorf("IsStreamReady = %v, want %v", + metadata.IsStreamReady(), test.wantReady) + } + switch { + case test.wantClose == nil && metadata.CloseStreamDurationMS != nil: + t.Errorf( + "CloseStreamDurationMS = %d, want absent", + *metadata.CloseStreamDurationMS, + ) + case test.wantClose != nil && metadata.CloseStreamDurationMS == nil: + t.Errorf("CloseStreamDurationMS absent, want %d", *test.wantClose) + case test.wantClose != nil && + *metadata.CloseStreamDurationMS != *test.wantClose: + t.Errorf( + "CloseStreamDurationMS = %d, want %d", + *metadata.CloseStreamDurationMS, + *test.wantClose, + ) + } + }) + } +}