frostfs-sdk-go/pool/parts_buffer_pool.go

package pool

import (
	"fmt"
	"sync"
)

type PartBuffer struct {
	Buffer []byte
	len    uint64
}

type PartsBufferPool struct {
	syncPool      *sync.Pool
	limit         uint64
	maxObjectSize uint64

	mu        sync.Mutex
	available uint64
}

func NewPartBufferPool(limit uint64, maxObjectSize uint64) *PartsBufferPool {
	return &PartsBufferPool{
		limit:         limit,
		maxObjectSize: maxObjectSize,
		available:     limit,
		syncPool:      &sync.Pool{New: func() any { return make([]byte, maxObjectSize) }},
	}
}

func (p *PartsBufferPool) ParBufferSize() uint64 {
	return p.maxObjectSize
}

func (p *PartsBufferPool) GetBuffer() (*PartBuffer, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if p.maxObjectSize > p.available {
		return nil, fmt.Errorf("requested buffer size %d is greater than available: %d", p.maxObjectSize, p.available)
	}

	p.available -= p.maxObjectSize

	return &PartBuffer{
		Buffer: p.syncPool.Get().([]byte),
		len:    p.maxObjectSize,
	}, nil
}

func (p *PartsBufferPool) FreeBuffer(buff *PartBuffer) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	used := p.limit - p.available
	if buff.len > used {
		return fmt.Errorf("buffer size %d to free is greater than used: %d", buff.len, used)
	}

	p.available += buff.len
	p.syncPool.Put(buff.Buffer)

	return nil
}