From bcbe393af32bc6c6a388ce5e203e559e18f32aa3 Mon Sep 17 00:00:00 2001
From: Nick Craig-Wood <nick@craig-wood.com>
Date: Fri, 27 Nov 2020 17:25:57 +0000
Subject: [PATCH] sftp: implement Shutdown method

---
 backend/sftp/sftp.go      | 24 ++++++++++++++++++++++++
 fstest/fstests/fstests.go |  9 +++++++++
 2 files changed, 33 insertions(+)

diff --git a/backend/sftp/sftp.go b/backend/sftp/sftp.go
index a06207386..fdf41cc0a 100644
--- a/backend/sftp/sftp.go
+++ b/backend/sftp/sftp.go
@@ -412,6 +412,23 @@ func (f *Fs) putSftpConnection(pc **conn, err error) {
 	f.poolMu.Unlock()
 }
 
+// Drain the pool of any connections
+func (f *Fs) drainPool(ctx context.Context) (err error) {
+	f.poolMu.Lock()
+	defer f.poolMu.Unlock()
+	for i, c := range f.pool {
+		if cErr := c.closed(); cErr == nil {
+			cErr = c.close()
+			if cErr != nil {
+				err = cErr
+			}
+		}
+		f.pool[i] = nil
+	}
+	f.pool = nil
+	return err
+}
+
 // NewFs creates a new Fs object from the name and root. It connects to
 // the host specified in the config file.
 func NewFs(ctx context.Context, name, root string, m configmap.Mapper) (fs.Fs, error) {
@@ -1057,6 +1074,12 @@ func (f *Fs) About(ctx context.Context) (*fs.Usage, error) {
 	return usage, nil
 }
 
+// Shutdown the backend, closing any background tasks and any
+// cached connections.
+func (f *Fs) Shutdown(ctx context.Context) error {
+	return f.drainPool(ctx)
+}
+
 // Fs is the filesystem this remote sftp file object is located within
 func (o *Object) Fs() fs.Info {
 	return o.fs
@@ -1407,5 +1430,6 @@ var (
 	_ fs.Mover       = &Fs{}
 	_ fs.DirMover    = &Fs{}
 	_ fs.Abouter     = &Fs{}
+	_ fs.Shutdowner  = &Fs{}
 	_ fs.Object      = &Object{}
 )
diff --git a/fstest/fstests/fstests.go b/fstest/fstests/fstests.go
index 622825afd..641a67f40 100644
--- a/fstest/fstests/fstests.go
+++ b/fstest/fstests/fstests.go
@@ -1932,6 +1932,15 @@ func Run(t *testing.T, opt *Opt) {
 		_ = operations.Purge(ctx, f, "")
 	}
 
+	t.Run("FsShutdown", func(t *testing.T) {
+		do := f.Features().Shutdown
+		if do == nil {
+			t.Skip("Shutdown method not supported")
+		}
+		require.NoError(t, do(ctx))
+		require.NoError(t, do(ctx), "must be able to call Shutdown twice")
+	})
+
 	// Remove the local directory so we don't clutter up /tmp
 	if strings.HasPrefix(remoteName, "/") {
 		t.Log("remoteName", remoteName)