forked from TrueCloudLab/frostfs-sdk-go
64 lines
1.3 KiB
Go
64 lines
1.3 KiB
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
|
|
}
|