src/main/java/com/google/devtools/build/lib/remote/RemoteActionInputFetcher.java

// Copyright 2019 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.remote;


import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.RequestMetadata;
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionOutputDirectoryHelper;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.FileArtifactValue;
import com.google.devtools.build.lib.actions.FileStateType;
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
import com.google.devtools.build.lib.cmdline.LabelConstants;
import com.google.devtools.build.lib.events.Reporter;
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.remote.util.TracingMetadataUtils;
import com.google.devtools.build.lib.util.TempPathGenerator;
import com.google.devtools.build.lib.vfs.OutputPermissions;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.PathFragment;
import com.google.devtools.build.lib.vfs.Symlinks;
import java.io.IOException;
import java.util.Collection;
import javax.annotation.Nullable;

/**
 * Stages output files that are stored remotely to the local filesystem.
 *
 * <p>This is used to ensure that the inputs to a local action are present, even when they are
 * provided by a remote action when building without the bytes, or by an external repository when
 * building with a remote repository cache enabled.
 */
public class RemoteActionInputFetcher extends AbstractActionInputPrefetcher {

  private final String buildRequestId;
  private final String commandId;
  private final CombinedCache combinedCache;
  private final ConcurrentArtifactPathTrie rewoundActionOutputs = new ConcurrentArtifactPathTrie();

  RemoteActionInputFetcher(
      Reporter reporter,
      String buildRequestId,
      String commandId,
      CombinedCache combinedCache,
      Supplier<Path> execRootSupplier,
      Path outputBase,
      TempPathGenerator tempPathGenerator,
      RemoteOutputChecker remoteOutputChecker,
      @Nullable ActionOutputDirectoryHelper outputDirectoryHelper,
      OutputPermissions outputPermissions) {
    super(
        reporter,
        execRootSupplier,
        outputBase,
        tempPathGenerator,
        remoteOutputChecker,
        outputDirectoryHelper,
        outputPermissions);
    this.buildRequestId = Preconditions.checkNotNull(buildRequestId);
    this.commandId = Preconditions.checkNotNull(commandId);
    this.combinedCache = Preconditions.checkNotNull(combinedCache);
  }

  @Override
  public boolean requiresTreeMetadataWhenTreeFileIsInput() {
    return true;
  }

  @Override
  protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOException {
    input.atomicallyWriteRelativeTo(execRoot());
  }

  @Override
  protected boolean canDownloadFile(Path path, FileArtifactValue metadata) {
    // Only files and directories have remote-only content that can be downloaded.
    if (metadata.getType() != FileStateType.REGULAR_FILE
        && metadata.getType() != FileStateType.DIRECTORY) {
      return false;
    }
    // When action rewinding is enabled, an action that had remote metadata at some point during the
    // build may have been re-executed locally to regenerate lost inputs, but may then be rewound
    // again and thus have its (now local) outputs deleted. In this case, we need to download the
    // outputs again, even if they are now considered local.
    return metadata.isRemote() || (forceRefetch(path) && !path.exists(Symlinks.NOFOLLOW));
  }

  @Override
  protected boolean forceRefetch(Path path) {
    // External repo paths are never rewound action outputs. Check this first (comparing fragments,
    // since the path may be on the host file system while the output base is on an overlay) to avoid
    // resolving the exec root during external repo materialization, before it is known.
    if (path.asFragment()
        .startsWith(
            outputBase.getRelative(LabelConstants.EXTERNAL_REPOSITORY_LOCATION).asFragment())) {
      return false;
    }
    // Caches for download operations and output directory creation need to be disregarded for the
    // outputs of rewound actions as they may have been deleted after they were first created.
    Path execRoot = execRoot();
    return path.startsWith(execRoot) && rewoundActionOutputs.contains(path.relativeTo(execRoot));
  }

  @Override
  protected ListenableFuture<Void> doDownloadFile(
      @Nullable ActionExecutionMetadata action,
      Reporter reporter,
      Path tempPath,
      PathFragment execPath,
      FileArtifactValue metadata,
      Priority priority,
      Reason reason)
      throws IOException {
    RequestMetadata requestMetadata =
        TracingMetadataUtils.buildMetadata(
            buildRequestId,
            commandId,
            switch (reason) {
              case INPUTS -> "input";
              case OUTPUTS -> "output";
            },
            action);
    RemoteActionExecutionContext context = RemoteActionExecutionContext.create(requestMetadata);

    Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());

    return combinedCache.downloadFile(
        context,
        execPath.getPathString(),
        execPath,
        tempPath,
        digest,
        new CombinedCache.DownloadProgressReporter(
            progress -> {
              if (action != null) {
                progress.postTo(reporter, action);
              }
            },
            execPath.toString(),
            digest.getSizeBytes()));
  }

  public void handleRewoundActionOutputs(Collection<Artifact> outputs) {
    // SkyframeActionExecutor#prepareForRewinding does *not* call this method because the
    // RemoteActionFileSystem corresponds to an ActionFileSystemType with inMemoryFileSystem() ==
    // true. While it is true that resetting outputDirectoryHelper isn't necessary to undo the
    // caching of output directory creation during action preparation, we still need to reset here
    // since outputDirectoryHelper is also used by AbstractActionInputPrefetcher.
    if (outputDirectoryHelper != null) {
      outputDirectoryHelper.invalidateTreeArtifactDirectoryCreation(outputs);
    }
    for (Artifact output : outputs) {
      // Action templates have TreeFileArtifacts as outputs, which isn't supported by the trie. We
      // only need to track the tree artifacts themselves.
      if (output instanceof Artifact.TreeFileArtifact) {
        rewoundActionOutputs.add(output.getParent());
      } else {
        rewoundActionOutputs.add(output);
      }
    }
  }
}