external store: use separate ctx for GCS clients (#60402) (#60502)

close #60155

Author: ti-chi-bot

br/pkg/storage/gcs.go

@@ -109,8 +109,9 @@ type GCSStorage struct {
 	clientCnt int64
 	clientOps []option.ClientOption
 
-	handles []*storage.BucketHandle
-	clients []*storage.Client
+	handles      []*storage.BucketHandle
+	clients      []*storage.Client
+	clientCancel context.CancelFunc
 }
 
 // GetBucketHandle gets the handle to the GCS API on the bucket.
@@ -334,6 +335,7 @@ func (s *GCSStorage) Rename(ctx context.Context, oldFileName, newFileName string
 
 // Close implements ExternalStorage interface.
 func (s *GCSStorage) Close() {
+	s.clientCancel()
 	for _, client := range s.clients {
 		if err := client.Close(); err != nil {
 			log.Warn("failed to close gcs client", zap.Error(err))
@@ -415,39 +417,73 @@ skipHandleCred:
 	return ret, nil
 }
 
-// Reset resets the GCS storage.
+// Reset resets the GCS storage. Reset should not be used concurrently with
+// Close.
 func (s *GCSStorage) Reset(ctx context.Context) error {
 	logutil.Logger(ctx).Info("resetting gcs storage")
 
-	for _, client := range s.clients {
-		_ = client.Close()
-	}
-
-	s.clients = make([]*storage.Client, gcsClientCnt)
-	eg, egCtx := util.NewErrorGroupWithRecoverWithCtx(ctx)
-	for i := range s.clients {
-		eg.Go(func() error {
-			client, err := storage.NewClient(egCtx, s.clientOps...)
+	s.cancelAndCloseGCSClients()
+
+	clientCtx, clientCancel := context.WithCancel(context.Background())
+	s.clients = make([]*storage.Client, 0, gcsClientCnt)
+	wg := util.WaitGroupWrapper{}
+	cliCh := make(chan *storage.Client)
+	wg.RunWithLog(func() {
+		for range gcsClientCnt {
+			select {
+			case cli := <-cliCh:
+				s.clients = append(s.clients, cli)
+			case <-ctx.Done():
+				clientCancel()
+				return
+			case <-clientCtx.Done():
+				return
+			}
+		}
+	})
+	firstErr := atomic.NewError(nil)
+	for range gcsClientCnt {
+		wg.RunWithLog(func() {
+			client, err := storage.NewClient(clientCtx, s.clientOps...)
 			if err != nil {
-				return errors.Trace(err)
+				firstErr.CompareAndSwap(nil, err)
+				clientCancel()
+				return
 			}
 			client.SetRetry(storage.WithErrorFunc(shouldRetry), storage.WithPolicy(storage.RetryAlways))
-			s.clients[i] = client
-			return nil
+			select {
+			case cliCh <- client:
+			case <-clientCtx.Done():
+			}
 		})
 	}
-	err := eg.Wait()
-	if err != nil {
+	wg.Wait()
+	if err := firstErr.Load(); err != nil {
+		s.cancelAndCloseGCSClients()
 		return errors.Trace(err)
 	}
 
+	s.clientCancel = clientCancel
 	s.handles = make([]*storage.BucketHandle, gcsClientCnt)
 	for i := range s.handles {
 		s.handles[i] = s.clients[i].Bucket(s.gcs.Bucket)
 	}
 	return nil
 }
 
+func (s *GCSStorage) cancelAndCloseGCSClients() {
+	if s.clientCancel != nil {
+		s.clientCancel()
+		s.clientCancel = nil
+	}
+
+	for _, client := range s.clients {
+		if client != nil {
+			_ = client.Close()
+		}
+	}
+}
+
 func shouldRetry(err error) bool {
 	if storage.ShouldRetry(err) {
 		return true

br/pkg/storage/gcs_test.go

@@ -9,6 +9,8 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"testing"
 	"time"
@@ -570,3 +572,30 @@ func TestSpeedReadManyFiles(t *testing.T) {
 	require.NoError(t, eg.Wait())
 	t.Logf("read %d large files cost %v", len(testFiles), time.Since(now))
 }
+
+func TestCtxUsage(t *testing.T) {
+	httpSvr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer httpSvr.Close()
+
+	ctx := context.Background()
+	gcs := &backuppb.GCS{
+		Endpoint:      httpSvr.URL,
+		Bucket:        "test",
+		Prefix:        "prefix",
+		StorageClass:  "NEARLINE",
+		PredefinedAcl: "private",
+		CredentialsBlob: fmt.Sprintf(`
+{
+	"type":"external_account",
+	"audience":"//iam.googleapis.com/projects/1234567890123/locations/global/workloadIdentityPools/my-pool/providers/my-provider",
+	"subject_token_type":"urn:ietf:params:oauth:token-type:access_token",
+	"credential_source":{"url":"%s"}
+}`, httpSvr.URL),
+	}
+	stg, err := NewGCSStorage(ctx, gcs, &ExternalStorageOptions{})
+	require.NoError(t, err)
+
+	_, err = stg.FileExists(ctx, "key")
+	// before the fix, it's context canceled error
+	require.ErrorContains(t, err, "invalid_request")
+}