diff --git a/backend/smb/filepool.go b/backend/smb/filepool.go deleted file mode 100644 index 0061ec1ff..000000000 --- a/backend/smb/filepool.go +++ /dev/null @@ -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() -} diff --git a/backend/smb/filepool_test.go b/backend/smb/filepool_test.go deleted file mode 100644 index 019e688f4..000000000 --- a/backend/smb/filepool_test.go +++ /dev/null @@ -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) -} diff --git a/backend/smb/smb.go b/backend/smb/smb.go index c97a8e391..b179bcabb 100644 --- a/backend/smb/smb.go +++ b/backend/smb/smb.go @@ -13,6 +13,7 @@ import ( "sync/atomic" "time" + "github.com/cloudsoda/go-smb2" "github.com/rclone/rclone/fs" "github.com/rclone/rclone/fs/config" "github.com/rclone/rclone/fs/config/configmap" @@ -21,6 +22,7 @@ import ( "github.com/rclone/rclone/lib/bucket" "github.com/rclone/rclone/lib/encoder" "github.com/rclone/rclone/lib/env" + "github.com/rclone/rclone/lib/filepool" "github.com/rclone/rclone/lib/pacer" "github.com/rclone/rclone/lib/readers" ) @@ -514,8 +516,15 @@ func (f *Fs) About(ctx context.Context) (_ *fs.Usage, err error) { 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 { - pool *filePool + fs *Fs + pool *filepool.Pool[*file] closed bool closeMu sync.Mutex wg sync.WaitGroup @@ -531,13 +540,13 @@ func (w *smbWriterAt) WriteAt(p []byte, off int64) (int, error) { w.closeMu.Unlock() defer w.wg.Done() - f, err := w.pool.get() + f, err := w.pool.Get() if err != nil { return 0, fmt.Errorf("failed to get file from pool: %w", err) } n, writeErr := f.WriteAt(p, off) - w.pool.put(f, writeErr) + w.pool.Put(f, writeErr) if writeErr != nil { 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 // 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)) } // Remove session - w.pool.fs.removeSession() + w.fs.removeSession() if len(errs) > 0 { return errors.Join(errs...) @@ -575,6 +584,32 @@ func (w *smbWriterAt) Close() error { 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 // // 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() 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 } diff --git a/lib/filepool/filepool.go b/lib/filepool/filepool.go new file mode 100644 index 000000000..0152a3b63 --- /dev/null +++ b/lib/filepool/filepool.go @@ -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() +} diff --git a/lib/filepool/filepool_test.go b/lib/filepool/filepool_test.go new file mode 100644 index 000000000..6e8fb36b3 --- /dev/null +++ b/lib/filepool/filepool_test.go @@ -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) +}