Let Remote Output Service use root_directory_digest

Even though bb_clientd should continue to make use of Tree objects, we
can optimize access to the root directory of output directories a bit
faster by making use of its digest. This allows us to extract it out of
the Tree message while slicing it.

pkg/cas/tree_directory_walker.go

@@ -18,12 +18,23 @@ type treeDirectoryWalker struct {
 // all Directory messages are stored in a single Tree object in the
 // Content Addressable Storage (CAS). This is the case for output
 // directories of build actions.
-func NewTreeDirectoryWalker(fetcher cas.DirectoryFetcher, treeDigest digest.Digest) cas.DirectoryWalker {
-	return &treeRootDirectoryWalker{
-		treeDirectoryWalker: treeDirectoryWalker{
-			fetcher:    fetcher,
-			treeDigest: treeDigest,
-		},
+func NewTreeDirectoryWalker(fetcher cas.DirectoryFetcher, treeDigest digest.Digest, rootDirectoryDigest *digest.Digest) cas.DirectoryWalker {
+	dw := treeDirectoryWalker{
+		fetcher:    fetcher,
+		treeDigest: treeDigest,
+	}
+	if rootDirectoryDigest == nil {
+		// Digest of the root directory contained in the Tree
+		// message is unknown. Fall back to calling
+		// DirectoryFetcher.GetTreeRootDirectory(), which is far
+		// less efficient.
+		return &treeRootDirectoryWalker{
+			treeDirectoryWalker: dw,
+		}
+	}
+	return &treeChildDirectoryWalker{
+		treeDirectoryWalker: &dw,
+		childDigest:         *rootDirectoryDigest,
 	}
 }
 

pkg/cas/tree_directory_walker_test.go

@@ -21,7 +21,7 @@ func TestTreeDirectoryWalker(t *testing.T) {
 
 	directoryFetcher := mock.NewMockDirectoryFetcher(ctrl)
 	treeDigest := digest.MustNewDigest("example", remoteexecution.DigestFunction_MD5, "6884a9e20905b512d1122a2b1ad8ba16", 123)
-	rootDirectoryWalker := cas.NewTreeDirectoryWalker(directoryFetcher, treeDigest)
+	rootDirectoryWalker := cas.NewTreeDirectoryWalker(directoryFetcher, treeDigest, nil)
 
 	exampleRootDirectory := &remoteexecution.Directory{
 		Directories: []*remoteexecution.DirectoryNode{

pkg/filesystem/virtual/remote_output_service_directory.go

@@ -495,17 +495,25 @@ func (d *RemoteOutputServiceDirectory) BatchCreate(ctx context.Context, request
 	for _, entry := range request.Directories {
 		childDigest, err := buildState.digestFunction.NewDigestFromProto(entry.TreeDigest)
 		if err != nil {
-			return nil, util.StatusWrapf(err, "Invalid digest for directory %#v", entry.Path)
+			return nil, util.StatusWrapf(err, "Invalid tree digest for directory %#v", entry.Path)
 		}
 		if sizeBytes := childDigest.GetSizeBytes(); sizeBytes > d.maximumTreeSizeBytes {
 			return nil, status.Errorf(codes.InvalidArgument, "Directory %#v is %d bytes in size, which exceeds the permitted maximum of %d bytes", entry.Path, sizeBytes, d.maximumTreeSizeBytes)
 		}
+		var rootDirectoryDigest *digest.Digest
+		if entry.RootDirectoryDigest != nil {
+			d, err := buildState.digestFunction.NewDigestFromProto(entry.RootDirectoryDigest)
+			if err != nil {
+				return nil, util.StatusWrapf(err, "Invalid root directory digest for directory %#v", entry.Path)
+			}
+			rootDirectoryDigest = &d
+		}
 		if err := prefixCreator.createChild(
 			entry.Path,
 			virtual.InitialNode{}.FromDirectory(
 				virtual.NewCASInitialContentsFetcher(
 					context.Background(),
-					cd_cas.NewTreeDirectoryWalker(d.directoryFetcher, childDigest),
+					cd_cas.NewTreeDirectoryWalker(d.directoryFetcher, childDigest, rootDirectoryDigest),
 					outputPathState.casFileFactory,
 					d.symlinkFactory,
 					buildState.digestFunction))); err != nil {