feat(triple): cut per-message marshal allocations with a pooled MarshalAppend codec path (#3718)
* feat(triple): add marshalAppender optional codec interface
Today the proto codec marshals via proto.Marshal, which always builds a fresh
[]byte from scratch, so every unary call allocates the whole payload even
though the request buffer pool already holds reusable capacity. The marshaler
should instead be able to serialize into a caller-provided buffer.
This commit adds marshalAppender, an optional Codec extension whose
MarshalAppend serializes the message into a caller-supplied dst slice.
protoBinaryCodec implements it by delegating to
proto.MarshalOptions{}.MarshalAppend, and tripleServerCodecSession implements
it too, wrapping the generic response before forwarding when needed, so both
IDL and non-IDL server calls expose the new capability. Codecs without the
extension (hessian2, json, msgpack, third-party codecs) are untouched and
keep their exact previous behavior, so the change is purely additive.
Signed-off-by: lizining <lizining1231@outlook.com>
* feat(triple): add pooled marshal pipeline via marshalToPool
Adds marshalToPool, the shared borrow-marshal-adopt dance that the envelope
writer and the unary marshaler both need. It borrows a buffer from the pool,
lets the codec's MarshalAppend fill its spare capacity, and when MarshalAppend
had to grow the backing array, swaps the larger array back into the pooled
buffer so the capacity is recycled instead of dropped. On the happy path
nothing is allocated and the buffer is handed to the caller ready to be
wrapped or written.
Signed-off-by: lizining <lizining1231@outlook.com>
* feat(triple): use marshalAppend fast path in envelope writer
The envelope writer used to marshal with codec.Marshal and then copy the
freshly allocated slice into the outgoing envelope, so every message paid one
allocation plus an extra copy. It now probes the codec for marshalAppender
and, when present, runs marshalAndWrite: marshalToPool fills a pooled buffer
in place and the bytes are wrapped straight into the envelope.
The previous logic is kept as marshalWithFallback and serves codecs without
the extension, so the produced bytes stay identical on both paths.
Signed-off-by: lizining <lizining1231@outlook.com>
* feat(triple): add marshalAppend fast path and shared compression tail
tripleUnaryMarshaler.Marshal used to serialize every unary payload with
codec.Marshal, which allocates the whole []byte from scratch, and then ran the
compression and write tail inline at the end of the function.
Marshal now probes the codec for marshalAppender: capable codecs take the new
marshalAndWrite path, which marshals into a pooled buffer via marshalToPool so
the common case allocates nothing, while codecs without the extension keep the
old path. To keep the two paths from ever drifting, the compression tail was
extracted from Marshal into one shared compressAndWrite that both paths call,
so the threshold check, pooled gzip, sendMaxBytes enforcement, compression
header and final write always behave the same.
Signed-off-by: lizining <lizining1231@outlook.com>
* test(triple): add marshal fast-path A/B benchmarks
Adds A/B benchmarks that run the pooled MarshalAppend fast path against the
plain codec.Marshal slow path, for both the unary marshaler and the envelope
writer, with and without gzip. The fast path is expected to allocate nothing
per op while the slow-path benches keep measuring the old baseline, so a
regression in the fast path shows up immediately as an allocation or
throughput change in the bench output.
Signed-off-by: lizining <lizining1231@outlook.com>
* test(triple): add regression tests for pooled marshal fast path
Adds nine TestMarshalPerf* regression cases that lock the pooled marshal
behavior in place so the fast and slow paths cannot silently diverge.
WireParity asserts both paths emit byte-identical bodies and headers for every
payload size, compressed or not; PoolInvariants and LargeBufferDropped check
that pooled buffers come back clean and oversized ones are released;
BackupCodecFallback, TypeGuard and ErrorGuard keep the previous fallback,
type-guard and error semantics; CompressionAndMaxBytes and ConcurrentSend
verify compression, sendMaxBytes and pool safety under concurrency; and
ServerSessionFastPath drives tripleServerCodecSession through the fast path
end to end.
Signed-off-by: lizining <lizining1231@outlook.com>
* style(triple): modernize regression test syntax for CI fmt check
* style(triple): apply modernize and gofmt to marshal perf regression test
Signed-off-by: lizining <lizining1231@outlook.com>
* style(triple): align test comments with repo style and clarify semantics
Signed-off-by: lizining <lizining1231@outlook.com>
* test(triple): strengthen concurrent marshal regression tests
Replace the concurrent send test's length-only assertion with per-frame
content verification: each goroutine marshals globally unique messages and
the output is decoded frame by frame, so cross-goroutine buffer reuse that
corrupts payloads fails even without -race.
Add a start gate so the shared buffer pool is actually contended, plus a
fastPathProbeCodec-based test asserting envelopeWriter.Marshal takes the
MarshalAppend fast path. Clarify that syncBuffer must not be shared by
concurrent writers.
Signed-off-by: lizining <lizining1231@outlook.com>
---------
Signed-off-by: lizining <lizining1231@outlook.com>diff --git a/protocol/triple/triple_protocol/buffer_pool.go b/protocol/triple/triple_protocol/buffer_pool.go
index 51dadb8..d34948b 100644
--- a/protocol/triple/triple_protocol/buffer_pool.go
+++ b/protocol/triple/triple_protocol/buffer_pool.go
@@ -24,7 +24,7 @@
const (
initialBufferSize = 512
- maxRecycleBufferSize = 8 * 1024 * 1024 // if >8MiB, don't hold onto a buffer
+ maxRecycleBufferSize = 8 * 1024 * 1024 // Don't recycle buffers larger than this.
)
type bufferPool struct {
@@ -55,3 +55,27 @@
buffer.Reset()
b.Pool.Put(buffer)
}
+
+// marshalToPool serializes message with appender into a *bytes.Buffer drawn
+// from pool and returns it. If the pooled array is too small and the appender
+// grows the slice, the larger array is swapped in so it can be recycled once
+// the caller returns the buffer. On failure the buffer is put back and a
+// CodeInternal error is returned.
+func marshalToPool(pool *bufferPool, appender marshalAppender, message any) (*bytes.Buffer, *Error) {
+ buffer := pool.Get()
+ raw, err := appender.MarshalAppend(buffer.Bytes(), message)
+ if err != nil {
+ pool.Put(buffer)
+ return nil, errorf(CodeInternal, "marshal message: %w", err)
+ }
+ if cap(raw) > buffer.Cap() {
+ // MarshalAppend grew the slice: swap the larger array in so it can be
+ // recycled next time.
+ *buffer = *bytes.NewBuffer(raw)
+ } else {
+ // No reallocation occurred; adopt the bytes the appender already wrote
+ // into the pooled array.
+ buffer.Write(raw)
+ }
+ return buffer, nil
+}
diff --git a/protocol/triple/triple_protocol/codec.go b/protocol/triple/triple_protocol/codec.go
index 8665b7c..b3ca089 100644
--- a/protocol/triple/triple_protocol/codec.go
+++ b/protocol/triple/triple_protocol/codec.go
@@ -98,6 +98,17 @@
IsBinary() bool
}
+// marshalAppender is an extension to Codec for serializing into a caller-provided
+// buffer. Codecs that implement it can serialize with zero allocation when the
+// buffer's capacity is sufficient (protobuf fast path); otherwise the appender
+// itself grows the buffer. MarshalAppend must produce output byte-identical to
+// Marshal for the same message — the extension only changes where the bytes are
+// written. Codecs that do not implement it are unaffected.
+type marshalAppender interface {
+ // MarshalAppend marshals the given message, appending the result to dst.
+ MarshalAppend(dst []byte, message any) ([]byte, error)
+}
+
// protoBinaryCodec handles standard protobuf binary serialization for IDL
// calls. Non-IDL (Java Dubbo Triple generic call) wrapper handling on the
// server side is handled by tripleServerCodecSession, which delegates to
@@ -116,6 +127,14 @@
return proto.Marshal(protoMessage)
}
+func (c *protoBinaryCodec) MarshalAppend(dst []byte, message any) ([]byte, error) {
+ protoMessage, ok := message.(proto.Message)
+ if !ok {
+ return nil, errNotProto(message)
+ }
+ return proto.MarshalOptions{}.MarshalAppend(dst, protoMessage)
+}
+
func (c *protoBinaryCodec) Unmarshal(data []byte, message any) error {
protoMessage, ok := message.(proto.Message)
if !ok {
@@ -636,6 +655,43 @@
}
// Non-IDL: wrap the response in a TripleResponseWrapper whose Data is
// serialized with the inner codec resolved from the request's SerializeType.
+ wrapper, err := s.responseWrapper(message)
+ if err != nil {
+ return nil, err
+ }
+ return proto.Marshal(wrapper)
+}
+
+// MarshalAppend extends Marshal for the marshalAppender fast path. The IDL
+// proto leg forwards to the delegate's appender (zero allocation when the
+// pooled buffer's cap is sufficient); the Non-IDL leg still serializes the
+// inner payload with the inner codec, but appends the outer
+// TripleResponseWrapper into the caller-provided buffer instead of allocating
+// a fresh slice.
+func (s *tripleServerCodecSession) MarshalAppend(dst []byte, message any) ([]byte, error) {
+ if pm, isProto := message.(proto.Message); isProto {
+ if appender, ok := s.delegate.(marshalAppender); ok {
+ return appender.MarshalAppend(dst, pm)
+ }
+ // A custom proto codec without the appender extension must still
+ // round-trip; fall back to a plain marshal plus append. The output is
+ // byte-for-byte what Marshal would produce.
+ raw, err := s.delegate.Marshal(pm)
+ if err != nil {
+ return nil, err
+ }
+ return append(dst, raw...), nil
+ }
+ wrapper, err := s.responseWrapper(message)
+ if err != nil {
+ return nil, err
+ }
+ return proto.MarshalOptions{}.MarshalAppend(dst, wrapper)
+}
+
+// responseWrapper builds the TripleResponseWrapper for a Non-IDL response,
+// resolving the inner codec from the SerializeType captured by Unmarshal.
+func (s *tripleServerCodecSession) responseWrapper(message any) (*interoperability.TripleResponseWrapper, error) {
inner, err := resolveInnerCodec(s.serializeType)
if err != nil {
return nil, fmt.Errorf("marshal triple wrapper response: %w", err)
@@ -667,10 +723,10 @@
}
// Use inner.Name() instead of s.serializeType so that an absent SerializeType
// (defaulted to hessian2 by resolveInnerCodec) is normalized on the wire.
- return proto.Marshal(&interoperability.TripleResponseWrapper{
+ return &interoperability.TripleResponseWrapper{
SerializeType: inner.Name(),
Data: data,
- })
+ }, nil
}
// unmarshalWrapperRequestArgs decodes the inner args of a TripleRequestWrapper
diff --git a/protocol/triple/triple_protocol/envelope.go b/protocol/triple/triple_protocol/envelope.go
index a139092..2280652 100644
--- a/protocol/triple/triple_protocol/envelope.go
+++ b/protocol/triple/triple_protocol/envelope.go
@@ -76,15 +76,37 @@
}
return nil
}
- raw, err := w.codec.Marshal(message)
- if err != nil {
- if w.backupCodec != nil && w.codec.Name() != w.backupCodec.Name() {
- logger.Debugf("[Triple][Codec] failed to marshal message with primary codec %s, trying fallback codec %s", w.codec.Name(), w.backupCodec.Name())
- raw, err = w.backupCodec.Marshal(message)
- }
+ // Fast path: if the codec can append to a caller-provided buffer, marshal
+ // directly into a pooled buffer to avoid the per-request allocation.
+ if appender, ok := w.codec.(marshalAppender); ok {
+ buffer, err := marshalToPool(w.bufferPool, appender, message)
if err != nil {
- return errorf(CodeInternal, "marshal message: %w", err)
+ // Preserve the backup-codec fallback semantics of the slow path.
+ // Only a marshal failure may fall back; write failures (I/O,
+ // compression, or the sendMaxBytes limit) are terminal and must
+ // never retry through the backup codec.
+ if w.backupCodec != nil && w.codec.Name() != w.backupCodec.Name() {
+ logger.Debugf("[Triple][Codec] failed to marshal message with primary codec %s, trying fallback codec %s", w.codec.Name(), w.backupCodec.Name())
+ return w.marshalWithFallback(w.backupCodec, message)
+ }
+ return err
}
+ defer w.bufferPool.Put(buffer)
+ return w.Write(&envelope{Data: buffer})
+ }
+ return w.marshalWithFallback(w.codec, message)
+}
+
+// marshalWithFallback is the slow path: Marshal, falling back to the backup
+// codec on error, then wrap the result in a fresh buffer.
+func (w *envelopeWriter) marshalWithFallback(codec Codec, message any) *Error {
+ raw, err := codec.Marshal(message)
+ if err != nil && w.backupCodec != nil && codec.Name() != w.backupCodec.Name() {
+ logger.Debugf("[Triple][Codec] failed to marshal message with primary codec %s, trying fallback codec %s", codec.Name(), w.backupCodec.Name())
+ raw, err = w.backupCodec.Marshal(message)
+ }
+ if err != nil {
+ return errorf(CodeInternal, "marshal message: %w", err)
}
// We can't avoid allocating the byte slice, so we may as well reuse it once
// we're done with it.
diff --git a/protocol/triple/triple_protocol/marshal_perf_bench_test.go b/protocol/triple/triple_protocol/marshal_perf_bench_test.go
new file mode 100644
index 0000000..e5f3efa
--- /dev/null
+++ b/protocol/triple/triple_protocol/marshal_perf_bench_test.go
@@ -0,0 +1,169 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package triple_protocol
+
+import (
+ "compress/gzip"
+ "fmt"
+ "io"
+ "net/http"
+ "strings"
+ "testing"
+)
+
+import (
+ pingv1 "dubbo.apache.org/dubbo-go/v3/protocol/triple/triple_protocol/internal/gen/proto/connect/ping/v1"
+)
+
+// noAppenderCodec embeds a Codec without the marshalAppender extension, so
+// the marshaler's appender type assertion fails and marshaling takes the
+// codec.Marshal slow path — the baseline that the fast path is measured
+// against.
+type noAppenderCodec struct{ Codec }
+
+// marshalPerfPayloadSizes covers fixed-overhead-dominated small messages
+// through bandwidth-dominated large messages.
+var marshalPerfPayloadSizes = []int{128, 1024, 16 * 1024, 1024 * 1024}
+
+func marshalPerfSizeLabel(size int) string {
+ if size == 1024*1024 {
+ return "1MiB"
+ }
+ return fmt.Sprintf("%dB", size)
+}
+
+func newMarshalPerfMessage(size int) *pingv1.PingRequest {
+ return &pingv1.PingRequest{Text: strings.Repeat("a", size)}
+}
+
+// benchMarshalConfig selects which marshal branches a benchmark drives: gzip
+// compression and the sendMaxBytes limit checks.
+type benchMarshalConfig struct {
+ compress bool
+ sendMaxBytes int
+}
+
+// newBenchGzipPool returns a gzip compression pool for benchmarks. The
+// initial io.Discard sink is replaced by the compressionPool with Reset on
+// each use.
+func newBenchGzipPool() *compressionPool {
+ return newCompressionPool(
+ func() Decompressor { return &gzip.Reader{} },
+ func() Compressor { return gzip.NewWriter(io.Discard) },
+ )
+}
+
+// benchTripleUnaryMarshaler drives tripleUnaryMarshaler.Marshal with the given
+// codec. protoBinaryCodec exercises the MarshalAppend fast path;
+// noAppenderCodec exercises the codec.Marshal slow path.
+func benchTripleUnaryMarshaler(b *testing.B, codec Codec, message *pingv1.PingRequest, cfg benchMarshalConfig) {
+ b.Helper()
+ m := &tripleUnaryMarshaler{
+ writer: io.Discard,
+ codec: codec,
+ bufferPool: newBufferPool(),
+ sendMaxBytes: cfg.sendMaxBytes,
+ }
+ if cfg.compress {
+ m.compressionPool = newBenchGzipPool()
+ m.compressionName = compressionGzip
+ // setHeaderCanonical writes into header on the compressed path.
+ m.header = make(http.Header)
+ }
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ if err := m.Marshal(message); err != nil {
+ b.Fatal(err)
+ }
+ }
+}
+
+func runUnaryMarshalerBench(b *testing.B, codec Codec, cfg benchMarshalConfig) {
+ for _, size := range marshalPerfPayloadSizes {
+ b.Run(marshalPerfSizeLabel(size), func(b *testing.B) {
+ benchTripleUnaryMarshaler(b, codec, newMarshalPerfMessage(size), cfg)
+ })
+ }
+}
+
+func BenchmarkUnaryMarshalerFastPath(b *testing.B) {
+ runUnaryMarshalerBench(b, &protoBinaryCodec{}, benchMarshalConfig{})
+}
+
+func BenchmarkUnaryMarshalerSlowPath(b *testing.B) {
+ runUnaryMarshalerBench(b, &noAppenderCodec{&protoBinaryCodec{}}, benchMarshalConfig{})
+}
+
+// BenchmarkUnaryMarshaler*Compressed drive the gzip-compressed fast and slow
+// paths with a generous sendMaxBytes that is exercised but never tripped.
+func BenchmarkUnaryMarshalerFastPathCompressed(b *testing.B) {
+ runUnaryMarshalerBench(b, &protoBinaryCodec{}, benchMarshalConfig{compress: true, sendMaxBytes: 1 << 30})
+}
+
+func BenchmarkUnaryMarshalerSlowPathCompressed(b *testing.B) {
+ runUnaryMarshalerBench(b, &noAppenderCodec{&protoBinaryCodec{}}, benchMarshalConfig{compress: true, sendMaxBytes: 1 << 30})
+}
+
+// benchEnvelopeWriter drives envelopeWriter.Marshal with the given codec.
+func benchEnvelopeWriter(b *testing.B, codec Codec, message *pingv1.PingRequest, cfg benchMarshalConfig) {
+ b.Helper()
+ w := &envelopeWriter{
+ writer: io.Discard,
+ codec: codec,
+ bufferPool: newBufferPool(),
+ sendMaxBytes: cfg.sendMaxBytes,
+ }
+ if cfg.compress {
+ w.compressionPool = newBenchGzipPool()
+ }
+ b.ReportAllocs()
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ if err := w.Marshal(message); err != nil {
+ b.Fatal(err)
+ }
+ }
+}
+
+func runEnvelopeWriterBench(b *testing.B, codec Codec, cfg benchMarshalConfig) {
+ for _, size := range marshalPerfPayloadSizes {
+ b.Run(marshalPerfSizeLabel(size), func(b *testing.B) {
+ benchEnvelopeWriter(b, codec, newMarshalPerfMessage(size), cfg)
+ })
+ }
+}
+
+func BenchmarkEnvelopeWriterFastPath(b *testing.B) {
+ runEnvelopeWriterBench(b, &protoBinaryCodec{}, benchMarshalConfig{})
+}
+
+func BenchmarkEnvelopeWriterSlowPath(b *testing.B) {
+ runEnvelopeWriterBench(b, &noAppenderCodec{&protoBinaryCodec{}}, benchMarshalConfig{})
+}
+
+// BenchmarkEnvelopeWriter*Compressed drive the gzip-compressed branch of
+// envelopeWriter.Write (compression into a second pooled buffer plus the
+// compressed-size limit check).
+func BenchmarkEnvelopeWriterFastPathCompressed(b *testing.B) {
+ runEnvelopeWriterBench(b, &protoBinaryCodec{}, benchMarshalConfig{compress: true, sendMaxBytes: 1 << 30})
+}
+
+func BenchmarkEnvelopeWriterSlowPathCompressed(b *testing.B) {
+ runEnvelopeWriterBench(b, &noAppenderCodec{&protoBinaryCodec{}}, benchMarshalConfig{compress: true, sendMaxBytes: 1 << 30})
+}
diff --git a/protocol/triple/triple_protocol/marshal_perf_regression_test.go b/protocol/triple/triple_protocol/marshal_perf_regression_test.go
new file mode 100644
index 0000000..24bb769
--- /dev/null
+++ b/protocol/triple/triple_protocol/marshal_perf_regression_test.go
@@ -0,0 +1,937 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package triple_protocol
+
+import (
+ "bytes"
+ "compress/gzip"
+ "encoding/binary"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "strings"
+ "sync"
+ "testing"
+)
+
+import (
+ pingv1 "dubbo.apache.org/dubbo-go/v3/protocol/triple/triple_protocol/internal/gen/proto/connect/ping/v1"
+)
+
+// syncBuffer is an io.Writer whose individual Write calls are safe for
+// concurrent use. Each goroutine in the concurrency test writes to its own
+// syncBuffer: emitting an envelope frame takes two Write calls (5-byte prefix
+// then payload), so giving each writer its own buffer prevents frames from
+// interleaving.
+type syncBuffer struct {
+ mu sync.Mutex
+ buf bytes.Buffer
+}
+
+func (s *syncBuffer) Write(p []byte) (int, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return s.buf.Write(p)
+}
+
+func (s *syncBuffer) Bytes() []byte {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return append([]byte(nil), s.buf.Bytes()...)
+}
+
+func (s *syncBuffer) Len() int {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+ return s.buf.Len()
+}
+
+// regressionCodecList returns the two codecs under test: the fast path
+// (implements marshalAppender) and the slow path (does not).
+func regressionCodecList() []struct {
+ name string
+ codec Codec
+} {
+ return []struct {
+ name string
+ codec Codec
+ }{
+ {name: "fast", codec: &protoBinaryCodec{}},
+ {name: "slow", codec: &noAppenderCodec{&protoBinaryCodec{}}},
+ }
+}
+
+func newTestGzipCompressionPool() *compressionPool {
+ return newCompressionPool(
+ func() Decompressor { return &gzip.Reader{} },
+ func() Compressor { return gzip.NewWriter(io.Discard) },
+ )
+}
+
+// newEnvelopeWriterForTest builds an envelopeWriter (gRPC/Triple wire) with
+// optional gzip compression.
+func newEnvelopeWriterForTest(t *testing.T, codec Codec, compress bool, compressMinBytes, sendMaxBytes int) (*envelopeWriter, *syncBuffer) {
+ t.Helper()
+ out := &syncBuffer{}
+ w := &envelopeWriter{
+ writer: out,
+ codec: codec,
+ compressMinBytes: compressMinBytes,
+ bufferPool: newBufferPool(),
+ sendMaxBytes: sendMaxBytes,
+ }
+ if compress {
+ w.compressionPool = newTestGzipCompressionPool()
+ }
+ return w, out
+}
+
+// newTripleMarshalerForTest builds a tripleUnaryMarshaler (Triple HTTP body
+// wire) with optional gzip compression.
+func newTripleMarshalerForTest(
+ t *testing.T,
+ codec Codec,
+ compress bool,
+ compressMinBytes, sendMaxBytes int,
+) (*tripleUnaryMarshaler, *syncBuffer, http.Header) {
+ t.Helper()
+ out := &syncBuffer{}
+ header := make(http.Header)
+ m := &tripleUnaryMarshaler{
+ writer: out,
+ codec: codec,
+ compressMinBytes: compressMinBytes,
+ bufferPool: newBufferPool(),
+ header: header,
+ sendMaxBytes: sendMaxBytes,
+ }
+ if compress {
+ m.compressionPool = newTestGzipCompressionPool()
+ m.compressionName = compressionGzip
+ }
+ return m, out, header
+}
+
+func regressionPayloadSizes() []int {
+ mi := 1024 * 1024
+ return []int{0, 1, 511, 512, 513, 1024, 8*mi - 1, 8 * mi, 8*mi + 1}
+}
+
+func regressionSizeLabel(size int) string {
+ switch size {
+ case 0:
+ return "empty"
+ case 1024:
+ return "1KiB"
+ case 8*1024*1024 - 1:
+ return "8MiB-1"
+ case 8 * 1024 * 1024:
+ return "8MiB"
+ case 8*1024*1024 + 1:
+ return "8MiB+1"
+ default:
+ return fmt.Sprintf("%dB", size)
+ }
+}
+
+// envelopePrefix parses the 5-byte gRPC/Triple envelope prefix written by
+// envelopeWriter: byte 0 is flags, bytes 1..4 (wire[1:5]) are the big-endian
+// payload length.
+func envelopePrefix(t *testing.T, wire []byte) (flags byte, length int) {
+ t.Helper()
+ if len(wire) < 5 {
+ t.Fatalf("wire output shorter than envelope prefix: got %d bytes", len(wire))
+ }
+ return wire[0], int(binary.BigEndian.Uint32(wire[1:5]))
+}
+
+func wantCode(t *testing.T, err *Error, want Code) {
+ t.Helper()
+ if err == nil {
+ t.Fatalf("expected *Error with code %v, got nil", want)
+ }
+ if got := err.Code(); got != want {
+ t.Fatalf("expected code %v, got %v (message: %v)", want, got, err)
+ }
+}
+
+// TestMarshalPerfWireParity verifies that the MarshalAppend fast path and the
+// codec.Marshal slow path emit byte-identical wire output on both wire types.
+func TestMarshalPerfWireParity(t *testing.T) {
+ for _, compress := range []bool{false, true} {
+ t.Run(fmt.Sprintf("compressed=%v", compress), func(t *testing.T) {
+ for _, size := range regressionPayloadSizes() {
+ t.Run(regressionSizeLabel(size), func(t *testing.T) {
+ t.Parallel()
+ msg := newMarshalPerfMessage(size)
+
+ // triple HTTP body wire
+ t.Run("triple", func(t *testing.T) {
+ var fastWire, slowWire []byte
+ var fastHeader, slowHeader http.Header
+ for _, c := range regressionCodecList() {
+ m, out, header := newTripleMarshalerForTest(t, c.codec, compress, 0, 0)
+ if err := m.Marshal(msg); err != nil {
+ t.Fatalf("%s path triple marshal: %v", c.name, err)
+ }
+ switch c.name {
+ case "fast":
+ fastWire, fastHeader = out.Bytes(), header
+ case "slow":
+ slowWire, slowHeader = out.Bytes(), header
+ }
+ }
+ if !bytes.Equal(fastWire, slowWire) {
+ t.Fatalf("triple wire mismatch: fast %d bytes vs slow %d bytes", len(fastWire), len(slowWire))
+ }
+ if compress {
+ if fastHeader.Get(tripleUnaryHeaderCompression) != compressionGzip ||
+ slowHeader.Get(tripleUnaryHeaderCompression) != compressionGzip {
+ t.Fatalf("compression header not set identically: fast=%q slow=%q",
+ fastHeader.Get(tripleUnaryHeaderCompression), slowHeader.Get(tripleUnaryHeaderCompression))
+ }
+ }
+ })
+
+ // gRPC/Triple envelope wire (5-byte prefix framing)
+ t.Run("envelope", func(t *testing.T) {
+ var fastWire, slowWire []byte
+ for _, c := range regressionCodecList() {
+ w, out := newEnvelopeWriterForTest(t, c.codec, compress, 0, 0)
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("%s path envelope marshal: %v", c.name, err)
+ }
+ switch c.name {
+ case "fast":
+ fastWire = out.Bytes()
+ case "slow":
+ slowWire = out.Bytes()
+ }
+ }
+ if !bytes.Equal(fastWire, slowWire) {
+ t.Fatalf("envelope wire mismatch: fast %d bytes vs slow %d bytes", len(fastWire), len(slowWire))
+ }
+ })
+ })
+ }
+ })
+ }
+}
+
+// TestMarshalPerfPoolInvariants verifies buffer pool reuse invariants and
+// envelope framing for empty and nil messages.
+func TestMarshalPerfPoolInvariants(t *testing.T) {
+ t.Parallel()
+
+ t.Run("get returns empty buffer", func(t *testing.T) {
+ pool := newBufferPool()
+ buf := pool.Get()
+ defer pool.Put(buf)
+ if buf.Len() != 0 {
+ t.Fatalf("bufferPool.Get() returned buffer with Len()==%d, want 0", buf.Len())
+ }
+ })
+
+ t.Run("put resets length", func(t *testing.T) {
+ pool := newBufferPool()
+ buf := pool.Get()
+ if _, err := buf.WriteString("junk to reset"); err != nil {
+ t.Fatal(err)
+ }
+ pool.Put(buf)
+ again := pool.Get()
+ defer pool.Put(again)
+ if again.Len() != 0 {
+ t.Fatalf("reused buffer not reset: Len()==%d, want 0", again.Len())
+ }
+ })
+
+ t.Run("nil message writes nothing", func(t *testing.T) {
+ for _, c := range regressionCodecList() {
+ w, out := newEnvelopeWriterForTest(t, c.codec, false, 0, 0)
+ if err := w.Marshal(nil); err != nil {
+ t.Fatalf("%s path nil message: %v", c.name, err)
+ }
+ if out.Len() != 0 {
+ t.Fatalf("%s path nil message wrote %d bytes, want 0", c.name, out.Len())
+ }
+ }
+ })
+
+ t.Run("empty proto yields zero-length envelope", func(t *testing.T) {
+ // compressMinBytes=1: an empty payload (0 < 1) must not be compressed,
+ // so the envelope is exactly a 5-byte prefix announcing length 0.
+ for _, c := range regressionCodecList() {
+ w, out := newEnvelopeWriterForTest(t, c.codec, true, 1, 0)
+ msg := &pingv1.PingRequest{Text: ""}
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("%s path: %v", c.name, err)
+ }
+ flags, length := envelopePrefix(t, out.Bytes())
+ if flags != 0 {
+ t.Fatalf("%s path: empty payload compressed (flags=%#x), want 0", c.name, flags)
+ }
+ if length != 0 {
+ t.Fatalf("%s path: envelope length %d, want 0", c.name, length)
+ }
+ if out.Len() != 5 {
+ t.Fatalf("%s path: wire length %d, want exactly 5 (prefix only)", c.name, out.Len())
+ }
+ }
+ })
+
+ t.Run("compressMinBytes boundary does not panic", func(t *testing.T) {
+ // The compression decision uses the codec's encoded output length: a
+ // threshold equal to the encoded length compresses, one byte more does not.
+ codec := &protoBinaryCodec{}
+ for _, size := range []int{511, 512, 513} {
+ msg := newMarshalPerfMessage(size)
+ raw, err := codec.Marshal(msg)
+ if err != nil {
+ t.Fatal(err)
+ }
+ encodedLen := len(raw)
+ for _, minBytes := range []int{encodedLen, encodedLen + 1} {
+ wantCompressed := minBytes <= encodedLen
+ for _, c := range regressionCodecList() {
+ w, out := newEnvelopeWriterForTest(t, c.codec, true, minBytes, 0)
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("%s path size %d minBytes %d: %v", c.name, size, minBytes, err)
+ }
+ flags, _ := envelopePrefix(t, out.Bytes())
+ gotCompressed := flags&flagEnvelopeCompressed != 0
+ if gotCompressed != wantCompressed {
+ t.Fatalf("%s path size %d (encoded %d) minBytes %d: compressed=%v, want %v",
+ c.name, size, encodedLen, minBytes, gotCompressed, wantCompressed)
+ }
+ }
+ }
+ }
+ })
+}
+
+// errorCodec implements Codec and marshalAppender, always failing. It forces
+// the envelopeWriter fast path (MarshalAppend) and slow path (Marshal) into
+// their backup-codec fallback branches.
+type errorCodec struct{}
+
+var _ Codec = errorCodec{}
+var _ marshalAppender = errorCodec{}
+
+func (errorCodec) Name() string { return "always-error" }
+
+func (errorCodec) Marshal(any) ([]byte, error) {
+ return nil, errors.New("primary codec marshal failed")
+}
+
+func (errorCodec) MarshalAppend(dst []byte, _ any) ([]byte, error) {
+ return nil, errors.New("primary codec marshal failed")
+}
+
+func (errorCodec) Unmarshal([]byte, any) error {
+ return errors.New("primary codec unmarshal failed")
+}
+
+// TestMarshalPerfBackupCodecFallback verifies the backup-codec fallback on the
+// fast (MarshalAppend) and slow (Marshal) paths.
+func TestMarshalPerfBackupCodecFallback(t *testing.T) {
+ msg := &pingv1.PingRequest{Text: "fallback"}
+ // Reference wire bytes produced by the healthy codec.
+ var refBytes []byte
+ {
+ ref, out := newEnvelopeWriterForTest(t, &protoBinaryCodec{}, false, 0, 0)
+ if err := ref.Marshal(msg); err != nil {
+ t.Fatal(err)
+ }
+ refBytes = out.Bytes()
+ }
+
+ t.Run("primary fails fast path falls back", func(t *testing.T) {
+ // errorCodec implements marshalAppender, so the fast path is taken;
+ // on failure it must fall back to the backup codec and produce the
+ // same wire bytes as the healthy codec.
+ w, out := newEnvelopeWriterForTest(t, errorCodec{}, false, 0, 0)
+ w.backupCodec = &protoBinaryCodec{}
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("fast fallback: %v", err)
+ }
+ if got := out.Bytes(); !bytes.Equal(got, refBytes) {
+ t.Fatalf("fast fallback wire mismatch: got %d bytes, want %d", len(got), len(refBytes))
+ }
+ })
+
+ t.Run("primary fails slow path falls back", func(t *testing.T) {
+ w, out := newEnvelopeWriterForTest(t, &noAppenderCodec{errorCodec{}}, false, 0, 0)
+ w.backupCodec = &protoBinaryCodec{}
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("slow fallback: %v", err)
+ }
+ if got := out.Bytes(); !bytes.Equal(got, refBytes) {
+ t.Fatalf("slow fallback wire mismatch: got %d bytes, want %d", len(got), len(refBytes))
+ }
+ })
+
+ t.Run("write failure with healthy appender does not fall back", func(t *testing.T) {
+ // The fast path falls back only on a marshal failure: a write error
+ // must surface as-is, without consulting the backup codec or retrying.
+ backup := &countingNamedCodec{Codec: &protoBinaryCodec{}, name: "hessian2"}
+ out := &failingWriter{}
+ w := &envelopeWriter{
+ writer: out,
+ codec: &protoBinaryCodec{}, // implements marshalAppender
+ backupCodec: backup,
+ bufferPool: newBufferPool(),
+ }
+ wantCode(t, w.Marshal(msg), CodeUnknown)
+ if out.writes != 1 {
+ t.Fatalf("writer called %d times, want 1 (write failure must not retry)", out.writes)
+ }
+ if backup.marshalCalls != 0 {
+ t.Fatalf("backup codec Marshal called %d times on a write failure, want 0", backup.marshalCalls)
+ }
+ })
+
+ t.Run("same name does not double fallback", func(t *testing.T) {
+ // A backup sharing the codec's name must not be consulted, so the error
+ // surfaces on both the fast and slow legs.
+ legs := []struct {
+ name string
+ codec Codec
+ }{
+ {name: "fast", codec: errorCodec{}},
+ {name: "slow", codec: &noAppenderCodec{errorCodec{}}},
+ }
+ for _, leg := range legs {
+ w, _ := newEnvelopeWriterForTest(t, leg.codec, false, 0, 0)
+ w.backupCodec = &failingNamedCodec{name: "always-error"}
+ if err := w.Marshal(msg); err == nil {
+ t.Fatalf("%s path: expected error when backup shares codec name, got nil", leg.name)
+ } else if asErr, ok := asError(err); ok && asErr.Code() != CodeInternal {
+ t.Fatalf("%s path: expected CodeInternal, got %v", leg.name, asErr.Code())
+ }
+ }
+ })
+
+ t.Run("nil backup returns internal error", func(t *testing.T) {
+ for _, c := range regressionCodecList() {
+ w, _ := newEnvelopeWriterForTest(t, errorCodec{}, false, 0, 0)
+ if c.name == "slow" {
+ w.codec = &noAppenderCodec{errorCodec{}}
+ }
+ w.backupCodec = nil
+ err := w.Marshal(msg)
+ if err == nil {
+ t.Fatalf("%s path: expected error with nil backup, got nil", c.name)
+ }
+ wantCode(t, err, CodeInternal)
+ }
+ })
+}
+
+// failingNamedCodec lets tests control the Name() seen by the fallback check.
+type failingNamedCodec struct {
+ errorCodec
+ name string
+}
+
+func (c *failingNamedCodec) Name() string { return c.name }
+
+// failingWriter is an io.Writer that records how many times it was called and
+// always fails, letting tests observe spurious write retries.
+type failingWriter struct {
+ writes int
+}
+
+func (w *failingWriter) Write(p []byte) (int, error) {
+ w.writes++
+ return 0, errors.New("write failed")
+}
+
+// countingNamedCodec wraps a codec under a distinct Name and counts Marshal
+// calls, letting tests observe whether the backup-codec fallback ran.
+type countingNamedCodec struct {
+ Codec
+ name string
+ marshalCalls int
+}
+
+func (c *countingNamedCodec) Name() string { return c.name }
+
+func (c *countingNamedCodec) Marshal(message any) ([]byte, error) {
+ c.marshalCalls++
+ return c.Codec.Marshal(message)
+}
+
+// TestMarshalPerfCompressionAndMaxBytes verifies the sendMaxBytes limit and
+// compression headers on both paths.
+func TestMarshalPerfCompressionAndMaxBytes(t *testing.T) {
+ t.Parallel()
+
+ t.Run("sendMaxBytes exceeded returns ResourceExhausted on both paths", func(t *testing.T) {
+ msg := newMarshalPerfMessage(1024)
+ for _, wire := range []string{"envelope", "triple"} {
+ for _, c := range regressionCodecList() {
+ switch wire {
+ case "envelope":
+ w, _ := newEnvelopeWriterForTest(t, c.codec, false, 0, 100)
+ wantCode(t, w.Marshal(msg), CodeResourceExhausted)
+ case "triple":
+ m, _, _ := newTripleMarshalerForTest(t, c.codec, false, 0, 100)
+ wantCode(t, m.Marshal(msg), CodeResourceExhausted)
+ }
+ }
+ }
+ })
+
+ t.Run("compressed over-limit also returns ResourceExhausted", func(t *testing.T) {
+ // sendMaxBytes=1 with gzip enabled: even the compressed form of any
+ // non-empty message exceeds 1 byte, so both paths must reject it.
+ msg := newMarshalPerfMessage(64)
+ for _, wire := range []string{"envelope", "triple"} {
+ for _, c := range regressionCodecList() {
+ switch wire {
+ case "envelope":
+ w, _ := newEnvelopeWriterForTest(t, c.codec, true, 0, 1)
+ wantCode(t, w.Marshal(msg), CodeResourceExhausted)
+ case "triple":
+ m, _, _ := newTripleMarshalerForTest(t, c.codec, true, 0, 1)
+ wantCode(t, m.Marshal(msg), CodeResourceExhausted)
+ }
+ }
+ }
+ })
+
+ t.Run("compression header set identically on triple wire", func(t *testing.T) {
+ msg := newMarshalPerfMessage(1024)
+ for _, c := range regressionCodecList() {
+ m, _, header := newTripleMarshalerForTest(t, c.codec, true, 512, 0)
+ if err := m.Marshal(msg); err != nil {
+ t.Fatalf("%s path: %v", c.name, err)
+ }
+ if got := header.Get(tripleUnaryHeaderCompression); got != compressionGzip {
+ t.Fatalf("%s path: compression header %q, want %q", c.name, got, compressionGzip)
+ }
+ }
+ })
+}
+
+// TestMarshalPerfLargeBufferDropped verifies that buffers larger than
+// maxRecycleBufferSize are dropped rather than recycled.
+func TestMarshalPerfLargeBufferDropped(t *testing.T) {
+ t.Parallel()
+
+ t.Run("pool refuses to recycle oversized buffer", func(t *testing.T) {
+ pool := newBufferPool()
+ big := pool.Get()
+ data := make([]byte, 8*1024*1024+1)
+ if _, err := big.Write(data); err != nil {
+ t.Fatal(err)
+ }
+ if big.Cap() <= maxRecycleBufferSize {
+ t.Fatalf("test setup: buffer cap %d not larger than %d", big.Cap(), maxRecycleBufferSize)
+ }
+ pool.Put(big)
+ next := pool.Get()
+ defer pool.Put(next)
+ if next.Cap() > maxRecycleBufferSize {
+ t.Fatalf("pool recycled an oversized buffer: cap %d > max %d", next.Cap(), maxRecycleBufferSize)
+ }
+ if next.Len() != 0 {
+ t.Fatalf("pool returned non-empty buffer: Len()==%d", next.Len())
+ }
+ })
+
+ t.Run("large message does not poison subsequent small message", func(t *testing.T) {
+ large := newMarshalPerfMessage(8*1024*1024 + 1)
+ small := newMarshalPerfMessage(1)
+ for _, wire := range []string{"envelope", "triple"} {
+ for _, c := range regressionCodecList() {
+ // newMarshalTo returns a marshal function whose marshaler is
+ // bound to the given buffer pool, so callers control sharing.
+ newMarshalTo := func(pool *bufferPool) func(*pingv1.PingRequest) ([]byte, error) {
+ out := &bytes.Buffer{}
+ switch wire {
+ case "envelope":
+ w := &envelopeWriter{writer: out, codec: c.codec, bufferPool: pool}
+ return func(msg *pingv1.PingRequest) ([]byte, error) {
+ out.Reset()
+ if err := w.Marshal(msg); err != nil {
+ return nil, err
+ }
+ return out.Bytes(), nil
+ }
+ default: // triple
+ m := &tripleUnaryMarshaler{writer: out, codec: c.codec, bufferPool: pool}
+ return func(msg *pingv1.PingRequest) ([]byte, error) {
+ out.Reset()
+ if err := m.Marshal(msg); err != nil {
+ return nil, err
+ }
+ return out.Bytes(), nil
+ }
+ }
+ }
+ // Marshal the huge message first, then the tiny one through the
+ // same pool; recycled residue would leak into the small output.
+ shared := newMarshalTo(newBufferPool())
+ if _, err := shared(large); err != nil {
+ t.Fatalf("%s/%s large marshal: %v", wire, c.name, err)
+ }
+ afterLarge, err := shared(small)
+ if err != nil {
+ t.Fatalf("%s/%s small marshal after large: %v", wire, c.name, err)
+ }
+ // The tiny message's wire bytes must be exactly the same as if
+ // it had never been preceded by an 8MiB+1 message. The reference
+ // is marshaled on an independent, uncontaminated pool.
+ expected, err := newMarshalTo(newBufferPool())(small)
+ if err != nil {
+ t.Fatalf("%s/%s reference small marshal: %v", wire, c.name, err)
+ }
+ if !bytes.Equal(afterLarge, expected) {
+ t.Fatalf("%s/%s small message output changed after 8MiB+1 message: got %d bytes, want %d",
+ wire, c.name, len(afterLarge), len(expected))
+ }
+ if len(afterLarge) == 0 {
+ t.Fatalf("%s/%s small message produced empty output", wire, c.name)
+ }
+ }
+ }
+ })
+}
+
+// TestMarshalPerfConcurrentSend verifies bufferPool safety under concurrent use.
+func TestMarshalPerfConcurrentSend(t *testing.T) {
+ for _, c := range regressionCodecList() {
+ t.Run(c.name, func(t *testing.T) {
+ // Drive concurrent envelopeWriters sharing one bufferPool. Each
+ // goroutine sends a globally unique payload per message, then
+ // decodes every emitted frame and verifies its content. Distinct
+ // content plus per-frame decode prevents cross-goroutine corruption
+ // from going unnoticed without -race: overwritten frames keep the
+ // same length, which a length-only check would not catch.
+ sharedPool := newBufferPool()
+ const goroutines = 32
+ const iters = 100
+ var wg sync.WaitGroup
+ errCh := make(chan error, goroutines)
+ // Start gate: every goroutine blocks until all 32 are ready. This
+ // keeps the pool actually contended and prevents the outcome from
+ // depending on the scheduler interleaving the goroutines.
+ var ready sync.WaitGroup
+ ready.Add(goroutines)
+ start := make(chan struct{})
+ for g := range goroutines {
+ wg.Go(func() {
+ ready.Done()
+ <-start
+ out := &syncBuffer{}
+ w := &envelopeWriter{
+ writer: out,
+ codec: c.codec,
+ bufferPool: sharedPool,
+ }
+ // Number uniquely identifies (goroutine, iteration) across
+ // the whole test, so content swapped between goroutines is
+ // reported, not silently accepted.
+ want := make([]*pingv1.PingRequest, 0, iters)
+ for i := range iters {
+ msg := &pingv1.PingRequest{
+ Number: int64(g*iters + i),
+ Text: fmt.Sprintf("goroutine-%02d-message-%03d", g, i),
+ }
+ want = append(want, msg)
+ if err := w.Marshal(msg); err != nil {
+ errCh <- fmt.Errorf("goroutine %d marshal message %d: %w", g, i, err)
+ return
+ }
+ }
+ if err := verifyEnvelopeFrames(out.Bytes(), c.codec, want); err != nil {
+ errCh <- fmt.Errorf("goroutine %d: %w", g, err)
+ }
+ })
+ }
+ ready.Wait()
+ close(start)
+ wg.Wait()
+ close(errCh)
+ for err := range errCh {
+ t.Fatal(err)
+ }
+ })
+ }
+}
+
+// verifyEnvelopeFrames walks the gRPC/Triple envelope stream written by
+// envelopeWriter and checks that every frame unmarshals to the corresponding
+// expected message. The frames must be uncompressed (flags == 0) and exactly
+// account for the whole stream.
+func verifyEnvelopeFrames(wire []byte, codec Codec, want []*pingv1.PingRequest) error {
+ pos := 0
+ for n, expected := range want {
+ if len(wire)-pos < 5 {
+ return fmt.Errorf("frame %d: truncated envelope prefix (%d bytes left)", n, len(wire)-pos)
+ }
+ if flags := wire[pos]; flags != 0 {
+ return fmt.Errorf("frame %d: unexpected flags 0x%x (only uncompressed frames expected)", n, flags)
+ }
+ length := int(binary.BigEndian.Uint32(wire[pos+1 : pos+5]))
+ pos += 5
+ if len(wire)-pos < length {
+ return fmt.Errorf("frame %d: payload length %d exceeds remaining %d bytes", n, length, len(wire)-pos)
+ }
+ payload := wire[pos : pos+length]
+ pos += length
+ var got pingv1.PingRequest
+ if err := codec.Unmarshal(payload, &got); err != nil {
+ return fmt.Errorf("frame %d: unmarshal: %w", n, err)
+ }
+ if got.Number != expected.Number || got.Text != expected.Text {
+ return fmt.Errorf(
+ "frame %d: content mismatch: got Number=%d Text=%q, want Number=%d Text=%q",
+ n, got.Number, got.Text, expected.Number, expected.Text,
+ )
+ }
+ }
+ if pos != len(wire) {
+ return fmt.Errorf("%d trailing bytes after %d frames", len(wire)-pos, len(want))
+ }
+ return nil
+}
+
+// Type guard: only protoBinaryCodec and the tripleServerCodecSession wrapper
+// may expose the marshalAppender fast path; all other codecs must stay on
+// codec.Marshal.
+func TestMarshalPerfTypeGuard(t *testing.T) {
+ t.Parallel()
+
+ implements := func(c Codec) bool {
+ _, ok := c.(marshalAppender)
+ return ok
+ }
+
+ if !implements(&protoBinaryCodec{}) {
+ t.Fatal("protoBinaryCodec must implement marshalAppender")
+ }
+ // tripleServerCodecSession wraps the proto codec on the server response
+ // path; it must reach the fast path too, or the optimization is skipped
+ // for every native triple server response.
+ if !implements(&tripleServerCodecSession{delegate: &protoBinaryCodec{}}) {
+ t.Fatal("tripleServerCodecSession must implement marshalAppender (server response fast path)")
+ }
+
+ never := []struct {
+ name string
+ codec Codec
+ }{
+ {name: "noAppender wrapper", codec: &noAppenderCodec{&protoBinaryCodec{}}},
+ {name: "proto wrapper (hessian inner)", codec: newProtoWrapperCodec(&hessian2Codec{})},
+ {name: "hessian2", codec: &hessian2Codec{}},
+ {name: "msgpack", codec: &msgpackCodec{}},
+ {name: "json", codec: &protoJSONCodec{name: codecNameJSON}},
+ }
+ for _, tc := range never {
+ if implements(tc.codec) {
+ t.Fatalf("%s must NOT implement marshalAppender (fast path scope leak)", tc.name)
+ }
+ }
+}
+
+// TestMarshalPerfEnvelopeWriterFastPath verifies that envelopeWriter.Marshal
+// really takes the MarshalAppend fast path when the codec implements
+// marshalAppender. fastPathProbeCodec.Marshal always fails, so only the
+// appender branch can produce a successful write; the probe counts the
+// appender calls to make the branch observable.
+func TestMarshalPerfEnvelopeWriterFastPath(t *testing.T) {
+ probe := &fastPathProbeCodec{}
+ w, out := newEnvelopeWriterForTest(t, probe, false, 0, 0)
+ if err := w.Marshal(&pingv1.PingRequest{Text: "fast-path"}); err != nil {
+ t.Fatalf("envelopeWriter.Marshal with appender codec failed: %v", err)
+ }
+ if probe.appendCalls != 1 {
+ t.Fatalf("MarshalAppend called %d times, want 1 (fast path not taken)", probe.appendCalls)
+ }
+ if out.Len() == 0 {
+ t.Fatal("fast path wrote no bytes")
+ }
+}
+
+// TestMarshalPerfErrorGuard verifies that non-proto messages never panic:
+// MarshalAppend returns errNotProto, and the envelope and triple marshalers
+// surface CodeInternal on both the fast and slow paths.
+func TestMarshalPerfErrorGuard(t *testing.T) {
+ t.Parallel()
+
+ nonProto := any("definitely not a proto message")
+
+ t.Run("MarshalAppend rejects non-proto", func(t *testing.T) {
+ c := &protoBinaryCodec{}
+ _, err := c.MarshalAppend(make([]byte, 0, 16), nonProto)
+ if err == nil {
+ t.Fatal("expected errNotProto from MarshalAppend, got nil")
+ }
+ })
+
+ t.Run("marshalers return CodeInternal without panicking", func(t *testing.T) {
+ for _, wire := range []string{"envelope", "triple"} {
+ for _, c := range regressionCodecList() {
+ var err *Error
+ switch wire {
+ case "envelope":
+ w, _ := newEnvelopeWriterForTest(t, c.codec, false, 0, 0)
+ err = w.Marshal(nonProto)
+ case "triple":
+ m, _, _ := newTripleMarshalerForTest(t, c.codec, false, 0, 0)
+ err = m.Marshal(nonProto)
+ }
+ wantCode(t, err, CodeInternal)
+ }
+ }
+ })
+
+ t.Run("primary error with nil backup is CodeInternal", func(t *testing.T) {
+ // Already exercised in the fallback test; assert the plain error text
+ // does not swallow the underlying errNotProto cause.
+ c := &protoBinaryCodec{}
+ _, err := c.MarshalAppend(nil, nonProto)
+ if err == nil || !strings.Contains(err.Error(), "doesn't implement proto.Message") {
+ t.Fatalf("expected errNotProto naming proto.Message, got: %v", err)
+ }
+ })
+}
+
+// fastPathProbeCodec distinguishes the fast and slow paths: Marshal always
+// fails while MarshalAppend counts calls, so a slow-path leak fails the test.
+type fastPathProbeCodec struct {
+ appendCalls int
+}
+
+var _ Codec = (*fastPathProbeCodec)(nil)
+var _ marshalAppender = (*fastPathProbeCodec)(nil)
+
+func (c *fastPathProbeCodec) Name() string { return codecNameProto }
+func (c *fastPathProbeCodec) Marshal(any) ([]byte, error) {
+ return nil, errors.New("slow path taken: codec.Marshal must not run on the marshalAppender fast path")
+}
+func (c *fastPathProbeCodec) Unmarshal([]byte, any) error { return nil }
+func (c *fastPathProbeCodec) MarshalAppend(dst []byte, _ any) ([]byte, error) {
+ c.appendCalls++
+ return append(dst, "probe-payload"...), nil
+}
+
+// TestMarshalPerfServerSessionFastPath verifies that tripleServerCodecSession
+// responses reach the MarshalAppend fast path: IDL proto output stays
+// byte-identical to the naked codec, non-IDL MarshalAppend equals Marshal,
+// and delegates without the appender extension fall back gracefully.
+func TestMarshalPerfServerSessionFastPath(t *testing.T) {
+ t.Parallel()
+
+ msg := &pingv1.PingRequest{Text: "server-session"}
+
+ t.Run("IDL proto response matches naked codec wire", func(t *testing.T) {
+ // codecSession implements marshalAppender, so tripleUnaryMarshaler
+ // takes the MarshalAppend branch on the server response path.
+ fast, outFast, _ := newTripleMarshalerForTest(t, &tripleServerCodecSession{delegate: &protoBinaryCodec{}}, false, 0, 0)
+ if err := fast.Marshal(msg); err != nil {
+ t.Fatalf("session marshal: %v", err)
+ }
+ ref, outRef, _ := newTripleMarshalerForTest(t, &protoBinaryCodec{}, false, 0, 0)
+ if err := ref.Marshal(msg); err != nil {
+ t.Fatalf("reference marshal: %v", err)
+ }
+ if got, want := outFast.Bytes(), outRef.Bytes(); !bytes.Equal(got, want) {
+ t.Fatalf("session fast path wire mismatch: got %d bytes, want %d", len(got), len(want))
+ }
+ })
+
+ t.Run("Non-IDL wrapped MarshalAppend equals Marshal", func(t *testing.T) {
+ messages := []struct {
+ name string
+ msg any
+ }{
+ {name: "scalar payload", msg: "payload"},
+ {name: "one-result container", msg: []any{"payload"}},
+ {name: "void container", msg: []any{}},
+ }
+ for _, serializeType := range []string{codecNameHessian2, codecNameMsgPack} {
+ for _, mc := range messages {
+ session := &tripleServerCodecSession{delegate: &protoBinaryCodec{}, serializeType: serializeType}
+ got, err := session.MarshalAppend(nil, mc.msg)
+ if err != nil {
+ t.Fatalf("%s/%s MarshalAppend: %v", serializeType, mc.name, err)
+ }
+ want, err := session.Marshal(mc.msg)
+ if err != nil {
+ t.Fatalf("%s/%s Marshal: %v", serializeType, mc.name, err)
+ }
+ if !bytes.Equal(got, want) {
+ t.Fatalf("%s/%s MarshalAppend != Marshal (%d vs %d bytes)", serializeType, mc.name, len(got), len(want))
+ }
+ }
+ }
+ })
+
+ t.Run("custom proto delegate without appender falls back", func(t *testing.T) {
+ // A session over a proto codec lacking the appender extension must not
+ // break: MarshalAppend degrades to marshal-plus-append, byte-equal to
+ // the naked codec's output.
+ session := &tripleServerCodecSession{delegate: &noAppenderCodec{&protoBinaryCodec{}}}
+ got, err := session.MarshalAppend(nil, msg)
+ if err != nil {
+ t.Fatalf("fallback MarshalAppend: %v", err)
+ }
+ want, err := (&protoBinaryCodec{}).Marshal(msg)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if !bytes.Equal(got, want) {
+ t.Fatalf("fallback MarshalAppend mismatch: got %d bytes, want %d", len(got), len(want))
+ }
+ })
+
+ t.Run("server response must hit MarshalAppend, not the slow path", func(t *testing.T) {
+ // The probe's Marshal always fails, so a successful Marshal proves the
+ // MarshalAppend fast path was taken.
+ for _, compress := range []bool{false, true} {
+ probe := &fastPathProbeCodec{}
+ session := &tripleServerCodecSession{delegate: probe}
+ m, _, _ := newTripleMarshalerForTest(t, session, compress, 0, 0)
+ if err := m.Marshal(msg); err != nil {
+ t.Fatalf("compressed=%v Marshal: %v (slow path leaked into server responses?)", compress, err)
+ }
+ if probe.appendCalls != 1 {
+ t.Fatalf("compressed=%v MarshalAppend calls = %d, want 1 (fast path not taken)", compress, probe.appendCalls)
+ }
+ }
+ })
+
+ t.Run("envelope writer session must hit MarshalAppend", func(t *testing.T) {
+ // The session may also sit behind an envelopeWriter (gRPC/Triple
+ // envelope wire); the same no-slow-path guarantee must hold there.
+ probe := &fastPathProbeCodec{}
+ session := &tripleServerCodecSession{delegate: probe}
+ w, _ := newEnvelopeWriterForTest(t, session, false, 0, 0)
+ if err := w.Marshal(msg); err != nil {
+ t.Fatalf("envelope Marshal: %v (slow path leaked?)", err)
+ }
+ if probe.appendCalls != 1 {
+ t.Fatalf("envelope MarshalAppend calls = %d, want 1 (fast path not taken)", probe.appendCalls)
+ }
+ })
+}
diff --git a/protocol/triple/triple_protocol/protocol_triple.go b/protocol/triple/triple_protocol/protocol_triple.go
index d5c0938..e48a018 100644
--- a/protocol/triple/triple_protocol/protocol_triple.go
+++ b/protocol/triple/triple_protocol/protocol_triple.go
@@ -505,6 +505,11 @@
if message == nil {
return m.write(nil)
}
+ // Fast path: if the codec can append to a caller-provided buffer, marshal
+ // directly into a pooled buffer to avoid the per-request allocation.
+ if appender, ok := m.codec.(marshalAppender); ok {
+ return m.marshalAndWrite(message, appender)
+ }
data, err := m.codec.Marshal(message)
if err != nil {
return errorf(CodeInternal, "marshal message: %w", err)
@@ -512,15 +517,36 @@
// Can't avoid allocating the slice, but we can reuse it.
uncompressed := bytes.NewBuffer(data)
defer m.bufferPool.Put(uncompressed)
- if len(data) < m.compressMinBytes || m.compressionPool == nil {
- if m.sendMaxBytes > 0 && len(data) > m.sendMaxBytes {
- return NewError(CodeResourceExhausted, fmt.Errorf("message size %d exceeds sendMaxBytes %d", len(data), m.sendMaxBytes))
+ return m.compressAndWrite(uncompressed)
+}
+
+// marshalAndWrite serializes message into a pooled *bytes.Buffer, compressing if
+// necessary, and writes it.
+func (m *tripleUnaryMarshaler) marshalAndWrite(message any, appender marshalAppender) *Error {
+ buffer, err := marshalToPool(m.bufferPool, appender, message)
+ if err != nil {
+ return err
+ }
+ defer m.bufferPool.Put(buffer)
+ return m.compressAndWrite(buffer)
+}
+
+// compressAndWrite enforces the compression threshold and the sendMaxBytes
+// limit on the marshaled bytes in buffer, sets the compression header when the
+// payload is compressed, and writes the result. It is the shared tail of both
+// Marshal (slow path, bytes produced by codec.Marshal) and marshalAndWrite
+// (fast path, bytes produced by marshalAppender.MarshalAppend), so the
+// compression policy can never drift between the two paths.
+func (m *tripleUnaryMarshaler) compressAndWrite(buffer *bytes.Buffer) *Error {
+ if buffer.Len() < m.compressMinBytes || m.compressionPool == nil {
+ if m.sendMaxBytes > 0 && buffer.Len() > m.sendMaxBytes {
+ return NewError(CodeResourceExhausted, fmt.Errorf("message size %d exceeds sendMaxBytes %d", buffer.Len(), m.sendMaxBytes))
}
- return m.write(data)
+ return m.write(buffer.Bytes())
}
compressed := m.bufferPool.Get()
defer m.bufferPool.Put(compressed)
- if err := m.compressionPool.Compress(compressed, uncompressed); err != nil {
+ if err := m.compressionPool.Compress(compressed, buffer); err != nil {
return err
}
if m.sendMaxBytes > 0 && compressed.Len() > m.sendMaxBytes {
