From 59691fbabea0004be0414fa57df455578fc4f9c9 Mon Sep 17 00:00:00 2001
From: "Adam Jones (Codethink)" <adam.jones@codethink.co.uk>
Date: Wed, 11 Dec 2019 08:57:40 +0000
Subject: [PATCH] Ensure the AC correctly tracks the output blobs associated
 with a cached action

---
 cmd/bb_storage/main.go        |  2 +-
 pkg/ac/action_cache_server.go | 30 +++++++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/cmd/bb_storage/main.go b/cmd/bb_storage/main.go
index 1f76ecc3..40aca8d6 100644
--- a/cmd/bb_storage/main.go
+++ b/cmd/bb_storage/main.go
@@ -91,7 +91,7 @@ func main() {
 		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
 		grpc.StatsHandler(&ocgrpc.ServerHandler{}),
 	)
-	remoteexecution.RegisterActionCacheServer(s, ac.NewActionCacheServer(actionCache, allowActionCacheUpdatesForInstances))
+	remoteexecution.RegisterActionCacheServer(s, ac.NewActionCacheServer(actionCache, contentAddressableStorageBlobAccess, allowActionCacheUpdatesForInstances))
 	remoteexecution.RegisterContentAddressableStorageServer(s, cas.NewContentAddressableStorageServer(contentAddressableStorageBlobAccess))
 	bytestream.RegisterByteStreamServer(s, cas.NewByteStreamServer(contentAddressableStorageBlobAccess, 1<<16))
 	remoteexecution.RegisterCapabilitiesServer(s, buildQueue)
diff --git a/pkg/ac/action_cache_server.go b/pkg/ac/action_cache_server.go
index 6446f0c9..530f91b9 100644
--- a/pkg/ac/action_cache_server.go
+++ b/pkg/ac/action_cache_server.go
@@ -2,8 +2,10 @@ package ac
 
 import (
 	"context"
+	"fmt"
 
 	remoteexecution "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
+	"github.com/buildbarn/bb-storage/pkg/blobstore"
 	"github.com/buildbarn/bb-storage/pkg/util"
 
 	"google.golang.org/grpc/codes"
@@ -12,14 +14,16 @@ import (
 
 type actionCacheServer struct {
 	actionCache              ActionCache
+	blobAccess               blobstore.BlobAccess
 	allowUpdatesForInstances map[string]bool
 }
 
 // NewActionCacheServer creates a GRPC service for serving the contents
 // of a Bazel Action Cache (AC) to Bazel.
-func NewActionCacheServer(actionCache ActionCache, allowUpdatesForInstances map[string]bool) remoteexecution.ActionCacheServer {
+func NewActionCacheServer(actionCache ActionCache, blobAccess blobstore.BlobAccess, allowUpdatesForInstances map[string]bool) remoteexecution.ActionCacheServer {
 	return &actionCacheServer{
 		actionCache:              actionCache,
+		blobAccess:               blobAccess,
 		allowUpdatesForInstances: allowUpdatesForInstances,
 	}
 }
@@ -29,6 +33,30 @@ func (s *actionCacheServer) GetActionResult(ctx context.Context, in *remoteexecu
 	if err != nil {
 		return nil, err
 	}
+
+	actionResult, err := s.actionCache.GetActionResult(ctx, digest)
+	if err != nil {
+		return nil, err
+	}
+
+	var digests []*util.Digest
+	for _, output := range actionResult.OutputFiles {
+		if output.Digest != nil {
+			OutputDigest, err := util.NewDigest(in.InstanceName, output.Digest)
+
+			if err != nil {
+				fmt.Println(err)
+			} else {
+				digests = append(digests, OutputDigest)
+			}
+		}
+	}
+
+	// Check that we have no missing blobs
+	if missing, err := s.blobAccess.FindMissing(ctx, digests); missing != nil {
+		return nil, status.Errorf(codes.NotFound, "Blob not found")
+	}
+
 	return s.actionCache.GetActionResult(ctx, digest)
 }