shared_buffer_pool.go

/*
 *
 * Copyright 2023 gRPC authors.
 *
 * 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 grpc

import "sync"

// SharedBufferPool is a pool of buffers that can be shared, resulting in
// decreased memory allocation. Currently, in gRPC-go, it is only utilized
// for parsing incoming messages.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
type SharedBufferPool interface {
	// Get returns a buffer with specified length from the pool.
	//
	// The returned byte slice may be not zero initialized.
	Get(length int) []byte

	// Put returns a buffer to the pool.
	Put(*[]byte)
}

// NewSharedBufferPool creates a simple SharedBufferPool with buckets
// of different sizes to optimize memory usage. This prevents the pool from
// wasting large amounts of memory, even when handling messages of varying sizes.
//
// # Experimental
//
// Notice: This API is EXPERIMENTAL and may be changed or removed in a
// later release.
func NewSharedBufferPool() SharedBufferPool {
	return &simpleSharedBufferPool{
		pools: [poolArraySize]simpleSharedBufferChildPool{
			newBytesPool(level0PoolMaxSize),
			newBytesPool(level1PoolMaxSize),
			newBytesPool(level2PoolMaxSize),
			newBytesPool(level3PoolMaxSize),
			newBytesPool(level4PoolMaxSize),
			newBytesPool(0),
		},
	}
}

// simpleSharedBufferPool is a simple implementation of SharedBufferPool.
type simpleSharedBufferPool struct {
	pools [poolArraySize]simpleSharedBufferChildPool
}

func (p *simpleSharedBufferPool) Get(size int) []byte {
	return p.pools[p.poolIdx(size)].Get(size)
}

func (p *simpleSharedBufferPool) Put(bs *[]byte) {
	p.pools[p.poolIdx(cap(*bs))].Put(bs)
}

func (p *simpleSharedBufferPool) poolIdx(size int) int {
	switch {
	case size <= level0PoolMaxSize:
		return level0PoolIdx
	case size <= level1PoolMaxSize:
		return level1PoolIdx
	case size <= level2PoolMaxSize:
		return level2PoolIdx
	case size <= level3PoolMaxSize:
		return level3PoolIdx
	case size <= level4PoolMaxSize:
		return level4PoolIdx
	default:
		return levelMaxPoolIdx
	}
}

const (
	level0PoolMaxSize = 16                     //  16  B
	level1PoolMaxSize = level0PoolMaxSize * 16 // 256  B
	level2PoolMaxSize = level1PoolMaxSize * 16 //   4 KB
	level3PoolMaxSize = level2PoolMaxSize * 16 //  64 KB
	level4PoolMaxSize = level3PoolMaxSize * 16 //   1 MB
)

const (
	level0PoolIdx = iota
	level1PoolIdx
	level2PoolIdx
	level3PoolIdx
	level4PoolIdx
	levelMaxPoolIdx
	poolArraySize
)

type simpleSharedBufferChildPool interface {
	Get(size int) []byte
	Put(any)
}

type bufferPool struct {
	sync.Pool

	defaultSize int
}

func (p *bufferPool) Get(size int) []byte {
	bs := p.Pool.Get().(*[]byte)

	if cap(*bs) < size {
		p.Pool.Put(bs)

		return make([]byte, size)
	}

	return (*bs)[:size]
}

func newBytesPool(size int) simpleSharedBufferChildPool {
	return &bufferPool{
		Pool: sync.Pool{
			New: func() any {
				bs := make([]byte, size)
				return &bs
			},
		},
		defaultSize: size,
	}
}

// nopBufferPool is a buffer pool just makes new buffer without pooling.
type nopBufferPool struct {
}

func (nopBufferPool) Get(length int) []byte {
	return make([]byte, length)
}

func (nopBufferPool) Put(*[]byte) {
}