Repository: buchgr/bazel-remote
Pull request #345 (2020-09-07 05:57:59)

add new key format for S3 to avoid per-prefix rate limits
The new key format is not used by default, since it would invalidate
the contents of pre-existing caches.

* Add cache.Key for common usage
* Refactor disk cache to utilize cache.Key
* Refactor s3proxy to utilize cache.Key


Patch:
diff --git a/README.md b/README.md
index 24e4b5a..fb14aa7 100644
--- a/README.md
+++ b/README.md
@@ -137,6 +137,7 @@ GLOBAL OPTIONS:
    --s3.disable_ssl                   Whether to disable TLS/SSL when using the S3 cache backend. (default: false, ie enable TLS/SSL) [$BAZEL_REMOTE_S3_DISABLE_SSL]
    --s3.iam_role_endpoint value       Endpoint for using IAM security credentials. By default it will look for credentials in the standard locations for the AWS platform. [$BAZEL_REMOTE_S3_IAM_ROLE_ENDPOINT]
    --s3.region value                  The AWS region. Required when not specifying S3/minio access keys. [$BAZEL_REMOTE_S3_REGION]
+   --s3.key_version value             Set to 1 for the legacy flat key format, or 2 for the newer format that reduces the impact of S3 rate limits. (default: 1) [$BAZEL_REMOTE_S3_KEY_VERSION]
    --disable_http_ac_validation       Whether to disable ActionResult validation for HTTP requests. (default: false, ie enable validation) [$BAZEL_REMOTE_DISABLE_HTTP_AC_VALIDATION]
    --disable_grpc_ac_deps_check       Whether to disable ActionResult dependency checks for gRPC GetActionResult requests. (default: false, ie enable ActionCache dependency checks) [$BAZEL_REMOTE_DISABLE_GRPS_AC_DEPS_CHECK]
    --enable_ac_key_instance_mangling  Whether to enable mangling ActionCache keys with non-empty instance names. (default: false, ie disable mangling) [$BAZEL_REMOTE_ENABLE_AC_KEY_INSTANCE_MANGLING]
@@ -205,6 +206,7 @@ host: localhost
 #  access_key_id: EXAMPLE_ACCESS_KEY
 #  secret_access_key: EXAMPLE_SECRET_KEY
 #  disable_ssl: true
+#  key_version: 2
 #
 # Provide either access_key_id/secret_access_key, or iam_role_endpoint/region.
 # iam_role_endpoint can also be left empty, and figured out automatically.
@@ -246,7 +248,7 @@ If you want the docker container to run in the background pass the `-d` flag rig
 You can adjust the maximum cache size by appending `--max_size=N`, where N is
 the maximum size in Gibibytes.
 
-### Kubernetes note 
+### Kubernetes note
 
 Don't name your deployment `bazel-remote`!
 
@@ -323,3 +325,10 @@ To avoid leaking your password in log files, you can place this flag in a
 For more details, see Bazel's [remote
 caching](https://docs.bazel.build/versions/master/remote-caching.html#run-bazel-using-the-remote-cache)
 documentation.
+
+## AWS S3 note
+
+To avoid per-prefix rate limiting with Amazon S3, you may want to try using
+`--s3.key_format=2`, which stores blobs across a larger number of prefixes.
+Reference:
+[Optimizing Amazon S3 Performance](https://docs.aws.amazon.com/AmazonS3/latest/dev/optimizing-performance.html).
diff --git a/cache/cache.go b/cache/cache.go
index 12dfd11..49c7cc9 100644
--- a/cache/cache.go
+++ b/cache/cache.go
@@ -4,6 +4,7 @@ import (
 	"crypto/sha256"
 	"encoding/hex"
 	"io"
+	"path/filepath"
 )
 
 // EntryKind describes the kind of cache entry
@@ -85,3 +86,8 @@ func TransformActionCacheKey(key, instance string, logger Logger) string {
 
 	return newKey
 }
+
+// Key returns the proper cache key for an entry kind and hash.
+func Key(kind EntryKind, hash string) string {
+	return filepath.Join(kind.String(), hash[:2], hash)
+}
diff --git a/cache/disk/disk.go b/cache/disk/disk.go
index 1805112..40c4a6a 100644
--- a/cache/disk/disk.go
+++ b/cache/disk/disk.go
@@ -247,7 +247,7 @@ func (c *Cache) Put(kind cache.EntryKind, hash string, size int64, r io.Reader)
 		return nil
 	}
 
-	key := cacheKey(kind, hash)
+	key := cache.Key(kind, hash)
 
 	var tf *os.File // Tempfile.
 
@@ -486,7 +486,7 @@ func (c *Cache) Get(kind cache.EntryKind, hash string, size int64) (rc io.ReadCl
 	}
 
 	var err error
-	key := cacheKey(kind, hash)
+	key := cache.Key(kind, hash)
 
 	var tf *os.File // Tempfile we will write to.
 
@@ -588,7 +588,7 @@ func (c *Cache) Contains(kind cache.EntryKind, hash string, size int64) (bool, i
 	}
 
 	foundSize := int64(-1)
-	key := cacheKey(kind, hash)
+	key := cache.Key(kind, hash)
 
 	c.mu.Lock()
 	val, exists := c.lru.Get(key)
@@ -640,12 +640,8 @@ func ensureDirExists(path string) {
 	}
 }
 
-func cacheKey(kind cache.EntryKind, hash string) string {
-	return filepath.Join(kind.String(), hash[:2], hash)
-}
-
 func cacheFilePath(kind cache.EntryKind, cacheDir string, hash string) string {
-	return filepath.Join(cacheDir, cacheKey(kind, hash))
+	return filepath.Join(cacheDir, cache.Key(kind, hash))
 }
 
 // GetValidatedActionResult returns a valid ActionResult and its serialized
diff --git a/cache/s3proxy/BUILD.bazel b/cache/s3proxy/BUILD.bazel
index 193322b..14612eb 100644
--- a/cache/s3proxy/BUILD.bazel
+++ b/cache/s3proxy/BUILD.bazel
@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "go_default_library",
@@ -14,3 +14,10 @@ go_library(
         "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
     ],
 )
+
+go_test(
+    name = "go_default_test",
+    srcs = ["s3proxy_test.go"],
+    embed = [":go_default_library"],
+    deps = ["//cache:go_default_library"],
+)
diff --git a/cache/s3proxy/s3proxy.go b/cache/s3proxy/s3proxy.go
index b361c74..de424e5 100644
--- a/cache/s3proxy/s3proxy.go
+++ b/cache/s3proxy/s3proxy.go
@@ -29,6 +29,7 @@ type s3Cache struct {
 	mcore        *minio.Core
 	prefix       string
 	bucket       string
+	keyVersion   int
 	uploadQueue  chan<- uploadReq
 	accessLogger cache.Logger
 	errorLogger  cache.Logger
@@ -96,6 +97,7 @@ func New(s3Config *config.S3CloudStorageConfig, accessLogger cache.Logger,
 		mcore:        minioCore,
 		prefix:       s3Config.Prefix,
 		bucket:       s3Config.Bucket,
+		keyVersion:   s3Config.KeyVersion,
 		uploadQueue:  uploadQueue,
 		accessLogger: accessLogger,
 		errorLogger:  errorLogger,
@@ -113,10 +115,16 @@ func New(s3Config *config.S3CloudStorageConfig, accessLogger cache.Logger,
 }
 
 func (c *s3Cache) objectKey(hash string, kind cache.EntryKind) string {
+	baseKey := fmt.Sprintf("%s/%s", kind, hash)
+	if c.keyVersion == 2 {
+		baseKey = cache.Key(kind, hash)
+	}
+
 	if c.prefix == "" {
-		return fmt.Sprintf("%s/%s", kind, hash)
+		return baseKey
 	}
-	return fmt.Sprintf("%s/%s/%s", c.prefix, kind, hash)
+
+	return fmt.Sprintf("%s/%s", c.prefix, baseKey)
 }
 
 // Helper function for logging responses
diff --git a/config/config.go b/config/config.go
index 8b4e7bd..c9efc2d 100644
--- a/config/config.go
+++ b/config/config.go
@@ -20,6 +20,7 @@ type S3CloudStorageConfig struct {
 	DisableSSL      bool   `yaml:"disable_ssl"`
 	IAMRoleEndpoint string `yaml:"iam_role_endpoint"`
 	Region          string `yaml:"region"`
+	KeyVersion      int    `yaml:"key_version"`
 }
 
 // GoogleCloudStorageConfig stores the configuration of a GCS proxy backend.
@@ -184,6 +185,10 @@ func validateConfig(c *Config) error {
 		if c.S3CloudStorage.AccessKeyID != "" && c.S3CloudStorage.IAMRoleEndpoint != "" {
 			return errors.New("Expected either 's3.access_key_id' or 's3.iam_role_endpoint', found both")
 		}
+
+		if c.S3CloudStorage.KeyVersion < 1 || c.S3CloudStorage.KeyVersion > 2 {
+			return fmt.Errorf("s3.key_version must be either 1 or 2, found %d", c.S3CloudStorage.KeyVersion)
+		}
 	}
 
 	return nil
diff --git a/main.go b/main.go
index cc4c2a8..9887686 100644
--- a/main.go
+++ b/main.go
@@ -206,6 +206,13 @@ func main() {
 			Usage:   "The AWS region. Required when not specifying S3/minio access keys.",
 			EnvVars: []string{"BAZEL_REMOTE_S3_REGION"},
 		},
+		&cli.IntFlag{
+			Name:        "s3.key_version",
+			Usage:       "Set to 1 for the legacy flat key format, or 2 for the newer format that reduces the impact of S3 rate limits.",
+			Value:       1,
+			DefaultText: "1",
+			EnvVars:     []string{"BAZEL_REMOTE_S3_KEY_VERSION"},
+		},
 		&cli.BoolFlag{
 			Name:        "disable_http_ac_validation",
 			Usage:       "Whether to disable ActionResult validation for HTTP requests.",
@@ -256,6 +263,7 @@ func main() {
 					DisableSSL:      ctx.Bool("s3.disable_ssl"),
 					IAMRoleEndpoint: ctx.String("s3.iam_role_endpoint"),
 					Region:          ctx.String("s3.region"),
+					KeyVersion:      ctx.Int("s3.key_version"),
 				}
 			}
 			c, err = config.New(


Test patch:
diff --git a/cache/s3proxy/s3proxy_test.go b/cache/s3proxy/s3proxy_test.go
new file mode 100644
index 0000000..30a5245
--- /dev/null
+++ b/cache/s3proxy/s3proxy_test.go
@@ -0,0 +1,48 @@
+package s3proxy
+
+import (
+	"log"
+	"testing"
+
+	"github.com/buchgr/bazel-remote/cache"
+)
+
+func TestObjectKey(t *testing.T) {
+	result, expected := "", ""
+	logger := &log.Logger{}
+	tc := &s3Cache{
+		mcore:        nil,
+		bucket:       "test",
+		keyVersion:   1,
+		uploadQueue:  make(chan uploadReq, 1),
+		accessLogger: logger,
+		errorLogger:  logger,
+	}
+
+	// Legacy key format tests
+	tc.keyVersion = 1
+	tc.prefix = ""
+	result, expected = tc.objectKey("1234", cache.CAS), "cas/1234"
+	checkTestCase(t, result, expected, "legacy format without prefix")
+
+	tc.prefix = "test"
+	result, expected = tc.objectKey("1234", cache.CAS), "test/cas/1234"
+	checkTestCase(t, result, expected, "legacy format with prefix")
+
+	// New key format tests
+	tc.keyVersion = 2
+
+	tc.prefix = ""
+	result, expected = tc.objectKey("1234", cache.CAS), "cas/12/1234"
+	checkTestCase(t, result, expected, "new format with prefix")
+
+	tc.prefix = "test"
+	result, expected = tc.objectKey("1234", cache.CAS), "test/cas/12/1234"
+	checkTestCase(t, result, expected, "new format with prefix")
+}
+
+func checkTestCase(t *testing.T, result string, expected string, testCase string) {
+	if result != expected {
+		t.Errorf("%s objectKey did not match. (result: '%s' expected: '%s'", testCase, result, expected)
+	}
+}
diff --git a/config/config_test.go b/config/config_test.go
index 33e9308..a3b8604 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -160,6 +160,7 @@ s3_proxy:
   prefix: test-prefix
   access_key_id: EXAMPLE_ACCESS_KEY
   secret_access_key: EXAMPLE_SECRET_KEY
+  key_version: 2
 `
 	config, err := newFromYaml([]byte(yaml))
 	if err != nil {
@@ -177,6 +178,7 @@ s3_proxy:
 			Prefix:          "test-prefix",
 			AccessKeyID:     "EXAMPLE_ACCESS_KEY",
 			SecretAccessKey: "EXAMPLE_SECRET_KEY",
+			KeyVersion:      2,
 		},
 	}