148 lines
3.3 KiB
Go
148 lines
3.3 KiB
Go
// Copyright 2021 PingCAP, Inc.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package membuf
|
|
|
|
const bigValueSize = 1 << 16 // 64K
|
|
|
|
var allocBufLen = 1 << 20 // 1M
|
|
|
|
// Allocator is the abstract interface for allocating and freeing memory.
|
|
type Allocator interface {
|
|
Alloc(n int) []byte
|
|
Free([]byte)
|
|
}
|
|
|
|
type stdAllocator struct{}
|
|
|
|
func (stdAllocator) Alloc(n int) []byte {
|
|
return make([]byte, n)
|
|
}
|
|
|
|
func (stdAllocator) Free(_ []byte) {}
|
|
|
|
// Pool is like `sync.Pool`, which manages memory for all bytes buffers.
|
|
//
|
|
// NOTE: we don't used a `sync.Pool` because when will sync.Pool release is depending on the
|
|
// garbage collector which always release the memory so late. Use a fixed size chan to reuse
|
|
// can decrease the memory usage to 1/3 compare with sync.Pool.
|
|
type Pool struct {
|
|
allocator Allocator
|
|
recycleCh chan []byte
|
|
}
|
|
|
|
// NewPool creates a new pool.
|
|
func NewPool(size int, allocator Allocator) *Pool {
|
|
return &Pool{
|
|
allocator: allocator,
|
|
recycleCh: make(chan []byte, size),
|
|
}
|
|
}
|
|
|
|
func (p *Pool) acquire() []byte {
|
|
select {
|
|
case b := <-p.recycleCh:
|
|
return b
|
|
default:
|
|
return p.allocator.Alloc(allocBufLen)
|
|
}
|
|
}
|
|
|
|
func (p *Pool) release(b []byte) {
|
|
select {
|
|
case p.recycleCh <- b:
|
|
default:
|
|
p.allocator.Free(b)
|
|
}
|
|
}
|
|
|
|
// NewBuffer creates a new buffer in current pool.
|
|
func (p *Pool) NewBuffer() *Buffer {
|
|
return &Buffer{pool: p, bufs: make([][]byte, 0, 128), curBufIdx: -1}
|
|
}
|
|
|
|
var globalPool = NewPool(1024, stdAllocator{})
|
|
|
|
// NewBuffer creates a new buffer in global pool.
|
|
func NewBuffer() *Buffer { return globalPool.NewBuffer() }
|
|
|
|
// Buffer represents the reuse buffer.
|
|
type Buffer struct {
|
|
pool *Pool
|
|
bufs [][]byte
|
|
curBuf []byte
|
|
curIdx int
|
|
curBufIdx int
|
|
curBufLen int
|
|
}
|
|
|
|
// addBuf adds buffer to Buffer.
|
|
func (b *Buffer) addBuf() {
|
|
if b.curBufIdx < len(b.bufs)-1 {
|
|
b.curBufIdx++
|
|
b.curBuf = b.bufs[b.curBufIdx]
|
|
} else {
|
|
buf := b.pool.acquire()
|
|
b.bufs = append(b.bufs, buf)
|
|
b.curBuf = buf
|
|
b.curBufIdx = len(b.bufs) - 1
|
|
}
|
|
|
|
b.curBufLen = len(b.curBuf)
|
|
b.curIdx = 0
|
|
}
|
|
|
|
// Reset resets the buffer.
|
|
func (b *Buffer) Reset() {
|
|
if len(b.bufs) > 0 {
|
|
b.curBuf = b.bufs[0]
|
|
b.curBufLen = len(b.bufs[0])
|
|
b.curBufIdx = 0
|
|
b.curIdx = 0
|
|
}
|
|
}
|
|
|
|
// Destroy frees all buffers.
|
|
func (b *Buffer) Destroy() {
|
|
for _, buf := range b.bufs {
|
|
b.pool.release(buf)
|
|
}
|
|
b.bufs = nil
|
|
}
|
|
|
|
// TotalSize represents the total memory size of this Buffer.
|
|
func (b *Buffer) TotalSize() int64 {
|
|
return int64(len(b.bufs) * allocBufLen)
|
|
}
|
|
|
|
// AllocBytes allocates bytes with the given length.
|
|
func (b *Buffer) AllocBytes(n int) []byte {
|
|
if n > bigValueSize {
|
|
return make([]byte, n)
|
|
}
|
|
if b.curIdx+n > b.curBufLen {
|
|
b.addBuf()
|
|
}
|
|
idx := b.curIdx
|
|
b.curIdx += n
|
|
return b.curBuf[idx:b.curIdx:b.curIdx]
|
|
}
|
|
|
|
// AddBytes adds the bytes into this Buffer.
|
|
func (b *Buffer) AddBytes(bytes []byte) []byte {
|
|
buf := b.AllocBytes(len(bytes))
|
|
copy(buf, bytes)
|
|
return buf
|
|
}
|