Repository: buildbuddy-io/buildbuddy
Pull request #9116: Optimize distributed.Cache.Get when lookaside cache is enabled
Author: login: vanja-p
Created: 2025-04-29T17:02:51Z | Merged: 2025-05-21T22:30:04Z

There's 2 optimizations here:
1. When the entry doesn't belong in the lookaside cache (like AC responses), don't create a `teeReadCloser` in `lookasideWriter` since it's unnecessary overhead.
2. When we do write to the lookaside cache, presize the buffer

Benchmark results:
```goos: linux
goarch: amd64
cpu: AMD Ryzen 9 7900X 12-Core Processor            
                                          │    base     │               after                │
                                          │   sec/op    │   sec/op     vs base               │
GetSingle/LookasideDistPebble10/AC-24       102.2µ ± 1%   100.7µ ± 1%  -1.54% (p=0.003 n=15)
GetSingle/LookasideDistPebble10/CAS-24      1.315µ ± 1%   1.321µ ± 1%       ~ (p=0.467 n=15)
GetSingle/LookasideDistPebble100/AC-24      104.8µ ± 2%   102.3µ ± 1%  -2.44% (p=0.000 n=15)
GetSingle/LookasideDistPebble100/CAS-24     1.341µ ± 1%   1.342µ ± 3%       ~ (p=0.340 n=15)
GetSingle/LookasideDistPebble1000/AC-24     108.2µ ± 0%   105.7µ ± 0%  -2.32% (p=0.000 n=15)
GetSingle/LookasideDistPebble1000/CAS-24    107.9µ ± 1%   107.8µ ± 1%       ~ (p=0.806 n=15)
GetSingle/LookasideDistPebble10000/AC-24    144.2µ ± 1%   135.4µ ± 1%  -6.08% (p=0.000 n=15)
GetSingle/LookasideDistPebble10000/CAS-24   139.9µ ± 1%   138.0µ ± 0%  -1.37% (p=0.001 n=15)
geomean                                     38.11µ        37.47µ       -1.69%

                                          │     base     │                after                │
                                          │     B/s      │     B/s       vs base               │
GetSingle/LookasideDistPebble10/AC-24       97.66Ki ± 0%   97.66Ki ± 0%       ~ (p=0.483 n=15)
GetSingle/LookasideDistPebble10/CAS-24      7.248Mi ± 1%   7.219Mi ± 1%       ~ (p=0.479 n=15)
GetSingle/LookasideDistPebble100/AC-24      927.7Ki ± 1%   957.0Ki ± 1%  +3.16% (p=0.000 n=15)
GetSingle/LookasideDistPebble100/CAS-24     71.11Mi ± 1%   71.06Mi ± 3%       ~ (p=0.340 n=15)
GetSingle/LookasideDistPebble1000/AC-24     8.812Mi ± 0%   9.022Mi ± 0%  +2.38% (p=0.000 n=15)
GetSingle/LookasideDistPebble1000/CAS-24    8.831Mi ± 1%   8.850Mi ± 1%       ~ (p=0.783 n=15)
GetSingle/LookasideDistPebble10000/AC-24    66.16Mi ± 1%   70.44Mi ± 1%  +6.47% (p=0.000 n=15)
GetSingle/LookasideDistPebble10000/CAS-24   68.17Mi ± 1%   69.11Mi ± 0%  +1.39% (p=0.000 n=15)
geomean                                     7.929Mi        8.058Mi       +1.62%

                                          │     base      │                after                 │
                                          │     B/op      │     B/op      vs base                │
GetSingle/LookasideDistPebble10/AC-24        30.58Ki ± 0%   29.87Ki ± 0%   -2.32% (p=0.000 n=15)
GetSingle/LookasideDistPebble10/CAS-24       1.385Ki ± 0%   1.385Ki ± 0%        ~ (p=0.710 n=15)
GetSingle/LookasideDistPebble100/AC-24       30.96Ki ± 0%   30.26Ki ± 0%   -2.25% (p=0.000 n=15)
GetSingle/LookasideDistPebble100/CAS-24      1.480Ki ± 0%   1.480Ki ± 0%        ~ (p=1.000 n=15)
GetSingle/LookasideDistPebble1000/AC-24      36.46Ki ± 0%   34.21Ki ± 0%   -6.15% (p=0.000 n=15)
GetSingle/LookasideDistPebble1000/CAS-24     35.43Ki ± 0%   35.43Ki ± 0%        ~ (p=0.493 n=15)
GetSingle/LookasideDistPebble10000/AC-24    104.91Ki ± 0%   50.74Ki ± 0%  -51.63% (p=0.000 n=15)
GetSingle/LookasideDistPebble10000/CAS-24    61.12Ki ± 0%   61.09Ki ± 0%        ~ (p=0.533 n=15)
geomean                                      18.87Ki        17.00Ki        -9.93%

                                          │    base    │                after                │
                                          │ allocs/op  │ allocs/op   vs base                 │
GetSingle/LookasideDistPebble10/AC-24       482.0 ± 0%   477.0 ± 0%  -1.04% (p=0.000 n=15)
GetSingle/LookasideDistPebble10/CAS-24      23.00 ± 0%   23.00 ± 0%       ~ (p=1.000 n=15) ¹
GetSingle/LookasideDistPebble100/AC-24      482.0 ± 0%   477.0 ± 0%  -1.04% (p=0.000 n=15)
GetSingle/LookasideDistPebble100/CAS-24     23.00 ± 0%   23.00 ± 0%       ~ (p=1.000 n=15) ¹
GetSingle/LookasideDistPebble1000/AC-24     484.0 ± 0%   477.0 ± 0%  -1.45% (p=0.000 n=15)
GetSingle/LookasideDistPebble1000/CAS-24    488.0 ± 0%   488.0 ± 0%       ~ (p=1.000 n=15) ¹
GetSingle/LookasideDistPebble10000/AC-24    489.0 ± 0%   473.0 ± 0%  -3.27% (p=0.000 n=15)
GetSingle/LookasideDistPebble10000/CAS-24   484.0 ± 0%   484.0 ± 0%       ~ (p=1.000 n=15) ¹
geomean                                     226.3        224.3       -0.85%
```


Diff:
diff --git a/enterprise/server/backends/distributed/distributed.go b/enterprise/server/backends/distributed/distributed.go
index d7423be039b..31e36afd132 100644
--- a/enterprise/server/backends/distributed/distributed.go
+++ b/enterprise/server/backends/distributed/distributed.go
@@ -343,6 +343,7 @@ func (t *teeReadCloser) Read(p []byte) (int, error) {
 	}
 	return read, t.lastReadErr
 }
+
 func (t *teeReadCloser) Close() error {
 	err := t.rc.Close()
 	if err == nil && t.lastReadErr != nil && t.lastReadErr != io.EOF {
@@ -389,17 +390,21 @@ func (c *Cache) addLookasideEntry(r *rspb.ResourceName, data []byte) {
 		c.log.Debugf("Not setting lookaside entry for resource: %s", r)
 		return
 	}
+	c.setLookasideEntry(k, data)
+}
+
+func (c *Cache) setLookasideEntry(lookasideKey string, data []byte) {
 	entry := lookasideCacheEntry{
 		createdAtMillis: time.Now().UnixMilli(),
 		data:            data,
 	}
 
 	c.lookasideMu.Lock()
-	if !c.lookaside.Contains(k) {
-		c.lookaside.Add(k, entry)
+	if !c.lookaside.Contains(lookasideKey) {
+		c.lookaside.Add(lookasideKey, entry)
 	}
 	c.lookasideMu.Unlock()
-	c.log.Debugf("Set %q in lookaside cache", k)
+	c.log.Debugf("Set %q in lookaside cache", lookasideKey)
 }
 
 // getLookasideEntry returns the resource and if it was found in the lookaside
@@ -442,11 +447,11 @@ func (c *Cache) getLookasideEntry(r *rspb.ResourceName) ([]byte, bool) {
 	return nil, false
 }
 
-func (c *Cache) lookasideWriter(r *rspb.ResourceName) (interfaces.CommittedWriteCloser, error) {
-	buffer := new(bytes.Buffer)
+func (c *Cache) lookasideWriter(r *rspb.ResourceName, lookasideKey string) (interfaces.CommittedWriteCloser, error) {
+	buffer := bytes.NewBuffer(make([]byte, 0, r.GetDigest().GetSizeBytes()))
 	wc := ioutil.NewCustomCommitWriteCloser(buffer)
 	wc.CommitFn = func(int64) error {
-		c.addLookasideEntry(r, buffer.Bytes())
+		c.setLookasideEntry(lookasideKey, buffer.Bytes())
 		return nil
 	}
 	return wc, nil
@@ -463,7 +468,11 @@ func (c *Cache) teeReadCloser(r *rspb.ResourceName, rc io.ReadCloser) io.ReadClo
 	if r.GetDigest().GetSizeBytes() > *maxLookasideEntryBytes {
 		return rc
 	}
-	lwc, err := c.lookasideWriter(r)
+	k, ok := lookasideKey(r)
+	if !ok {
+		return rc
+	}
+	lwc, err := c.lookasideWriter(r, k)
 	if err != nil {
 		return rc
 	}
diff --git a/enterprise/server/test/performance/cache/cache_test.go b/enterprise/server/test/performance/cache/cache_test.go
index f769acb90b3..0864d948947 100644
--- a/enterprise/server/test/performance/cache/cache_test.go
+++ b/enterprise/server/test/performance/cache/cache_test.go
@@ -37,8 +37,8 @@ func init() {
 	log.Configure()
 }
 
-func getAnonContext(t testing.TB, te *testenv.TestEnv) context.Context {
-	ctx, err := prefix.AttachUserPrefixToContext(context.Background(), te.GetAuthenticator())
+func getAnonContext(t testing.TB, env environment.Env) context.Context {
+	ctx, err := prefix.AttachUserPrefixToContext(context.Background(), env.GetAuthenticator())
 	if err != nil {
 		t.Fatalf("error attaching user prefix: %v", err)
 	}
@@ -99,14 +99,15 @@ func getMigrationCache(t testing.TB, env environment.Env, src, dest interfaces.C
 	return migration_cache.NewMigrationCache(env, config, src, dest)
 }
 
-func getDistributedCache(t testing.TB, te *testenv.TestEnv, c interfaces.Cache) interfaces.Cache {
+func getDistributedCache(t testing.TB, te environment.Env, c interfaces.Cache, lookasideCacheSizeBytes int64) interfaces.Cache {
 	listenAddr := fmt.Sprintf("localhost:%d", testport.FindFree(t))
 	conf := distributed.CacheConfig{
-		ListenAddr:         listenAddr,
-		GroupName:          "default",
-		ReplicationFactor:  1,
-		Nodes:              []string{listenAddr},
-		DisableLocalLookup: true,
+		ListenAddr:              listenAddr,
+		GroupName:               "default",
+		ReplicationFactor:       1,
+		Nodes:                   []string{listenAddr},
+		DisableLocalLookup:      true,
+		LookasideCacheSizeBytes: lookasideCacheSizeBytes,
 	}
 	dc, err := distributed.NewDistributedCache(te, c, conf, te.GetHealthChecker())
 	if err != nil {
@@ -117,7 +118,7 @@ func getDistributedCache(t testing.TB, te *testenv.TestEnv, c interfaces.Cache)
 	return dc
 }
 
-func getPebbleCache(t testing.TB, te *testenv.TestEnv) interfaces.Cache {
+func getPebbleCache(t testing.TB, te environment.Env) interfaces.Cache {
 	testRootDir := testfs.MakeTempDir(t)
 	pc, err := pebble_cache.NewPebbleCache(te, &pebble_cache.Options{
 		Name:          testRootDir,
@@ -179,8 +180,8 @@ func benchmarkRead(ctx context.Context, c interfaces.Cache, digestSizeBytes int6
 	}
 }
 
-func benchmarkGet(ctx context.Context, c interfaces.Cache, digestSizeBytes int64, b *testing.B) {
-	digestBufs := makeDigests(b, numDigests, digestSizeBytes, rspb.CacheType_CAS)
+func benchmarkGet(ctx context.Context, c interfaces.Cache, digestSizeBytes int64, b *testing.B, cacheType rspb.CacheType) {
+	digestBufs := makeDigests(b, numDigests, digestSizeBytes, cacheType)
 	setDigestsInCache(b, ctx, c, digestBufs)
 	b.ReportAllocs()
 	b.SetBytes(digestSizeBytes)
@@ -238,12 +239,13 @@ type namedCache struct {
 	Name string
 }
 
-func getAllCaches(b *testing.B, te *testenv.TestEnv) []*namedCache {
+func getAllCaches(b *testing.B, te environment.Env) []*namedCache {
 	flags.Set(b, "cache.distributed_cache.consistent_hash_function", "SHA256")
 	dc := getDiskCache(b, te)
-	ddc := getDistributedCache(b, te, dc)
+	ddc := getDistributedCache(b, te, dc, 0)
 	pc := getPebbleCache(b, te)
-	dpc := getDistributedCache(b, te, pc)
+	dpc := getDistributedCache(b, te, pc, 0)
+	lpc := getDistributedCache(b, te, getPebbleCache(b, te), 100_000)
 
 	time.Sleep(100 * time.Millisecond)
 	caches := []*namedCache{
@@ -252,8 +254,9 @@ func getAllCaches(b *testing.B, te *testenv.TestEnv) []*namedCache {
 		{ddc, "DistDisk"},
 		{getPebbleCache(b, te), "LocalPebble"},
 		{dpc, "DistPebble"},
+		{lpc, "LookasideDistPebble"},
 		{getMigrationCache(b, te, getPebbleCache(b, te), getPebbleCache(b, te)), "LocalMigration"},
-		{getDistributedCache(b, te, getMigrationCache(b, te, getPebbleCache(b, te), getPebbleCache(b, te))), "DistMigration"},
+		{getDistributedCache(b, te, getMigrationCache(b, te, getPebbleCache(b, te), getPebbleCache(b, te)), 0), "DistMigration"},
 	}
 	return caches
 }
@@ -295,15 +298,18 @@ func BenchmarkRead(b *testing.B) {
 
 func BenchmarkGetSingle(b *testing.B) {
 	sizes := []int64{10, 100, 1000, 10000}
+	// sizes := []int64{10}
 	te := testenv.GetTestEnv(b)
 	ctx := getAnonContext(b, te)
 
 	for _, cache := range getAllCaches(b, te) {
 		for _, size := range sizes {
-			name := fmt.Sprintf("%s%d", cache.Name, size)
-			b.Run(name, func(b *testing.B) {
-				benchmarkGet(ctx, cache, size, b)
-			})
+			for _, cacheType := range []rspb.CacheType{rspb.CacheType_AC, rspb.CacheType_CAS} {
+				name := fmt.Sprintf("%s%d/%s", cache.Name, size, cacheType)
+				b.Run(name, func(b *testing.B) {
+					benchmarkGet(ctx, cache, size, b, cacheType)
+				})
+			}
 		}
 	}
 }


Reviews:
user: login: jdhollen
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T16:59:37Z
user: login: vanja-p
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T17:48:51Z
user: login: tylerwilliams
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-21T21:39:58Z

Review comments:
diff_hunk: @@ -1093,34 +1127,35 @@ func (c *Cache) distributedReader(ctx context.Context, rn *rspb.ResourceName, of
 	return nil, status.NotFoundErrorf("Exhausted all peers attempting to read %q.", rn.GetDigest().GetHash())
 }
 
-// Below, in Get(), this value is the max initial allocatable buffer size.
-// Set it somewhat conservatively so that we're not DOSed by someone crafting
-// remote_instance_names that match this just to use memory.
-const maxInitialByteBufferSize = (1024 * 1024 * 4)
-
-func (c *Cache) Get(ctx context.Context, rn *rspb.ResourceName) ([]byte, error) {
-	r, err := c.distributedReader(ctx, rn, 0, 0, "Get" /*=metricsLabel*/)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-
-	var buf *bytes.Buffer
+func bufferSize(rn *rspb.ResourceName) int {
 	if rn.GetCacheType() == rspb.CacheType_CAS {
 		// If this is a CAS object, size the buffer to fit exactly.
-		buf = bytes.NewBuffer(make([]byte, 0, int(rn.GetDigest().GetSizeBytes())))
+		// Clamp the size between 0 and 10MB, to protect from invalid and
+		// malicious requests.
+		return min(10_000_000, max(0, int(rn.GetDigest().GetSizeBytes())))
 	} else if strings.HasPrefix(rn.GetInstanceName(), content_addressable_storage_server.TreeCacheRemoteInstanceName) {
 		// If this is a TreeCache entry that we wrote; pull the size
 		// from the remote instance name.
 		parts := strings.Split(rn.GetInstanceName(), "/")
 		if s, err := strconv.Atoi(parts[len(parts)-1]); err == nil {
-			buf = bytes.NewBuffer(make([]byte, 0, min(s, maxInitialByteBufferSize)))
+			// Limit this to 4MiB so that we're not DOSed by someone crafting
+			// remote_instance_names that match this just to use memory.
+			return min(1024*1024*4, s)
 		} else {
-			buf = new(bytes.Buffer)
+			return 0
 		}
-	} else {
-		buf = new(bytes.Buffer)
 	}
+	// The median and average AC results are less than 4KiB: go/action-result-size
+	return 4 * 1024
+
+}
path: enterprise/server/backends/distributed/distributed.go
user: login: jdhollen
body: im obligated to nitpick your missing newline here because you added one elsewhere
created_at: 2025-05-09T16:06:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1093,34 +1127,35 @@ func (c *Cache) distributedReader(ctx context.Context, rn *rspb.ResourceName, of
 	return nil, status.NotFoundErrorf("Exhausted all peers attempting to read %q.", rn.GetDigest().GetHash())
 }
 
-// Below, in Get(), this value is the max initial allocatable buffer size.
-// Set it somewhat conservatively so that we're not DOSed by someone crafting
-// remote_instance_names that match this just to use memory.
-const maxInitialByteBufferSize = (1024 * 1024 * 4)
-
-func (c *Cache) Get(ctx context.Context, rn *rspb.ResourceName) ([]byte, error) {
-	r, err := c.distributedReader(ctx, rn, 0, 0, "Get" /*=metricsLabel*/)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-
-	var buf *bytes.Buffer
+func bufferSize(rn *rspb.ResourceName) int {
 	if rn.GetCacheType() == rspb.CacheType_CAS {
 		// If this is a CAS object, size the buffer to fit exactly.
-		buf = bytes.NewBuffer(make([]byte, 0, int(rn.GetDigest().GetSizeBytes())))
+		// Clamp the size between 0 and 10MB, to protect from invalid and
+		// malicious requests.
+		return min(10_000_000, max(0, int(rn.GetDigest().GetSizeBytes())))
path: enterprise/server/backends/distributed/distributed.go
user: login: jdhollen
body: this seems like it would cause us to create and throw out (large) buffers a few times for internal calls to this api for cas objects.  the benchmark you ran caps out at something smaller than this size, right?

i'm notabout other places where we're really doing this in practice, but we at least do it for very-large cached trees (look at the first line of `lookupCachedTreeNodeInCAS`)
created_at: 2025-05-09T16:59:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1093,34 +1127,35 @@ func (c *Cache) distributedReader(ctx context.Context, rn *rspb.ResourceName, of
 	return nil, status.NotFoundErrorf("Exhausted all peers attempting to read %q.", rn.GetDigest().GetHash())
 }
 
-// Below, in Get(), this value is the max initial allocatable buffer size.
-// Set it somewhat conservatively so that we're not DOSed by someone crafting
-// remote_instance_names that match this just to use memory.
-const maxInitialByteBufferSize = (1024 * 1024 * 4)
-
-func (c *Cache) Get(ctx context.Context, rn *rspb.ResourceName) ([]byte, error) {
-	r, err := c.distributedReader(ctx, rn, 0, 0, "Get" /*=metricsLabel*/)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-
-	var buf *bytes.Buffer
+func bufferSize(rn *rspb.ResourceName) int {
 	if rn.GetCacheType() == rspb.CacheType_CAS {
 		// If this is a CAS object, size the buffer to fit exactly.
-		buf = bytes.NewBuffer(make([]byte, 0, int(rn.GetDigest().GetSizeBytes())))
+		// Clamp the size between 0 and 10MB, to protect from invalid and
+		// malicious requests.
+		return min(10_000_000, max(0, int(rn.GetDigest().GetSizeBytes())))
 	} else if strings.HasPrefix(rn.GetInstanceName(), content_addressable_storage_server.TreeCacheRemoteInstanceName) {
 		// If this is a TreeCache entry that we wrote; pull the size
 		// from the remote instance name.
 		parts := strings.Split(rn.GetInstanceName(), "/")
 		if s, err := strconv.Atoi(parts[len(parts)-1]); err == nil {
-			buf = bytes.NewBuffer(make([]byte, 0, min(s, maxInitialByteBufferSize)))
+			// Limit this to 4MiB so that we're not DOSed by someone crafting
+			// remote_instance_names that match this just to use memory.
+			return min(1024*1024*4, s)
 		} else {
-			buf = new(bytes.Buffer)
+			return 0
 		}
-	} else {
-		buf = new(bytes.Buffer)
 	}
+	// The median and average AC results are less than 4KiB: go/action-result-size
+	return 4 * 1024
+
+}
path: enterprise/server/backends/distributed/distributed.go
user: login: vanja-p
body: Please do. Why is gofmt not catching this?
created_at: 2025-05-09T17:45:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1093,34 +1127,35 @@ func (c *Cache) distributedReader(ctx context.Context, rn *rspb.ResourceName, of
 	return nil, status.NotFoundErrorf("Exhausted all peers attempting to read %q.", rn.GetDigest().GetHash())
 }
 
-// Below, in Get(), this value is the max initial allocatable buffer size.
-// Set it somewhat conservatively so that we're not DOSed by someone crafting
-// remote_instance_names that match this just to use memory.
-const maxInitialByteBufferSize = (1024 * 1024 * 4)
-
-func (c *Cache) Get(ctx context.Context, rn *rspb.ResourceName) ([]byte, error) {
-	r, err := c.distributedReader(ctx, rn, 0, 0, "Get" /*=metricsLabel*/)
-	if err != nil {
-		return nil, err
-	}
-	defer r.Close()
-
-	var buf *bytes.Buffer
+func bufferSize(rn *rspb.ResourceName) int {
 	if rn.GetCacheType() == rspb.CacheType_CAS {
 		// If this is a CAS object, size the buffer to fit exactly.
-		buf = bytes.NewBuffer(make([]byte, 0, int(rn.GetDigest().GetSizeBytes())))
+		// Clamp the size between 0 and 10MB, to protect from invalid and
+		// malicious requests.
+		return min(10_000_000, max(0, int(rn.GetDigest().GetSizeBytes())))
path: enterprise/server/backends/distributed/distributed.go
user: login: vanja-p
body: Yep, the benchmarks are smaller than this cap. It will definitely make things slower when reads are bigger than the cap. I don't mind removing it, since we didn't have it before, but in https://github.com/buildbuddy-io/buildbuddy/pull/8887, we decided to add a cap to the pebble cache.
created_at: 2025-05-09T17:48:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: vanja-p
created_at: 2025-05-06T15:40:23Z
body: ping
author_association: CONTRIBUTOR
user: login: jdhollen
created_at: 2025-05-19T23:28:20Z
body: Sorry for being slow here--i'm theoretically fine with this change, but there are some conflicts that you'll need to sync past now.
author_association: CONTRIBUTOR
user: login: vanja-p
created_at: 2025-05-21T21:06:15Z
body: > Sorry for being slow here--i'm theoretically fine with this change, but there are some conflicts that you'll need to sync past now.

I handled the conflicts and simplified this change. Turns out that implementing WriterTo here doesn't help. Also, I reverted most of the buffer size changes, but I did presize the `lookasideWriter`, since it's a trivial change.

I also added benchmarks for the AC action type, so that the main benefit is clear. Note that `GetSingle/LookasideDistPebble10000/AC` allocates half the memory now.
author_association: CONTRIBUTOR