Repository: buchgr/bazel-remote
Pull request #530 (2022-03-09 23:40:37)

Improve GetValidatedActionResult() performance by validating action result blobs in parallel
Speed up `GetValidatedActionResult()` for proxy backends by validating
result blobs in parallel.

Previous behavior would iterate through result blobs sequentially. An
action with hundreds (or thousands) of result blobs to check at 10-20ms
each could easily run into timeouts for the caller.

Really, this is just reusing some of the functionality that already exists in `FindMissingCasBlobs()`.  As a result, this should also improve performance for the local cache lookups since it is using `findMissingLocalCAS()` to check the local cache in batches.

justfalter: I've got some tests to fix up in this first
justfalter: Tests added, should be good to go.
justfalter: Any thoughts @mostynb? This has made a pretty significant difference in my build environment (bazel-remote + aws s3). 

Patch:
diff --git a/cache/disk/disk.go b/cache/disk/disk.go
index d460f53..b0557ce 100644
--- a/cache/disk/disk.go
+++ b/cache/disk/disk.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -1116,13 +1117,12 @@ func (c *diskCache) GetValidatedActionResult(ctx context.Context, hash string) (
 		return nil, nil, err // Should we return "not found" instead of an error?
 	}
 
+	pendingValidations := []*pb.Digest{}
+
 	for _, f := range result.OutputFiles {
 		// f was validated in validate.ActionResult but blobs were not checked for existence
 		if len(f.Contents) == 0 {
-			found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
-			if !found {
-				return nil, nil, nil // aka "not found"
-			}
+			pendingValidations = append(pendingValidations, f.Digest)
 		}
 	}
 
@@ -1156,40 +1156,33 @@ func (c *diskCache) GetValidatedActionResult(ctx context.Context, hash string) (
 		}
 
 		for _, f := range tree.Root.GetFiles() {
-			if f.Digest == nil {
-				continue
-			}
-			found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
-			if !found {
-				return nil, nil, nil // aka "not found"
+			if f.Digest != nil {
+				pendingValidations = append(pendingValidations, f.Digest)
 			}
 		}
 
 		for _, child := range tree.GetChildren() {
 			for _, f := range child.GetFiles() {
-				if f.Digest == nil {
-					continue
-				}
-				found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
-				if !found {
-					return nil, nil, nil // aka "not found"
+				if f.Digest != nil {
+					pendingValidations = append(pendingValidations, f.Digest)
 				}
 			}
 		}
 	}
 
 	if result.StdoutDigest != nil {
-		found, _ := c.Contains(ctx, cache.CAS, result.StdoutDigest.Hash, result.StdoutDigest.SizeBytes)
-		if !found {
-			return nil, nil, nil // aka "not found"
-		}
+		pendingValidations = append(pendingValidations, result.StdoutDigest)
 	}
 
 	if result.StderrDigest != nil {
-		found, _ := c.Contains(ctx, cache.CAS, result.StderrDigest.Hash, result.StderrDigest.SizeBytes)
-		if !found {
-			return nil, nil, nil // aka "not found"
-		}
+		pendingValidations = append(pendingValidations, result.StderrDigest)
+	}
+
+	err = c.findMissingCasBlobsInternal(ctx, pendingValidations, true)
+	if errors.Is(err, errMissingBlob) {
+		return nil, nil, nil // aka "not found"
+	} else if err != nil {
+		return nil, nil, err
 	}
 
 	return result, acdata, nil
diff --git a/cache/disk/findmissing.go b/cache/disk/findmissing.go
index 92d2e8f..a099851 100644
--- a/cache/disk/findmissing.go
+++ b/cache/disk/findmissing.go
@@ -2,6 +2,7 @@ package disk
 
 import (
 	"context"
+	"errors"
 	"sync"
 
 	"github.com/buchgr/bazel-remote/cache"
@@ -13,19 +14,52 @@ import (
 )
 
 type proxyCheck struct {
-	wg     *sync.WaitGroup
-	digest **pb.Digest
-	ctx    context.Context
+	wg          *sync.WaitGroup
+	digest      **pb.Digest
+	ctx         context.Context
+	onProxyMiss func()
 }
 
+var errMissingBlob = errors.New("a blob could not be found")
+
 // Optimised implementation of FindMissingBlobs, which batches local index
 // lookups and performs concurrent proxy lookups for local cache misses.
 // Returns a slice with the blobs that are missing from the cache.
 //
 // Note that this modifies the input slice and returns a subset of it.
 func (c *diskCache) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest) ([]*pb.Digest, error) {
+	err := c.findMissingCasBlobsInternal(ctx, blobs, false)
+	if err != nil {
+		return nil, err
+	}
+	return filterNonNil(blobs), nil
+}
+
+// Identifies local and proxy cache misses for blobs. Modifies the input `blobs` slice such that found
+// blobs are replaced with nil, while the missing digests remain unchanged.
+//
+// When failFast is true and a blob could not be found in the local cache nor in the
+// proxy back end, the search will immediately terminate and errMissingBlob will be returned. Given that the
+// search is terminated early, the contents of blobs will only have partially been updated.
+func (c *diskCache) findMissingCasBlobsInternal(ctx context.Context, blobs []*pb.Digest, failFast bool) error {
+	// batchSize moderates how long the cache lock is held by findMissingLocalCAS.
 	const batchSize = 20
 
+	var cancelContextForFailFast context.CancelFunc = nil
+	cancelledDueToFailFast := false
+
+	if failFast && c.proxy != nil {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
+		defer cancel()
+
+		cancelContextForFailFast = func() {
+			// Indicate that we were canceled so that we can fail fast.
+			cancelledDueToFailFast = true
+			cancel()
+		}
+	}
+
 	var wg sync.WaitGroup
 
 	var chunk []*pb.Digest
@@ -34,7 +68,10 @@ func (c *diskCache) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest)
 	for len(remaining) > 0 {
 		select {
 		case <-ctx.Done():
-			return nil, status.Error(codes.Canceled, "Request was cancelled")
+			if cancelledDueToFailFast {
+				return errMissingBlob
+			}
+			return status.Error(codes.Canceled, "Request was cancelled")
 		default:
 		}
 
@@ -47,27 +84,71 @@ func (c *diskCache) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest)
 		}
 
 		numMissing := c.findMissingLocalCAS(chunk)
-		if numMissing > 0 && c.proxy != nil {
-			wg.Add(numMissing)
+		if numMissing == 0 {
+			continue
+		}
+
+		if c.proxy == nil && failFast {
+			// There's no proxy, there are missing blobs from the local cache, and we are failing fast.
+			return errMissingBlob
+		} else if c.proxy != nil {
 			for i := range chunk {
-				if chunk[i] != nil {
-					c.containsQueue <- proxyCheck{
-						wg:     &wg,
-						digest: &chunk[i],
-						ctx:    ctx,
+				if chunk[i] == nil {
+					continue
+				}
+
+				if chunk[i].SizeBytes > c.maxProxyBlobSize {
+					// The blob would exceed the limit, so skip it.
+					if failFast {
+						return errMissingBlob
 					}
+					continue
+				}
+
+				// Adding to the containsQueue channel may have blocked on a previous iteration,
+				// so check to see if the context has cancelled.
+				select {
+				case <-ctx.Done():
+					if cancelledDueToFailFast {
+						return errMissingBlob
+					}
+					return status.Error(codes.Canceled, "Request was cancelled")
+				default:
+				}
+
+				wg.Add(1)
+				c.containsQueue <- proxyCheck{
+					wg:     &wg,
+					digest: &chunk[i],
+					ctx:    ctx,
+					// When failFast is true, onProxyMiss will have been set to a function that
+					// will cancel the context, causing the remaining proxyChecks to short-circuit.
+					onProxyMiss: cancelContextForFailFast,
 				}
 			}
 		}
 	}
 
 	if c.proxy != nil {
-		wg.Wait()
+		// Adapt the waitgroup for select
+		waitCh := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(waitCh)
+		}()
+
+		// Wait for all proxyChecks to finish or a context cancellation.
+		select {
+		case <-ctx.Done():
+			if cancelledDueToFailFast {
+				return errMissingBlob
+			}
+			return status.Error(codes.Canceled, "Request was cancelled")
+		case <-waitCh: // Everything in the waitgroup has finished.
+		}
 	}
 
-	missingBlobs := filterNonNil(blobs)
-
-	return missingBlobs, nil
+	return nil
 }
 
 // Move all the non-nil items in the input slice to the
@@ -88,21 +169,27 @@ func filterNonNil(blobs []*pb.Digest) []*pb.Digest {
 // of missing blobs.
 func (c *diskCache) findMissingLocalCAS(blobs []*pb.Digest) int {
 	var exists bool
+	var item lruItem
 	var key string
 	missing := 0
 
 	c.mu.Lock()
 
 	for i := range blobs {
-		if blobs[i].SizeBytes == 0 {
+		if blobs[i].SizeBytes == 0 && blobs[i].Hash == emptySha256 {
 			c.accessLogger.Printf("GRPC CAS HEAD %s OK", blobs[i].Hash)
 			blobs[i] = nil
 			continue
 		}
 
+		foundSize := int64(-1)
 		key = cache.LookupKey(cache.CAS, blobs[i].Hash)
-		_, exists = c.lru.Get(key)
+		item, exists = c.lru.Get(key)
 		if exists {
+			foundSize = item.size
+		}
+
+		if exists && !isSizeMismatch(blobs[i].SizeBytes, foundSize) {
 			c.accessLogger.Printf("GRPC CAS HEAD %s OK", blobs[i].Hash)
 			blobs[i] = nil
 		} else {
@@ -118,6 +205,17 @@ func (c *diskCache) findMissingLocalCAS(blobs []*pb.Digest) int {
 func (c *diskCache) containsWorker() {
 	var ok bool
 	for req := range c.containsQueue {
+		if req.ctx != nil {
+			select {
+			case <-req.ctx.Done():
+				// Fast-fail if the context has already been cancelled.
+				c.accessLogger.Printf("GRPC CAS HEAD %s CANCELLED", (*req.digest).Hash)
+				req.wg.Done()
+				continue
+			default:
+			}
+		}
+
 		ok, _ = c.proxy.Contains(req.ctx, cache.CAS, (*req.digest).Hash)
 		if ok {
 			c.accessLogger.Printf("GRPC CAS HEAD %s OK", (*req.digest).Hash)
@@ -126,6 +224,9 @@ func (c *diskCache) containsWorker() {
 			*(req.digest) = nil
 		} else {
 			c.accessLogger.Printf("GRPC CAS HEAD %s NOT FOUND", (*req.digest).Hash)
+			if req.onProxyMiss != nil {
+				req.onProxyMiss()
+			}
 		}
 		req.wg.Done()
 	}
diff --git a/go.mod b/go.mod
index 3b2618f..855f4e2 100644
--- a/go.mod
+++ b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/google/uuid v1.3.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/klauspost/compress v1.14.1
-	github.com/minio/minio-go/v7 v7.0.20
+	github.com/minio/minio-go/v7 v7.0.24-0.20220302185830-f2d40be3e82f
 	github.com/mostynb/go-grpc-compression v1.1.16
 	github.com/mostynb/zstdpool-syncpool v0.0.12
 	github.com/prometheus/client_golang v1.11.0
diff --git a/go.sum b/go.sum
index 8652352..e9951dd 100644
--- a/go.sum
+++ b/go.sum
@@ -265,6 +265,8 @@ github.com/minio/md5-simd v1.1.2 h1:Gdi1DZK69+ZVMoNHRXJyNcxrMA4dSxoYHZSQbirFg34=
 github.com/minio/md5-simd v1.1.2/go.mod h1:MzdKDxYpY2BT9XQFocsiZf/NKVtR7nkE4RoEpN+20RM=
 github.com/minio/minio-go/v7 v7.0.20 h1:0+Xt1SkCKDgcx5cmo3UxXcJ37u5Gy+/2i/+eQYqmYJw=
 github.com/minio/minio-go/v7 v7.0.20/go.mod h1:ei5JjmxwHaMrgsMrn4U/+Nmg+d8MKS1U2DAn1ou4+Do=
+github.com/minio/minio-go/v7 v7.0.24-0.20220302185830-f2d40be3e82f h1:5+NaPkyAjRvzl81ZN3N7VgMYsQwbYwnj3NQ2CJpxwX0=
+github.com/minio/minio-go/v7 v7.0.24-0.20220302185830-f2d40be3e82f/go.mod h1:ei5JjmxwHaMrgsMrn4U/+Nmg+d8MKS1U2DAn1ou4+Do=
 github.com/minio/sha256-simd v0.1.1/go.mod h1:B5e1o+1/KgNmWrSQK08Y6Z1Vb5pwIktudl0J58iy0KM=
 github.com/minio/sha256-simd v1.0.0 h1:v1ta+49hkWZyvaKwrQB8elexRqm6Y0aMLjCNsrYxo6g=
 github.com/minio/sha256-simd v1.0.0/go.mod h1:OuYzVNI5vcoYIAmbIvHPl3N3jUzVedXbKy5RFepssQM=
diff --git a/utils/BUILD.bazel b/utils/BUILD.bazel
index cdeaa62..8c280b1 100644
--- a/utils/BUILD.bazel
+++ b/utils/BUILD.bazel
@@ -5,4 +5,7 @@ go_library(
     srcs = ["testutils.go"],
     importpath = "github.com/buchgr/bazel-remote/utils",
     visibility = ["//visibility:public"],
+    deps = [
+        "//genproto/build/bazel/remote/execution/v2:go_default_library",
+    ],
 )


Test patch:
diff --git a/cache/disk/findmissing_test.go b/cache/disk/findmissing_test.go
index 5d1e851..6c6c4f5 100644
--- a/cache/disk/findmissing_test.go
+++ b/cache/disk/findmissing_test.go
@@ -1,13 +1,19 @@
 package disk
 
 import (
+	"bytes"
 	"context"
+	"errors"
+	"fmt"
 	"io"
+	"io/ioutil"
+	"os"
 	"sync"
 	"testing"
 
 	"github.com/buchgr/bazel-remote/cache"
 	testutils "github.com/buchgr/bazel-remote/utils"
+	"google.golang.org/protobuf/proto"
 
 	pb "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/execution/v2"
 )
@@ -148,3 +154,337 @@ func TestContainsWorker(t *testing.T) {
 		t.Error("Expected digests[1] to not be found in the proxy and left as-is")
 	}
 }
+
+type proxyAdapter struct {
+	cache Cache
+}
+
+func NewProxyAdapter(cache Cache) (*proxyAdapter, error) {
+	if cache == nil {
+		return nil, fmt.Errorf("cache cannot be nil")
+	}
+	return &proxyAdapter{
+		cache: cache,
+	}, nil
+}
+
+func (p *proxyAdapter) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
+	err := p.cache.Put(kind, hash, size, rc)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func (p *proxyAdapter) Get(ctx context.Context, kind cache.EntryKind, hash string) (rc io.ReadCloser, size int64, err error) {
+	return p.cache.Get(ctx, kind, hash, size, 0)
+}
+
+func (p *proxyAdapter) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {
+	return p.cache.Contains(ctx, kind, hash, -1)
+}
+
+func TestFindMissingCasBlobsWithProxy(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cacheDir := tempDir(t)
+	defer os.RemoveAll(cacheDir)
+	proxyCacheDir := tempDir(t)
+	defer os.RemoveAll(proxyCacheDir)
+
+	cacheForProxy, err := New(proxyCacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	proxy, err := NewProxyAdapter(cacheForProxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testCache, err := New(cacheDir, 10*1024, WithProxyBackend(proxy), WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	data1, digest1 := testutils.RandomDataAndDigest(100)
+	_, digest2 := testutils.RandomDataAndDigest(200)
+	data3, digest3 := testutils.RandomDataAndDigest(300)
+	_, digest4 := testutils.RandomDataAndDigest(400)
+
+	proxy.Put(cache.CAS, digest1.Hash, digest1.SizeBytes, ioutil.NopCloser(bytes.NewReader(data1)))
+	proxy.Put(cache.CAS, digest3.Hash, digest3.SizeBytes, ioutil.NopCloser(bytes.NewReader(data3)))
+
+	missing, err := testCache.FindMissingCasBlobs(ctx, []*pb.Digest{
+		&digest1,
+		&digest2,
+		&digest3,
+		&digest4,
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(missing) != 2 {
+		t.Fatalf("Expected missing array to have exactly two entries, got %d", len(missing))
+	}
+
+	if !proto.Equal(missing[0], &digest2) {
+		t.Fatalf("Expected missing[0] == digest2, got: %+v", missing[0])
+	}
+
+	if !proto.Equal(missing[1], &digest4) {
+		t.Fatalf("Expected missing[1] == digest4, got: %+v", missing[1])
+	}
+}
+
+func TestFindMissingCasBlobsWithProxyFailFast(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cacheDir := tempDir(t)
+	defer os.RemoveAll(cacheDir)
+	proxyCacheDir := tempDir(t)
+	defer os.RemoveAll(proxyCacheDir)
+
+	cacheForProxy, err := New(proxyCacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	proxy, err := NewProxyAdapter(cacheForProxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Explicitly avoid using WithProxyBackEnd, as we want to control the workers.
+	testCacheI, err := New(cacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	actualDiskCache := testCacheI.(*diskCache)
+	actualDiskCache.proxy = proxy
+	actualDiskCache.containsQueue = make(chan proxyCheck, 4)
+	defer func() {
+		close(actualDiskCache.containsQueue)
+	}()
+	// Spawn a single worker.
+	go actualDiskCache.containsWorker()
+
+	data1, digest1 := testutils.RandomDataAndDigest(100)
+	_, digest2 := testutils.RandomDataAndDigest(200)
+	data3, digest3 := testutils.RandomDataAndDigest(300)
+	_, digest4 := testutils.RandomDataAndDigest(400)
+
+	proxy.Put(cache.CAS, digest1.Hash, digest1.SizeBytes, ioutil.NopCloser(bytes.NewReader(data1)))
+	proxy.Put(cache.CAS, digest3.Hash, digest3.SizeBytes, ioutil.NopCloser(bytes.NewReader(data3)))
+
+	blobs := []*pb.Digest{
+		&digest1,
+		&digest2,
+		&digest3,
+		&digest4,
+	}
+	err = actualDiskCache.findMissingCasBlobsInternal(ctx, blobs, true)
+
+	if !errors.Is(err, errMissingBlob) {
+		t.Fatalf("Expected err to be errMissingBlob, got: %s", err)
+	}
+
+	if proto.Equal(blobs[0], &digest1) {
+		t.Fatalf("Expected blobs[0] to equal digest1, got: %+v", blobs[0])
+	}
+}
+
+func TestFindMissingCasBlobsWithProxyFailFastNoneMissing(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cacheDir := tempDir(t)
+	defer os.RemoveAll(cacheDir)
+	proxyCacheDir := tempDir(t)
+	defer os.RemoveAll(proxyCacheDir)
+
+	cacheForProxy, err := New(proxyCacheDir, 40*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	proxy, err := NewProxyAdapter(cacheForProxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Explicitly avoid using WithProxyBackEnd, as we want to control the workers.
+	testCacheI, err := New(cacheDir, 40*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	actualDiskCache := testCacheI.(*diskCache)
+	actualDiskCache.proxy = proxy
+	actualDiskCache.containsQueue = make(chan proxyCheck, 4)
+	defer func() {
+		close(actualDiskCache.containsQueue)
+	}()
+	// Spawn a single worker.
+	go actualDiskCache.containsWorker()
+
+	data1, digest1 := testutils.RandomDataAndDigest(100)
+	data2, digest2 := testutils.RandomDataAndDigest(200)
+	data3, digest3 := testutils.RandomDataAndDigest(300)
+	data4, digest4 := testutils.RandomDataAndDigest(400)
+
+	proxy.Put(cache.CAS, digest1.Hash, digest1.SizeBytes, ioutil.NopCloser(bytes.NewReader(data1)))
+	proxy.Put(cache.CAS, digest2.Hash, digest2.SizeBytes, ioutil.NopCloser(bytes.NewReader(data2)))
+	proxy.Put(cache.CAS, digest3.Hash, digest3.SizeBytes, ioutil.NopCloser(bytes.NewReader(data3)))
+	proxy.Put(cache.CAS, digest4.Hash, digest4.SizeBytes, ioutil.NopCloser(bytes.NewReader(data4)))
+
+	blobs := []*pb.Digest{
+		&digest1,
+		&digest2,
+		&digest3,
+		&digest4,
+	}
+
+	err = actualDiskCache.findMissingCasBlobsInternal(ctx, blobs, true)
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if blobs[0] != nil {
+		t.Fatalf("Expected blobs[0] to be nil, got: %+v", blobs[0])
+	}
+
+	if blobs[1] != nil {
+		t.Fatalf("Expected blobs[1] to be nil, got: %+v", blobs[1])
+	}
+
+	if blobs[2] != nil {
+		t.Fatalf("Expected blobs[3] to be nil, got: %+v", blobs[2])
+	}
+
+	if blobs[3] != nil {
+		t.Fatalf("Expected blobs[3] to be nil, got: %+v", blobs[3])
+	}
+}
+
+func TestFindMissingCasBlobsWithProxyFailFastMaxProxyBlobSize(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cacheDir := tempDir(t)
+	defer os.RemoveAll(cacheDir)
+	proxyCacheDir := tempDir(t)
+	defer os.RemoveAll(proxyCacheDir)
+
+	cacheForProxy, err := New(proxyCacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	proxy, err := NewProxyAdapter(cacheForProxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Explicitly avoid using WithProxyBackEnd, as we want to control the workers.
+	testCacheI, err := New(cacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()), WithProxyMaxBlobSize(300))
+	if err != nil {
+		t.Fatal(err)
+	}
+	actualDiskCache := testCacheI.(*diskCache)
+	actualDiskCache.proxy = proxy
+	actualDiskCache.containsQueue = make(chan proxyCheck, 4)
+	defer func() {
+		close(actualDiskCache.containsQueue)
+	}()
+	// Spawn a single worker.
+	go actualDiskCache.containsWorker()
+
+	data1, digest1 := testutils.RandomDataAndDigest(100)
+	data2, digest2 := testutils.RandomDataAndDigest(200)
+	data3, digest3 := testutils.RandomDataAndDigest(300) // We expect this blob to not be found.
+
+	// Put blobs directly into proxy backend, where it will not be filtered out.
+	proxy.Put(cache.CAS, digest1.Hash, digest1.SizeBytes, ioutil.NopCloser(bytes.NewReader(data1)))
+	proxy.Put(cache.CAS, digest2.Hash, digest2.SizeBytes, ioutil.NopCloser(bytes.NewReader(data2)))
+	proxy.Put(cache.CAS, digest3.Hash, digest3.SizeBytes, ioutil.NopCloser(bytes.NewReader(data3)))
+
+	blobs := []*pb.Digest{
+		&digest1,
+		&digest2,
+		&digest3,
+	}
+	err = actualDiskCache.findMissingCasBlobsInternal(ctx, blobs, true)
+
+	if !errors.Is(err, errMissingBlob) {
+		t.Fatalf("Expected err to be errMissingBlob, got: %s", err)
+	}
+
+	if blobs[0] == nil {
+		t.Fatalf("Expected blobs[0] to be nil, got: %+v", blobs[0])
+	}
+
+	if blobs[1] == nil {
+		t.Fatalf("Expected blobs[1] to be nil, got: %+v", blobs[1])
+	}
+
+	if !proto.Equal(blobs[2], &digest3) {
+		t.Fatalf("Expected blobs[2] == digest3, got %+v", blobs[2])
+	}
+}
+
+func TestFindMissingCasBlobsWithProxyMaxProxyBlobSize(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	cacheDir := tempDir(t)
+	defer os.RemoveAll(cacheDir)
+	proxyCacheDir := tempDir(t)
+	defer os.RemoveAll(proxyCacheDir)
+
+	cacheForProxy, err := New(proxyCacheDir, 10*1024, WithAccessLogger(testutils.NewSilentLogger()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	proxy, err := NewProxyAdapter(cacheForProxy)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testCache, err := New(cacheDir, 10*1024, WithProxyBackend(proxy), WithAccessLogger(testutils.NewSilentLogger()), WithProxyMaxBlobSize(500))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	data1, digest1 := testutils.RandomDataAndDigest(100)
+	data2, digest2 := testutils.RandomDataAndDigest(600)
+
+	proxy.Put(cache.CAS, digest1.Hash, digest1.SizeBytes, ioutil.NopCloser(bytes.NewReader(data1)))
+	proxy.Put(cache.CAS, digest2.Hash, digest2.SizeBytes, ioutil.NopCloser(bytes.NewReader(data2)))
+
+	missing, err := testCache.FindMissingCasBlobs(ctx, []*pb.Digest{
+		&digest1,
+		&digest2,
+	})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(missing) != 1 {
+		t.Fatalf("Expected missing array to have exactly one entry, got %d", len(missing))
+	}
+
+	if !proto.Equal(missing[0], &digest2) {
+		t.Fatalf("Expected missing[0] == digest2, got %+v", missing[0])
+	}
+}
diff --git a/utils/testutils.go b/utils/testutils.go
index 73074a4..e06d28d 100644
--- a/utils/testutils.go
+++ b/utils/testutils.go
@@ -7,6 +7,8 @@ import (
 	"io/ioutil"
 	"log"
 	"testing"
+
+	pb "github.com/buchgr/bazel-remote/genproto/build/bazel/remote/execution/v2"
 )
 
 // TempDir creates a temporary directory and returns its name. If an error
@@ -38,6 +40,14 @@ func RandomDataAndHash(size int64) ([]byte, string) {
 	return data, hashStr
 }
 
+func RandomDataAndDigest(size int64) ([]byte, pb.Digest) {
+	data, hash := RandomDataAndHash(size)
+	return data, pb.Digest{
+		Hash:      hash,
+		SizeBytes: size,
+	}
+}
+
 // NewSilentLogger returns a cheap logger that doesn't print anything, useful
 // for tests.
 func NewSilentLogger() *log.Logger {