// Copyright 2020 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.devtools.build.lib.worker.WorkerProtocol.WorkRequest;
import com.google.devtools.build.lib.worker.WorkerProtocol.WorkResponse;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.sun.management.OperatingSystemMXBean;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintStream;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.management.ManagementFactory;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;

/**
 * A helper class that handles WorkRequests (https://bazel.build/docs/persistent-workers), including
 * multiplex workers (https://bazel.build/docs/multiplex-worker).
 */
public class WorkRequestHandler implements AutoCloseable {
  /** Contains the logic for reading {@link WorkRequest}s and writing {@link WorkResponse}s. */
  public interface WorkerMessageProcessor {
    /** Reads the next incoming request from this worker's stdin. */
    WorkRequest readWorkRequest() throws IOException;

    /**
     * Writes the provided {@link WorkResponse} to this worker's stdout. This function is also
     * responsible for flushing the stdout.
     */
    void writeWorkResponse(WorkResponse workResponse) throws IOException;

    /** Clean up. */
    void close() throws IOException;
  }

  /** Holds information necessary to properly handle a request, especially for cancellation. */
  static class RequestInfo {
    /** The thread handling the request. */
    final Thread thread;
    /** If true, we have received a cancel request for this request. */
    private final AtomicBoolean cancelled = new AtomicBoolean(false);
    /**
     * The builder for the response to this request. Since only one response must be sent per
     * request, this builder must be accessed through takeBuilder(), which zeroes this field and
     * returns the builder.
     */
    private WorkResponse.Builder responseBuilder = WorkResponse.newBuilder();

    RequestInfo(Thread thread) {
      this.thread = thread;
    }

    /** Sets whether this request has been cancelled. */
    void setCancelled() {
      cancelled.set(true);
    }

    /** Returns true if this request has been cancelled. */
    boolean isCancelled() {
      return cancelled.get();
    }

    /**
     * Returns the response builder. If called more than once on the same instance, subsequent calls
     * will return {@code null}.
     */
    synchronized Optional<WorkResponse.Builder> takeBuilder() {
      WorkResponse.Builder b = responseBuilder;
      responseBuilder = null;
      return Optional.ofNullable(b);
    }

    /**
     * Adds {@code s} as output to when the response eventually gets built. Does nothing if the
     * response has already been taken. There is no guarantee that the response hasn't already been
     * taken, making this call a no-op. This may be called multiple times. No delimiters are added
     * between strings from multiple calls.
     */
    synchronized void addOutput(String s) {
      if (responseBuilder != null) {
        responseBuilder.setOutput(responseBuilder.getOutput() + s);
      }
    }
  }

  /** Requests that are currently being processed. Visible for testing. */
  final ConcurrentMap<Integer, RequestInfo> activeRequests = new ConcurrentHashMap<>();

  /** The function to be called after each {@link WorkRequest} is read. */
  private final WorkRequestCallback callback;

  /** This worker's stderr. */
  private final PrintStream stderr;

  final WorkerMessageProcessor messageProcessor;

  private final BiConsumer<Integer, Thread> cancelCallback;
  /**
   * A scheduler that runs garbage collection after a certain amount of CPU time has passed. In our
   * experience, explicit GC reclaims much more than implicit GC. This scheduler helps make sure
   * very busy workers don't grow ridiculously large.
   */
  private final CpuTimeBasedGcScheduler gcScheduler;
  /**
   * A scheduler that runs garbage collection after a certain amount of time without any activity.
   * In our experience, explicit GC reclaims much more than implicit GC. This scheduler helps make
   * sure workers don't hang on to excessive memory after they are done working.
   */
  private final IdleGcScheduler idleGcScheduler;

  /**
   * If set, this worker will stop handling requests and shut itself down. This can happen if
   * something throws an {@link Error}.
   */
  private final AtomicBoolean shutdownWorker = new AtomicBoolean(false);

  /**
   * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
   * received.
   *
   * @param callback Callback method for executing a single WorkRequest in a thread. The first
   *     argument to {@code callback} is the set of command-line arguments, the second is where all
   *     error messages and other user-oriented messages should be written to. The callback must
   *     return an exit code indicating success (zero) or failure (nonzero).
   * @param stderr Stream that log messages should be written to, typically the process' stderr.
   * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server and
   *     writing {@code WorkResponses} to the server.
   */
  @Deprecated
  public WorkRequestHandler(
      BiFunction<List<String>, PrintWriter, Integer> callback,
      PrintStream stderr,
      WorkerMessageProcessor messageProcessor) {
    this(callback, stderr, messageProcessor, Duration.ZERO, null);
  }

  /**
   * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
   * received.
   *
   * @param callback Callback method for executing a single WorkRequest in a thread. The first
   *     argument to {@code callback} is the set of command-line arguments, the second is where all
   *     error messages and other user-oriented messages should be written to. The callback must
   *     return an exit code indicating success (zero) or failure (nonzero).
   * @param stderr Stream that log messages should be written to, typically the process' stderr.
   * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server and
   *     writing {@code WorkResponses} to the server.
   * @param cpuUsageBeforeGc The minimum amount of CPU time between explicit garbage collection
   *     calls. Pass Duration.ZERO to not do explicit garbage collection.
   * @deprecated Use WorkRequestHandlerBuilder instead.
   */
  @Deprecated()
  public WorkRequestHandler(
      BiFunction<List<String>, PrintWriter, Integer> callback,
      PrintStream stderr,
      WorkerMessageProcessor messageProcessor,
      Duration cpuUsageBeforeGc) {
    this(callback, stderr, messageProcessor, cpuUsageBeforeGc, null);
  }

  /**
   * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
   * received. Only used for the Builder.
   *
   * @deprecated Use WorkRequestHandlerBuilder instead.
   */
  @Deprecated
  private WorkRequestHandler(
      BiFunction<List<String>, PrintWriter, Integer> callback,
      PrintStream stderr,
      WorkerMessageProcessor messageProcessor,
      Duration cpuUsageBeforeGc,
      BiConsumer<Integer, Thread> cancelCallback) {
    this(
        new WorkRequestCallback((request, pw) -> callback.apply(request.getArgumentsList(), pw)),
        stderr,
        messageProcessor,
        cpuUsageBeforeGc,
        cancelCallback,
        Duration.ZERO);
  }

  /**
   * Creates a {@code WorkRequestHandler} that will call {@code callback} for each WorkRequest
   * received. Only used for the Builder.
   *
   * @param callback WorkRequestCallback object with Callback method for executing a single
   *     WorkRequest in a thread. The first argument to {@code callback} is the WorkRequest, the
   *     second is where all error messages and other user-oriented messages should be written to.
   *     The callback must return an exit code indicating success (zero) or failure (nonzero).
   */
  private WorkRequestHandler(
      WorkRequestCallback callback,
      PrintStream stderr,
      WorkerMessageProcessor messageProcessor,
      Duration cpuUsageBeforeGc,
      BiConsumer<Integer, Thread> cancelCallback,
      Duration idleTimeBeforeGc) {
    this.callback = callback;
    this.stderr = stderr;
    this.messageProcessor = messageProcessor;
    this.gcScheduler = new CpuTimeBasedGcScheduler(cpuUsageBeforeGc);
    this.cancelCallback = cancelCallback;
    this.idleGcScheduler = new IdleGcScheduler(idleTimeBeforeGc);
  }

  /** A wrapper class for the callback BiFunction */
  public static class WorkRequestCallback {

    /**
     * Callback method for executing a single WorkRequest in a thread. The first argument to {@code
     * callback} is the WorkRequest, the second is where all error messages and other user-oriented
     * messages should be written to. The callback must return an exit code indicating success
     * (zero) or failure (nonzero).
     */
    private final BiFunction<WorkRequest, PrintWriter, Integer> callback;

    public WorkRequestCallback(BiFunction<WorkRequest, PrintWriter, Integer> callback) {
      this.callback = callback;
    }

    public Integer apply(WorkRequest workRequest, PrintWriter printWriter)
        throws InterruptedException {
      Integer result = callback.apply(workRequest, printWriter);
      if (Thread.interrupted()) {
        throw new InterruptedException("Work request interrupted: " + workRequest.getRequestId());
      }
      return result;
    }
  }

  /** Builder class for WorkRequestHandler. Required parameters are passed to the constructor. */
  public static class WorkRequestHandlerBuilder {
    private final WorkRequestCallback callback;
    private final PrintStream stderr;
    private final WorkerMessageProcessor messageProcessor;
    private Duration cpuUsageBeforeGc = Duration.ZERO;
    private BiConsumer<Integer, Thread> cancelCallback;
    private Duration idleTimeBeforeGc = Duration.ZERO;

    /**
     * Creates a {@code WorkRequestHandlerBuilder}.
     *
     * @param callback Callback method for executing a single WorkRequest in a thread. The first
     *     argument to {@code callback} is the set of command-line arguments, the second is where
     *     all error messages and other user-oriented messages should be written to. The callback
     *     must return an exit code indicating success (zero) or failure (nonzero).
     * @param stderr Stream that log messages should be written to, typically the process' stderr.
     * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server
     *     and writing {@code WorkResponses} to the server.
     * @deprecated use WorkRequestHandlerBuilder with WorkRequestCallback instead
     */
    @Deprecated
    public WorkRequestHandlerBuilder(
        BiFunction<List<String>, PrintWriter, Integer> callback,
        PrintStream stderr,
        WorkerMessageProcessor messageProcessor) {
      this(
          new WorkRequestCallback((request, pw) -> callback.apply(request.getArgumentsList(), pw)),
          stderr,
          messageProcessor);
    }

    /**
     * Creates a {@code WorkRequestHandlerBuilder}.
     *
     * @param callback WorkRequestCallback object with Callback method for executing a single
     *     WorkRequest in a thread. The first argument to {@code callback} is the WorkRequest, the
     *     second is where all error messages and other user-oriented messages should be written to.
     *     The callback must return an exit code indicating success (zero) or failure (nonzero).
     * @param stderr Stream that log messages should be written to, typically the process' stderr.
     * @param messageProcessor Object responsible for parsing {@code WorkRequest}s from the server
     *     and writing {@code WorkResponses} to the server.
     */
    public WorkRequestHandlerBuilder(
        WorkRequestCallback callback, PrintStream stderr, WorkerMessageProcessor messageProcessor) {
      this.callback = callback;
      this.stderr = stderr;
      this.messageProcessor = messageProcessor;
    }

    /**
     * Sets the minimum amount of CPU time between explicit garbage collection calls. Pass
     * Duration.ZERO to not do explicit garbage collection (the default).
     */
    @CanIgnoreReturnValue
    public WorkRequestHandlerBuilder setCpuUsageBeforeGc(Duration cpuUsageBeforeGc) {
      this.cpuUsageBeforeGc = cpuUsageBeforeGc;
      return this;
    }

    /**
     * Sets a callback will be called when a cancellation message has been received. The callback
     * will be call with the request ID and the thread executing the request.
     */
    @CanIgnoreReturnValue
    public WorkRequestHandlerBuilder setCancelCallback(BiConsumer<Integer, Thread> cancelCallback) {
      this.cancelCallback = cancelCallback;
      return this;
    }

    /** Sets the time without any work that should elapse before forcing a GC. */
    @CanIgnoreReturnValue
    public WorkRequestHandlerBuilder setIdleTimeBeforeGc(Duration idleTimeBeforeGc) {
      this.idleTimeBeforeGc = idleTimeBeforeGc;
      return this;
    }

    /** Returns a WorkRequestHandler instance with the values in this Builder. */
    public WorkRequestHandler build() {
      return new WorkRequestHandler(
          callback, stderr, messageProcessor, cpuUsageBeforeGc, cancelCallback, idleTimeBeforeGc);
    }
  }

  /**
   * Runs an infinite loop of reading {@link WorkRequest} from {@code in}, running the callback,
   * then writing the corresponding {@link WorkResponse} to {@code out}. If there is an error
   * reading or writing the requests or responses, it writes an error message on {@code err} and
   * returns. If {@code in} reaches EOF, it also returns.
   *
   * <p>This function also wraps the system streams in a {@link WorkerIO} instance that prevents the
   * underlying tool from writing to {@link System#out} or reading from {@link System#in}, which
   * would corrupt the worker worker protocol. When the while loop exits, the original system
   * streams will be swapped back into {@link System}.
   */
  public void processRequests() throws IOException {
    // Wrap the system streams into a WorkerIO instance to prevent unexpected reads and writes on
    // stdin/stdout.
    WorkerIO workerIO = WorkerIO.capture();

    try {
      while (!shutdownWorker.get()) {
        WorkRequest request = messageProcessor.readWorkRequest();
        idleGcScheduler.markActivity(true);
        if (request == null) {
          break;
        }
        if (request.getCancel()) {
          respondToCancelRequest(request);
        } else {
          startResponseThread(workerIO, request);
        }
      }
    } catch (IOException e) {
      stderr.println("Error reading next WorkRequest: " + e);
      e.printStackTrace(stderr);
    } finally {
      idleGcScheduler.stop();
      // TODO(b/220878242): Give the outstanding requests a chance to send a "shutdown" response,
      // but also try to kill stuck threads. For now, we just interrupt the remaining threads.
      // We considered doing System.exit here, but that is hard to test and would deny the callers
      // of this method a chance to clean up. Instead, we initiate the cleanup of our resources here
      // and the caller can decide whether to wait for an orderly shutdown or now.
      for (RequestInfo ri : activeRequests.values()) {
        if (ri.thread.isAlive()) {
          try {
            ri.thread.interrupt();
          } catch (RuntimeException e) {
            // If we can't interrupt, we can't do much else.
          }
        }
      }

      try {
        // Unwrap the system streams placing the original streams back
        workerIO.close();
      } catch (Exception e) {
        stderr.println(e.getMessage());
      }
    }
  }

  /** Starts a thread for the given request. */
  void startResponseThread(WorkerIO workerIO, WorkRequest request) {
    Thread currentThread = Thread.currentThread();
    String threadName =
        request.getRequestId() > 0
            ? "multiplex-request-" + request.getRequestId()
            : "singleplex-request";
    // TODO(larsrc): See if this can be handled with a queue instead, without introducing more
    // race conditions.
    if (request.getRequestId() == 0) {
      while (activeRequests.containsKey(request.getRequestId())) {
        // b/194051480: Previous singleplex requests can still be in activeRequests for a bit after
        // the response has been sent. We need to wait for them to vanish.
        try {
          Thread.sleep(1);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
          return;
        }
      }
    }
    Thread t =
        new Thread(
            () -> {
              RequestInfo requestInfo = activeRequests.get(request.getRequestId());
              if (requestInfo == null) {
                // Already cancelled
                idleGcScheduler.markActivity(!activeRequests.isEmpty());
                return;
              }
              try {
                respondToRequest(workerIO, request, requestInfo);
              } catch (IOException e) {
                // IOExceptions here means a problem talking to the server, so we must shut down.
                if (!shutdownWorker.compareAndSet(false, true)) {
                  stderr.println("Error communicating with server, shutting down worker.");
                  e.printStackTrace(stderr);
                  currentThread.interrupt();
                }
              } finally {
                activeRequests.remove(request.getRequestId());
                idleGcScheduler.markActivity(!activeRequests.isEmpty());
              }
            },
            threadName);
    t.setUncaughtExceptionHandler(
        (t1, e) -> {
          // Shut down the worker in case of severe issues. We don't handle RuntimeException here,
          // as those are not serious enough to merit shutting down the worker.
          if (e instanceof Error && shutdownWorker.compareAndSet(false, true)) {
            stderr.println("Error thrown by worker thread, shutting down worker.");
            e.printStackTrace(stderr);
            currentThread.interrupt();
            idleGcScheduler.stop();
            System.exit(1);
          }
        });
    RequestInfo previous = activeRequests.putIfAbsent(request.getRequestId(), new RequestInfo(t));
    if (previous != null) {
      // Kill worker since this shouldn't happen: server didn't follow the worker protocol
      throw new IllegalStateException("Request still active: " + request.getRequestId());
    }
    t.start();
  }

  /**
   * Handles and responds to the given {@link WorkRequest}.
   *
   * @throws IOException if there is an error talking to the server. Errors from calling the {@link
   *     #callback} are reported with exit code 1.
   */
  @VisibleForTesting
  void respondToRequest(WorkerIO workerIO, WorkRequest request, RequestInfo requestInfo)
      throws IOException {
    int exitCode;
    StringWriter sw = new StringWriter();
    try (PrintWriter pw = new PrintWriter(sw)) {
      try {
        exitCode = callback.apply(request, pw);
      } catch (InterruptedException e) {
        exitCode = 1;
      } catch (RuntimeException e) {
        e.printStackTrace(pw);
        exitCode = 1;
      }

      try {
        // Read out the captured string for the final WorkResponse output
        String captured = workerIO.readCapturedAsUtf8String().trim();
        if (!captured.isEmpty()) {
          pw.write(captured);
        }
      } catch (IOException e) {
        stderr.println(e.getMessage());
      }
    }
    Optional<WorkResponse.Builder> optBuilder = requestInfo.takeBuilder();
    if (optBuilder.isPresent()) {
      WorkResponse.Builder builder = optBuilder.get();
      builder.setRequestId(request.getRequestId());
      if (requestInfo.isCancelled()) {
        builder.setWasCancelled(true);
      } else {
        builder.setOutput(builder.getOutput() + sw).setExitCode(exitCode);
      }
      WorkResponse response = builder.build();
      synchronized (this) {
        messageProcessor.writeWorkResponse(response);
      }
    }
    gcScheduler.maybePerformGc();
  }

  /**
   * Marks the given request as cancelled and uses {@link #cancelCallback} to request cancellation.
   *
   * <p>For simplicity, and to avoid blocking in {@link #cancelCallback}, response to cancellation
   * is still handled by {@link #respondToRequest} once the canceled request aborts (or finishes).
   */
  void respondToCancelRequest(WorkRequest request) {
    // Theoretically, we could have gotten two singleplex requests, and we can't tell those apart.
    // However, that's a violation of the protocol, so we don't try to handle it (not least because
    // handling it would be quite error-prone).
    RequestInfo ri = activeRequests.get(request.getRequestId());

    if (ri == null) {
      return;
    }
    if (cancelCallback == null) {
      ri.setCancelled();
      // This is either an error on the server side or a version mismatch between the server setup
      // and the binary. It's better to wait for the regular work to finish instead of breaking the
      // build, but we should inform the user about the bad setup.
      ri.addOutput(
          String.format(
              "Cancellation request received for worker request %d, but this worker does not"
                  + " support cancellation.\n",
              request.getRequestId()));
    } else {
      if (ri.thread.isAlive() && !ri.isCancelled()) {
        ri.setCancelled();
        Thread t =
            new Thread(
                // Response will be sent from request thread once request handler returns.
                // We can ignore any exceptions in cancel callback since it's best effort.
                () -> cancelCallback.accept(request.getRequestId(), ri.thread));
        t.start();
      }
    }
  }

  @Override
  public void close() throws IOException {
    messageProcessor.close();
  }

  /** Schedules GC when the worker has been idle for a while */
  private static class IdleGcScheduler {
    private Instant lastActivity = Instant.EPOCH;
    private Instant lastGc = Instant.EPOCH;
    /** Minimum duration from the end of activity until we perform an idle GC. */
    private final Duration idleTimeBeforeGc;

    private final ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1);
    private ScheduledFuture<?> futureGc = null;

    /**
     * Creates a new scheduler.
     *
     * @param idleTimeBeforeGc The time from the last activity until attempting GC.
     */
    public IdleGcScheduler(Duration idleTimeBeforeGc) {
      this.idleTimeBeforeGc = idleTimeBeforeGc;
    }

    synchronized void start() {
      if (!idleTimeBeforeGc.isZero()) {
        futureGc =
            executor.schedule(this::maybeDoGc, idleTimeBeforeGc.toMillis(), TimeUnit.MILLISECONDS);
      }
    }

    /**
     * Should be called whenever there is some sort of activity starting or ending. Better to call
     * too often.
     */
    synchronized void markActivity(boolean anythingActive) {
      lastActivity = Instant.now();
      if (futureGc != null) {
        futureGc.cancel(false);
        futureGc = null;
      }
      if (!anythingActive) {
        start();
      }
    }

    private void maybeDoGc() {
      if (lastGc.isBefore(lastActivity)
          && lastActivity.isBefore(Instant.now().minus(idleTimeBeforeGc))) {
        System.gc();
        lastGc = Instant.now();
      } else {
        start();
      }
    }

    synchronized void stop() {
      if (futureGc != null) {
        futureGc.cancel(false);
        futureGc = null;
      }
      executor.shutdown();
    }
  }

  /**
   * Class that performs GC occasionally, based on how much CPU time has passed. This strikes a
   * compromise between blindly doing GC after e.g. every request, which takes too much CPU, and not
   * doing explicit GC at all, which causes poor garbage collection in some cases.
   */
  private static class CpuTimeBasedGcScheduler {
    /**
     * After this much CPU time has elapsed, we may force a GC run. Set to {@link Duration#ZERO} to
     * disable.
     */
    private final Duration cpuUsageBeforeGc;

    /** The total process CPU time at the last GC run (or from the start of the worker). */
    private final AtomicReference<Duration> cpuTimeAtLastGc;

    /** Used to get the CPU time used by this process. */
    private static final OperatingSystemMXBean bean =
        (OperatingSystemMXBean) ManagementFactory.getOperatingSystemMXBean();

    /**
     * Creates a new {@link CpuTimeBasedGcScheduler} that may perform GC after {@code
     * cpuUsageBeforeGc} amount of CPU time has been used.
     */
    public CpuTimeBasedGcScheduler(Duration cpuUsageBeforeGc) {
      this.cpuUsageBeforeGc = cpuUsageBeforeGc;
      this.cpuTimeAtLastGc = new AtomicReference<>(getCpuTime());
    }

    private Duration getCpuTime() {
      return !cpuUsageBeforeGc.isZero()
          ? Duration.ofNanos(bean.getProcessCpuTime())
          : Duration.ZERO;
    }

    /** Call occasionally to perform a GC if enough CPU time has been used. */
    private void maybePerformGc() {
      if (!cpuUsageBeforeGc.isZero()) {
        Duration currentCpuTime = getCpuTime();
        Duration lastCpuTime = cpuTimeAtLastGc.get();
        // Do GC when enough CPU time has been used, but only if nobody else beat us to it.
        if (currentCpuTime.minus(lastCpuTime).compareTo(cpuUsageBeforeGc) > 0
            && cpuTimeAtLastGc.compareAndSet(lastCpuTime, currentCpuTime)) {
          System.gc();
          // Avoid counting GC CPU time against CPU time before next GC.
          cpuTimeAtLastGc.compareAndSet(currentCpuTime, getCpuTime());
        }
      }
    }
  }

  /**
   * Class that wraps the standard {@link System#in}, {@link System#out}, and {@link System#err}
   * with our own ByteArrayOutputStream that allows {@link WorkRequestHandler} to safely capture
   * outputs that can't be directly captured by the PrintStream associated with the work request.
   *
   * <p>This is most useful when integrating JVM tools that write exceptions and logs directly to
   * {@link System#out} and {@link System#err}, which would corrupt the persistent worker protocol.
   * We also redirect {@link System#in}, just in case a tool should attempt to read it.
   *
   * <p>WorkerIO implements {@link AutoCloseable} and will swap the original streams back into
   * {@link System} once close has been called.
   */
  public static class WorkerIO implements AutoCloseable {
    private final InputStream originalInputStream;
    private final PrintStream originalOutputStream;
    private final PrintStream originalErrorStream;
    private final ByteArrayOutputStream capturedStream;
    private final AutoCloseable restore;

    /**
     * Creates a new {@link WorkerIO} that allows {@link WorkRequestHandler} to capture standard
     * output and error streams that can't be directly captured by the PrintStream associated with
     * the work request.
     */
    @VisibleForTesting
    WorkerIO(
        InputStream originalInputStream,
        PrintStream originalOutputStream,
        PrintStream originalErrorStream,
        ByteArrayOutputStream capturedStream,
        AutoCloseable restore) {
      this.originalInputStream = originalInputStream;
      this.originalOutputStream = originalOutputStream;
      this.originalErrorStream = originalErrorStream;
      this.capturedStream = capturedStream;
      this.restore = restore;
    }

    /** Wraps the standard System streams and WorkerIO instance */
    public static WorkerIO capture() {
      // Save the original streams
      InputStream originalInputStream = System.in;
      PrintStream originalOutputStream = System.out;
      PrintStream originalErrorStream = System.err;

      // Replace the original streams with our own instances
      ByteArrayOutputStream capturedStream = new ByteArrayOutputStream();
      PrintStream outputBuffer = new PrintStream(capturedStream, true);
      ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(new byte[0]);
      System.setIn(byteArrayInputStream);
      System.setOut(outputBuffer);
      System.setErr(outputBuffer);

      return new WorkerIO(
          originalInputStream,
          originalOutputStream,
          originalErrorStream,
          capturedStream,
          () -> {
            System.setIn(originalInputStream);
            System.setOut(originalOutputStream);
            System.setErr(originalErrorStream);
            outputBuffer.close();
            byteArrayInputStream.close();
          });
    }

    /** Returns the original input stream most commonly provided by {@link System#in} */
    @VisibleForTesting
    InputStream getOriginalInputStream() {
      return originalInputStream;
    }

    /** Returns the original output stream most commonly provided by {@link System#out} */
    @VisibleForTesting
    PrintStream getOriginalOutputStream() {
      return originalOutputStream;
    }

    /** Returns the original error stream most commonly provided by {@link System#err} */
    @VisibleForTesting
    PrintStream getOriginalErrorStream() {
      return originalErrorStream;
    }

    /** Returns the captured outputs as a UTF-8 string */
    @VisibleForTesting
    String readCapturedAsUtf8String() throws IOException {
      capturedStream.flush();
      String captureOutput = capturedStream.toString(StandardCharsets.UTF_8);
      capturedStream.reset();
      return captureOutput;
    }

    @Override
    public void close() throws Exception {
      restore.close();
    }
  }
}