github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/file/s3file/bucketcache.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache-2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package s3file
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"time"
    11  
    12  	"github.com/aws/aws-sdk-go/aws"
    13  	"github.com/aws/aws-sdk-go/aws/credentials"
    14  	awsrequest "github.com/aws/aws-sdk-go/aws/request"
    15  	"github.com/aws/aws-sdk-go/aws/session"
    16  	"github.com/aws/aws-sdk-go/service/s3"
    17  	"github.com/aws/aws-sdk-go/service/s3/s3iface"
    18  	"github.com/aws/aws-sdk-go/service/s3/s3manager"
    19  	"github.com/Schaudge/grailbase/file"
    20  	"github.com/Schaudge/grailbase/sync/loadingcache"
    21  )
    22  
    23  // bucketRegionCacheDuration is chosen fairly arbitrarily. We expect region changes to be
    24  // extremely rare (deleting a bucket, then recreating elsewhere) so a long time seems fine.
    25  const bucketRegionCacheDuration = time.Hour
    26  
    27  // FindBucketRegion locates the AWS region in which bucket is located.
    28  // The lookup is cached internally.
    29  //
    30  // It assumes the region is in the "aws" partition, not other partitions like "aws-us-gov".
    31  // See: https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingBucket.html
    32  func FindBucketRegion(ctx context.Context, bucket string) (string, error) {
    33  	return globalBucketRegionCache.locate(ctx, bucket)
    34  }
    35  
    36  type bucketRegionCache struct {
    37  	cache loadingcache.Map
    38  	// getBucketRegionWithClient indirectly references s3manager.GetBucketRegionWithClient to
    39  	// allow unit testing.
    40  	getBucketRegionWithClient func(ctx aws.Context, svc s3iface.S3API, bucket string, opts ...awsrequest.Option) (string, error)
    41  }
    42  
    43  var (
    44  	globalBucketRegionCache = bucketRegionCache{
    45  		getBucketRegionWithClient: s3manager.GetBucketRegionWithClient,
    46  	}
    47  	bucketRegionClient = s3.New(
    48  		session.Must(session.NewSessionWithOptions(session.Options{
    49  			Config: aws.Config{
    50  				// This client is only used for looking up bucket locations, which doesn't
    51  				// require any credentials.
    52  				Credentials: credentials.AnonymousCredentials,
    53  				// Note: This region is just used to infer the relevant AWS partition (group of
    54  				// regions). This would fail for, say, "aws-us-gov", but we only use "aws".
    55  				// See: https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingBucket.html
    56  				Region: aws.String("us-west-2"),
    57  			},
    58  			SharedConfigState: session.SharedConfigDisable,
    59  		})),
    60  	)
    61  )
    62  
    63  func (c *bucketRegionCache) locate(ctx context.Context, bucket string) (string, error) {
    64  	var region string
    65  	err := c.cache.
    66  		GetOrCreate(bucket).
    67  		GetOrLoad(ctx, &region, func(ctx context.Context, opts *loadingcache.LoadOpts) (err error) {
    68  			opts.CacheFor(bucketRegionCacheDuration)
    69  			policy := newBackoffPolicy([]s3iface.S3API{bucketRegionClient}, file.Opts{})
    70  			for {
    71  				var ids s3RequestIDs
    72  				region, err = c.getBucketRegionWithClient(ctx,
    73  					bucketRegionClient, bucket, ids.captureOption())
    74  				if err == nil {
    75  					return nil
    76  				}
    77  				if !policy.shouldRetry(ctx, err, fmt.Sprintf("locate region: %s", bucket)) {
    78  					return annotate(err, ids, &policy)
    79  				}
    80  			}
    81  		})
    82  	if err != nil {
    83  		return "", err
    84  	}
    85  	return region, nil
    86  }