1 // Package buffer implements a buffer for serialization, consisting of a chain of []byte-s to
2 // reduce copying and to allow reuse of individual chunks.
10 // PoolConfig contains configuration for the allocation and reuse strategy.
11 type PoolConfig struct {
12 StartSize int // Minimum chunk size that is allocated.
13 PooledSize int // Minimum chunk size that is reused, reusing chunks too small will result in overhead.
14 MaxSize int // Maximum chunk size that will be allocated.
17 var config = PoolConfig{
23 // Reuse pool: chunk size -> pool.
24 var buffers = map[int]*sync.Pool{}
27 for l := config.PooledSize; l <= config.MaxSize; l *= 2 {
28 buffers[l] = new(sync.Pool)
36 // Init sets up a non-default pooling and allocation strategy. Should be run before serialization is done.
37 func Init(cfg PoolConfig) {
42 // putBuf puts a chunk to reuse pool if it can be reused.
43 func putBuf(buf []byte) {
45 if size < config.PooledSize {
48 if c := buffers[size]; c != nil {
53 // getBuf gets a chunk from reuse pool or creates a new one if reuse failed.
54 func getBuf(size int) []byte {
55 if size < config.PooledSize {
56 return make([]byte, 0, size)
59 if c := buffers[size]; c != nil {
65 return make([]byte, 0, size)
68 // Buffer is a buffer optimized for serialization without extra copying.
71 // Buf is the current chunk that can be used for serialization.
78 // EnsureSpace makes sure that the current chunk contains at least s free bytes,
79 // possibly creating a new chunk.
80 func (b *Buffer) EnsureSpace(s int) {
81 if cap(b.Buf)-len(b.Buf) >= s {
86 if cap(b.toPool) != cap(b.Buf) {
87 // Chunk was reallocated, toPool can be pooled.
91 b.bufs = make([][]byte, 0, 8)
93 b.bufs = append(b.bufs, b.Buf)
99 if l > config.MaxSize {
106 // AppendByte appends a single byte to buffer.
107 func (b *Buffer) AppendByte(data byte) {
108 if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
111 b.Buf = append(b.Buf, data)
114 // AppendBytes appends a byte slice to buffer.
115 func (b *Buffer) AppendBytes(data []byte) {
117 if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
121 sz := cap(b.Buf) - len(b.Buf)
126 b.Buf = append(b.Buf, data[:sz]...)
131 // AppendBytes appends a string to buffer.
132 func (b *Buffer) AppendString(data string) {
134 if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
138 sz := cap(b.Buf) - len(b.Buf)
143 b.Buf = append(b.Buf, data[:sz]...)
148 // Size computes the size of a buffer by adding sizes of every chunk.
149 func (b *Buffer) Size() int {
151 for _, buf := range b.bufs {
157 // DumpTo outputs the contents of a buffer to a writer and resets the buffer.
158 func (b *Buffer) DumpTo(w io.Writer) (written int, err error) {
160 for _, buf := range b.bufs {
162 n, err = w.Write(buf)
169 n, err = w.Write(b.Buf)
181 // BuildBytes creates a single byte slice with all the contents of the buffer. Data is
182 // copied if it does not fit in a single chunk. You can optionally provide one byte
183 // slice as argument that it will try to reuse.
184 func (b *Buffer) BuildBytes(reuse ...[]byte) []byte {
185 if len(b.bufs) == 0 {
195 // If we got a buffer as argument and it is big enought, reuse it.
196 if len(reuse) == 1 && cap(reuse[0]) >= size {
199 ret = make([]byte, 0, size)
201 for _, buf := range b.bufs {
202 ret = append(ret, buf...)
206 ret = append(ret, b.Buf...)
216 type readCloser struct {
221 func (r *readCloser) Read(p []byte) (n int, err error) {
222 for _, buf := range r.bufs {
223 // Copy as much as we can.
224 x := copy(p[n:], buf[r.offset:])
225 n += x // Increment how much we filled.
227 // Did we empty the whole buffer?
228 if r.offset+x == len(buf) {
229 // On to the next buffer.
233 // We can release this buffer.
243 // No buffers left or nothing read?
244 if len(r.bufs) == 0 {
250 func (r *readCloser) Close() error {
251 // Release all remaining buffers.
252 for _, buf := range r.bufs {
255 // In case Close gets called multiple times.
261 // ReadCloser creates an io.ReadCloser with all the contents of the buffer.
262 func (b *Buffer) ReadCloser() io.ReadCloser {
263 ret := &readCloser{0, append(b.bufs, b.Buf)}