restic/internal/archiver/buffer.go

90 lines
1.8 KiB
Go
Raw Normal View History

2018-03-30 20:43:18 +00:00
package archiver
import (
"context"
"sync"
)
// Buffer is a reusable buffer. After the buffer has been used, Release should
// be called so the underlying slice is put back into the pool.
type Buffer struct {
Data []byte
2018-04-29 13:34:41 +00:00
Put func(*Buffer)
2018-03-30 20:43:18 +00:00
}
// Release puts the buffer back into the pool it came from.
2018-04-29 13:34:41 +00:00
func (b *Buffer) Release() {
2018-03-30 20:43:18 +00:00
if b.Put != nil {
2018-04-29 13:34:41 +00:00
b.Put(b)
2018-03-30 20:43:18 +00:00
}
}
// BufferPool implements a limited set of reusable buffers.
type BufferPool struct {
2018-04-29 13:34:41 +00:00
ch chan *Buffer
2018-03-30 20:43:18 +00:00
chM sync.Mutex
defaultSize int
clearOnce sync.Once
}
// NewBufferPool initializes a new buffer pool. When the context is cancelled,
// all buffers are released. The pool stores at most max items. New buffers are
// created with defaultSize, buffers that are larger are released and not put
// back.
func NewBufferPool(ctx context.Context, max int, defaultSize int) *BufferPool {
b := &BufferPool{
2018-04-29 13:34:41 +00:00
ch: make(chan *Buffer, max),
2018-03-30 20:43:18 +00:00
defaultSize: defaultSize,
}
go func() {
<-ctx.Done()
b.clear()
}()
return b
}
// Get returns a new buffer, either from the pool or newly allocated.
2018-04-29 13:34:41 +00:00
func (pool *BufferPool) Get() *Buffer {
2018-03-30 20:43:18 +00:00
pool.chM.Lock()
defer pool.chM.Unlock()
select {
case buf := <-pool.ch:
2018-04-29 13:34:41 +00:00
return buf
2018-03-30 20:43:18 +00:00
default:
2018-04-29 13:34:41 +00:00
}
b := &Buffer{
Put: pool.Put,
Data: make([]byte, pool.defaultSize),
2018-03-30 20:43:18 +00:00
}
return b
}
2018-04-29 13:34:41 +00:00
// Put returns a buffer to the pool for reuse.
func (pool *BufferPool) Put(b *Buffer) {
if cap(b.Data) > pool.defaultSize {
return
}
2018-03-30 20:43:18 +00:00
pool.chM.Lock()
defer pool.chM.Unlock()
select {
case pool.ch <- b:
default:
}
}
// clear empties the buffer so that all items can be garbage collected.
func (pool *BufferPool) clear() {
pool.clearOnce.Do(func() {
ch := pool.ch
pool.chM.Lock()
pool.ch = nil
pool.chM.Unlock()
close(ch)
for range ch {
}
})
}