134 lines
3.4 KiB
Go
134 lines
3.4 KiB
Go
// Copyright 2017 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,
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package bytespool
|
|
|
|
import (
|
|
"bytes"
|
|
"sync"
|
|
)
|
|
|
|
// BytesPool maintains large bytes pools, used for reducing memory allocation.
|
|
// It has a slice of pools which handle different size of bytes.
|
|
// Can be safely used concurrently.
|
|
type BytesPool struct {
|
|
buckets []sync.Pool
|
|
}
|
|
|
|
const (
|
|
kilo = 1024
|
|
mega = kilo * kilo
|
|
baseSize = kilo
|
|
numBuckets = 18
|
|
maxSize = 128 * mega
|
|
)
|
|
|
|
// DefaultPool is a default BytesBool instance.
|
|
var DefaultPool = NewBytesPool()
|
|
|
|
// NewBytesPool creates a new bytes pool.
|
|
func NewBytesPool() *BytesPool {
|
|
bp := new(BytesPool)
|
|
bp.buckets = make([]sync.Pool, numBuckets)
|
|
for i := uint(0); i < numBuckets; i++ {
|
|
bp.buckets[i].New = makeNewFunc(i)
|
|
}
|
|
return bp
|
|
}
|
|
|
|
func makeNewFunc(shift uint) func() interface{} {
|
|
return func() interface{} {
|
|
return make([]byte, baseSize<<shift)
|
|
}
|
|
}
|
|
|
|
// Alloc allocates a bytes which has the size of power of two.
|
|
// The caller should keep the origin bytes and use the returned data.
|
|
// When finished using, the origin bytes should be freed to the pool.
|
|
// The allocated data may not have zero value.
|
|
func (bp *BytesPool) Alloc(size int) (origin, data []byte) {
|
|
if size > maxSize {
|
|
return nil, make([]byte, size)
|
|
}
|
|
i := bucketIdx(size)
|
|
origin = bp.buckets[i].Get().([]byte)
|
|
data = origin[:size]
|
|
return
|
|
}
|
|
|
|
// Free frees the data which should be the original bytes return by Alloc.
|
|
// It returns the bucket index of the data. returns -1 means the data is not returned to the pool.
|
|
func (bp *BytesPool) Free(origin []byte) int {
|
|
originLen := len(origin)
|
|
if originLen > maxSize || originLen < baseSize || !isPowerOfTwo(originLen) {
|
|
return -1
|
|
}
|
|
i := bucketIdx(originLen)
|
|
bp.buckets[i].Put(origin)
|
|
return i
|
|
}
|
|
|
|
func isPowerOfTwo(x int) bool {
|
|
return x&(x-1) == 0
|
|
}
|
|
|
|
func bucketIdx(size int) (i int) {
|
|
for size > baseSize {
|
|
size = (size + 1) >> 1
|
|
i++
|
|
}
|
|
return
|
|
}
|
|
|
|
// ReadCloser frees the origin bytes when Close is called.
|
|
type ReadCloser struct {
|
|
pool *BytesPool
|
|
origin []byte
|
|
buffer *bytes.Buffer
|
|
closed bool
|
|
}
|
|
|
|
// Close implements the io.ReadCloser interface.
|
|
// It frees the origin bytes allocated from the pool
|
|
func (r *ReadCloser) Close() error {
|
|
if r.closed {
|
|
return nil
|
|
}
|
|
r.pool.Free(r.origin)
|
|
r.closed = true
|
|
return nil
|
|
}
|
|
|
|
// Read implements the io.ReadCloser interface.
|
|
func (r *ReadCloser) Read(b []byte) (n int, err error) {
|
|
return r.buffer.Read(b)
|
|
}
|
|
|
|
// SharedBytes returns shared bytes of the underlying buffer.
|
|
// The returned bytes may not be used after Close is called.
|
|
// And it may not be used at the same time with Read.
|
|
func (r *ReadCloser) SharedBytes() []byte {
|
|
return r.buffer.Bytes()
|
|
}
|
|
|
|
// NewReadCloser creates a ReadCloser.
|
|
// origin should be allocated from the pool, data can be any slice of the origin.
|
|
func NewReadCloser(pool *BytesPool, origin, data []byte) *ReadCloser {
|
|
r := &ReadCloser{
|
|
pool: pool,
|
|
origin: origin,
|
|
buffer: bytes.NewBuffer(data),
|
|
}
|
|
return r
|
|
}
|