github.com/pawelgaczynski/gain@v0.4.0-alpha.0.20230821120126-41f1e60a18da/pkg/pool/ringbuffer/ringbuffer.go

// Copyright (c) 2023 Paweł Gaczyński
// Copyright (c) 2019 Andy Pan
// Copyright (c) 2016 Aliaksandr Valialkin, VertaMedia
//
// 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.
//
// Use of this source code is governed by a MIT license that can be found
// at https://github.com/valyala/bytebufferpool/blob/master/LICENSE
package ringbuffer

import (
	"sort"
	"sync/atomic"

	"github.com/pawelgaczynski/gain/pkg/buffer/magicring"
	"github.com/pawelgaczynski/gain/pkg/pool/sync"
)

const (
	minBitSize = 12 // 2**6=64 is a CPU cache line size
	steps      = 15

	minSize = 1 << minBitSize

	calibrateCallsThreshold = 42000
	maxPercentile           = 0.95
)

// RingBuffer is the alias of ring.Buffer.
type RingBuffer = magicring.RingBuffer

// Pool represents ring-buffer pool.
//
// Distinct pools may be used for distinct types of byte buffers.
// Properly determined byte buffer types with their own pools may help to reduce
// memory waste.
type Pool struct {
	calls       [steps]uint64
	calibrating uint64

	defaultSize uint64
	maxSize     uint64

	pool sync.Pool[*RingBuffer]
}

var builtinPool = NewRingBufferPool()

// Get returns an empty byte buffer from the pool.
//
// Got byte buffer may be returned to the pool via Put call.
// This reduces the number of memory allocations required for byte buffer
// management.
func Get() *RingBuffer { return builtinPool.Get() }

// Get returns new byte buffer with zero length.
//
// The byte buffer may be returned to the pool via Put after the use
// in order to minimize GC overhead.
func (p *Pool) Get() *RingBuffer {
	v := p.pool.Get()
	if v != nil {
		return v
	}

	buffer := magicring.NewMagicBuffer(int(atomic.LoadUint64(&p.defaultSize)))

	return buffer
}

// Put returns byte buffer to the pool.
//
// RingBuffer mustn't be touched after returning it to the pool,
// otherwise, data races will occur.
func Put(b *RingBuffer) {
	b.Zeroes()
	builtinPool.Put(b)
}

// Put releases byte buffer obtained via Get to the pool.
//
// The buffer mustn't be accessed after returning to the pool.
func (p *Pool) Put(buffer *RingBuffer) {
	idx := indexRingBufferPool(buffer.Cap())
	if atomic.AddUint64(&p.calls[idx], 1) > calibrateCallsThreshold {
		p.calibrate()
	}

	maxSize := int(atomic.LoadUint64(&p.maxSize))
	if maxSize == 0 || buffer.Cap() <= maxSize {
		buffer.Reset()
		p.pool.Put(buffer)
	}
}

func (p *Pool) calibrate() {
	if !atomic.CompareAndSwapUint64(&p.calibrating, 0, 1) {
		return
	}

	callData := make(callSizes, 0, steps)

	var callsSum uint64

	for i := uint64(0); i < steps; i++ {
		calls := atomic.SwapUint64(&p.calls[i], 0)
		callsSum += calls
		callData = append(callData, callSize{
			calls: calls,
			size:  minSize << i,
		})
	}
	sort.Sort(callData)

	defaultSize := callData[0].size
	maxSize := defaultSize

	maxSum := uint64(float64(callsSum) * maxPercentile)
	callsSum = 0

	for i := 0; i < steps; i++ {
		if callsSum > maxSum {
			break
		}
		callsSum += callData[i].calls

		size := callData[i].size
		if size > maxSize {
			maxSize = size
		}
	}
	atomic.StoreUint64(&p.defaultSize, defaultSize)
	atomic.StoreUint64(&p.maxSize, maxSize)

	atomic.StoreUint64(&p.calibrating, 0)
}

func NewRingBufferPool() Pool {
	p := Pool{
		pool: sync.NewPool[*RingBuffer](),
	}
	atomic.StoreUint64(&p.defaultSize, uint64(magicring.DefaultMagicBufferSize))

	return p
}

type callSize struct {
	calls uint64
	size  uint64
}

type callSizes []callSize

func (ci callSizes) Len() int {
	return len(ci)
}

func (ci callSizes) Less(i, j int) bool {
	return ci[i].calls > ci[j].calls
}

func (ci callSizes) Swap(i, j int) {
	ci[i], ci[j] = ci[j], ci[i]
}

func indexRingBufferPool(n int) int {
	n--
	n >>= minBitSize
	idx := 0

	for n > 0 {
		n >>= 1
		idx++
	}

	if idx >= steps {
		idx = steps - 1
	}

	return idx
}