feat(lib/backend): add functionality to close clients (#451)

* feat(lib/backend): add functionality to close clients

What?
This PR aims to add functionality to close the clients and it's
resources in case the resources have closers.

Why?
There are use cases in for gcsbacknd using `storage.transfermanager` to
close the single downloader it creates. If we don't close it, there's a
possibility of leaking resources.

How?
Add `Close` to client interface and implement the method for each
client.\
Then `backend.Manager` will use them to expose a `Close()` method that
can be used by build-index and origin to close them.

* add close methods to tests for proper cleanups

* resolved comments

Author: sambhav-jain-16

build-index/cmd/cmd.go

@@ -144,6 +144,7 @@ func Run(flags *Flags, opts ...Option) {
 	if err != nil {
 		log.Fatalf("Error creating backend manager: %s", err)
 	}
+	defer closers.Close(backends)
 
 	tls, err := config.TLS.BuildClient()
 	if err != nil {

lib/backend/client.go

@@ -65,4 +65,8 @@ type Client interface {
 
 	// List lists entries whose names start with prefix.
 	List(prefix string, opts ...ListOption) (*ListResult, error)
+
+	// Close can be used to close any resources held by the client.
+	// This operation should be idempotent.
+	Close() error
 }

lib/backend/gcsbackend/client.go

@@ -70,10 +70,11 @@ func (f *factory) Create(
 
 // Client implements a backend.Client for GCS.
 type Client struct {
-	config Config
-	pather namepath.Pather
-	stats  tally.Scope
-	gcs    GCS
+	config  Config
+	pather  namepath.Pather
+	stats   tally.Scope
+	gcs     GCS
+	sClient *storage.Client
 }
 
 // Option allows setting optional Client parameters.
@@ -111,7 +112,13 @@ func NewClient(
 
 	if len(opts) > 0 {
 		// For mock.
-		client := &Client{config, pather, stats, nil}
+		client := &Client{
+			config:  config,
+			pather:  pather,
+			stats:   stats,
+			gcs:     nil,
+			sClient: nil,
+		}
 		for _, opt := range opts {
 			opt(client)
 		}
@@ -125,8 +132,13 @@ func NewClient(
 		return nil, fmt.Errorf("invalid gcs credentials: %s", err)
 	}
 
-	client := &Client{config, pather, stats,
-		NewGCS(ctx, sClient.Bucket(config.Bucket), &config)}
+	client := &Client{
+		config:  config,
+		pather:  pather,
+		stats:   stats,
+		gcs:     NewGCS(ctx, sClient.Bucket(config.Bucket), &config),
+		sClient: sClient,
+	}
 
 	log.Infof("Initalized GCS backend with config: %s", config)
 	return client, nil
@@ -216,6 +228,14 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 	return result, nil
 }
 
+// Close closes the storage client
+func (c *Client) Close() error {
+	if c.sClient == nil {
+		return nil
+	}
+	return c.sClient.Close()
+}
+
 // isObjectNotFound is helper function for identify non-existing object error.
 func isObjectNotFound(err error) bool {
 	return err == storage.ErrObjectNotExist || err == storage.ErrBucketNotExist

lib/backend/gcsbackend/client_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/uber/kraken/core"
 	"github.com/uber/kraken/lib/backend"
 	mockgcsbackend "github.com/uber/kraken/mocks/lib/backend/gcsbackend"
+	"github.com/uber/kraken/utils/closers"
 	"github.com/uber/kraken/utils/mockutil"
 	"github.com/uber/kraken/utils/randutil"
 	"github.com/uber/kraken/utils/rwutil"
@@ -98,6 +99,7 @@ func TestClientStat(t *testing.T) {
 	defer cleanup()
 
 	client := mocks.new()
+	defer closers.Close(client)
 
 	var objectAttrs storage.ObjectAttrs
 	objectAttrs.Size = 100
@@ -116,6 +118,7 @@ func TestClientDownload(t *testing.T) {
 	defer cleanup()
 
 	client := mocks.new()
+	defer closers.Close(client)
 	data := randutil.Text(32)
 
 	mocks.gcs.EXPECT().Download(
@@ -194,6 +197,7 @@ func TestClientList(t *testing.T) {
 	defer cleanup()
 
 	client := mocks.new()
+	defer closers.Close(client)
 
 	contToken := ""
 	mocks.gcs.EXPECT().GetObjectIterator(

lib/backend/hdfsbackend/client.go

@@ -274,3 +274,9 @@ func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListR
 		Names: files,
 	}, nil
 }
+
+// Close closes the client and releases any held resources.
+func (c *Client) Close() error {
+	// No resources to close for HDFS client
+	return nil
+}

lib/backend/hdfsbackend/client_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/uber/kraken/core"
 	"github.com/uber/kraken/lib/backend/hdfsbackend/webhdfs"
 	mockwebhdfs "github.com/uber/kraken/mocks/lib/backend/hdfsbackend/webhdfs"
+	"github.com/uber/kraken/utils/closers"
 	"github.com/uber/kraken/utils/mockutil"
 	"github.com/uber/kraken/utils/randutil"
 
@@ -76,6 +77,7 @@ func TestClientStat(t *testing.T) {
 	defer cleanup()
 
 	client := mocks.new()
+	defer closers.Close(client)
 
 	mocks.webhdfs.EXPECT().GetFileStatus("/root/test").Return(webhdfs.FileStatus{Length: 32}, nil)
 
@@ -149,6 +151,7 @@ func TestClientList(t *testing.T) {
 			defer cleanup()
 
 			client := mocks.new()
+			defer closers.Close(client)
 
 			mocks.webhdfs.EXPECT().ListFileStatus("/root").Return([]webhdfs.FileStatus{{
 				PathSuffix: "foo",
@@ -214,6 +217,7 @@ func TestClientListErrorDoesNotLeakGoroutines(t *testing.T) {
 	defer cleanup()
 
 	client := mocks.new()
+	defer closers.Close(client)
 
 	initDirectoryTree(mocks, "/root", 10, 3) // 1000 nodes.
 

lib/backend/httpbackend/http.go

@@ -123,3 +123,9 @@ func (c *Client) Upload(namespace, name string, src io.Reader) error {
 func (c *Client) List(prefix string, opts ...backend.ListOption) (*backend.ListResult, error) {
 	return nil, errors.New("not supported")
 }
+
+// Close closes the client and releases any held resources.
+func (c *Client) Close() error {
+	// No resources to close for http client
+	return nil
+}

lib/backend/httpbackend/http_test.go

@@ -22,6 +22,7 @@ import (
 	"github.com/uber-go/tally"
 	"github.com/uber/kraken/core"
 	"github.com/uber/kraken/lib/backend/backenderrors"
+	"github.com/uber/kraken/utils/closers"
 	"github.com/uber/kraken/utils/memsize"
 	"github.com/uber/kraken/utils/randutil"
 	"github.com/uber/kraken/utils/testutil"
@@ -56,6 +57,7 @@ func TestHttpDownloadSuccess(t *testing.T) {
 	config := Config{DownloadURL: "http://" + addr + "/data/%s"}
 	client, err := NewClient(config, tally.NoopScope)
 	require.NoError(err)
+	defer closers.Close(client)
 
 	var b bytes.Buffer
 	require.NoError(client.Download(core.NamespaceFixture(), "data", &b))
@@ -77,6 +79,7 @@ func TestHttpDownloadFileNotFound(t *testing.T) {
 	config := Config{DownloadURL: "http://" + addr + "/data/%s"}
 	client, err := NewClient(config, tally.NoopScope)
 	require.NoError(err)
+	defer closers.Close(client)
 
 	var b bytes.Buffer
 	require.Equal(backenderrors.ErrBlobNotFound, client.Download(core.NamespaceFixture(), "data", &b))
@@ -92,6 +95,7 @@ func TestDownloadMalformedURLThrowsError(t *testing.T) {
 	config := Config{DownloadURL: "http://" + addr + "/data"}
 	client, err := NewClient(config, tally.NoopScope)
 	require.NoError(err)
+	defer closers.Close(client)
 
 	var b bytes.Buffer
 	require.Error(client.Download(core.NamespaceFixture(), "data", &b))

lib/backend/manager.go

@@ -168,3 +168,16 @@ func (m *Manager) CheckReadiness() error {
 	}
 	return nil
 }
+
+func (m *Manager) Close() error {
+	totalErrors := make([]error, 0)
+	for _, b := range m.backends {
+		if err := b.client.Close(); err != nil {
+			totalErrors = append(totalErrors, fmt.Errorf("closing backend for namespace '%s': %s", b.regexp.String(), err))
+		}
+	}
+	if len(totalErrors) > 0 {
+		return errors.Join(totalErrors...)
+	}
+	return nil
+}

lib/backend/noop.go

@@ -48,3 +48,9 @@ func (c NoopClient) Download(namespace, name string, dst io.Writer) error {
 func (c NoopClient) List(prefix string, opts ...ListOption) (*ListResult, error) {
 	return nil, nil
 }
+
+// Close closes the client and releases any held resources.
+func (c NoopClient) Close() error {
+	// No resources to close for noop client
+	return nil
+}