cubefs

1
// Copyright 2014 The Go Authors. All rights reserved.
2
// Use of this source code is governed by a BSD-style
3
// license that can be found in the LICENSE file.
4

5
package http2
6

7
import (
8
	"errors"
9
	"fmt"
10
	"sync"
11
)
12

13
// Buffer chunks are allocated from a pool to reduce pressure on GC.
14
// The maximum wasted space per dataBuffer is 2x the largest size class,
15
// which happens when the dataBuffer has multiple chunks and there is
16
// one unread byte in both the first and last chunks. We use a few size
17
// classes to minimize overheads for servers that typically receive very
18
// small request bodies.
19
//
20
// TODO: Benchmark to determine if the pools are necessary. The GC may have
21
// improved enough that we can instead allocate chunks like this:
22
// make([]byte, max(16<<10, expectedBytesRemaining))
23
var (
24
	dataChunkSizeClasses = []int{
25
		1 << 10,
26
		2 << 10,
27
		4 << 10,
28
		8 << 10,
29
		16 << 10,
30
	}
31
	dataChunkPools = [...]sync.Pool{
32
		{New: func() interface{} { return make([]byte, 1<<10) }},
33
		{New: func() interface{} { return make([]byte, 2<<10) }},
34
		{New: func() interface{} { return make([]byte, 4<<10) }},
35
		{New: func() interface{} { return make([]byte, 8<<10) }},
36
		{New: func() interface{} { return make([]byte, 16<<10) }},
37
	}
38
)
39

40
func getDataBufferChunk(size int64) []byte {
41
	i := 0
42
	for ; i < len(dataChunkSizeClasses)-1; i++ {
43
		if size <= int64(dataChunkSizeClasses[i]) {
44
			break
45
		}
46
	}
47
	return dataChunkPools[i].Get().([]byte)
48
}
49

50
func putDataBufferChunk(p []byte) {
51
	for i, n := range dataChunkSizeClasses {
52
		if len(p) == n {
53
			dataChunkPools[i].Put(p)
54
			return
55
		}
56
	}
57
	panic(fmt.Sprintf("unexpected buffer len=%v", len(p)))
58
}
59

60
// dataBuffer is an io.ReadWriter backed by a list of data chunks.
61
// Each dataBuffer is used to read DATA frames on a single stream.
62
// The buffer is divided into chunks so the server can limit the
63
// total memory used by a single connection without limiting the
64
// request body size on any single stream.
65
type dataBuffer struct {
66
	chunks   [][]byte
67
	r        int   // next byte to read is chunks[0][r]
68
	w        int   // next byte to write is chunks[len(chunks)-1][w]
69
	size     int   // total buffered bytes
70
	expected int64 // we expect at least this many bytes in future Write calls (ignored if <= 0)
71
}
72

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

75
// Read copies bytes from the buffer into p.
76
// It is an error to read when no data is available.
77
func (b *dataBuffer) Read(p []byte) (int, error) {
78
	if b.size == 0 {
79
		return 0, errReadEmpty
80
	}
81
	var ntotal int
82
	for len(p) > 0 && b.size > 0 {
83
		readFrom := b.bytesFromFirstChunk()
84
		n := copy(p, readFrom)
85
		p = p[n:]
86
		ntotal += n
87
		b.r += n
88
		b.size -= n
89
		// If the first chunk has been consumed, advance to the next chunk.
90
		if b.r == len(b.chunks[0]) {
91
			putDataBufferChunk(b.chunks[0])
92
			end := len(b.chunks) - 1
93
			copy(b.chunks[:end], b.chunks[1:])
94
			b.chunks[end] = nil
95
			b.chunks = b.chunks[:end]
96
			b.r = 0
97
		}
98
	}
99
	return ntotal, nil
100
}
101

102
func (b *dataBuffer) bytesFromFirstChunk() []byte {
103
	if len(b.chunks) == 1 {
104
		return b.chunks[0][b.r:b.w]
105
	}
106
	return b.chunks[0][b.r:]
107
}
108

109
// Len returns the number of bytes of the unread portion of the buffer.
110
func (b *dataBuffer) Len() int {
111
	return b.size
112
}
113

114
// Write appends p to the buffer.
115
func (b *dataBuffer) Write(p []byte) (int, error) {
116
	ntotal := len(p)
117
	for len(p) > 0 {
118
		// If the last chunk is empty, allocate a new chunk. Try to allocate
119
		// enough to fully copy p plus any additional bytes we expect to
120
		// receive. However, this may allocate less than len(p).
121
		want := int64(len(p))
122
		if b.expected > want {
123
			want = b.expected
124
		}
125
		chunk := b.lastChunkOrAlloc(want)
126
		n := copy(chunk[b.w:], p)
127
		p = p[n:]
128
		b.w += n
129
		b.size += n
130
		b.expected -= int64(n)
131
	}
132
	return ntotal, nil
133
}
134

135
func (b *dataBuffer) lastChunkOrAlloc(want int64) []byte {
136
	if len(b.chunks) != 0 {
137
		last := b.chunks[len(b.chunks)-1]
138
		if b.w < len(last) {
139
			return last
140
		}
141
	}
142
	chunk := getDataBufferChunk(want)
143
	b.chunks = append(b.chunks, chunk)
144
	b.w = 0
145
	return chunk
146
}
147