69 lines
1.6 KiB
Go
69 lines
1.6 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 {
|
|
// We have to use pointer (even for slices), see https://staticcheck.dev/docs/checks/#SA6002
|
|
// It's based on interfaces implementation in 2016, so maybe something has changed since then.
|
|
// We can use no pointer for multi-kilobyte slices though https://github.com/golang/go/issues/16323#issuecomment-254401036
|
|
buff := make([]byte, maxObjectSize)
|
|
return &buff
|
|
}},
|
|
}
|
|
}
|
|
|
|
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()
|
|
|
|
if buff.len+p.available > p.limit {
|
|
return fmt.Errorf("buffer size %d to free is too large, available: %d, limit: %d", buff.len, p.available, p.limit)
|
|
}
|
|
|
|
p.available += buff.len
|
|
p.syncPool.Put(&buff.Buffer)
|
|
|
|
return nil
|
|
}
|