src/test/java/com/google/devtools/build/lib/worker/ExampleWorkerMultiplexer.java

// Copyright 2015 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 static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.devtools.build.lib.worker.ExampleWorkerMultiplexerOptions.ExampleWorkMultiplexerOptions;
import com.google.devtools.build.lib.worker.WorkerProtocol.Input;
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
import com.google.devtools.common.options.OptionsParser;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/** An example implementation of a worker process that is used for integration tests. */
public class ExampleWorkerMultiplexer {

  static final Pattern FLAG_FILE_PATTERN = Pattern.compile("(?:@|--?flagfile=)(.+)");

  // Creating Executor Service with a thread pool of Size 3.
  static final int CONCURRENT_THREAD_NUMBER = 3;

  // A UUID that uniquely identifies this running worker process.
  static final UUID WORKER_UUID = UUID.randomUUID();

  // A counter that increases with each work unit processed.
  static int workUnitCounter = 1;

  static int counterOutput = workUnitCounter;

  static Semaphore protectResponse = new Semaphore(1);

  // Keep state across multiple builds.
  static final LinkedHashMap<String, String> inputs = new LinkedHashMap<>();

  private ExampleWorkerMultiplexer() {}

  public static void main(String[] args) throws Exception {
    if (ImmutableSet.copyOf(args).contains("--persistent_worker")) {
      OptionsParser parser =
          OptionsParser.builder()
              .optionsClasses(ExampleWorkerMultiplexerOptions.class)
              .allowResidue(false)
              .build();
      parser.parse(args);
      ExampleWorkerMultiplexerOptions workerOptions =
          parser.getOptions(ExampleWorkerMultiplexerOptions.class);
      Preconditions.checkState(workerOptions.persistentWorker);

      runPersistentWorker(workerOptions);
    } else {
      // This is a single invocation of the example that exits after it processed the request.
      processRequest(parserHelper(ImmutableList.copyOf(args)));
    }
  }

  private static void runPersistentWorker(ExampleWorkerMultiplexerOptions workerOptions)
      throws IOException, ExecutionException, InterruptedException {
    PrintStream originalStdOut = System.out;
    PrintStream originalStdErr = System.err;

    ExecutorService executorService = Executors.newFixedThreadPool(CONCURRENT_THREAD_NUMBER);
    List<Future<?>> results = new ArrayList<>();

    while (true) {
      try {
        WorkRequest request = WorkRequest.parseDelimitedFrom(System.in);
        int requestId = request.getRequestId();
        if (request == null) {
          break;
        }

        inputs.clear();
        for (Input input : request.getInputsList()) {
          inputs.put(input.getPath(), input.getDigest().toStringUtf8());
        }

        // If true, returns corrupt responses instead of correct protobufs.
        boolean poisoned = false;
        if (workerOptions.poisonAfter > 0 && workUnitCounter > workerOptions.poisonAfter) {
          poisoned = true;
        }

        if (poisoned && workerOptions.hardPoison) {
          System.err.println("I'm a very poisoned worker and will just crash.");
          System.exit(1);
        } else {
          int exitCode = 0;
          try {
            OptionsParser parser = parserHelper(request.getArgumentsList());
            ExampleWorkMultiplexerOptions options =
                parser.getOptions(ExampleWorkMultiplexerOptions.class);
            if (options.writeCounter) {
              counterOutput = workUnitCounter++;
            }
            results.add(
                executorService.submit(
                    createTask(originalStdOut, originalStdErr, requestId, parser, poisoned)));
          } catch (Exception e) {
            e.printStackTrace();
            exitCode = 1;
            WorkResponse.newBuilder()
                .setRequestId(requestId)
                .setOutput(new ByteArrayOutputStream().toString())
                .setExitCode(exitCode)
                .build()
                .writeDelimitedTo(System.out);
          }
        }

        if (workerOptions.exitAfter > 0 && workUnitCounter > workerOptions.exitAfter) {
          System.in.close();
        }
      } finally {
        // Be a good worker process and consume less memory when idle.
        System.gc();
      }
    }

    for (Future<?> result : results) {
      result.get();
    }
  }

  private static OptionsParser parserHelper(List<String> args) throws Exception {
    ImmutableList.Builder<String> expandedArgs = ImmutableList.builder();
    for (String arg : args) {
      Matcher flagFileMatcher = FLAG_FILE_PATTERN.matcher(arg);
      if (flagFileMatcher.matches()) {
        expandedArgs.addAll(Files.readAllLines(Paths.get(flagFileMatcher.group(1)), UTF_8));
      } else {
        expandedArgs.add(arg);
      }
    }

    OptionsParser parser =
        OptionsParser.builder()
            .optionsClasses(ExampleWorkMultiplexerOptions.class)
            .allowResidue(true)
            .build();
    parser.parse(expandedArgs.build());

    return parser;
  }

  private static Runnable createTask(
      PrintStream originalStdOut,
      PrintStream originalStdErr,
      int requestId,
      OptionsParser parser,
      boolean poisoned) {
    return () -> {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();
      int exitCode = 0;

      try {
        try (PrintStream ps = new PrintStream(baos)) {
          System.setOut(ps);
          System.setErr(ps);

          if (poisoned) {
            System.out.println("I'm a poisoned worker and this is not a protobuf.");
            System.out.println("Here's a fake stack trace for you:");
            System.out.println("    at com.example.Something(Something.java:83)");
            System.out.println("    at java.lang.Thread.run(Thread.java:745)");
            System.out.print("And now, 8k of random bytes: ");
            byte[] b = new byte[8192];
            new Random().nextBytes(b);
            System.out.write(b);
          } else {
            try {
              processRequest(parser);
            } catch (Exception e) {
              e.printStackTrace();
              exitCode = 1;
            }
          }
        } finally {
          System.setOut(originalStdOut);
          System.setErr(originalStdErr);
        }

        if (poisoned) {
          baos.writeTo(System.out);
        } else {
          protectResponse.acquire();
          WorkResponse.newBuilder()
              .setRequestId(requestId)
              .setOutput(baos.toString())
              .setExitCode(exitCode)
              .build()
              .writeDelimitedTo(System.out);
          protectResponse.release();
        }
        System.out.flush();
      } catch (IOException | InterruptedException e) {
        throw new IllegalStateException(e);
      }
    };
  }

  private static void processRequest(OptionsParser parser) throws Exception {
    ExampleWorkMultiplexerOptions options = parser.getOptions(ExampleWorkMultiplexerOptions.class);

    List<String> outputs = new ArrayList<>();

    if (options.delay) {
      Integer randomDelay = new Random().nextInt(200) + 100;
      TimeUnit.MILLISECONDS.sleep(randomDelay);
      outputs.add("DELAY " + randomDelay + " MILLISECONDS");
    }

    if (options.writeUUID) {
      outputs.add("UUID " + WORKER_UUID.toString());
    }

    if (options.writeCounter) {
      outputs.add("COUNTER " + counterOutput);
    }

    String residueStr = Joiner.on(' ').join(parser.getResidue());
    if (options.uppercase) {
      residueStr = residueStr.toUpperCase();
    }
    outputs.add(residueStr);

    if (options.printInputs) {
      for (Map.Entry<String, String> input : inputs.entrySet()) {
        outputs.add("INPUT " + input.getKey() + " " + input.getValue());
      }
    }

    if (options.printEnv) {
      for (Map.Entry<String, String> entry : System.getenv().entrySet()) {
        outputs.add(entry.getKey() + "=" + entry.getValue());
      }
    }

    String outputStr = Joiner.on('\n').join(outputs);
    if (options.outputFile.isEmpty()) {
      System.out.println(outputStr);
    } else {
      try (PrintStream outputFile = new PrintStream(options.outputFile)) {
        outputFile.println(outputStr);
      }
    }
  }
}