146 lines
3.9 KiB
Go
146 lines
3.9 KiB
Go
// Copyright 2014 The Go Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style
|
|
// license that can be found in the LICENSE file.
|
|
|
|
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
|
|
b.size -= 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)
|
|
for len(p) > 0 {
|
|
// 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
|
|
}
|