Source: databuffer.go in package golang.org/x/net/http2


package http2

import (
	"errors"
	"fmt"
	"sync"
)

Buffer chunks are allocated from a pool to reduce pressure on GC. The maximum wasted space per dataBuffer is 2x the largest size class, which happens when the dataBuffer has multiple chunks and there is one unread byte in both the first and last chunks. We use a few size classes to minimize overheads for servers that typically receive very small request bodies. TODO: Benchmark to determine if the pools are necessary. The GC may have improved enough that we can instead allocate chunks like this: make([]byte, max(16<<10, expectedBytesRemaining))

var (
	dataChunkSizeClasses = []int{
		1 << 10,
		2 << 10,
		4 << 10,
		8 << 10,
		16 << 10,
	}
	dataChunkPools = [...]sync.Pool{
		{New: func() interface{} { return make([]byte, 1<<10) }},
		{New: func() interface{} { return make([]byte, 2<<10) }},
		{New: func() interface{} { return make([]byte, 4<<10) }},
		{New: func() interface{} { return make([]byte, 8<<10) }},
		{New: func() interface{} { return make([]byte, 16<<10) }},
	}
)

func getDataBufferChunk(size int64) []byte {
	i := 0
	for ; i < len(dataChunkSizeClasses)-1; i++ {
		if size <= int64(dataChunkSizeClasses[i]) {
			break
		}
	}
	return dataChunkPools[i].Get().([]byte)
}

func putDataBufferChunk(p []byte) {
	for i, n := range dataChunkSizeClasses {
		if len(p) == n {
			dataChunkPools[i].Put(p)
			return
		}
	}
	panic(fmt.Sprintf("unexpected buffer len=%v", len(p)))
}

dataBuffer is an io.ReadWriter backed by a list of data chunks. Each dataBuffer is used to read DATA frames on a single stream. The buffer is divided into chunks so the server can limit the total memory used by a single connection without limiting the request body size on any single stream.

type dataBuffer struct {
	chunks   [][]byte
	r        int   // next byte to read is chunks[0][r]
	w        int   // next byte to write is chunks[len(chunks)-1][w]
	size     int   // total buffered bytes
	expected int64 // we expect at least this many bytes in future Write calls (ignored if <= 0)
}

var errReadEmpty = errors.New("read from empty dataBuffer")

Read copies bytes from the buffer into p. It is an error to read when no data is available.

func (b *dataBuffer) Read(p []byte) (int, error) {
	if b.size == 0 {
		return 0, errReadEmpty
	}
	var ntotal int
	for len(p) > 0 && b.size > 0 {
		readFrom := b.bytesFromFirstChunk()
		n := copy(p, readFrom)
		p = p[n:]
		ntotal += n
		b.r += n

If the first chunk has been consumed, advance to the next chunk.

		if b.r == len(b.chunks[0]) {
			putDataBufferChunk(b.chunks[0])
			end := len(b.chunks) - 1
			copy(b.chunks[:end], b.chunks[1:])
			b.chunks[end] = nil
			b.chunks = b.chunks[:end]
			b.r = 0
		}
	}
	return ntotal, nil
}

func (b *dataBuffer) bytesFromFirstChunk() []byte {
	if len(b.chunks) == 1 {
		return b.chunks[0][b.r:b.w]
	}
	return b.chunks[0][b.r:]
}

Len returns the number of bytes of the unread portion of the buffer.

func (b *dataBuffer) Len() int {
	return b.size
}

Write appends p to the buffer.

func (b *dataBuffer) Write(p []byte) (int, error) {
	ntotal := len(p)

If the last chunk is empty, allocate a new chunk. Try to allocate enough to fully copy p plus any additional bytes we expect to receive. However, this may allocate less than len(p).

		want := int64(len(p))
		if b.expected > want {
			want = b.expected
		}
		chunk := b.lastChunkOrAlloc(want)
		n := copy(chunk[b.w:], p)
		p = p[n:]
		b.w += n
		b.size += n
		b.expected -= int64(n)
	}
	return ntotal, nil
}

func (b *dataBuffer) lastChunkOrAlloc(want int64) []byte {
	if len(b.chunks) != 0 {
		last := b.chunks[len(b.chunks)-1]
		if b.w < len(last) {
			return last
		}
	}
	chunk := getDataBufferChunk(want)
	b.chunks = append(b.chunks, chunk)
	b.w = 0
	return chunk