rclone/lib/pool/pool.go

// Package pool implements a memory pool similar in concept to
// sync.Pool but with more determinism.
package pool

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"time"

	"github.com/rclone/rclone/fs"
	"github.com/rclone/rclone/lib/mmap"
	"golang.org/x/sync/semaphore"
)

const (
	// BufferSize is the page size of the Global() pool
	BufferSize = 1024 * 1024
	// BufferCacheSize is the max number of buffers to keep in the cache for the Global() pool
	BufferCacheSize = 64
	// BufferCacheFlushTime is the max time to keep buffers in the Global() pool
	BufferCacheFlushTime = 5 * time.Second
)

// Pool of internal buffers
//
// We hold buffers in cache. Every time we Get or Put we update
// minFill which is the minimum len(cache) seen.
//
// Every flushTime we remove minFill buffers from the cache as they
// were not used in the previous flushTime interval.
type Pool struct {
	mu           sync.Mutex
	cache        [][]byte
	minFill      int // the minimum fill of the cache
	bufferSize   int
	poolSize     int
	timer        *time.Timer
	inUse        int
	alloced      int
	flushTime    time.Duration
	flushPending bool
	alloc        func(int) ([]byte, error)
	free         func([]byte) error
}

// totalMemory is a semaphore used to control total buffer usage of
// all Pools. It may be nil in which case the total buffer usage
// will not be controlled. It counts memory in active use, it does not
// count memory cached in the pool.
var totalMemory *semaphore.Weighted

// Make sure we initialise the totalMemory semaphore once
var totalMemoryInit sync.Once

// 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([][]byte, 0, poolSize),
		poolSize:   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
		}
	}

	// Initialise total memory limit if required
	totalMemoryInit.Do(func() {
		ci := fs.GetConfig(context.Background())

		// Set max buffer memory limiter
		if ci.MaxBufferMemory > 0 {
			totalMemory = semaphore.NewWeighted(int64(ci.MaxBufferMemory))
		}
	})

	bp.timer = time.AfterFunc(flushTime, bp.flushAged)
	return bp
}

// get gets the last buffer in bp.cache
//
// Call with mu held
func (bp *Pool) get() []byte {
	n := len(bp.cache) - 1
	buf := bp.cache[n]
	bp.cache[n] = nil // clear buffer pointer from bp.cache
	bp.cache = bp.cache[:n]
	return buf
}

// getN gets the last n buffers in bp.cache
//
// will panic if you ask for too many buffers
//
// Call with mu held
func (bp *Pool) getN(n int) [][]byte {
	i := len(bp.cache) - n
	bufs := slices.Clone(bp.cache[i:])
	bp.cache = slices.Delete(bp.cache, i, len(bp.cache))
	return bufs
}

// put puts the buffer on the end of bp.cache
//
// Call with mu held
func (bp *Pool) put(buf []byte) {
	bp.cache = append(bp.cache, buf)
}

// put puts the bufs on the end of bp.cache
//
// Call with mu held
func (bp *Pool) putN(bufs [][]byte) {
	bp.cache = append(bp.cache, bufs...)
}

// buffers returns the number of buffers in bp.ache
//
// Call with mu held
func (bp *Pool) buffers() int {
	return len(bp.cache)
}

// flush n entries from the entire buffer pool
// Call with mu held
func (bp *Pool) flush(n int) {
	for range n {
		bp.freeBuffer(bp.get())
	}
	bp.minFill = len(bp.cache)
}

// Flush the entire buffer pool
func (bp *Pool) Flush() {
	bp.mu.Lock()
	bp.flush(len(bp.cache))
	bp.mu.Unlock()
}

// Remove bp.minFill buffers
func (bp *Pool) flushAged() {
	bp.mu.Lock()
	bp.flushPending = false
	bp.flush(bp.minFill)
	// If there are still items in the cache, schedule another flush
	if len(bp.cache) != 0 {
		bp.kickFlusher()
	}
	bp.mu.Unlock()
}

// InUse returns the number of buffers in use which haven't been
// returned to the pool
func (bp *Pool) InUse() int {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	return bp.inUse
}

// InPool returns the number of buffers in the pool
func (bp *Pool) InPool() int {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	return len(bp.cache)
}

// Alloced returns the number of buffers allocated and not yet freed
func (bp *Pool) Alloced() int {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	return bp.alloced
}

// starts or resets the buffer flusher timer - call with mu held
func (bp *Pool) kickFlusher() {
	if bp.flushPending {
		return
	}
	bp.flushPending = true
	bp.timer.Reset(bp.flushTime)
}

// Make sure minFill is correct - call with mu held
func (bp *Pool) updateMinFill() {
	if len(bp.cache) < bp.minFill {
		bp.minFill = len(bp.cache)
	}
}

// acquire mem bytes of memory for the user
func (bp *Pool) acquire(mem int64) error {
	if totalMemory == nil {
		return nil
	}
	ctx := context.Background()
	return totalMemory.Acquire(ctx, mem)
}

// release mem bytes of memory from the user
func (bp *Pool) release(mem int64) {
	if totalMemory == nil {
		return
	}
	totalMemory.Release(mem)
}

// Get a buffer from the pool or allocate one
func (bp *Pool) Get() []byte {
	return bp.GetN(1)[0]
}

// GetN get n buffers atomically from the pool or allocate them
func (bp *Pool) GetN(n int) [][]byte {
	bp.mu.Lock()
	var (
		waitTime = time.Millisecond // retry time if allocation failed
		err      error              // allocation error
		buf      []byte             // allocated buffer
		bufs     [][]byte           // bufs so far
		have     int                // have this many buffers in bp.cache
		want     int                // want this many extra buffers
		acquired bool               // whether we have acquired the memory or not
	)
	for {
		acquired = false
		bp.mu.Unlock()
		err = bp.acquire(int64(bp.bufferSize) * int64(n))
		bp.mu.Lock()
		if err != nil {
			goto FAIL
		}
		acquired = true
		have = min(bp.buffers(), n)
		want = n - have
		bufs = bp.getN(have) // get as many buffers as we have from the cache
		for range want {
			buf, err = bp.alloc(bp.bufferSize)
			if err != nil {
				goto FAIL
			}
			bp.alloced++
			bufs = append(bufs, buf)
		}
		break
	FAIL:
		// Release the buffers and the allocation if it succeeded
		bp.putN(bufs)
		if acquired {
			bp.release(int64(bp.bufferSize) * int64(n))
		}
		fs.Logf(nil, "Failed to get memory for buffer, waiting for %v: %v", waitTime, err)
		bp.mu.Unlock()
		time.Sleep(waitTime)
		bp.mu.Lock()
		waitTime *= 2
		clear(bufs)
		bufs = nil
	}
	bp.inUse += n
	bp.updateMinFill()
	bp.mu.Unlock()
	return bufs
}

// freeBuffer returns mem to the os if required - call with lock held
func (bp *Pool) freeBuffer(mem []byte) {
	err := bp.free(mem)
	if err != nil {
		fs.Logf(nil, "Failed to free memory: %v", err)
	}
	bp.alloced--
}

// _put returns the buffer to the buffer cache or frees it
//
// call with lock held
//
// Note that if you try to return a buffer of the wrong size it will
// panic.
func (bp *Pool) _put(buf []byte) {
	buf = buf[0:cap(buf)]
	if len(buf) != bp.bufferSize {
		panic(fmt.Sprintf("Returning buffer sized %d but expecting %d", len(buf), bp.bufferSize))
	}
	if len(bp.cache) < bp.poolSize {
		bp.put(buf)
	} else {
		bp.freeBuffer(buf)
	}
	bp.release(int64(bp.bufferSize))
}

// 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(buf []byte) {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	bp._put(buf)
	bp.inUse--
	bp.updateMinFill()
	bp.kickFlusher()
}

// PutN returns the buffers to the buffer cache or frees it,
//
// Note that if you try to return a buffer of the wrong size to PutN it
// will panic.
func (bp *Pool) PutN(bufs [][]byte) {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	for _, buf := range bufs {
		bp._put(buf)
	}
	bp.inUse -= len(bufs)
	bp.updateMinFill()
	bp.kickFlusher()
}

// bufferPool is a global pool of buffers
var bufferPool *Pool
var bufferPoolOnce sync.Once

// Global gets a global pool of BufferSize, BufferCacheSize, BufferCacheFlushTime.
func Global() *Pool {
	bufferPoolOnce.Do(func() {
		// Initialise the buffer pool when used
		ci := fs.GetConfig(context.Background())
		bufferPool = New(BufferCacheFlushTime, BufferSize, BufferCacheSize, ci.UseMmap)
	})
	return bufferPool
}