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