118 lines
2.5 KiB
Go
118 lines
2.5 KiB
Go
// Package pool implements a memory pool similar in concept to
|
|
// sync.Pool but with more determinism.
|
|
package pool
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/ncw/rclone/lib/mmap"
|
|
)
|
|
|
|
// Pool of internal buffers
|
|
type Pool struct {
|
|
cache chan []byte
|
|
bufferSize int
|
|
timer *time.Timer
|
|
inUse int32
|
|
flushTime time.Duration
|
|
alloc func(int) ([]byte, error)
|
|
free func([]byte) error
|
|
}
|
|
|
|
// New makes a buffer pool
|
|
//
|
|
// flushTime is the interval the buffer pools is flushed
|
|
// bufferSize is the size of the allocations
|
|
// poolSize is the maximum number of free buffers in the pool
|
|
// useMmap should be set to use mmap allocations
|
|
func New(flushTime time.Duration, bufferSize, poolSize int, useMmap bool) *Pool {
|
|
bp := &Pool{
|
|
cache: make(chan []byte, poolSize),
|
|
flushTime: flushTime,
|
|
bufferSize: bufferSize,
|
|
}
|
|
if useMmap {
|
|
bp.alloc = mmap.Alloc
|
|
bp.free = mmap.Free
|
|
} else {
|
|
bp.alloc = func(size int) ([]byte, error) {
|
|
return make([]byte, size), nil
|
|
}
|
|
bp.free = func([]byte) error {
|
|
return nil
|
|
}
|
|
}
|
|
bp.timer = time.AfterFunc(flushTime, bp.Flush)
|
|
return bp
|
|
}
|
|
|
|
// Flush the entire buffer pool
|
|
func (bp *Pool) Flush() {
|
|
for {
|
|
select {
|
|
case b := <-bp.cache:
|
|
bp.freeBuffer(b)
|
|
default:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// InUse returns the approximate number of buffers in use which
|
|
// haven't been returned to the pool.
|
|
func (bp *Pool) InUse() int {
|
|
return int(atomic.LoadInt32(&bp.inUse))
|
|
}
|
|
|
|
// starts or resets the buffer flusher timer
|
|
func (bp *Pool) kickFlusher() {
|
|
bp.timer.Reset(bp.flushTime)
|
|
}
|
|
|
|
// Get a buffer from the pool or allocate one
|
|
func (bp *Pool) Get() []byte {
|
|
select {
|
|
case b := <-bp.cache:
|
|
return b
|
|
default:
|
|
}
|
|
mem, err := bp.alloc(bp.bufferSize)
|
|
if err != nil {
|
|
log.Printf("Failed to get memory for buffer, waiting for a freed one: %v", err)
|
|
return <-bp.cache
|
|
}
|
|
atomic.AddInt32(&bp.inUse, 1)
|
|
return mem
|
|
}
|
|
|
|
// freeBuffer returns mem to the os if required
|
|
func (bp *Pool) freeBuffer(mem []byte) {
|
|
err := bp.free(mem)
|
|
if err != nil {
|
|
log.Printf("Failed to free memory: %v", err)
|
|
} else {
|
|
atomic.AddInt32(&bp.inUse, -1)
|
|
}
|
|
}
|
|
|
|
// Put returns the buffer to the buffer cache or frees it
|
|
//
|
|
// Note that if you try to return a buffer of the wrong size to Put it
|
|
// will panic.
|
|
func (bp *Pool) Put(mem []byte) {
|
|
mem = mem[0:cap(mem)]
|
|
if len(mem) != bp.bufferSize {
|
|
panic(fmt.Sprintf("Returning buffer sized %d but expecting %d", len(mem), bp.bufferSize))
|
|
}
|
|
select {
|
|
case bp.cache <- mem:
|
|
bp.kickFlusher()
|
|
return
|
|
default:
|
|
}
|
|
bp.freeBuffer(mem)
|
|
mem = nil
|
|
}
|