Repository: buchgr/bazel-remote
Pull request #454 (2021-07-03 21:31:25)

plumb context.Context through to proxy backends


Patch:
diff --git a/cache/cache.go b/cache/cache.go
index 7c5281e..0f47784 100644
--- a/cache/cache.go
+++ b/cache/cache.go
@@ -1,6 +1,7 @@
 package cache
 
 import (
+	"context"
 	"crypto/sha256"
 	"encoding/hex"
 	"io"
@@ -63,10 +64,10 @@ func (e *Error) Error() string {
 // Implementations are expected to be safe for concurrent use.
 type Proxy interface {
 
-	// Put makes a reasonable effort to upload the cache item identified by
-	// `hash` with logical size `size`, whose data is readable from `rc` to
-	// the proxy backend. The data available in `rc` is in the same format
-	// as used by the disk.Cache instance.
+	// Put makes a reasonable effort to asynchronously upload the cache
+	// item identified by `hash` with logical size `size`, whose data is
+	// readable from `rc` to the proxy backend. The data available in
+	// `rc` is in the same format as used by the disk.Cache instance.
 	//
 	// This is allowed to fail silently (for example when under heavy load).
 	Put(kind EntryKind, hash string, size int64, rc io.ReadCloser)
@@ -75,12 +76,12 @@ type Proxy interface {
 	// `hash` can be read, its logical size, and an error if something went
 	// wrong. The data available from `rc` is in the same format as used by
 	// the disk.Cache instance.
-	Get(kind EntryKind, hash string) (rc io.ReadCloser, size int64, err error)
+	Get(ctx context.Context, kind EntryKind, hash string) (rc io.ReadCloser, size int64, err error)
 
 	// Contains returns whether or not the cache item exists on the
 	// remote end, and the size if it exists (and -1 if the size is
 	// unknown).
-	Contains(kind EntryKind, hash string) (bool, int64)
+	Contains(ctx context.Context, kind EntryKind, hash string) (bool, int64)
 }
 
 // TransformActionCacheKey takes an ActionCache key and an instance name
diff --git a/cache/disk/disk.go b/cache/disk/disk.go
index 7f8d41b..4e50b67 100644
--- a/cache/disk/disk.go
+++ b/cache/disk/disk.go
@@ -2,6 +2,7 @@ package disk
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha256"
 	"fmt"
 	"io"
@@ -794,18 +795,18 @@ var errOnlyCompressedCAS = &cache.Error{
 // item is not found, the io.ReadCloser will be nil. If some error occurred
 // when processing the request, then it is returned. Callers should provide
 // the `size` of the item to be retrieved, or -1 if unknown.
-func (c *Cache) Get(kind cache.EntryKind, hash string, size int64, offset int64) (rc io.ReadCloser, s int64, rErr error) {
-	return c.get(kind, hash, size, offset, false)
+func (c *Cache) Get(ctx context.Context, kind cache.EntryKind, hash string, size int64, offset int64) (rc io.ReadCloser, s int64, rErr error) {
+	return c.get(ctx, kind, hash, size, offset, false)
 }
 
 // GetZstd is just like Get, except the data available from rc is zstandard
 // compressed. Note that the returned `s` value still refers to the amount
 // of data once it has been decompressed.
-func (c *Cache) GetZstd(hash string, size int64, offset int64) (rc io.ReadCloser, s int64, rErr error) {
-	return c.get(cache.CAS, hash, size, offset, true)
+func (c *Cache) GetZstd(ctx context.Context, hash string, size int64, offset int64) (rc io.ReadCloser, s int64, rErr error) {
+	return c.get(ctx, cache.CAS, hash, size, offset, true)
 }
 
-func (c *Cache) get(kind cache.EntryKind, hash string, size int64, offset int64, zstd bool) (rc io.ReadCloser, s int64, rErr error) {
+func (c *Cache) get(ctx context.Context, kind cache.EntryKind, hash string, size int64, offset int64, zstd bool) (rc io.ReadCloser, s int64, rErr error) {
 	// The hash format is checked properly in the http/grpc code.
 	// Just perform a simple/fast check here, to catch bad tests.
 	if len(hash) != sha256HashStrSize {
@@ -882,7 +883,7 @@ func (c *Cache) get(kind cache.EntryKind, hash string, size int64, offset int64,
 		return nil, -1, nil
 	}
 
-	r, foundSize, err := c.proxy.Get(kind, hash)
+	r, foundSize, err := c.proxy.Get(ctx, kind, hash)
 	if r != nil {
 		defer r.Close()
 	}
@@ -958,7 +959,7 @@ func (c *Cache) get(kind cache.EntryKind, hash string, size int64, offset int64,
 // one) will be checked.
 //
 // Callers should provide the `size` of the item, or -1 if unknown.
-func (c *Cache) Contains(kind cache.EntryKind, hash string, size int64) (bool, int64) {
+func (c *Cache) Contains(ctx context.Context, kind cache.EntryKind, hash string, size int64) (bool, int64) {
 
 	// The hash format is checked properly in the http/grpc code.
 	// Just perform a simple/fast check here, to catch bad tests.
@@ -985,7 +986,7 @@ func (c *Cache) Contains(kind cache.EntryKind, hash string, size int64) (bool, i
 	}
 
 	if c.proxy != nil {
-		exists, foundSize = c.proxy.Contains(kind, hash)
+		exists, foundSize = c.proxy.Contains(ctx, kind, hash)
 		if exists && !isSizeMismatch(size, foundSize) {
 			return true, foundSize
 		}
@@ -1026,9 +1027,9 @@ func ensureDirExists(path string) {
 // value from the CAS if it and all its dependencies are also available. If
 // not, nil values are returned. If something unexpected went wrong, return
 // an error.
-func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte, error) {
+func (c *Cache) GetValidatedActionResult(ctx context.Context, hash string) (*pb.ActionResult, []byte, error) {
 
-	rc, sizeBytes, err := c.Get(cache.AC, hash, -1, 0)
+	rc, sizeBytes, err := c.Get(ctx, cache.AC, hash, -1, 0)
 	if rc != nil {
 		defer rc.Close()
 	}
@@ -1053,7 +1054,7 @@ func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte,
 
 	for _, f := range result.OutputFiles {
 		if len(f.Contents) == 0 {
-			found, _ := c.Contains(cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
+			found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
 			if !found {
 				return nil, nil, nil // aka "not found"
 			}
@@ -1061,7 +1062,7 @@ func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte,
 	}
 
 	for _, d := range result.OutputDirectories {
-		r, size, err := c.Get(cache.CAS, d.TreeDigest.Hash, d.TreeDigest.SizeBytes, 0)
+		r, size, err := c.Get(ctx, cache.CAS, d.TreeDigest.Hash, d.TreeDigest.SizeBytes, 0)
 		if r == nil {
 			return nil, nil, err // aka "not found", or an err if non-nil
 		}
@@ -1092,7 +1093,7 @@ func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte,
 			if f.Digest == nil {
 				continue
 			}
-			found, _ := c.Contains(cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
+			found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
 			if !found {
 				return nil, nil, nil // aka "not found"
 			}
@@ -1103,7 +1104,7 @@ func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte,
 				if f.Digest == nil {
 					continue
 				}
-				found, _ := c.Contains(cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
+				found, _ := c.Contains(ctx, cache.CAS, f.Digest.Hash, f.Digest.SizeBytes)
 				if !found {
 					return nil, nil, nil // aka "not found"
 				}
@@ -1112,14 +1113,14 @@ func (c *Cache) GetValidatedActionResult(hash string) (*pb.ActionResult, []byte,
 	}
 
 	if result.StdoutDigest != nil {
-		found, _ := c.Contains(cache.CAS, result.StdoutDigest.Hash, result.StdoutDigest.SizeBytes)
+		found, _ := c.Contains(ctx, cache.CAS, result.StdoutDigest.Hash, result.StdoutDigest.SizeBytes)
 		if !found {
 			return nil, nil, nil // aka "not found"
 		}
 	}
 
 	if result.StderrDigest != nil {
-		found, _ := c.Contains(cache.CAS, result.StderrDigest.Hash, result.StderrDigest.SizeBytes)
+		found, _ := c.Contains(ctx, cache.CAS, result.StderrDigest.Hash, result.StderrDigest.SizeBytes)
 		if !found {
 			return nil, nil, nil // aka "not found"
 		}
diff --git a/cache/disk/findmissing.go b/cache/disk/findmissing.go
index 8645710..77634ec 100644
--- a/cache/disk/findmissing.go
+++ b/cache/disk/findmissing.go
@@ -15,6 +15,7 @@ import (
 type proxyCheck struct {
 	wg     *sync.WaitGroup
 	digest **pb.Digest
+	ctx    context.Context
 }
 
 // Optimised implementation of FindMissingBlobs, which batches local index
@@ -53,6 +54,7 @@ func (c *Cache) FindMissingCasBlobs(ctx context.Context, blobs []*pb.Digest) ([]
 					c.containsQueue <- proxyCheck{
 						wg:     &wg,
 						digest: &chunk[i],
+						ctx:    ctx,
 					}
 				}
 			}
@@ -116,7 +118,7 @@ func (c *Cache) findMissingLocalCAS(blobs []*pb.Digest) int {
 func (c *Cache) containsWorker() {
 	var ok bool
 	for req := range c.containsQueue {
-		ok, _ = c.proxy.Contains(cache.CAS, (*req.digest).Hash)
+		ok, _ = c.proxy.Contains(req.ctx, cache.CAS, (*req.digest).Hash)
 		if ok {
 			c.accessLogger.Printf("GRPC CAS HEAD %s OK", (*req.digest).Hash)
 			// The blob exists on the proxy, remove it from the
diff --git a/cache/httpproxy/httpproxy.go b/cache/httpproxy/httpproxy.go
index baee112..5387243 100644
--- a/cache/httpproxy/httpproxy.go
+++ b/cache/httpproxy/httpproxy.go
@@ -3,6 +3,7 @@
 package httpproxy
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -57,14 +58,24 @@ func (r *remoteHTTPProxyCache) uploadFile(item uploadReq) {
 
 	url := r.requestURL(item.hash, item.kind)
 
-	rsp, err := r.remote.Head(url)
+	req, err := http.NewRequestWithContext(context.Background(), http.MethodHead, url, nil)
+	if err != nil {
+		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)
+		item.rc.Close()
+		return
+	}
+
+	rsp, err := r.remote.Do(req)
 	if err == nil && rsp.StatusCode == http.StatusOK {
 		r.accessLogger.Printf("SKIP UPLOAD %s", item.hash)
+		item.rc.Close()
 		return
 	}
 
-	req, err := http.NewRequest(http.MethodPut, url, item.rc)
+	req, err = http.NewRequestWithContext(context.Background(), http.MethodPut, url, item.rc)
 	if err != nil {
+		r.errorLogger.Printf("INTERNAL ERROR, FAILED TO SETUP HTTP PROXY UPLOAD %s: %s", url, err)
+
 		// item.rc will be closed if we call req.Do(), but not if we
 		// return earlier.
 		item.rc.Close()
@@ -160,9 +171,16 @@ func (r *remoteHTTPProxyCache) Put(kind cache.EntryKind, hash string, size int64
 	}
 }
 
-func (r *remoteHTTPProxyCache) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
+func (r *remoteHTTPProxyCache) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
 	url := r.requestURL(hash, kind)
-	rsp, err := r.remote.Get(url)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		cacheMisses.Inc()
+		return nil, -1, err
+	}
+
+	rsp, err := r.remote.Do(req)
 	if err != nil {
 		cacheMisses.Inc()
 		return nil, -1, err
@@ -216,11 +234,16 @@ func (r *remoteHTTPProxyCache) Get(kind cache.EntryKind, hash string) (io.ReadCl
 	return rsp.Body, sizeBytes, nil
 }
 
-func (r *remoteHTTPProxyCache) Contains(kind cache.EntryKind, hash string) (bool, int64) {
+func (r *remoteHTTPProxyCache) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {
 
 	url := r.requestURL(hash, kind)
 
-	rsp, err := r.remote.Head(url)
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
+	if err != nil {
+		return false, -1
+	}
+
+	rsp, err := r.remote.Do(req)
 	if err == nil && rsp.StatusCode == http.StatusOK {
 		if kind != cache.CAS {
 			return true, rsp.ContentLength
diff --git a/cache/s3proxy/s3proxy.go b/cache/s3proxy/s3proxy.go
index 0220b59..440b8eb 100644
--- a/cache/s3proxy/s3proxy.go
+++ b/cache/s3proxy/s3proxy.go
@@ -216,10 +216,10 @@ func (c *s3Cache) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadC
 	}
 }
 
-func (c *s3Cache) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
+func (c *s3Cache) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
 
 	rc, info, _, err := c.mcore.GetObject(
-		context.Background(),
+		ctx,
 		c.bucket,                 // bucketName
 		c.objectKey(hash, kind),  // objectName
 		minio.GetObjectOptions{}, // opts
@@ -245,12 +245,12 @@ func (c *s3Cache) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64,
 	return rc, info.Size, nil
 }
 
-func (c *s3Cache) Contains(kind cache.EntryKind, hash string) (bool, int64) {
+func (c *s3Cache) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {
 	size := int64(-1)
 	exists := false
 
 	s, err := c.mcore.StatObject(
-		context.Background(),
+		ctx,
 		c.bucket,                  // bucketName
 		c.objectKey(hash, kind),   // objectName
 		minio.StatObjectOptions{}, // opts
diff --git a/server/grpc_ac.go b/server/grpc_ac.go
index e68e5d1..4ab95f6 100644
--- a/server/grpc_ac.go
+++ b/server/grpc_ac.go
@@ -57,7 +57,7 @@ func (s *grpcServer) GetActionResult(ctx context.Context,
 	if !s.depsCheck {
 		logPrefix = "GRPC AC GET NODEPSCHECK"
 
-		rdr, sizeBytes, err := s.cache.Get(cache.AC, req.ActionDigest.Hash, unknownActionResultSize, 0)
+		rdr, sizeBytes, err := s.cache.Get(ctx, cache.AC, req.ActionDigest.Hash, unknownActionResultSize, 0)
 		if err != nil {
 			s.accessLogger.Printf("%s %s %s", logPrefix, req.ActionDigest.Hash, err)
 			return nil, status.Error(codes.Unknown, err.Error())
@@ -86,7 +86,7 @@ func (s *grpcServer) GetActionResult(ctx context.Context,
 		return result, nil
 	}
 
-	result, _, err := s.cache.GetValidatedActionResult(req.ActionDigest.Hash)
+	result, _, err := s.cache.GetValidatedActionResult(ctx, req.ActionDigest.Hash)
 	if err != nil {
 		s.accessLogger.Printf("%s %s %s", logPrefix, req.ActionDigest.Hash, err)
 		return nil, status.Error(codes.Unknown, err.Error())
@@ -102,14 +102,14 @@ func (s *grpcServer) GetActionResult(ctx context.Context,
 
 	var inlinedSoFar int64
 
-	err = s.maybeInline(req.InlineStdout,
+	err = s.maybeInline(ctx, req.InlineStdout,
 		&result.StdoutRaw, &result.StdoutDigest, &inlinedSoFar)
 	if err != nil {
 		s.accessLogger.Printf("%s %s %s", logPrefix, req.ActionDigest.Hash, err)
 		return nil, status.Error(codes.Unknown, err.Error())
 	}
 
-	err = s.maybeInline(req.InlineStderr,
+	err = s.maybeInline(ctx, req.InlineStderr,
 		&result.StderrRaw, &result.StderrDigest, &inlinedSoFar)
 	if err != nil {
 		s.accessLogger.Printf("%s %s %s", logPrefix, req.ActionDigest.Hash, err)
@@ -122,7 +122,7 @@ func (s *grpcServer) GetActionResult(ctx context.Context,
 	}
 	for _, of := range result.GetOutputFiles() {
 		_, ok := inlinableFiles[of.Path]
-		err = s.maybeInline(ok, &of.Contents, &of.Digest, &inlinedSoFar)
+		err = s.maybeInline(ctx, ok, &of.Contents, &of.Digest, &inlinedSoFar)
 		if err != nil {
 			s.accessLogger.Printf("%s %s %s", logPrefix, req.ActionDigest.Hash, err)
 			return nil, status.Error(codes.Unknown, err.Error())
@@ -134,7 +134,7 @@ func (s *grpcServer) GetActionResult(ctx context.Context,
 	return result, nil
 }
 
-func (s *grpcServer) maybeInline(inline bool, slice *[]byte, digest **pb.Digest, inlinedSoFar *int64) error {
+func (s *grpcServer) maybeInline(ctx context.Context, inline bool, slice *[]byte, digest **pb.Digest, inlinedSoFar *int64) error {
 
 	if (*inlinedSoFar + int64(len(*slice))) > maxInlineSize {
 		inline = false
@@ -156,7 +156,7 @@ func (s *grpcServer) maybeInline(inline bool, slice *[]byte, digest **pb.Digest,
 			}
 		}
 
-		found, _ := s.cache.Contains(cache.CAS, (*digest).Hash, (*digest).SizeBytes)
+		found, _ := s.cache.Contains(ctx, cache.CAS, (*digest).Hash, (*digest).SizeBytes)
 		if !found {
 			err := s.cache.Put(cache.CAS, (*digest).Hash, (*digest).SizeBytes,
 				bytes.NewReader(*slice))
@@ -181,7 +181,7 @@ func (s *grpcServer) maybeInline(inline bool, slice *[]byte, digest **pb.Digest,
 
 	// Otherwise, attempt to inline.
 	if (*digest).SizeBytes > 0 {
-		data, err := s.getBlobData((*digest).Hash, (*digest).SizeBytes)
+		data, err := s.getBlobData(ctx, (*digest).Hash, (*digest).SizeBytes)
 		if err != nil {
 			return err
 		}
diff --git a/server/grpc_asset.go b/server/grpc_asset.go
index 8b1bb63..06dfc83 100644
--- a/server/grpc_asset.go
+++ b/server/grpc_asset.go
@@ -64,14 +64,14 @@ func (s *grpcServer) FetchBlob(ctx context.Context, req *asset.FetchBlobRequest)
 
 			sha256Str = hex.EncodeToString(decoded)
 
-			found, size := s.cache.Contains(cache.CAS, sha256Str, -1)
+			found, size := s.cache.Contains(ctx, cache.CAS, sha256Str, -1)
 			if !found {
 				continue
 			}
 
 			if size < 0 {
 				// We don't know the size yet (bad http backend?).
-				r, actualSize, err := s.cache.Get(cache.CAS, sha256Str, -1, 0)
+				r, actualSize, err := s.cache.Get(ctx, cache.CAS, sha256Str, -1, 0)
 				if r != nil {
 					defer r.Close()
 				}
@@ -107,7 +107,7 @@ func (s *grpcServer) FetchBlob(ctx context.Context, req *asset.FetchBlobRequest)
 	// See if we can download one of the URIs.
 
 	for _, uri := range req.GetUris() {
-		ok, actualHash, size := s.fetchItem(uri, sha256Str)
+		ok, actualHash, size := s.fetchItem(ctx, uri, sha256Str)
 		if ok {
 			return &asset.FetchBlobResponse{
 				Status: &status.Status{Code: int32(codes.OK)},
@@ -127,7 +127,7 @@ func (s *grpcServer) FetchBlob(ctx context.Context, req *asset.FetchBlobRequest)
 	}, nil
 }
 
-func (s *grpcServer) fetchItem(uri string, expectedHash string) (bool, string, int64) {
+func (s *grpcServer) fetchItem(ctx context.Context, uri string, expectedHash string) (bool, string, int64) {
 	u, err := url.Parse(uri)
 	if err != nil {
 		s.errorLogger.Printf("unable to parse URI: %s err: %v", uri, err)
diff --git a/server/grpc_cas.go b/server/grpc_cas.go
index ec13974..5cab000 100644
--- a/server/grpc_cas.go
+++ b/server/grpc_cas.go
@@ -94,7 +94,7 @@ func (s *grpcServer) BatchUpdateBlobs(ctx context.Context,
 // Return the data for a blob, or an error.  If the blob was not
 // found, the returned error is errBlobNotFound. Only use this
 // function when it's OK to buffer the entire blob in memory.
-func (s *grpcServer) getBlobData(hash string, size int64) ([]byte, error) {
+func (s *grpcServer) getBlobData(ctx context.Context, hash string, size int64) ([]byte, error) {
 	if size < 0 {
 		return []byte{}, errBadSize
 	}
@@ -103,7 +103,7 @@ func (s *grpcServer) getBlobData(hash string, size int64) ([]byte, error) {
 		return []byte{}, nil
 	}
 
-	rdr, sizeBytes, err := s.cache.Get(cache.CAS, hash, size, 0)
+	rdr, sizeBytes, err := s.cache.Get(ctx, cache.CAS, hash, size, 0)
 	if err != nil {
 		rdr.Close()
 		return []byte{}, err
@@ -127,10 +127,10 @@ func (s *grpcServer) getBlobData(hash string, size int64) ([]byte, error) {
 	return data, rdr.Close()
 }
 
-func (s *grpcServer) getBlobResponse(digest *pb.Digest) *pb.BatchReadBlobsResponse_Response {
+func (s *grpcServer) getBlobResponse(ctx context.Context, digest *pb.Digest) *pb.BatchReadBlobsResponse_Response {
 	r := pb.BatchReadBlobsResponse_Response{Digest: digest}
 
-	data, err := s.getBlobData(digest.Hash, digest.SizeBytes)
+	data, err := s.getBlobData(ctx, digest.Hash, digest.SizeBytes)
 	if err == errBlobNotFound {
 		s.accessLogger.Printf("GRPC CAS GET %s NOT FOUND", digest.Hash)
 		r.Status = &status.Status{Code: int32(code.Code_NOT_FOUND)}
@@ -166,7 +166,7 @@ func (s *grpcServer) BatchReadBlobs(ctx context.Context,
 		if err != nil {
 			return nil, err
 		}
-		resp.Responses = append(resp.Responses, s.getBlobResponse(digest))
+		resp.Responses = append(resp.Responses, s.getBlobResponse(ctx, digest))
 	}
 
 	return &resp, nil
@@ -184,7 +184,7 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 		return err
 	}
 
-	data, err := s.getBlobData(in.RootDigest.Hash, in.RootDigest.SizeBytes)
+	data, err := s.getBlobData(stream.Context(), in.RootDigest.Hash, in.RootDigest.SizeBytes)
 	if err == errBlobNotFound {
 		s.accessLogger.Printf("GRPC CAS GETTREEREQUEST %s NOT FOUND",
 			in.RootDigest.Hash)
@@ -202,7 +202,7 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 		return grpc_status.Error(codes.DataLoss, err.Error())
 	}
 
-	err = s.fillDirectories(&resp, &dir, errorPrefix)
+	err = s.fillDirectories(stream.Context(), &resp, &dir, errorPrefix)
 	if err != nil {
 		return err
 	}
@@ -217,7 +217,7 @@ func (s *grpcServer) GetTree(in *pb.GetTreeRequest,
 
 // Attempt to populate `resp`. Return errors for invalid requests, but
 // otherwise attempt to return as many blobs as possible.
-func (s *grpcServer) fillDirectories(resp *pb.GetTreeResponse, dir *pb.Directory, errorPrefix string) error {
+func (s *grpcServer) fillDirectories(ctx context.Context, resp *pb.GetTreeResponse, dir *pb.Directory, errorPrefix string) error {
 
 	// Add this dir.
 	resp.Directories = append(resp.Directories, dir)
@@ -230,7 +230,7 @@ func (s *grpcServer) fillDirectories(resp *pb.GetTreeResponse, dir *pb.Directory
 			return err
 		}
 
-		data, err := s.getBlobData(dirNode.Digest.Hash, dirNode.Digest.SizeBytes)
+		data, err := s.getBlobData(ctx, dirNode.Digest.Hash, dirNode.Digest.SizeBytes)
 		if err == errBlobNotFound {
 			s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s NOT FOUND",
 				dirNode.Digest.Hash)
@@ -251,7 +251,7 @@ func (s *grpcServer) fillDirectories(resp *pb.GetTreeResponse, dir *pb.Directory
 		s.accessLogger.Printf("GRPC GETTREEREQUEST BLOB %s ADDED OK",
 			dirNode.Digest.Hash)
 
-		err = s.fillDirectories(resp, &dirMsg, errorPrefix)
+		err = s.fillDirectories(ctx, resp, &dirMsg, errorPrefix)
 		if err != nil {
 			return err
 		}
diff --git a/server/http.go b/server/http.go
index a5634cc..d64cbdb 100644
--- a/server/http.go
+++ b/server/http.go
@@ -115,7 +115,7 @@ func parseRequestURL(url string, validateAC bool) (kind cache.EntryKind, hash st
 	return cache.RAW, hash, instance, nil
 }
 func (h *httpCache) handleContainsValidAC(w http.ResponseWriter, r *http.Request, hash string) {
-	_, data, err := h.cache.GetValidatedActionResult(hash)
+	_, data, err := h.cache.GetValidatedActionResult(r.Context(), hash)
 	if err != nil {
 		http.Error(w, "Not found", http.StatusNotFound)
 		h.logResponse(http.StatusNotFound, r)
@@ -134,7 +134,7 @@ func (h *httpCache) handleContainsValidAC(w http.ResponseWriter, r *http.Request
 }
 
 func (h *httpCache) handleGetValidAC(w http.ResponseWriter, r *http.Request, hash string) {
-	_, data, err := h.cache.GetValidatedActionResult(hash)
+	_, data, err := h.cache.GetValidatedActionResult(r.Context(), hash)
 	if err != nil {
 		http.Error(w, "Not found", http.StatusNotFound)
 		h.logResponse(http.StatusNotFound, r)
@@ -220,10 +220,10 @@ func (h *httpCache) CacheHandler(w http.ResponseWriter, r *http.Request) {
 
 		zstdCompressed := false
 		if kind == cache.CAS && strings.Contains(r.Header.Get("Accept-Encoding"), "zstd") {
-			rdr, sizeBytes, err = h.cache.GetZstd(hash, -1, 0)
+			rdr, sizeBytes, err = h.cache.GetZstd(r.Context(), hash, -1, 0)
 			zstdCompressed = true
 		} else {
-			rdr, sizeBytes, err = h.cache.Get(kind, hash, -1, 0)
+			rdr, sizeBytes, err = h.cache.Get(r.Context(), kind, hash, -1, 0)
 		}
 		if err != nil {
 			if e, ok := err.(*cache.Error); ok {
@@ -397,7 +397,7 @@ func (h *httpCache) CacheHandler(w http.ResponseWriter, r *http.Request) {
 
 		// Unvalidated path:
 
-		ok, size := h.cache.Contains(kind, hash, -1)
+		ok, size := h.cache.Contains(r.Context(), kind, hash, -1)
 		if !ok {
 			http.Error(w, "Not found", http.StatusNotFound)
 			h.logResponse(http.StatusNotFound, r)


Test patch:
diff --git a/cache/disk/disk_test.go b/cache/disk/disk_test.go
index 821712e..e587899 100644
--- a/cache/disk/disk_test.go
+++ b/cache/disk/disk_test.go
@@ -2,6 +2,7 @@ package disk
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha256"
 	"encoding/hex"
 	"fmt"
@@ -41,6 +42,9 @@ const contentsHash = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e7304336293
 const contentsLength = int64(len(contents))
 
 func TestCacheBasics(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := tempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -62,7 +66,7 @@ func TestCacheBasics(t *testing.T) {
 	data, hash := testutils.RandomDataAndHash(itemSize)
 
 	// Non-existing item.
-	rdr, _, err := testCache.Get(cache.CAS, hash, itemSize, 0)
+	rdr, _, err := testCache.Get(ctx, cache.CAS, hash, itemSize, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -78,7 +82,7 @@ func TestCacheBasics(t *testing.T) {
 	}
 
 	// Get the item back.
-	rdr, sizeBytes, err := testCache.Get(cache.CAS, hash, itemSize, 0)
+	rdr, sizeBytes, err := testCache.Get(ctx, cache.CAS, hash, itemSize, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -123,6 +127,8 @@ func TestCachePutWrongSize(t *testing.T) {
 }
 
 func TestCacheGetContainsWrongSize(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	cacheDir := tempDir(t)
 	defer os.RemoveAll(cacheDir)
@@ -139,28 +145,30 @@ func TestCacheGetContainsWrongSize(t *testing.T) {
 		t.Fatal("Expected success", err)
 	}
 
-	found, _ = testCache.Contains(cache.CAS, contentsHash, contentsLength+1)
+	found, _ = testCache.Contains(ctx, cache.CAS, contentsHash, contentsLength+1)
 	if found {
 		t.Error("Expected not found, due to size being different")
 	}
 
-	rdr, _, _ = testCache.Get(cache.CAS, contentsHash, contentsLength+1, 0)
+	rdr, _, _ = testCache.Get(ctx, cache.CAS, contentsHash, contentsLength+1, 0)
 	if rdr != nil {
 		t.Error("Expected not found, due to size being different")
 	}
 
-	found, _ = testCache.Contains(cache.CAS, contentsHash, -1)
+	found, _ = testCache.Contains(ctx, cache.CAS, contentsHash, -1)
 	if !found {
 		t.Error("Expected found, when unknown size")
 	}
 
-	rdr, _, _ = testCache.Get(cache.CAS, contentsHash, -1, 0)
+	rdr, _, _ = testCache.Get(ctx, cache.CAS, contentsHash, -1, 0)
 	if rdr == nil {
 		t.Error("Expected found, when unknown size")
 	}
 }
 
 func TestCacheGetContainsWrongSizeWithProxy(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	cacheDir := tempDir(t)
 	defer os.RemoveAll(cacheDir)
@@ -179,12 +187,12 @@ func TestCacheGetContainsWrongSizeWithProxy(t *testing.T) {
 			testCache.lru.Len())
 	}
 
-	found, _ = testCache.Contains(cache.CAS, contentsHash, contentsLength+1)
+	found, _ = testCache.Contains(ctx, cache.CAS, contentsHash, contentsLength+1)
 	if found {
 		t.Fatal("Expected not found, due to size being different")
 	}
 
-	rdr, _, _ = testCache.Get(cache.CAS, contentsHash, contentsLength+1, 0)
+	rdr, _, _ = testCache.Get(ctx, cache.CAS, contentsHash, contentsLength+1, 0)
 	if rdr != nil {
 		t.Fatal("Expected not found, due to size being different")
 	}
@@ -194,12 +202,12 @@ func TestCacheGetContainsWrongSizeWithProxy(t *testing.T) {
 			testCache.lru.Len())
 	}
 
-	found, _ = testCache.Contains(cache.CAS, contentsHash, -1)
+	found, _ = testCache.Contains(ctx, cache.CAS, contentsHash, -1)
 	if !found {
 		t.Fatal("Expected found, when unknown size")
 	}
 
-	rdr, _, _ = testCache.Get(cache.CAS, contentsHash, -1, 0)
+	rdr, _, _ = testCache.Get(ctx, cache.CAS, contentsHash, -1, 0)
 	if rdr == nil {
 		t.Fatal("Expected found, when unknown size")
 	}
@@ -218,7 +226,7 @@ func (d proxyStub) Put(kind cache.EntryKind, hash string, size int64, rc io.Read
 	// Not implemented.
 }
 
-func (d proxyStub) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
+func (d proxyStub) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
 	if hash != contentsHash || kind != cache.CAS {
 		return nil, -1, nil
 	}
@@ -245,7 +253,7 @@ func (d proxyStub) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64,
 	return readme, contentsLength, nil
 }
 
-func (d proxyStub) Contains(kind cache.EntryKind, hash string) (bool, int64) {
+func (d proxyStub) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {
 	if hash != contentsHash || kind != cache.CAS {
 		return false, -1
 	}
@@ -273,11 +281,11 @@ func expectContentEquals(rdr io.ReadCloser, sizeBytes int64, expectedContent []b
 	return nil
 }
 
-func putGetCompare(kind cache.EntryKind, hash string, content string, testCache *Cache) error {
-	return putGetCompareBytes(kind, hash, []byte(content), testCache)
+func putGetCompare(ctx context.Context, kind cache.EntryKind, hash string, content string, testCache *Cache) error {
+	return putGetCompareBytes(ctx, kind, hash, []byte(content), testCache)
 }
 
-func putGetCompareBytes(kind cache.EntryKind, hash string, data []byte, testCache *Cache) error {
+func putGetCompareBytes(ctx context.Context, kind cache.EntryKind, hash string, data []byte, testCache *Cache) error {
 
 	r := bytes.NewReader(data)
 
@@ -286,10 +294,11 @@ func putGetCompareBytes(kind cache.EntryKind, hash string, data []byte, testCach
 		return err
 	}
 
-	rdr, sizeBytes, err := testCache.Get(kind, hash, int64(len(data)), 0)
+	rdr, sizeBytes, err := testCache.Get(ctx, kind, hash, int64(len(data)), 0)
 	if err != nil {
 		return err
 	}
+
 	// Get the item back
 	return expectContentEquals(rdr, sizeBytes, data)
 }
@@ -301,6 +310,9 @@ func hashStr(content string) string {
 
 // Make sure that we can overwrite items if we upload the same key again.
 func TestOverwrite(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := tempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -309,35 +321,38 @@ func TestOverwrite(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = putGetCompare(cache.CAS, hashStr("hello"), "hello", testCache)
+	err = putGetCompare(ctx, cache.CAS, hashStr("hello"), "hello", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = putGetCompare(cache.CAS, hashStr("hello"), "hello", testCache)
+	err = putGetCompare(ctx, cache.CAS, hashStr("hello"), "hello", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = putGetCompare(cache.AC, hashStr("world"), "world1", testCache)
+	err = putGetCompare(ctx, cache.AC, hashStr("world"), "world1", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = putGetCompare(cache.AC, hashStr("world"), "world2", testCache)
+	err = putGetCompare(ctx, cache.AC, hashStr("world"), "world2", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = putGetCompare(cache.RAW, hashStr("world"), "world3", testCache)
+	err = putGetCompare(ctx, cache.RAW, hashStr("world"), "world3", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = putGetCompare(cache.RAW, hashStr("world"), "world4", testCache)
+	err = putGetCompare(ctx, cache.RAW, hashStr("world"), "world4", testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
 }
 
 func TestCacheExistingFiles(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := tempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -445,7 +460,7 @@ func TestCacheExistingFiles(t *testing.T) {
 		break // First item evicted as expected.
 	}
 
-	found, _ := testCache.Contains(cache.CAS, items[0].hash, contentsLength)
+	found, _ := testCache.Contains(ctx, cache.CAS, items[0].hash, contentsLength)
 	if found {
 		t.Fatalf("%s should have been evicted", items[0].file)
 	}
@@ -545,6 +560,9 @@ func createRandomCASFile(dir string, size int64) (string, error) {
 }
 
 func TestMigrateFromOldDirectoryStructure(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -577,23 +595,26 @@ func TestMigrateFromOldDirectoryStructure(t *testing.T) {
 	}
 
 	var found bool
-	found, _ = testCache.Contains(cache.AC, acHash, 512)
+	found, _ = testCache.Contains(ctx, cache.AC, acHash, 512)
 	if !found {
 		t.Fatalf("Expected cache to contain AC entry '%s'", acHash)
 	}
 
-	found, _ = testCache.Contains(cache.CAS, casHash1, 1024)
+	found, _ = testCache.Contains(ctx, cache.CAS, casHash1, 1024)
 	if !found {
 		t.Fatalf("Expected cache to contain CAS entry '%s'", casHash1)
 	}
 
-	found, _ = testCache.Contains(cache.CAS, casHash2, 1024)
+	found, _ = testCache.Contains(ctx, cache.CAS, casHash2, 1024)
 	if !found {
 		t.Fatalf("Expected cache to contain CAS entry '%s'", casHash2)
 	}
 }
 
 func TestLoadExistingEntries(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	// Test that loading existing items works
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
@@ -674,28 +695,31 @@ func TestLoadExistingEntries(t *testing.T) {
 
 	var found bool
 
-	found, _ = testCache.Contains(cache.AC, acHash, blobSize)
+	found, _ = testCache.Contains(ctx, cache.AC, acHash, blobSize)
 	if !found {
 		t.Fatalf("Expected cache to contain AC entry '%s'", acHash)
 	}
 
-	found, _ = testCache.Contains(cache.CAS, casHash, blobSize)
+	found, _ = testCache.Contains(ctx, cache.CAS, casHash, blobSize)
 	if !found {
 		t.Fatalf("Expected cache to contain CAS entry '%s'", casHash)
 	}
 
-	found, _ = testCache.Contains(cache.CAS, casV1Hash, blobSize)
+	found, _ = testCache.Contains(ctx, cache.CAS, casV1Hash, blobSize)
 	if !found {
 		t.Fatalf("Expected cache to contain CAS V1 entry '%s'", casV1Hash)
 	}
 
-	found, _ = testCache.Contains(cache.RAW, rawHash, blobSize)
+	found, _ = testCache.Contains(ctx, cache.RAW, rawHash, blobSize)
 	if !found {
 		t.Fatalf("Expected cache to contain RAW entry '%s'", rawHash)
 	}
 }
 
 func TestDistinctKeyspaces(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -715,17 +739,17 @@ func TestDistinctKeyspaces(t *testing.T) {
 	// keyspaces, and verify that we have exactly three items in
 	// the cache.
 
-	err = putGetCompareBytes(cache.CAS, casHash, blob, testCache)
+	err = putGetCompareBytes(ctx, cache.CAS, casHash, blob, testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = putGetCompareBytes(cache.AC, casHash, blob, testCache)
+	err = putGetCompareBytes(ctx, cache.AC, casHash, blob, testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	err = putGetCompareBytes(cache.RAW, casHash, blob, testCache)
+	err = putGetCompareBytes(ctx, cache.RAW, casHash, blob, testCache)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -803,6 +827,8 @@ func newTestServer(t *testing.T) *testServer {
 }
 
 func TestHttpProxyBackend(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	backend := newTestServer(t)
 	url, err := url.Parse(backend.srv.URL)
@@ -833,7 +859,7 @@ func TestHttpProxyBackend(t *testing.T) {
 	blob, casHash := testutils.RandomDataAndHash(blobSize)
 
 	// Non-existing item
-	r, _, err := testCache.Get(cache.CAS, casHash, blobSize, 0)
+	r, _, err := testCache.Get(ctx, cache.CAS, casHash, blobSize, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -870,12 +896,12 @@ func TestHttpProxyBackend(t *testing.T) {
 	// Confirm that it does not contain the item we added to the
 	// first testCache and the proxy backend.
 
-	found, _ := testCache.Contains(cache.CAS, casHash, blobSize)
+	found, _ := testCache.Contains(ctx, cache.CAS, casHash, blobSize)
 	if found {
 		t.Fatalf("Expected the cache not to contain %s", casHash)
 	}
 
-	r, _, err = testCache.Get(cache.CAS, casHash, blobSize, 0)
+	r, _, err = testCache.Get(ctx, cache.CAS, casHash, blobSize, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -886,13 +912,13 @@ func TestHttpProxyBackend(t *testing.T) {
 	// Add the proxy backend and check that we can Get the item.
 	testCache.proxy = proxy
 
-	found, _ = testCache.Contains(cache.CAS, casHash, blobSize)
+	found, _ = testCache.Contains(ctx, cache.CAS, casHash, blobSize)
 	if !found {
 		t.Fatalf("Expected the cache to contain %s (via the proxy)",
 			casHash)
 	}
 
-	r, fetchedSize, err := testCache.Get(cache.CAS, casHash, blobSize, 0)
+	r, fetchedSize, err := testCache.Get(ctx, cache.CAS, casHash, blobSize, 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -921,6 +947,9 @@ func TestHttpProxyBackend(t *testing.T) {
 // Store an ActionResult with an output directory, then confirm that
 // GetValidatedActionResult returns the original item.
 func TestGetValidatedActionResult(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -1076,7 +1105,7 @@ func TestGetValidatedActionResult(t *testing.T) {
 	// to assume that the value should be returned unchanged by the cache
 	// layer.
 
-	rAR, rData, err := testCache.GetValidatedActionResult(arDataHashStr)
+	rAR, rData, err := testCache.GetValidatedActionResult(ctx, arDataHashStr)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1091,6 +1120,9 @@ func TestGetValidatedActionResult(t *testing.T) {
 }
 
 func TestGetWithOffset(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	cacheDir := testutils.TempDir(t)
 	defer os.RemoveAll(cacheDir)
 
@@ -1109,7 +1141,7 @@ func TestGetWithOffset(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	rc, foundSize, err := testCache.Get(cache.CAS, hash, int64(len(data)), 0)
+	rc, foundSize, err := testCache.Get(ctx, cache.CAS, hash, int64(len(data)), 0)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -1130,7 +1162,7 @@ func TestGetWithOffset(t *testing.T) {
 
 	// Now try some partial reads.
 	for _, offset := range []int64{42, 1023, 1024, 1025, 2048, 2303} {
-		rc, foundSize, err = testCache.Get(cache.CAS, hash, int64(len(data)), offset)
+		rc, foundSize, err = testCache.Get(ctx, cache.CAS, hash, int64(len(data)), offset)
 		if err != nil {
 			t.Fatal(err)
 		}
diff --git a/cache/disk/findmissing_test.go b/cache/disk/findmissing_test.go
index f92aa57..7398873 100644
--- a/cache/disk/findmissing_test.go
+++ b/cache/disk/findmissing_test.go
@@ -1,6 +1,7 @@
 package disk
 
 import (
+	"context"
 	"io"
 	"sync"
 	"testing"
@@ -84,11 +85,12 @@ type testCWProxy struct {
 	blob string
 }
 
-func (p *testCWProxy) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {}
-func (p *testCWProxy) Get(kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
+func (p *testCWProxy) Put(kind cache.EntryKind, hash string, size int64, rc io.ReadCloser) {
+}
+func (p *testCWProxy) Get(ctx context.Context, kind cache.EntryKind, hash string) (io.ReadCloser, int64, error) {
 	return nil, -1, nil
 }
-func (p *testCWProxy) Contains(kind cache.EntryKind, hash string) (bool, int64) {
+func (p *testCWProxy) Contains(ctx context.Context, kind cache.EntryKind, hash string) (bool, int64) {
 	if kind == cache.CAS && hash == p.blob {
 		return true, 42
 	}
diff --git a/cache/httpproxy/httpproxy_test.go b/cache/httpproxy/httpproxy_test.go
index 98c045b..edece8a 100644
--- a/cache/httpproxy/httpproxy_test.go
+++ b/cache/httpproxy/httpproxy_test.go
@@ -2,6 +2,7 @@ package httpproxy
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -88,6 +89,9 @@ func newTestServer() *testServer {
 }
 
 func TestEverything(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	s := newTestServer()
 	defer s.srv.Close()
 
@@ -204,7 +208,7 @@ func TestEverything(t *testing.T) {
 	var found bool
 	var size int64
 
-	found, size = diskCache.Contains(cache.AC, hash, int64(len(acData)))
+	found, size = diskCache.Contains(ctx, cache.AC, hash, int64(len(acData)))
 	if !found {
 		t.Fatalf("Expected to find AC item %s", hash)
 	}
@@ -213,7 +217,7 @@ func TestEverything(t *testing.T) {
 			len(acData), size)
 	}
 
-	found, size = diskCache.Contains(cache.CAS, hash, int64(len(casData)))
+	found, size = diskCache.Contains(ctx, cache.CAS, hash, int64(len(casData)))
 	if !found {
 		t.Fatalf("Expected to find CAS item %s", hash)
 	}
@@ -227,7 +231,7 @@ func TestEverything(t *testing.T) {
 	var data []byte
 	var rc io.ReadCloser
 
-	rc, size, err = diskCache.Get(cache.AC, hash, int64(len(acData)), 0)
+	rc, size, err = diskCache.Get(ctx, cache.AC, hash, int64(len(acData)), 0)
 	if err != nil {
 		t.Error(err)
 	}
@@ -247,7 +251,7 @@ func TestEverything(t *testing.T) {
 	}
 	rc.Close()
 
-	rc, size, err = diskCache.Get(cache.CAS, hash, int64(len(casData)), 0)
+	rc, size, err = diskCache.Get(ctx, cache.CAS, hash, int64(len(casData)), 0)
 	if err != nil {
 		t.Error(err)
 	}
@@ -285,7 +289,7 @@ func TestEverything(t *testing.T) {
 
 	// Confirm that we can HEAD both values successfully.
 
-	found, size = diskCache.Contains(cache.AC, hash, int64(len(acData)))
+	found, size = diskCache.Contains(ctx, cache.AC, hash, int64(len(acData)))
 	if !found {
 		t.Fatalf("Expected to find AC item %s", hash)
 	}
@@ -294,7 +298,7 @@ func TestEverything(t *testing.T) {
 			len(acData), size)
 	}
 
-	found, size = diskCache.Contains(cache.CAS, hash, int64(len(casData)))
+	found, size = diskCache.Contains(ctx, cache.CAS, hash, int64(len(casData)))
 	if !found {
 		t.Fatalf("Expected to find CAS item %s", hash)
 	}
@@ -305,7 +309,7 @@ func TestEverything(t *testing.T) {
 
 	// Confirm that we can GET both values successfully.
 
-	rc, size, err = diskCache.Get(cache.AC, hash, int64(len(acData)), 0)
+	rc, size, err = diskCache.Get(ctx, cache.AC, hash, int64(len(acData)), 0)
 	if err != nil {
 		t.Error(err)
 	}
@@ -325,7 +329,7 @@ func TestEverything(t *testing.T) {
 	}
 	rc.Close()
 
-	rc, size, err = diskCache.Get(cache.CAS, hash, int64(len(casData)), 0)
+	rc, size, err = diskCache.Get(ctx, cache.CAS, hash, int64(len(casData)), 0)
 	if err != nil {
 		t.Error(err)
 	}
diff --git a/server/grpc_bytestream.go b/server/grpc_bytestream.go
index f66a578..0380653 100644
--- a/server/grpc_bytestream.go
+++ b/server/grpc_bytestream.go
@@ -97,9 +97,9 @@ func (s *grpcServer) Read(req *bytestream.ReadRequest,
 	var foundSize int64
 
 	if cmp == casblob.Zstandard {
-		rc, foundSize, err = s.cache.GetZstd(hash, size, req.ReadOffset)
+		rc, foundSize, err = s.cache.GetZstd(resp.Context(), hash, size, req.ReadOffset)
 	} else {
-		rc, foundSize, err = s.cache.Get(cache.CAS, hash, size, req.ReadOffset)
+		rc, foundSize, err = s.cache.Get(resp.Context(), cache.CAS, hash, size, req.ReadOffset)
 	}
 
 	if rc != nil {
@@ -423,7 +423,7 @@ func (s *grpcServer) Write(srv bytestream.ByteStream_WriteServer) error {
 					return
 				}
 
-				exists, _ := s.cache.Contains(cache.CAS, hash, size)
+				exists, _ := s.cache.Contains(srv.Context(), cache.CAS, hash, size)
 				if exists {
 					// Blob already exists, return without writing anything.
 					resp.CommittedSize = size
diff --git a/server/grpc_test.go b/server/grpc_test.go
index fa84e66..d69c264 100644
--- a/server/grpc_test.go
+++ b/server/grpc_test.go
@@ -510,6 +510,9 @@ func TestGrpcAcRequestInlinedBlobs(t *testing.T) {
 }
 
 func TestGrpcByteStreamDeadline(t *testing.T) {
+	testCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	testBlobSize := int64(16)
 	testBlob, testBlobHash := testutils.RandomDataAndHash(testBlobSize)
 	testBlobDigest := pb.Digest{
@@ -588,7 +591,7 @@ func TestGrpcByteStreamDeadline(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	_, sz, err := diskCache.Get(cache.CAS, testBlobHash, testBlobSize, 0)
+	_, sz, err := diskCache.Get(testCtx, cache.CAS, testBlobHash, testBlobSize, 0)
 	if err != nil {
 		t.Fatalf("get error: %v\n", err)
 	}