pool byte buffers on deserialize + string arena

codegen/smithy-go-codegen/src/main/java/software/amazon/smithy/go/codegen/integration/HttpProtocolUnitTestResponseGenerator.java

@@ -235,9 +235,10 @@ public void generateSerdBenchmarkIteration(GoWriter writer, String clientName) {
         writer.writeGoTemplate("""
                 resp := &smithyhttp.Response{
                     Response: &http.Response{
-                        StatusCode: c.StatusCode,
-                        Header:     c.Header.Clone(),
-                        Body:       io.NopCloser(bytes.NewReader(c.Body)),
+                        StatusCode:    c.StatusCode,
+                        Header:        c.Header.Clone(),
+                        ContentLength: int64(len(c.Body)),
+                        Body:          io.NopCloser(bytes.NewReader(c.Body)),
                     },
                 }
                 output := &$outputSymbol:T{}

internal/serde/read_payload_blob.go

@@ -0,0 +1,34 @@
+package serde
+
+import (
+	"bytes"
+	"io"
+)
+
+// maxPresize bounds how large a buffer we will allocate up front from a
+// Content-Length header. Past this point we allocate maxPresize and let the read
+// grow the buffer to fit whatever actually arrives.
+//
+// We arbitrarily choose 512K.
+const maxPresize = 512 * 1024
+
+// ReadPayloadBlob consumes the given reader into a buffer that does not come
+// from (and will never be returned to) a pool, sizing it from contentLength
+// when that is known and within bounds.
+func ReadPayloadBlob(r io.Reader, contentLength int64) (*bytes.Buffer, error) {
+	buf := &bytes.Buffer{}
+	if contentLength > 0 {
+		presize := contentLength
+		if presize > maxPresize {
+			presize = maxPresize
+		}
+
+		// + bytes.MinRead prevents pointless doubling on EOF
+		buf.Grow(int(presize) + bytes.MinRead)
+	}
+
+	if _, err := buf.ReadFrom(r); err != nil {
+		return nil, err
+	}
+	return buf, nil
+}

internal/serde/read_payload_blob_test.go

@@ -0,0 +1,38 @@
+package serde
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestReadPayloadBlobHugeContentLength(t *testing.T) {
+	body := []byte("short")
+
+	buf, err := ReadPayloadBlob(bytes.NewReader(body), 64<<30)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(buf.Bytes(), body) {
+		t.Fatalf("got %q, want %q", buf.Bytes(), body)
+	}
+
+	if limit := int64(maxPresize) + bytes.MinRead + 64<<10; int64(buf.Cap()) > limit {
+		t.Errorf("capacity %d exceeds the presize cap %d", buf.Cap(), maxPresize)
+	}
+}
+
+func TestReadPayloadBlobWrongContentLength(t *testing.T) {
+	body := make([]byte, 8192)
+	for i := range body {
+		body[i] = byte(i)
+	}
+
+	buf, err := ReadPayloadBlob(bytes.NewReader(body), 16)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(buf.Bytes(), body) {
+		t.Fatal("body truncated when Content-Length understated it")
+	}
+}

internal/serde/string_arena.go

@@ -0,0 +1,66 @@
+package serde
+
+import "unsafe"
+
+// ArenaPayloadFactor is the ratio at which we size string arenas to the payload
+// (clamped by StringArenaMin/Max)
+//
+// We limit arenas to 4k but smaller payloads likely don't need it, so we size
+// down (somewhat arbitrarily) until we hit the 512-byte floor at which we just
+// don't arena.
+const ArenaPayloadFactor = 4
+
+// The size of a string arena is bounded, but not strictly by [min, max]:
+//   - Attempting to set a capacity below MinArenaSize will have no effect as
+//     beneath this limit it is not worth it. Calls to String with a capacity below
+//     this boundary will simply allocate.
+//   - An arena WILL be capped at MaxArenaSize.
+const (
+	MinArenaSize = 512
+	MaxArenaSize = 4096
+)
+
+// StringArena batches many small string allocations into one block.
+//
+// You MUST call Reset with the desired capacity before use. Failure to do so
+// will mean any calls to String will simply allocate and you won't get the
+// performance benefit of the arena.
+type StringArena struct {
+	buf []byte
+	cap int
+}
+
+// Reset prepares the arena for a new document, sizing its next block to cap.
+func (a *StringArena) Reset(cap int) {
+	if cap > MaxArenaSize {
+		cap = MaxArenaSize
+	}
+	a.buf = nil
+	a.cap = cap
+}
+
+// String attempts to intern a copy of b in the arena, returning a string slice
+// for it.
+//
+// If the arena is at capacity, it will just allocate a new string.
+func (a *StringArena) String(b []byte) string {
+	if len(b) == 0 {
+		return ""
+	}
+
+	if a.buf == nil { // lazy init, there might be no strings at all
+		if a.cap < MinArenaSize || len(b) > a.cap {
+			return string(b) // "not worth it" fallback
+		}
+
+		a.buf = make([]byte, 0, a.cap)
+	}
+
+	if len(b) > cap(a.buf)-len(a.buf) {
+		return string(b) // "over budget" fallback
+	}
+
+	off := len(a.buf)
+	a.buf = append(a.buf, b...)
+	return unsafe.String(&a.buf[off], len(b))
+}

internal/serde/string_arena_test.go

@@ -0,0 +1,107 @@
+package serde
+
+import (
+	"strings"
+	"testing"
+	"unsafe"
+)
+
+func TestStringArena_Basic(t *testing.T) {
+	var a StringArena
+	a.Reset(MaxArenaSize)
+	inputs := []string{"a", "", "hello", strings.Repeat("x", 100), "item", "amount"}
+	got := make([]string, len(inputs))
+	for i, in := range inputs {
+		got[i] = a.String([]byte(in))
+	}
+	for i, in := range inputs {
+		if got[i] != in {
+			t.Errorf("[%d] got %q, want %q", i, got[i], in)
+		}
+	}
+}
+
+func TestStringArena_SourceMutation(t *testing.T) {
+	var a StringArena
+	a.Reset(MaxArenaSize)
+	src := []byte("original")
+	s := a.String(src)
+
+	// the byte buffer this came from goes back into a pool and then is used
+	// for another response
+	copy(src, "MUTATED!")
+
+	if s != "original" {
+		t.Errorf("string changed with its source: got %q", s)
+	}
+}
+
+func TestStringArena_Offsets(t *testing.T) {
+	var a StringArena
+	a.Reset(MaxArenaSize)
+	first := a.String([]byte("first"))
+	second := a.String([]byte("second"))
+
+	if unsafe.StringData(first) == unsafe.StringData(second) {
+		t.Fatal("distinct strings share a start pointer")
+	}
+
+	// 2 should be right after 1, in literal memory
+	gap := uintptr(unsafe.Pointer(unsafe.StringData(second))) - uintptr(unsafe.Pointer(unsafe.StringData(first)))
+	if gap != uintptr(len(first)) {
+		t.Errorf("strings not contiguous in one block: gap %d, want %d", gap, len(first))
+	}
+}
+
+func TestStringArena_Reset(t *testing.T) {
+	var a StringArena
+	a.Reset(MaxArenaSize)
+	old := make([]string, 0, 64)
+	for i := 0; i < 64; i++ {
+		old = append(old, a.String([]byte(strings.Repeat("c", i%16+1))))
+	}
+
+	a.Reset(MaxArenaSize)
+	for i := 0; i < 200; i++ {
+		_ = a.String([]byte(strings.Repeat("z", i%16+1)))
+	}
+
+	// the zs shouldn't bleed into the cs
+	for i, s := range old {
+		if want := strings.Repeat("c", i%16+1); s != want {
+			t.Fatalf("[%d] string corrupted across Reset: got %q, want %q", i, s, want)
+		}
+	}
+}
+
+func TestStringArena_MinArenaSize(t *testing.T) {
+	var a StringArena
+	a.Reset(MinArenaSize - 1) // should just not arena
+
+	first := a.String([]byte("first"))
+	second := a.String([]byte("second"))
+
+	if first != "first" || second != "second" {
+		t.Fatalf("got %q, %q", first, second)
+	}
+	if a.buf != nil {
+		t.Errorf("arena allocated a block despite cap %d < MinArenaSize %d", MinArenaSize-1, MinArenaSize)
+	}
+}
+
+func TestStringArena_MaxArenaSize(t *testing.T) {
+	var a StringArena
+	a.Reset(MaxArenaSize + 1) // should clamp
+
+	huge := strings.Repeat("h", MaxArenaSize-1) // goes into arena
+	if got := a.String([]byte(huge)); got != huge {
+		t.Error("oversized string not returned intact")
+	}
+	hugest := strings.Repeat("hh", MaxArenaSize+1) // fallback
+	if got := a.String([]byte(hugest)); got != hugest {
+		t.Error("oversizeder string not returned intact")
+	}
+	if a.cap-len(a.buf) != 1 {
+		t.Errorf("arena should've been maxed out minus one byte")
+	}
+}

internal/sync/buffer_pool.go

@@ -0,0 +1,52 @@
+package sync
+
+import (
+	"bytes"
+	"io"
+	"sync"
+)
+
+// maxBufferSize is the largest buffer that will be returned to the pool.
+//
+// A pooled buffer retains the capacity its largest occupant forced it to grow
+// to, for the life of the pool. Dropping oversized buffers bounds that
+// retention, at the cost of re-growing after an unusually large response.
+//
+// We arbitrarily choose 64K.
+const maxBufferSize = 64 * 1024
+
+func newBuffer() any {
+	return new(bytes.Buffer)
+}
+
+// BufferPool pools bytes.Buffers for reading response bodies during protocol
+// deserialization.
+type BufferPool struct {
+	pool sync.Pool
+}
+
+// NewBufferPool returns a buffer pool ready for use.
+func NewBufferPool() *BufferPool {
+	return &BufferPool{
+		pool: sync.Pool{New: newBuffer},
+	}
+}
+
+// Get consumes the given reader into a buffer from the pool, returning that
+// buffer.
+func (p *BufferPool) Get(r io.Reader) (*bytes.Buffer, error) {
+	buf := p.pool.Get().(*bytes.Buffer)
+	if _, err := buf.ReadFrom(r); err != nil {
+		return nil, err
+	}
+
+	return buf, nil
+}
+
+// Put returns the buffer if it is within the cap limit.
+func (p *BufferPool) Put(buf *bytes.Buffer) {
+	if buf.Cap() <= maxBufferSize {
+		buf.Reset()
+		p.pool.Put(buf)
+	}
+}

internal/sync/buffer_pool_test.go

@@ -0,0 +1,41 @@
+package sync
+
+import (
+	"bytes"
+	"sync"
+	"testing"
+)
+
+// exercises the pool the way protocols use it, with mixed sizes including
+// oversized that will drop
+func TestBufferPoolConcurrent(t *testing.T) {
+	p := NewBufferPool()
+
+	sizes := []int{16, 4096, 64 * 1024, maxBufferSize + 1}
+
+	var wg sync.WaitGroup
+	for g := 0; g < 16; g++ {
+		wg.Add(1)
+		go func(g int) {
+			defer wg.Done()
+			for i := 0; i < 200; i++ {
+				size := sizes[(g+i)%len(sizes)]
+				fill := byte('a' + g%26)
+				want := bytes.Repeat([]byte{fill}, size)
+
+				buf, err := p.Get(bytes.NewReader(want))
+				if err != nil {
+					t.Errorf("goroutine %d: %v", g, err)
+					return
+				}
+				if !bytes.Equal(buf.Bytes(), want) {
+					t.Errorf("goroutine %d iter %d: content mismatch for size %d", g, i, size)
+					p.Put(buf)
+					return
+				}
+				p.Put(buf)
+			}
+		}(g)
+	}
+	wg.Wait()
+}

schema_ext.go

@@ -9,13 +9,14 @@ import (
 // (JSON, CBOR, etc.) uses a distinct slot to cache precomputed data.
 type ExtensionID int
 
-const numExtensionSlots = 4
+const numExtensionSlots = 5
 
 const (
-	ExtJSON  ExtensionID = iota // transport/http/protocol/internal/json
-	ExtCBOR                     // transport/http/protocol/internal/cbor
-	ExtXML                      // transport/http/protocol/internal/xml
-	ExtQuery                    // transport/http/protocol/internal/query
+	ExtJSON        ExtensionID = iota // transport/http/protocol/internal/json
+	ExtCBOR                           // transport/http/protocol/internal/cbor
+	ExtXML                            // transport/http/protocol/internal/xml
+	ExtQuery                          // transport/http/protocol/internal/query
+	ExtHTTPBinding                    // transport/http/protocol/internal/httpbinding
 )
 
 // SchemaExtension retrieves or lazily computes the extension for the given