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.
3package buffer
4
5import (
6	"io"
7	"sync"
8)
9
10// PoolConfig contains configuration for the allocation and reuse strategy.
11type 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.
15}
16
17var config = PoolConfig{
18	StartSize:  128,
19	PooledSize: 512,
20	MaxSize:    32768,
21}
22
23// Reuse pool: chunk size -> pool.
24var buffers = map[int]*sync.Pool{}
25
26func initBuffers() {
27	for l := config.PooledSize; l <= config.MaxSize; l *= 2 {
28		buffers[l] = new(sync.Pool)
29	}
30}
31
32func init() {
33	initBuffers()
34}
35
36// Init sets up a non-default pooling and allocation strategy. Should be run before serialization is done.
37func Init(cfg PoolConfig) {
38	config = cfg
39	initBuffers()
40}
41
42// putBuf puts a chunk to reuse pool if it can be reused.
43func putBuf(buf []byte) {
44	size := cap(buf)
45	if size < config.PooledSize {
46		return
47	}
48	if c := buffers[size]; c != nil {
49		c.Put(buf[:0])
50	}
51}
52
53// getBuf gets a chunk from reuse pool or creates a new one if reuse failed.
54func getBuf(size int) []byte {
55	if size < config.PooledSize {
56		return make([]byte, 0, size)
57	}
58
59	if c := buffers[size]; c != nil {
60		v := c.Get()
61		if v != nil {
62			return v.([]byte)
63		}
64	}
65	return make([]byte, 0, size)
66}
67
68// Buffer is a buffer optimized for serialization without extra copying.
69type Buffer struct {
70
71	// Buf is the current chunk that can be used for serialization.
72	Buf []byte
73
74	toPool []byte
75	bufs   [][]byte
76}
77
78// EnsureSpace makes sure that the current chunk contains at least s free bytes,
79// possibly creating a new chunk.
80func (b *Buffer) EnsureSpace(s int) {
81	if cap(b.Buf)-len(b.Buf) >= s {
82		return
83	}
84	l := len(b.Buf)
85	if l > 0 {
86		if cap(b.toPool) != cap(b.Buf) {
87			// Chunk was reallocated, toPool can be pooled.
88			putBuf(b.toPool)
89		}
90		if cap(b.bufs) == 0 {
91			b.bufs = make([][]byte, 0, 8)
92		}
93		b.bufs = append(b.bufs, b.Buf)
94		l = cap(b.toPool) * 2
95	} else {
96		l = config.StartSize
97	}
98
99	if l > config.MaxSize {
100		l = config.MaxSize
101	}
102	b.Buf = getBuf(l)
103	b.toPool = b.Buf
104}
105
106// AppendByte appends a single byte to buffer.
107func (b *Buffer) AppendByte(data byte) {
108	if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
109		b.EnsureSpace(1)
110	}
111	b.Buf = append(b.Buf, data)
112}
113
114// AppendBytes appends a byte slice to buffer.
115func (b *Buffer) AppendBytes(data []byte) {
116	for len(data) > 0 {
117		if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
118			b.EnsureSpace(1)
119		}
120
121		sz := cap(b.Buf) - len(b.Buf)
122		if sz > len(data) {
123			sz = len(data)
124		}
125
126		b.Buf = append(b.Buf, data[:sz]...)
127		data = data[sz:]
128	}
129}
130
131// AppendBytes appends a string to buffer.
132func (b *Buffer) AppendString(data string) {
133	for len(data) > 0 {
134		if cap(b.Buf) == len(b.Buf) { // EnsureSpace won't be inlined.
135			b.EnsureSpace(1)
136		}
137
138		sz := cap(b.Buf) - len(b.Buf)
139		if sz > len(data) {
140			sz = len(data)
141		}
142
143		b.Buf = append(b.Buf, data[:sz]...)
144		data = data[sz:]
145	}
146}
147
148// Size computes the size of a buffer by adding sizes of every chunk.
149func (b *Buffer) Size() int {
150	size := len(b.Buf)
151	for _, buf := range b.bufs {
152		size += len(buf)
153	}
154	return size
155}
156
157// DumpTo outputs the contents of a buffer to a writer and resets the buffer.
158func (b *Buffer) DumpTo(w io.Writer) (written int, err error) {
159	var n int
160	for _, buf := range b.bufs {
161		if err == nil {
162			n, err = w.Write(buf)
163			written += n
164		}
165		putBuf(buf)
166	}
167
168	if err == nil {
169		n, err = w.Write(b.Buf)
170		written += n
171	}
172	putBuf(b.toPool)
173
174	b.bufs = nil
175	b.Buf = nil
176	b.toPool = nil
177
178	return
179}
180
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.
183func (b *Buffer) BuildBytes() []byte {
184	if len(b.bufs) == 0 {
185
186		ret := b.Buf
187		b.toPool = nil
188		b.Buf = nil
189
190		return ret
191	}
192
193	ret := make([]byte, 0, b.Size())
194	for _, buf := range b.bufs {
195		ret = append(ret, buf...)
196		putBuf(buf)
197	}
198
199	ret = append(ret, b.Buf...)
200	putBuf(b.toPool)
201
202	b.bufs = nil
203	b.toPool = nil
204	b.Buf = nil
205
206	return ret
207}
208