fix(storage): stop BucketExist scan after first result (#1031)

Signed-off-by: Xuewei Wang <xueweiwang.sallery@gmail.com>

Author: Sallery-X

internal/storage/minio.go

@@ -322,28 +322,31 @@ func (m *MinioClient) ListPrefix(ctx context.Context, prefix string, recursive b
 
 // BucketExist checks if the bucket exists by listing a single object.
 // We use ListObjects instead of BucketExists (HEAD bucket) to minimize
-// required S3 permissions. MaxKeys=1 ensures only one object is fetched
-// to avoid iterating the entire bucket, which can timeout on large buckets.
+// required S3 permissions. MaxKeys=1 avoids iterating the entire bucket,
+// which can timeout on large buckets, while preserving caller context
+// propagation.
 func (m *MinioClient) BucketExist(ctx context.Context, prefix string) (bool, error) {
-	opts := minio.ListObjectsOptions{Prefix: prefix, MaxKeys: 1}
-	// MaxKeys=1 ensures at most one result, so the channel closes naturally.
-	// We must drain the channel to avoid goroutine leaks (see minio-go docs).
-	exists := true
-	var retErr error
-	for obj := range m.cli.ListObjects(ctx, m.cfg.Bucket, opts) {
-		if obj.Err != nil {
-			if minio.ToErrorResponse(obj.Err).Code == "NoSuchBucket" {
-				exists = false
-			} else {
-				retErr = fmt.Errorf("storage: %s list objects %w", m.cfg.Provider, obj.Err)
-			}
-		}
+	subCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	objCh := m.cli.ListObjects(subCtx, m.cfg.Bucket, minio.ListObjectsOptions{
+		Prefix:  prefix,
+		MaxKeys: 1,
+	})
+
+	obj, ok := <-objCh
+	if !ok {
+		return true, nil
 	}
 
-	if retErr != nil {
-		return false, retErr
+	if obj.Err != nil {
+		if minio.ToErrorResponse(obj.Err).Code == "NoSuchBucket" {
+			return false, nil
+		}
+		return false, fmt.Errorf("storage: %s list objects %w", m.cfg.Provider, obj.Err)
 	}
-	return exists, nil
+
+	return true, nil
 }
 
 func (m *MinioClient) CreateBucket(ctx context.Context) error {

internal/storage/minio_test.go

@@ -1,13 +1,19 @@
 package storage
 
 import (
+	"context"
 	"errors"
+	"io"
 	"net/http"
+	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/minio/minio-go/v7"
+	"github.com/minio/minio-go/v7/pkg/credentials"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSplitIntoParts(t *testing.T) {
@@ -89,3 +95,148 @@ func TestIsDeleteSuccessful(t *testing.T) {
 		})
 	}
 }
+
+func TestBucketExistStopsAfterFirstResult(t *testing.T) {
+	var listReqCount atomic.Int32
+	var firstReqPath string
+	var firstReqMaxKeys string
+	var firstReqContinuationToken string
+	cli, err := newInternalMinio(Config{
+		Provider: "s3",
+		Endpoint: "example.com",
+		UseSSL:   true,
+		Bucket:   "test-bucket",
+		Credential: Credential{
+			Type: Static,
+			AK:   "ak",
+			SK:   "sk",
+		},
+	}, &minio.Options{
+		Secure: true,
+		Creds:  credentials.NewStaticV4("ak", "sk", ""),
+		Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+			if _, ok := r.URL.Query()["location"]; ok {
+				return &http.Response{
+					StatusCode: http.StatusOK,
+					Header:     http.Header{"Content-Type": []string{"application/xml"}},
+					Body:       io.NopCloser(strings.NewReader(`<LocationConstraint xmlns="http://s3.amazonaws.com/doc/2006-03-01/"></LocationConstraint>`)),
+					Request:    r,
+				}, nil
+			}
+
+			if r.URL.Query().Get("list-type") != "2" {
+				return nil, errors.New("unexpected request")
+			}
+
+			if token := r.URL.Query().Get("continuation-token"); token != "" {
+				<-r.Context().Done()
+				return nil, r.Context().Err()
+			}
+
+			listReqCount.Add(1)
+			firstReqPath = r.URL.Path
+			firstReqMaxKeys = r.URL.Query().Get("max-keys")
+			firstReqContinuationToken = r.URL.Query().Get("continuation-token")
+
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Header:     http.Header{"Content-Type": []string{"application/xml"}},
+				Body: io.NopCloser(strings.NewReader(`<?xml version="1.0" encoding="UTF-8"?>
+<ListBucketResult>
+  <Name>test-bucket</Name>
+  <Prefix></Prefix>
+  <MaxKeys>1</MaxKeys>
+  <IsTruncated>true</IsTruncated>
+  <Contents>
+    <Key>first-object</Key>
+    <Size>1</Size>
+  </Contents>
+  <NextContinuationToken>next-page</NextContinuationToken>
+</ListBucketResult>`)),
+				Request: r,
+			}, nil
+		}),
+	})
+	require.NoError(t, err)
+
+	exists, err := cli.BucketExist(context.Background(), "")
+	require.NoError(t, err)
+	assert.True(t, exists)
+	assert.EqualValues(t, 1, listReqCount.Load())
+	assert.Equal(t, "/test-bucket/", firstReqPath)
+	assert.Equal(t, "1", firstReqMaxKeys)
+	assert.Empty(t, firstReqContinuationToken)
+}
+
+func TestBucketExistReturnsFalseForNoSuchBucket(t *testing.T) {
+	cli, err := newInternalMinio(Config{
+		Provider: "s3",
+		Endpoint: "example.com",
+		UseSSL:   true,
+		Bucket:   "missing-bucket",
+		Credential: Credential{
+			Type: Static,
+			AK:   "ak",
+			SK:   "sk",
+		},
+	}, &minio.Options{
+		Secure: true,
+		Creds:  credentials.NewStaticV4("ak", "sk", ""),
+		Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: http.StatusNotFound,
+				Header:     http.Header{"Content-Type": []string{"application/xml"}},
+				Body: io.NopCloser(strings.NewReader(`<?xml version="1.0" encoding="UTF-8"?>
+<Error>
+  <Code>NoSuchBucket</Code>
+  <Message>The specified bucket does not exist</Message>
+  <BucketName>missing-bucket</BucketName>
+  <Resource>/missing-bucket/</Resource>
+  <RequestId>req</RequestId>
+  <HostId>host</HostId>
+</Error>`)),
+				Request: r,
+			}, nil
+		}),
+	})
+	require.NoError(t, err)
+
+	exists, err := cli.BucketExist(context.Background(), "")
+	require.NoError(t, err)
+	assert.False(t, exists)
+}
+
+func TestBucketExistPropagatesContextCancellation(t *testing.T) {
+	cli, err := newInternalMinio(Config{
+		Provider: "s3",
+		Endpoint: "example.com",
+		UseSSL:   true,
+		Bucket:   "test-bucket",
+		Credential: Credential{
+			Type: Static,
+			AK:   "ak",
+			SK:   "sk",
+		},
+	}, &minio.Options{
+		Secure: true,
+		Creds:  credentials.NewStaticV4("ak", "sk", ""),
+		Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
+			return nil, context.Canceled
+		}),
+	})
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	exists, err := cli.BucketExist(ctx, "")
+	assert.False(t, exists)
+	require.Error(t, err)
+	assert.ErrorIs(t, err, context.Canceled)
+}
+
+type roundTripperFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req)
+}