diff --git a/lib/pool/pool.go b/lib/pool/pool.go index 5491f992d..654ae444f 100644 --- a/lib/pool/pool.go +++ b/lib/pool/pool.go @@ -5,6 +5,7 @@ package pool import ( "context" "fmt" + "slices" "sync" "time" @@ -46,7 +47,8 @@ type Pool struct { // totalMemory is a semaphore used to control total buffer usage of // all Pools. It it may be nil in which case the total buffer usage -// will not be controlled. +// 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 @@ -102,6 +104,18 @@ func (bp *Pool) get() []byte { 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 @@ -109,6 +123,20 @@ 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) { @@ -175,7 +203,7 @@ func (bp *Pool) updateMinFill() { } } -// acquire mem bytes of memory +// acquire mem bytes of memory for the user func (bp *Pool) acquire(mem int64) error { if totalMemory == nil { return nil @@ -184,7 +212,7 @@ func (bp *Pool) acquire(mem int64) error { return totalMemory.Acquire(ctx, mem) } -// release mem bytes of memory +// release mem bytes of memory from the user func (bp *Pool) release(mem int64) { if totalMemory == nil { return @@ -192,84 +220,62 @@ func (bp *Pool) release(mem int64) { totalMemory.Release(mem) } -// Reserve buffers for use. Blocks until they are free. -// -// Doesn't allocate any memory. -// -// Must be released by calling GetReserved() which releases 1 buffer or -// Release() to release any number of buffers. -func (bp *Pool) Reserve(buffers int) { - waitTime := time.Millisecond - for { - err := bp.acquire(int64(buffers) * int64(bp.bufferSize)) - if err == nil { - break - } - fs.Logf(nil, "Failed to get reservation for buffer, waiting for %v: %v", waitTime, err) - time.Sleep(waitTime) - waitTime *= 2 - } -} - -// Release previously Reserved buffers. -// -// Doesn't free any memory. -func (bp *Pool) Release(buffers int) { - bp.release(int64(buffers) * int64(bp.bufferSize)) -} - -// Get a buffer from the pool or allocate one -func (bp *Pool) getBlock(reserved bool) []byte { - bp.mu.Lock() - var buf []byte - waitTime := time.Millisecond - for { - if len(bp.cache) > 0 { - buf = bp.get() - if reserved { - // If got reserved memory from the cache we - // can release the reservation of one buffer. - bp.release(int64(bp.bufferSize)) - } - break - } else { - var err error - if !reserved { - bp.mu.Unlock() - err = bp.acquire(int64(bp.bufferSize)) - bp.mu.Lock() - } - if err == nil { - buf, err = bp.alloc(bp.bufferSize) - if err == nil { - bp.alloced++ - break - } - if !reserved { - bp.release(int64(bp.bufferSize)) - } - } - 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 - } - } - bp.inUse++ - bp.updateMinFill() - bp.mu.Unlock() - return buf -} - // Get a buffer from the pool or allocate one func (bp *Pool) Get() []byte { - return bp.getBlock(false) + return bp.GetN(1)[0] } -// GetReserved gets a reserved buffer from the pool or allocates one. -func (bp *Pool) GetReserved() []byte { - return bp.getBlock(true) +// 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 @@ -277,19 +283,17 @@ func (bp *Pool) freeBuffer(mem []byte) { err := bp.free(mem) if err != nil { fs.Logf(nil, "Failed to free memory: %v", err) - } else { - bp.release(int64(bp.bufferSize)) } bp.alloced-- } -// Put returns the buffer to the buffer cache or frees it +// _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() +// 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)) @@ -299,11 +303,37 @@ func (bp *Pool) Put(buf []byte) { } 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 diff --git a/lib/pool/pool_test.go b/lib/pool/pool_test.go index 425acbd7a..254ac01f2 100644 --- a/lib/pool/pool_test.go +++ b/lib/pool/pool_test.go @@ -53,17 +53,27 @@ func testGetPut(t *testing.T, useMmap bool, unreliable bool) { assert.Equal(t, 0, bp.InPool()) assert.Equal(t, 3, bp.Alloced()) + bs := bp.GetN(3) + assert.Equal(t, 6, bp.InUse()) + assert.Equal(t, 0, bp.InPool()) + assert.Equal(t, 6, bp.Alloced()) + bp.Put(b1) - assert.Equal(t, 2, bp.InUse()) + assert.Equal(t, 5, bp.InUse()) assert.Equal(t, 1, bp.InPool()) - assert.Equal(t, 3, bp.Alloced()) + assert.Equal(t, 6, bp.Alloced()) bp.Put(b2) - assert.Equal(t, 1, bp.InUse()) + assert.Equal(t, 4, bp.InUse()) assert.Equal(t, 2, bp.InPool()) - assert.Equal(t, 3, bp.Alloced()) + assert.Equal(t, 6, bp.Alloced()) bp.Put(b3) + assert.Equal(t, 3, bp.InUse()) + assert.Equal(t, 2, bp.InPool()) + assert.Equal(t, 5, bp.Alloced()) + + bp.PutN(bs) assert.Equal(t, 0, bp.InUse()) assert.Equal(t, 2, bp.InPool()) assert.Equal(t, 2, bp.Alloced()) @@ -89,6 +99,18 @@ func testGetPut(t *testing.T, useMmap bool, unreliable bool) { assert.Equal(t, 2, bp.InPool()) assert.Equal(t, 2, bp.Alloced()) + bsa := bp.GetN(3) + assert.Equal(t, addr(b1), addr(bsa[1])) + assert.Equal(t, addr(b2), addr(bsa[0])) + assert.Equal(t, 3, bp.InUse()) + assert.Equal(t, 0, bp.InPool()) + assert.Equal(t, 3, bp.Alloced()) + + bp.PutN(bsa) + assert.Equal(t, 0, bp.InUse()) + assert.Equal(t, 2, bp.InPool()) + assert.Equal(t, 2, bp.Alloced()) + assert.Panics(t, func() { bp.Put(make([]byte, 1)) }) @@ -265,11 +287,19 @@ func TestPoolMaxBufferMemory(t *testing.T) { wg.Add(1) go func() { defer wg.Done() - buf := bp.Get() - countBuf(1) - time.Sleep(100 * time.Millisecond) - bp.Put(buf) - countBuf(-1) + if i < 4 { + buf := bp.GetN(i + 1) + countBuf(i + 1) + time.Sleep(100 * time.Millisecond) + bp.PutN(buf) + countBuf(-(i + 1)) + } else { + buf := bp.Get() + countBuf(1) + time.Sleep(100 * time.Millisecond) + bp.Put(buf) + countBuf(-1) + } }() } diff --git a/lib/pool/reader_writer.go b/lib/pool/reader_writer.go index ac266b37c..3a14af2c0 100644 --- a/lib/pool/reader_writer.go +++ b/lib/pool/reader_writer.go @@ -36,7 +36,7 @@ type RW struct { out int // offset we are reading from reads int // count how many times the data has been read - reserved int // number of buffers reserved + reserved [][]byte // reserved buffers } var ( @@ -63,15 +63,16 @@ func NewRW(pool *Pool) *RW { // Reserve bytes of memory. // -// Reserve, but don't allocate n bytes of memory. +// This allocates n bytes of memory for later use. // // This is rounded up to the nearest buffer page size. +// +// Only safe to call once. func (rw *RW) Reserve(n int64) *RW { rw.mu.Lock() defer rw.mu.Unlock() buffers := int((n + int64(rw.pool.bufferSize) - 1) / int64(rw.pool.bufferSize)) - rw.pool.Reserve(buffers) - rw.reserved += buffers + rw.reserved = rw.pool.GetN(buffers) return rw } @@ -216,9 +217,12 @@ func (rw *RW) writePage() (page []byte) { if len(rw.pages) > 0 && rw.lastOffset < rw.pool.bufferSize { return rw.pages[len(rw.pages)-1][rw.lastOffset:] } - if rw.reserved > 0 { - page = rw.pool.GetReserved() - rw.reserved-- + if len(rw.reserved) > 0 { + // Get reserved pages if available + i := len(rw.reserved) - 1 + page = rw.reserved[i] + rw.reserved[i] = nil + rw.reserved = rw.reserved[:i] } else { page = rw.pool.Get() } @@ -338,14 +342,12 @@ func (rw *RW) Close() error { rw.mu.Lock() defer rw.mu.Unlock() rw.signalWrite() // signal more data available - for _, page := range rw.pages { - rw.pool.Put(page) - } + rw.pool.PutN(rw.pages) + clear(rw.pages) rw.pages = nil - if rw.reserved > 0 { - rw.pool.Release(rw.reserved) - rw.reserved = 0 - } + rw.pool.PutN(rw.reserved) + clear(rw.reserved) + rw.reserved = nil return nil }