content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.beam.examples.subprocess.kernel;
  
  import java.io.IOException;
  import java.lang.ProcessBuilder.Redirect;
  import java.nio.file.Files;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Stream;
  import org.apache.beam.examples.subprocess.configuration.SubProcessConfiguration;
  import org.apache.beam.examples.subprocess.utils.CallingSubProcessUtils;
  import org.apache.beam.examples.subprocess.utils.FileUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * This is the process kernel which deals with exec of the subprocess. It also deals with all I/O.
   */
  @SuppressWarnings({
    "nullness" // TODO(https://github.com/apache/beam/issues/20497)
  })
  public class SubProcessKernel {
  
    private static final Logger LOG = LoggerFactory.getLogger(SubProcessKernel.class);
  
    private static final int MAX_SIZE_COMMAND_LINE_ARGS = 128 * 1024;
  
    SubProcessConfiguration configuration;
    ProcessBuilder processBuilder;
  
    private SubProcessKernel() {}
  
    /**
     * Creates the SubProcess Kernel ready for execution. Will deal with all input and outputs to the
     * SubProcess
     *
     * @param options
     * @param binaryName
     */
    public SubProcessKernel(SubProcessConfiguration options, String binaryName) {
      this.configuration = options;
      this.processBuilder = new ProcessBuilder(binaryName);
    }
  
    public List<String> exec(SubProcessCommandLineArgs commands) throws Exception {
      try (CallingSubProcessUtils.Permit permit =
          new CallingSubProcessUtils.Permit(processBuilder.command().get(0))) {
  
        List<String> results = null;
  
        try (SubProcessIOFiles outputFiles = new SubProcessIOFiles(configuration.getWorkerPath())) {
  
          try {
            Process process = execBinary(processBuilder, commands, outputFiles);
            results = collectProcessResults(process, processBuilder, outputFiles);
          } catch (Exception ex) {
            LOG.error("Error running executable ", ex);
            throw ex;
          }
        } catch (IOException ex) {
          LOG.error(
              "Unable to delete the outputfiles. This can lead to performance issues and failure",
              ex);
        }
        return results;
      }
    }
  
    public byte[] execBinaryResult(SubProcessCommandLineArgs commands) throws Exception {
      try (CallingSubProcessUtils.Permit permit =
          new CallingSubProcessUtils.Permit(processBuilder.command().get(0))) {
  
        try (SubProcessIOFiles outputFiles = new SubProcessIOFiles(configuration.getWorkerPath())) {
  
          try {
            Process process = execBinary(processBuilder, commands, outputFiles);
            return collectProcessResultsBytes(process, processBuilder, outputFiles);
          } catch (Exception ex) {
            LOG.error("Error running executable ", ex);
            throw ex;
          }
        } catch (IOException ex) {
          LOG.error(
              "Unable to delete the outputfiles. This can lead to performance issues and failure",
              ex);
        }
        return new byte[0];
      }
    }
  
    private ProcessBuilder prepareBuilder(
        ProcessBuilder builder, SubProcessCommandLineArgs commands, SubProcessIOFiles outPutFiles)
        throws IllegalStateException {
  
      // Check we are not over the max size of command line parameters
      if (getTotalCommandBytes(commands) > MAX_SIZE_COMMAND_LINE_ARGS) {
        throw new IllegalStateException("Command is over 2MB in size");
      }
  
      appendExecutablePath(builder);
  
      // Add the result file path to the builder at position 1, 0 is reserved for the process itself
      builder.command().add(1, outPutFiles.resultFile.toString());
  
      // Shift commands by 2 ordinal positions and load into the builder
      for (SubProcessCommandLineArgs.Command s : commands.getParameters()) {
        builder.command().add(s.ordinalPosition + 2, s.value);
      }
  
      builder.redirectError(Redirect.appendTo(outPutFiles.errFile.toFile()));
      builder.redirectOutput(Redirect.appendTo(outPutFiles.outFile.toFile()));
  
      return builder;
    }
  
    /**
     * Add up the total bytes used by the process.
     *
     * @param commands
     * @return
     */
    private int getTotalCommandBytes(SubProcessCommandLineArgs commands) {
      int size = 0;
      for (SubProcessCommandLineArgs.Command c : commands.getParameters()) {
        size += c.value.length();
      }
      return size;
    }
  
    private Process execBinary(
        ProcessBuilder builder, SubProcessCommandLineArgs commands, SubProcessIOFiles outPutFiles)
        throws Exception {
      try {
  
        builder = prepareBuilder(builder, commands, outPutFiles);
        Process process = builder.start();
  
        boolean timeout = !process.waitFor(configuration.getWaitTime(), TimeUnit.SECONDS);
  
        if (timeout) {
          String log =
              String.format(
                  "Timeout waiting to run process with parameters %s . "
                      + "Check to see if your timeout is long enough. Currently set at %s.",
                  createLogEntryFromInputs(builder.command()), configuration.getWaitTime());
          throw new Exception(log);
        }
        return process;
  
      } catch (Exception ex) {
  
        LOG.error(
            String.format(
                "Error running process with parameters %s error was %s ",
                createLogEntryFromInputs(builder.command()), ex.getMessage()));
        throw new Exception(ex);
      }
    }
  
    /**
     * TODO clean up duplicate with byte[] version collectBinaryProcessResults.
     *
     * @param process
     * @param builder
     * @param outPutFiles
     * @return List of results
     * @throws Exception if process has non 0 value or no logs found then throw exception
     */
    private List<String> collectProcessResults(
        Process process, ProcessBuilder builder, SubProcessIOFiles outPutFiles) throws Exception {
  
      List<String> results = new ArrayList<>();
  
      try {
  
        LOG.debug(String.format("Executing process %s", createLogEntryFromInputs(builder.command())));
  
        // If process exit value is not 0 then subprocess failed, record logs
        if (process.exitValue() != 0) {
          outPutFiles.copyOutPutFilesToBucket(configuration, FileUtils.toStringParams(builder));
          String log = createLogEntryForProcessFailure(process, builder.command(), outPutFiles);
          throw new Exception(log);
        }
  
        // If no return file then either something went wrong or the binary is setup incorrectly for
        // the ret file either way throw error
        if (!Files.exists(outPutFiles.resultFile)) {
          String log = createLogEntryForProcessFailure(process, builder.command(), outPutFiles);
          outPutFiles.copyOutPutFilesToBucket(configuration, FileUtils.toStringParams(builder));
          throw new Exception(log);
        }
  
        // Everything looks healthy return values
        try (Stream<String> lines = Files.lines(outPutFiles.resultFile)) {
          for (String line : (Iterable<String>) lines::iterator) {
            results.add(line);
          }
        }
        return results;
      } catch (Exception ex) {
        String log =
            String.format(
                "Unexpected error runnng process. %s error message was %s",
                createLogEntryFromInputs(builder.command()), ex.getMessage());
        throw new Exception(log);
      }
    }
  
    /**
     * Used when the reault file contains binary data.
     *
     * @param process
     * @param builder
     * @param outPutFiles
     * @return Binary results
     * @throws Exception if process has non 0 value or no logs found then throw exception
     */
    private byte[] collectProcessResultsBytes(
        Process process, ProcessBuilder builder, SubProcessIOFiles outPutFiles) throws Exception {
  
      try {
  
        LOG.debug(String.format("Executing process %s", createLogEntryFromInputs(builder.command())));
  
        // If process exit value is not 0 then subprocess failed, record logs
        if (process.exitValue() != 0) {
          outPutFiles.copyOutPutFilesToBucket(configuration, FileUtils.toStringParams(builder));
          String log = createLogEntryForProcessFailure(process, builder.command(), outPutFiles);
          throw new Exception(log);
        }
  
        // If no return file then either something went wrong or the binary is setup incorrectly for
        // the ret file either way throw error
        if (!Files.exists(outPutFiles.resultFile)) {
          String log = createLogEntryForProcessFailure(process, builder.command(), outPutFiles);
          outPutFiles.copyOutPutFilesToBucket(configuration, FileUtils.toStringParams(builder));
          throw new Exception(log);
        }
  
        // Everything looks healthy return bytes
        return Files.readAllBytes(outPutFiles.resultFile);
  
      } catch (Exception ex) {
        String log =
            String.format(
                "Unexpected error runnng process. %s error message was %s",
                createLogEntryFromInputs(builder.command()), ex.getMessage());
        throw new Exception(log);
      }
    }
  
    private static String createLogEntryForProcessFailure(
        Process process, List<String> commands, SubProcessIOFiles files) {
  
      StringBuilder stringBuilder = new StringBuilder();
  
      // Highlight when no result file is found vs standard process error
      if (process.exitValue() == 0) {
        stringBuilder.append(String.format("%nProcess succeded but no result file was found %n"));
      } else {
        stringBuilder.append(
            String.format("%nProcess error failed with exit value of %s %n", process.exitValue()));
      }
  
      stringBuilder.append(
          String.format("Command info was %s %n", createLogEntryFromInputs(commands)));
  
      stringBuilder.append(
          String.format(
              "First line of error file is  %s %n", FileUtils.readLineOfLogFile(files.errFile)));
  
      stringBuilder.append(
          String.format(
              "First line of out file is %s %n", FileUtils.readLineOfLogFile(files.outFile)));
  
      stringBuilder.append(
          String.format(
              "First line of ret file is %s %n", FileUtils.readLineOfLogFile(files.resultFile)));
  
      return stringBuilder.toString();
    }
  
    private static String createLogEntryFromInputs(List<String> commands) {
      String params;
      if (commands != null) {
        params = String.join(",", commands);
      } else {
        params = "No-Commands";
      }
      return params;
    }
  
    // Pass the Path of the binary to the SubProcess in Command position 0
    private ProcessBuilder appendExecutablePath(ProcessBuilder builder) {
      String executable = builder.command().get(0);
      if (executable == null) {
        throw new IllegalArgumentException(
            "No executable provided to the Process Builder... we will do... nothing... ");
      }
      builder
          .command()
          .set(0, FileUtils.getFileResourceId(configuration.getWorkerPath(), executable).toString());
      return builder;
    }
  }
  