Allow a SpawnRunner to inject output metadata.

Adds a MetadataInjector type that can be accessed via the
SpawnExecutionContext.

src/main/java/com/google/devtools/build/lib/actions/cache/MetadataHandler.java

@@ -31,7 +31,7 @@
  * <p>Note that implementations of this interface call chmod on output files if {@link
  * #discardOutputMetadata} has been called.
  */
-public interface MetadataHandler extends MetadataProvider {
+public interface MetadataHandler extends MetadataProvider, MetadataInjector {
   @Override
   FileArtifactValue getMetadata(ActionInput actionInput) throws IOException;
 
@@ -56,9 +56,6 @@ public interface MetadataHandler extends MetadataProvider {
    */
   void injectDigest(ActionInput output, FileStatus statNoFollow, byte[] digest);
 
-  /** Injects a file that is only stored remotely. */
-  void injectRemoteFile(Artifact output, byte[] digest, long size, int locationIndex);
-
   /**
    * Marks an artifact as intentionally omitted. Acknowledges that this Artifact could have existed,
    * but was intentionally not saved, most likely as an optimization.

src/main/java/com/google/devtools/build/lib/actions/cache/MetadataInjector.java

@@ -0,0 +1,49 @@
+// Copyright 2018 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.actions.cache;
+
+import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.Artifact.SpecialArtifact;
+import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
+import com.google.devtools.build.lib.vfs.PathFragment;
+import java.util.Map;
+
+/**
+ * Allows to inject metadata of action outputs into skyframe.
+ */
+public interface MetadataInjector {
+
+  /**
+   * Injects metadata of a file that is stored remotely.
+   *
+   * @param file  a regular output file.
+   * @param digest  the digest of the file.
+   * @param sizeBytes the size of the file in bytes.
+   * @param locationIndex is only used in Blaze.
+   */
+  default void injectRemoteFile(Artifact file, byte[] digest, long sizeBytes, int locationIndex) {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Inject the metadata of a tree artifact whose contents are stored remotely.
+   *
+   * @param treeArtifact  an output directory.
+   * @param children  the metadata of the files stored in the directory.
+   */
+  default void injectRemoteDirectory(SpecialArtifact treeArtifact,
+      Map<PathFragment, RemoteFileArtifactValue> children) {
+    throw new UnsupportedOperationException();
+  }
+}

src/main/java/com/google/devtools/build/lib/exec/AbstractSpawnStrategy.java

@@ -32,6 +32,7 @@
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.SpawnResult.Status;
 import com.google.devtools.build.lib.actions.Spawns;
+import com.google.devtools.build.lib.actions.cache.MetadataHandler;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.exec.SpawnCache.CacheHandle;
@@ -193,6 +194,11 @@ public void prefetchInputs() throws IOException {
     public MetadataProvider getMetadataProvider() {
       return actionExecutionContext.getMetadataProvider();
     }
+
+    @Override
+    public MetadataHandler getMetadataInjector() {
+      return actionExecutionContext.getMetadataHandler();
+    }
 
     @Override
     public ArtifactExpander getArtifactExpander() {

src/main/java/com/google/devtools/build/lib/exec/SpawnRunner.java

@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.actions.MetadataProvider;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnResult;
+import com.google.devtools.build.lib.actions.cache.MetadataInjector;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.io.IOException;
@@ -199,6 +200,14 @@ SortedMap<PathFragment, ActionInput> getInputMapping(boolean expandTreeArtifacts
 
     /** Reports a progress update to the Spawn strategy. */
     void report(ProgressStatus state, String name);
+
+    /**
+     * Returns a {@link MetadataInjector} that allows a caller to inject metadata about spawn
+     * outputs that are stored remotely.
+     */
+    default MetadataInjector getMetadataInjector() {
+      throw new UnsupportedOperationException();
+    }
   }
 
   /**

src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java

@@ -16,6 +16,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -29,6 +30,7 @@
 import com.google.devtools.build.lib.actions.ArtifactFileMetadata;
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
+import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.actions.FileStateValue;
 import com.google.devtools.build.lib.actions.cache.Md5Digest;
 import com.google.devtools.build.lib.actions.cache.MetadataHandler;
@@ -479,6 +481,18 @@ public void injectRemoteFile(Artifact output, byte[] digest, long size, int loca
     store.injectRemoteFile(output, digest, size, locationIndex);
   }
 
+  @Override
+  public void injectRemoteDirectory(SpecialArtifact output,
+      Map<PathFragment, RemoteFileArtifactValue> children) {
+    ImmutableMap.Builder<TreeFileArtifact, FileArtifactValue> childFileValues = ImmutableMap.builder();
+    for (Map.Entry<PathFragment, RemoteFileArtifactValue> child : children.entrySet()) {
+      childFileValues.put(ActionInputHelper.treeFileArtifact(output, child.getKey()), child.getValue());
+    }
+
+    TreeArtifactValue treeArtifactValue = TreeArtifactValue.create(childFileValues.build());
+    store.putTreeArtifactData(output, treeArtifactValue);
+  }
+
   @Override
   public void markOmitted(ActionInput output) {
     Preconditions.checkState(executionMode.get());

src/test/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandlerTest.java

@@ -17,6 +17,7 @@
 import static org.junit.Assert.fail;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.actions.ActionInput;
 import com.google.devtools.build.lib.actions.ActionInputHelper;
 import com.google.devtools.build.lib.actions.ActionInputMap;
@@ -28,10 +29,12 @@
 import com.google.devtools.build.lib.actions.ArtifactPathResolver;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.actions.FileArtifactValue;
+import com.google.devtools.build.lib.actions.FileArtifactValue.RemoteFileArtifactValue;
 import com.google.devtools.build.lib.testutil.Scratch;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
 import java.io.FileNotFoundException;
+import java.util.Map;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -248,4 +251,66 @@ public void withUnknownOutputArtifactMissingDisallowedTreeArtifact() throws Exce
     } catch (IllegalStateException expected) {
     }
   }
+
+  @Test
+  public void injectRemoteArtifactMetadata() throws Exception {
+    PathFragment path = PathFragment.create("foo/bar");
+    Artifact artifact = new Artifact(path, outputRoot);
+    ActionMetadataHandler handler = new ActionMetadataHandler(
+        /* inputArtifactData= */ new ActionInputMap(0),
+        /* missingArtifactsAllowed= */ false,
+        /* outputs= */ ImmutableList.of(artifact),
+        /* tsgm= */ null,
+        ArtifactPathResolver.IDENTITY,
+        new OutputStore());
+    handler.discardOutputMetadata();
+
+    byte[] digest = new byte[]{1,2,3};
+    int size = 10;
+    handler.injectRemoteFile(artifact, digest, size, /* locationIndex= */ 1);
+
+    FileArtifactValue v = handler.getMetadata(artifact);
+    assertThat(v).isNotNull();
+    assertThat(v.getDigest()).isEqualTo(digest);
+    assertThat(v.getSize()).isEqualTo(size);
+    assertThat(v.isRemote()).isTrue();
+  }
+
+  @Test
+  public void injectRemoteTreeArtifactMetadata() throws Exception {
+    PathFragment path = PathFragment.create("bin/dir");
+    SpecialArtifact treeArtifact =
+        new SpecialArtifact(
+            outputRoot, path, ArtifactOwner.NullArtifactOwner.INSTANCE, SpecialArtifactType.TREE);
+    OutputStore store = new OutputStore();
+    ActionMetadataHandler handler = new ActionMetadataHandler(
+        /* inputArtifactData= */ new ActionInputMap(0),
+        /* missingArtifactsAllowed= */ false,
+        /* outputs= */ ImmutableList.of(treeArtifact),
+        /* tsgm= */ null,
+        ArtifactPathResolver.IDENTITY,
+        store);
+    handler.discardOutputMetadata();
+
+    RemoteFileArtifactValue fooValue = new RemoteFileArtifactValue(new byte[]{1,2,3}, 5, 1);
+    RemoteFileArtifactValue barValue = new RemoteFileArtifactValue(new byte[]{4,5,6}, 10, 1);
+    Map<PathFragment, RemoteFileArtifactValue> children =
+        ImmutableMap.<PathFragment, RemoteFileArtifactValue>builder()
+            .put(PathFragment.create("foo"), fooValue)
+            .put(PathFragment.create("bar"), barValue)
+            .build();
+
+    handler.injectRemoteDirectory(treeArtifact, children);
+
+    FileArtifactValue value = handler.getMetadata(treeArtifact);
+    assertThat(value).isNotNull();
+    TreeArtifactValue treeValue = store.getTreeArtifactData(treeArtifact);
+    assertThat(treeValue).isNotNull();
+    assertThat(treeValue.getDigest()).isEqualTo(value.getDigest());
+
+    assertThat(treeValue.getChildPaths()).containsExactly(PathFragment.create("foo"),
+        PathFragment.create("bar"));
+    assertThat(treeValue.getChildValues().values()).containsExactly(fooValue, barValue);
+    assertThat(treeValue.isRemote()).isTrue();
+  }
 }