filepool: add generic file handle pool in lib/filepool
Factor the connection-backed write handle pool out of the smb backend into a generic lib/filepool.Pool[T] with its own tests, so it can be reused by other backends that implement fs.OpenWriterAter over a connection pool. The smb backend keeps its behaviour, opening and releasing handles through small closures passed to the pool. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Splainte <r.wycke@hotmail.fr>
This commit is contained in:
committed by
Nick Craig-Wood
co-authored by
Claude Opus 4.8
parent
6a617a379b
commit
613b335962
@@ -1,99 +0,0 @@
|
|||||||
package smb
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"fmt"
|
|
||||||
"os"
|
|
||||||
"sync"
|
|
||||||
|
|
||||||
"github.com/cloudsoda/go-smb2"
|
|
||||||
"golang.org/x/sync/errgroup"
|
|
||||||
)
|
|
||||||
|
|
||||||
// FsInterface defines the methods that filePool needs from Fs
|
|
||||||
type FsInterface interface {
|
|
||||||
getConnection(ctx context.Context, share string) (*conn, error)
|
|
||||||
putConnection(pc **conn, err error)
|
|
||||||
removeSession()
|
|
||||||
}
|
|
||||||
|
|
||||||
type file struct {
|
|
||||||
*smb2.File
|
|
||||||
c *conn
|
|
||||||
}
|
|
||||||
|
|
||||||
type filePool struct {
|
|
||||||
ctx context.Context
|
|
||||||
fs FsInterface
|
|
||||||
share string
|
|
||||||
path string
|
|
||||||
|
|
||||||
mu sync.Mutex
|
|
||||||
pool []*file
|
|
||||||
}
|
|
||||||
|
|
||||||
func newFilePool(ctx context.Context, fs FsInterface, share, path string) *filePool {
|
|
||||||
return &filePool{
|
|
||||||
ctx: ctx,
|
|
||||||
fs: fs,
|
|
||||||
share: share,
|
|
||||||
path: path,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *filePool) get() (*file, error) {
|
|
||||||
p.mu.Lock()
|
|
||||||
if len(p.pool) > 0 {
|
|
||||||
f := p.pool[len(p.pool)-1]
|
|
||||||
p.pool = p.pool[:len(p.pool)-1]
|
|
||||||
p.mu.Unlock()
|
|
||||||
return f, nil
|
|
||||||
}
|
|
||||||
p.mu.Unlock()
|
|
||||||
|
|
||||||
c, err := p.fs.getConnection(p.ctx, p.share)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
|
|
||||||
fl, err := c.smbShare.OpenFile(p.path, os.O_WRONLY, 0o644)
|
|
||||||
if err != nil {
|
|
||||||
p.fs.putConnection(&c, err)
|
|
||||||
return nil, fmt.Errorf("failed to open: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
return &file{File: fl, c: c}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *filePool) put(f *file, err error) {
|
|
||||||
if f == nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
if err != nil {
|
|
||||||
_ = f.Close()
|
|
||||||
p.fs.putConnection(&f.c, err)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
p.mu.Lock()
|
|
||||||
p.pool = append(p.pool, f)
|
|
||||||
p.mu.Unlock()
|
|
||||||
}
|
|
||||||
|
|
||||||
func (p *filePool) drain() error {
|
|
||||||
p.mu.Lock()
|
|
||||||
files := p.pool
|
|
||||||
p.pool = nil
|
|
||||||
p.mu.Unlock()
|
|
||||||
|
|
||||||
g, _ := errgroup.WithContext(p.ctx)
|
|
||||||
for _, f := range files {
|
|
||||||
g.Go(func() error {
|
|
||||||
err := f.Close()
|
|
||||||
p.fs.putConnection(&f.c, err)
|
|
||||||
return err
|
|
||||||
})
|
|
||||||
}
|
|
||||||
return g.Wait()
|
|
||||||
}
|
|
||||||
@@ -1,228 +0,0 @@
|
|||||||
package smb
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
"sync"
|
|
||||||
"testing"
|
|
||||||
|
|
||||||
"github.com/cloudsoda/go-smb2"
|
|
||||||
"github.com/stretchr/testify/assert"
|
|
||||||
)
|
|
||||||
|
|
||||||
// Mock Fs that implements FsInterface
|
|
||||||
type mockFs struct {
|
|
||||||
mu sync.Mutex
|
|
||||||
putConnectionCalled bool
|
|
||||||
putConnectionErr error
|
|
||||||
getConnectionCalled bool
|
|
||||||
getConnectionErr error
|
|
||||||
getConnectionResult *conn
|
|
||||||
removeSessionCalled bool
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) putConnection(pc **conn, err error) {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
m.putConnectionCalled = true
|
|
||||||
m.putConnectionErr = err
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) getConnection(ctx context.Context, share string) (*conn, error) {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
m.getConnectionCalled = true
|
|
||||||
if m.getConnectionErr != nil {
|
|
||||||
return nil, m.getConnectionErr
|
|
||||||
}
|
|
||||||
if m.getConnectionResult != nil {
|
|
||||||
return m.getConnectionResult, nil
|
|
||||||
}
|
|
||||||
return &conn{}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) removeSession() {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
m.removeSessionCalled = true
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) isPutConnectionCalled() bool {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
return m.putConnectionCalled
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) getPutConnectionErr() error {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
return m.putConnectionErr
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *mockFs) isGetConnectionCalled() bool {
|
|
||||||
m.mu.Lock()
|
|
||||||
defer m.mu.Unlock()
|
|
||||||
return m.getConnectionCalled
|
|
||||||
}
|
|
||||||
|
|
||||||
func newMockFs() *mockFs {
|
|
||||||
return &mockFs{}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Helper function to create a mock file
|
|
||||||
func newMockFile() *file {
|
|
||||||
return &file{
|
|
||||||
File: &smb2.File{},
|
|
||||||
c: &conn{},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test filePool creation
|
|
||||||
func TestNewFilePool(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
share := "testshare"
|
|
||||||
path := "/test/path"
|
|
||||||
|
|
||||||
pool := newFilePool(ctx, fs, share, path)
|
|
||||||
|
|
||||||
assert.NotNil(t, pool)
|
|
||||||
assert.Equal(t, ctx, pool.ctx)
|
|
||||||
assert.Equal(t, fs, pool.fs)
|
|
||||||
assert.Equal(t, share, pool.share)
|
|
||||||
assert.Equal(t, path, pool.path)
|
|
||||||
assert.Empty(t, pool.pool)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test getting file from pool when pool has files
|
|
||||||
func TestFilePool_Get_FromPool(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
// Add a mock file to the pool
|
|
||||||
mockFile := newMockFile()
|
|
||||||
pool.pool = append(pool.pool, mockFile)
|
|
||||||
|
|
||||||
// Get file from pool
|
|
||||||
f, err := pool.get()
|
|
||||||
|
|
||||||
assert.NoError(t, err)
|
|
||||||
assert.NotNil(t, f)
|
|
||||||
assert.Equal(t, mockFile, f)
|
|
||||||
assert.Empty(t, pool.pool)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test getting file when pool is empty
|
|
||||||
func TestFilePool_Get_EmptyPool(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
|
|
||||||
// Set up the mock to return an error from getConnection
|
|
||||||
// This tests that the pool calls getConnection when empty
|
|
||||||
fs.getConnectionErr = errors.New("connection failed")
|
|
||||||
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "test/path")
|
|
||||||
|
|
||||||
// This should call getConnection and return the error
|
|
||||||
f, err := pool.get()
|
|
||||||
assert.Error(t, err)
|
|
||||||
assert.Nil(t, f)
|
|
||||||
assert.True(t, fs.isGetConnectionCalled())
|
|
||||||
assert.Equal(t, "connection failed", err.Error())
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test putting file successfully
|
|
||||||
func TestFilePool_Put_Success(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
mockFile := newMockFile()
|
|
||||||
|
|
||||||
pool.put(mockFile, nil)
|
|
||||||
|
|
||||||
assert.Len(t, pool.pool, 1)
|
|
||||||
assert.Equal(t, mockFile, pool.pool[0])
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test putting file with error
|
|
||||||
func TestFilePool_Put_WithError(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
mockFile := newMockFile()
|
|
||||||
|
|
||||||
pool.put(mockFile, errors.New("write error"))
|
|
||||||
|
|
||||||
// Should call putConnection with error
|
|
||||||
assert.True(t, fs.isPutConnectionCalled())
|
|
||||||
assert.Equal(t, errors.New("write error"), fs.getPutConnectionErr())
|
|
||||||
assert.Empty(t, pool.pool)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test putting nil file
|
|
||||||
func TestFilePool_Put_NilFile(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
// Should not panic
|
|
||||||
pool.put(nil, nil)
|
|
||||||
pool.put(nil, errors.New("some error"))
|
|
||||||
|
|
||||||
assert.Empty(t, pool.pool)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test draining pool with files
|
|
||||||
func TestFilePool_Drain_WithFiles(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
// Add mock files to pool
|
|
||||||
mockFile1 := newMockFile()
|
|
||||||
mockFile2 := newMockFile()
|
|
||||||
pool.pool = append(pool.pool, mockFile1, mockFile2)
|
|
||||||
|
|
||||||
// Before draining
|
|
||||||
assert.Len(t, pool.pool, 2)
|
|
||||||
|
|
||||||
_ = pool.drain()
|
|
||||||
assert.Empty(t, pool.pool)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test concurrent access to pool
|
|
||||||
func TestFilePool_ConcurrentAccess(t *testing.T) {
|
|
||||||
ctx := context.Background()
|
|
||||||
fs := newMockFs()
|
|
||||||
pool := newFilePool(ctx, fs, "testshare", "/test/path")
|
|
||||||
|
|
||||||
const numGoroutines = 10
|
|
||||||
for range numGoroutines {
|
|
||||||
mockFile := newMockFile()
|
|
||||||
pool.pool = append(pool.pool, mockFile)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Test concurrent get operations
|
|
||||||
done := make(chan bool, numGoroutines)
|
|
||||||
|
|
||||||
for range numGoroutines {
|
|
||||||
go func() {
|
|
||||||
defer func() { done <- true }()
|
|
||||||
|
|
||||||
f, err := pool.get()
|
|
||||||
if err == nil {
|
|
||||||
pool.put(f, nil)
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
for range numGoroutines {
|
|
||||||
<-done
|
|
||||||
}
|
|
||||||
|
|
||||||
// Pool should be in a consistent after the concurrence access
|
|
||||||
assert.Len(t, pool.pool, numGoroutines)
|
|
||||||
}
|
|
||||||
+42
-6
@@ -13,6 +13,7 @@ import (
|
|||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/cloudsoda/go-smb2"
|
||||||
"github.com/rclone/rclone/fs"
|
"github.com/rclone/rclone/fs"
|
||||||
"github.com/rclone/rclone/fs/config"
|
"github.com/rclone/rclone/fs/config"
|
||||||
"github.com/rclone/rclone/fs/config/configmap"
|
"github.com/rclone/rclone/fs/config/configmap"
|
||||||
@@ -21,6 +22,7 @@ import (
|
|||||||
"github.com/rclone/rclone/lib/bucket"
|
"github.com/rclone/rclone/lib/bucket"
|
||||||
"github.com/rclone/rclone/lib/encoder"
|
"github.com/rclone/rclone/lib/encoder"
|
||||||
"github.com/rclone/rclone/lib/env"
|
"github.com/rclone/rclone/lib/env"
|
||||||
|
"github.com/rclone/rclone/lib/filepool"
|
||||||
"github.com/rclone/rclone/lib/pacer"
|
"github.com/rclone/rclone/lib/pacer"
|
||||||
"github.com/rclone/rclone/lib/readers"
|
"github.com/rclone/rclone/lib/readers"
|
||||||
)
|
)
|
||||||
@@ -514,8 +516,15 @@ func (f *Fs) About(ctx context.Context) (_ *fs.Usage, err error) {
|
|||||||
return usage, nil
|
return usage, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// file is a pooled write handle together with the connection it lives on.
|
||||||
|
type file struct {
|
||||||
|
*smb2.File
|
||||||
|
c *conn
|
||||||
|
}
|
||||||
|
|
||||||
type smbWriterAt struct {
|
type smbWriterAt struct {
|
||||||
pool *filePool
|
fs *Fs
|
||||||
|
pool *filepool.Pool[*file]
|
||||||
closed bool
|
closed bool
|
||||||
closeMu sync.Mutex
|
closeMu sync.Mutex
|
||||||
wg sync.WaitGroup
|
wg sync.WaitGroup
|
||||||
@@ -531,13 +540,13 @@ func (w *smbWriterAt) WriteAt(p []byte, off int64) (int, error) {
|
|||||||
w.closeMu.Unlock()
|
w.closeMu.Unlock()
|
||||||
defer w.wg.Done()
|
defer w.wg.Done()
|
||||||
|
|
||||||
f, err := w.pool.get()
|
f, err := w.pool.Get()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return 0, fmt.Errorf("failed to get file from pool: %w", err)
|
return 0, fmt.Errorf("failed to get file from pool: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
n, writeErr := f.WriteAt(p, off)
|
n, writeErr := f.WriteAt(p, off)
|
||||||
w.pool.put(f, writeErr)
|
w.pool.Put(f, writeErr)
|
||||||
|
|
||||||
if writeErr != nil {
|
if writeErr != nil {
|
||||||
return n, fmt.Errorf("failed to write at offset %d: %w", off, writeErr)
|
return n, fmt.Errorf("failed to write at offset %d: %w", off, writeErr)
|
||||||
@@ -561,12 +570,12 @@ func (w *smbWriterAt) Close() error {
|
|||||||
var errs []error
|
var errs []error
|
||||||
|
|
||||||
// Drain the pool
|
// Drain the pool
|
||||||
if err := w.pool.drain(); err != nil {
|
if err := w.pool.Drain(); err != nil {
|
||||||
errs = append(errs, fmt.Errorf("failed to drain file pool: %w", err))
|
errs = append(errs, fmt.Errorf("failed to drain file pool: %w", err))
|
||||||
}
|
}
|
||||||
|
|
||||||
// Remove session
|
// Remove session
|
||||||
w.pool.fs.removeSession()
|
w.fs.removeSession()
|
||||||
|
|
||||||
if len(errs) > 0 {
|
if len(errs) > 0 {
|
||||||
return errors.Join(errs...)
|
return errors.Join(errs...)
|
||||||
@@ -575,6 +584,32 @@ func (w *smbWriterAt) Close() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// openPoolFile opens a fresh write handle on its own connection for share/path.
|
||||||
|
func (f *Fs) openPoolFile(share, path string) func(context.Context) (*file, error) {
|
||||||
|
return func(ctx context.Context) (*file, error) {
|
||||||
|
c, err := f.getConnection(ctx, share)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
fl, err := c.smbShare.OpenFile(path, os.O_WRONLY, 0o644)
|
||||||
|
if err != nil {
|
||||||
|
f.putConnection(&c, err)
|
||||||
|
return nil, fmt.Errorf("failed to open: %w", err)
|
||||||
|
}
|
||||||
|
return &file{File: fl, c: c}, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// releasePoolFile closes a pooled handle and returns its connection.
|
||||||
|
func (f *Fs) releasePoolFile(fl *file, err error) error {
|
||||||
|
closeErr := fl.Close()
|
||||||
|
if err == nil {
|
||||||
|
err = closeErr
|
||||||
|
}
|
||||||
|
f.putConnection(&fl.c, err)
|
||||||
|
return closeErr
|
||||||
|
}
|
||||||
|
|
||||||
// OpenWriterAt opens with a handle for random access writes
|
// OpenWriterAt opens with a handle for random access writes
|
||||||
//
|
//
|
||||||
// Pass in the remote desired and the size if known.
|
// Pass in the remote desired and the size if known.
|
||||||
@@ -624,7 +659,8 @@ func (f *Fs) OpenWriterAt(ctx context.Context, remote string, size int64) (fs.Wr
|
|||||||
o.fs.addSession()
|
o.fs.addSession()
|
||||||
|
|
||||||
return &smbWriterAt{
|
return &smbWriterAt{
|
||||||
pool: newFilePool(ctx, o.fs, share, smbPath),
|
fs: o.fs,
|
||||||
|
pool: filepool.New(ctx, o.fs.openPoolFile(share, smbPath), o.fs.releasePoolFile),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,78 @@
|
|||||||
|
// Package filepool keeps a set of reusable write handles open on a single
|
||||||
|
// remote path, one per connection, so several goroutines can write to the same
|
||||||
|
// file at once without sharing a handle.
|
||||||
|
//
|
||||||
|
// It is used by backends that implement fs.OpenWriterAter over a connection
|
||||||
|
// pool, where the core writes the chunks of a large file concurrently at
|
||||||
|
// non-overlapping offsets.
|
||||||
|
package filepool
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"golang.org/x/sync/errgroup"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Pool hands out handles of type T for a single file. Handles are reused when
|
||||||
|
// free and opened on demand otherwise. It is safe for concurrent use.
|
||||||
|
//
|
||||||
|
// The zero value is not usable; call New.
|
||||||
|
type Pool[T any] struct {
|
||||||
|
ctx context.Context
|
||||||
|
open func(context.Context) (T, error)
|
||||||
|
release func(handle T, err error) error
|
||||||
|
|
||||||
|
mu sync.Mutex
|
||||||
|
free []T
|
||||||
|
}
|
||||||
|
|
||||||
|
// New returns a Pool.
|
||||||
|
//
|
||||||
|
// open opens a fresh handle on its own connection. release closes a handle and
|
||||||
|
// returns its connection: err is the error that made the handle unusable (nil
|
||||||
|
// when the handle is simply being drained) and the returned error is the result
|
||||||
|
// of closing it.
|
||||||
|
func New[T any](ctx context.Context, open func(context.Context) (T, error), release func(handle T, err error) error) *Pool[T] {
|
||||||
|
return &Pool[T]{ctx: ctx, open: open, release: release}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get returns a free handle, opening a new one if none are free.
|
||||||
|
func (p *Pool[T]) Get() (T, error) {
|
||||||
|
p.mu.Lock()
|
||||||
|
if n := len(p.free); n > 0 {
|
||||||
|
h := p.free[n-1]
|
||||||
|
p.free = p.free[:n-1]
|
||||||
|
p.mu.Unlock()
|
||||||
|
return h, nil
|
||||||
|
}
|
||||||
|
p.mu.Unlock()
|
||||||
|
return p.open(p.ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Put returns a handle to the pool. If err is non-nil the write that used the
|
||||||
|
// handle failed, so the handle is released instead of being reused.
|
||||||
|
func (p *Pool[T]) Put(handle T, err error) {
|
||||||
|
if err != nil {
|
||||||
|
_ = p.release(handle, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
p.mu.Lock()
|
||||||
|
p.free = append(p.free, handle)
|
||||||
|
p.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Drain releases every free handle, closing them concurrently, and returns the
|
||||||
|
// first error encountered.
|
||||||
|
func (p *Pool[T]) Drain() error {
|
||||||
|
p.mu.Lock()
|
||||||
|
free := p.free
|
||||||
|
p.free = nil
|
||||||
|
p.mu.Unlock()
|
||||||
|
|
||||||
|
g := new(errgroup.Group)
|
||||||
|
for _, h := range free {
|
||||||
|
g.Go(func() error { return p.release(h, nil) })
|
||||||
|
}
|
||||||
|
return g.Wait()
|
||||||
|
}
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
package filepool
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
// handle is a fake pooled handle used by the tests.
|
||||||
|
type handle struct {
|
||||||
|
id int
|
||||||
|
released bool
|
||||||
|
relErr error
|
||||||
|
}
|
||||||
|
|
||||||
|
// harness wires a Pool up to counters so the tests can assert on the open and
|
||||||
|
// release calls without a real backend.
|
||||||
|
type harness struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
next int
|
||||||
|
opens int
|
||||||
|
openErr error
|
||||||
|
releases int
|
||||||
|
closeErr error
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *harness) open(context.Context) (*handle, error) {
|
||||||
|
h.mu.Lock()
|
||||||
|
defer h.mu.Unlock()
|
||||||
|
h.opens++
|
||||||
|
if h.openErr != nil {
|
||||||
|
return nil, h.openErr
|
||||||
|
}
|
||||||
|
h.next++
|
||||||
|
return &handle{id: h.next}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *harness) release(hd *handle, err error) error {
|
||||||
|
h.mu.Lock()
|
||||||
|
defer h.mu.Unlock()
|
||||||
|
h.releases++
|
||||||
|
hd.released = true
|
||||||
|
hd.relErr = err
|
||||||
|
return h.closeErr
|
||||||
|
}
|
||||||
|
|
||||||
|
func newPool(h *harness) *Pool[*handle] {
|
||||||
|
return New(context.Background(), h.open, h.release)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetOpensWhenEmpty(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 1, hd.id)
|
||||||
|
assert.Equal(t, 1, h.opens)
|
||||||
|
assert.Empty(t, p.free)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetReusesFreeHandle(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
p.Put(hd, nil)
|
||||||
|
assert.Len(t, p.free, 1)
|
||||||
|
|
||||||
|
got, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Same(t, hd, got, "a free handle should be reused instead of opening a new one")
|
||||||
|
assert.Equal(t, 1, h.opens)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetOpenError(t *testing.T) {
|
||||||
|
h := &harness{openErr: errors.New("connection failed")}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
assert.Error(t, err)
|
||||||
|
assert.Nil(t, hd)
|
||||||
|
assert.EqualError(t, err, "connection failed")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPutSuccessKeepsHandle(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
p.Put(hd, nil)
|
||||||
|
|
||||||
|
assert.Len(t, p.free, 1)
|
||||||
|
assert.Zero(t, h.releases, "a healthy handle must not be released")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestPutErrorReleasesHandle(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
writeErr := errors.New("write error")
|
||||||
|
p.Put(hd, writeErr)
|
||||||
|
|
||||||
|
assert.Empty(t, p.free, "a handle put back with an error must not be reused")
|
||||||
|
assert.Equal(t, 1, h.releases)
|
||||||
|
assert.True(t, hd.released)
|
||||||
|
assert.Equal(t, writeErr, hd.relErr, "release must receive the write error")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDrainReleasesEveryHandle(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
var handles []*handle
|
||||||
|
for range 3 {
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
handles = append(handles, hd)
|
||||||
|
}
|
||||||
|
for _, hd := range handles {
|
||||||
|
p.Put(hd, nil)
|
||||||
|
}
|
||||||
|
|
||||||
|
require.NoError(t, p.Drain())
|
||||||
|
assert.Empty(t, p.free)
|
||||||
|
assert.Equal(t, 3, h.releases)
|
||||||
|
for _, hd := range handles {
|
||||||
|
assert.True(t, hd.released)
|
||||||
|
assert.NoError(t, hd.relErr, "draining passes a nil error to release")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDrainReturnsCloseError(t *testing.T) {
|
||||||
|
h := &harness{closeErr: errors.New("close failed")}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
hd, err := p.Get()
|
||||||
|
require.NoError(t, err)
|
||||||
|
p.Put(hd, nil)
|
||||||
|
|
||||||
|
assert.EqualError(t, p.Drain(), "close failed")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestConcurrentGetPut(t *testing.T) {
|
||||||
|
h := &harness{}
|
||||||
|
p := newPool(h)
|
||||||
|
|
||||||
|
const workers = 10
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(workers)
|
||||||
|
for range workers {
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
for range 100 {
|
||||||
|
hd, err := p.Get()
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
p.Put(hd, nil)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
// Every handle handed out was returned, so draining must release them all
|
||||||
|
// with no leaks.
|
||||||
|
require.NoError(t, p.Drain())
|
||||||
|
assert.Empty(t, p.free)
|
||||||
|
assert.Equal(t, h.opens, h.releases)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user