src/main/java/com/google/devtools/build/lib/worker/WorkerParser.java

// Copyright 2021 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.worker;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.hash.HashCode;
import com.google.devtools.build.lib.actions.ExecException;
import com.google.devtools.build.lib.actions.ExecutionRequirements.WorkerProtocolFormat;
import com.google.devtools.build.lib.actions.Spawn;
import com.google.devtools.build.lib.actions.Spawns;
import com.google.devtools.build.lib.actions.UserExecException;
import com.google.devtools.build.lib.exec.BinTools;
import com.google.devtools.build.lib.exec.SpawnRunner.SpawnExecutionContext;
import com.google.devtools.build.lib.exec.local.LocalEnvProvider;
import com.google.devtools.build.lib.server.FailureDetails;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.PathFragment;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.SortedMap;
import java.util.regex.Pattern;

/**
 * A helper class to process a {@link Spawn} into a {@link WorkerKey}, which is used to select a
 * persistent worker process (actions with equal keys are allowed to use the same worker process),
 * and a separate list of flag files. The result is encapsulated as a {@link WorkerConfig}.
 */
class WorkerParser {
  private static final String ERROR_MESSAGE_PREFIX =
      "Worker strategy cannot execute this %s action, ";
  private static final String REASON_NO_FLAGFILE =
      "because the command-line arguments do not contain exactly one @flagfile or --flagfile=";
  private static final String REASON_EXCESS_FLAGFILE =
      "because the command-line arguments has a @flagfile or --flagfile= argument before the end";
  private static final String REASON_NO_FINAL_FLAGFILE =
      "because the command-line arguments does not end with a @flagfile or --flagfile= argument";

  /**
   * Pattern for @flagfile.txt and --flagfile=flagfile.txt. This doesn't handle @@-escapes, those
   * are checked for separately.
   */
  private static final Pattern FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");

  /**
   * Legacy pattern for @flagfile.txt and --flagfile=flagfile.txt. This doesn't handle @@-escapes.
   */
  private static final Pattern LEGACY_FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");

  /** The global execRoot. */
  private final Path execRoot;

  private final WorkerOptions workerOptions;
  private final LocalEnvProvider localEnvProvider;
  private final BinTools binTools;

  public WorkerParser(
      Path execRoot,
      WorkerOptions workerOptions,
      LocalEnvProvider localEnvProvider,
      BinTools binTools) {
    this.execRoot = execRoot;
    this.workerOptions = workerOptions;
    this.localEnvProvider = localEnvProvider;
    this.binTools = binTools;
  }

  /**
   * Processes the given {@link Spawn} and {@link SpawnExecutionContext} to compute the worker key.
   * This involves splitting the command line into the worker startup command and the separate list
   * of flag files. Returns a pair of the {@link WorkerKey} and list of flag files.
   */
  public WorkerConfig compute(Spawn spawn, SpawnExecutionContext context)
      throws ExecException, IOException, InterruptedException {
    // We assume that the spawn to be executed always gets at least one @flagfile.txt or
    // --flagfile=flagfile.txt argument, which contains the flags related to the work itself (as
    // opposed to start-up options for the executed tool). Thus, we can extract those elements from
    // its args and put them into the WorkRequest instead.
    List<String> flagFiles = new ArrayList<>();
    ImmutableList<String> workerArgs = splitSpawnArgsIntoWorkerArgsAndFlagFiles(spawn, flagFiles);
    ImmutableMap<String, String> env =
        localEnvProvider.rewriteLocalEnv(spawn.getEnvironment(), binTools, "/tmp");

    SortedMap<PathFragment, byte[]> workerFiles =
        WorkerFilesHash.getWorkerFilesWithDigests(
            spawn, context.getArtifactExpander(), context.getMetadataProvider());

    HashCode workerFilesCombinedHash = WorkerFilesHash.getCombinedHash(workerFiles);

    WorkerKey key =
        createWorkerKey(
            spawn,
            workerArgs,
            env,
            execRoot,
            workerFilesCombinedHash,
            workerFiles,
            workerOptions,
            context.speculating(),
            Spawns.getWorkerProtocolFormat(spawn));
    return new WorkerConfig(key, flagFiles);
  }

  /**
   * This method handles the logic of creating a WorkerKey (e.g., if sandboxing should be enabled or
   * not, when to use multiplex-workers)
   */
  @VisibleForTesting
  static WorkerKey createWorkerKey(
      Spawn spawn,
      ImmutableList<String> workerArgs,
      ImmutableMap<String, String> env,
      Path execRoot,
      HashCode workerFilesCombinedHash,
      SortedMap<PathFragment, byte[]> workerFiles,
      WorkerOptions options,
      boolean dynamic,
      WorkerProtocolFormat protocolFormat) {
    boolean multiplex = options.workerMultiplex && Spawns.supportsMultiplexWorkers(spawn);
    if (dynamic && !(Spawns.supportsMultiplexSandboxing(spawn) && options.multiplexSandboxing)) {
      multiplex = false;
    }
    boolean sandboxed;
    if (multiplex) {
      sandboxed =
          Spawns.supportsMultiplexSandboxing(spawn) && (options.multiplexSandboxing || dynamic);
    } else {
      sandboxed = options.workerSandboxing || dynamic;
    }
    return new WorkerKey(
        workerArgs,
        env,
        execRoot,
        Spawns.getWorkerKeyMnemonic(spawn),
        workerFilesCombinedHash,
        workerFiles,
        sandboxed,
        multiplex,
        Spawns.supportsWorkerCancellation(spawn),
        protocolFormat);
  }

  private static boolean isFlagFileArg(String arg) {
    return FLAG_FILE_PATTERN.matcher(arg).matches() && !arg.startsWith("@@");
  }

  private static boolean isLegacyFlagFileArg(String arg) {
    return LEGACY_FLAG_FILE_PATTERN.matcher(arg).matches();
  }

  /**
   * Splits the command-line arguments of the {@code Spawn} into the part that is used to start the
   * persistent worker ({@code workerArgs}) and the part that goes into the {@code WorkRequest}
   * protobuf ({@code flagFiles}).
   */
  @VisibleForTesting
  ImmutableList<String> splitSpawnArgsIntoWorkerArgsAndFlagFiles(
      Spawn spawn, List<String> flagFiles) throws UserExecException {
    ImmutableList.Builder<String> workerArgs = ImmutableList.builder();
    ImmutableList<String> args = spawn.getArguments();
    if (args.isEmpty()) {
      throwFlagFileFailure(REASON_NO_FLAGFILE, spawn);
    }
    if (workerOptions.strictFlagfiles) {
      if (!isFlagFileArg(Iterables.getLast(args))) {
        throwFlagFileFailure(REASON_NO_FINAL_FLAGFILE, spawn);
      }
      flagFiles.add(Iterables.getLast(args));
      for (int i = 0; i < args.size() - 1; i++) {
        if (isFlagFileArg(args.get(i))) {
          throwFlagFileFailure(REASON_EXCESS_FLAGFILE, spawn);
        } else {
          workerArgs.add(args.get(i));
        }
      }
    } else {
      for (String arg : args) {
        if (isLegacyFlagFileArg(arg)) {
          flagFiles.add(arg);
        } else {
          workerArgs.add(arg);
        }
      }
      if (flagFiles.isEmpty()) {
        throwFlagFileFailure(REASON_NO_FLAGFILE, spawn);
      }
    }

    ImmutableList.Builder<String> mnemonicFlags = ImmutableList.builder();

    workerOptions.workerExtraFlags.stream()
        .filter(entry -> entry.getKey().equals(spawn.getMnemonic()))
        .forEach(entry -> mnemonicFlags.add(entry.getValue()));

    return workerArgs.add("--persistent_worker").addAll(mnemonicFlags.build()).build();
  }

  private void throwFlagFileFailure(String reason, Spawn spawn) throws UserExecException {
    String message =
        String.format(
            ERROR_MESSAGE_PREFIX + reason + "%n%s", spawn.getMnemonic(), spawn.getArguments());
    throw new UserExecException(
        FailureDetails.FailureDetail.newBuilder()
            .setMessage(message)
            .setWorker(
                FailureDetails.Worker.newBuilder().setCode(FailureDetails.Worker.Code.NO_FLAGFILE))
            .build());
  }

  /** A pair of the {@link WorkerKey} and the list of flag files. */
  public static class WorkerConfig {
    private final WorkerKey workerKey;
    private final List<String> flagFiles;

    public WorkerConfig(WorkerKey workerKey, List<String> flagFiles) {
      this.workerKey = workerKey;
      this.flagFiles = ImmutableList.copyOf(flagFiles);
    }

    public WorkerKey getWorkerKey() {
      return workerKey;
    }

    public List<String> getFlagFiles() {
      return flagFiles;
    }
  }
}