com.hubspot.singularity.runner.base.shared.ProcessUtils Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of SingularityRunnerBase Show documentation
The newest version!
package com.hubspot.singularity.runner.base.shared;

import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.io.CharStreams;
import com.google.common.io.Closeables;
import com.hubspot.mesos.JavaUtils;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.util.Optional;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class ProcessUtils {
  private static final Logger LOG = LoggerFactory.getLogger(ProcessUtils.class);

  private final Optional log;

  public ProcessUtils() {
    this(null);
  }

  @SuppressFBWarnings("NP_PARAMETER_MUST_BE_NONNULL_BUT_MARKED_AS_NULLABLE")
  public ProcessUtils(@Nullable Logger log) {
    this.log = Optional.ofNullable(log);
  }

  public static class ProcessResult {
    private final int exitCode;
    private final String output;

    public ProcessResult(int exitCode, String output) {
      this.exitCode = exitCode;
      this.output = output;
    }

    public int getExitCode() {
      return exitCode;
    }

    public String getOutput() {
      return output;
    }

    @Override
    public String toString() {
      return "ProcessResult [exitCode=" + exitCode + ", output=" + output + "]";
    }
  }

  public ProcessResult sendSignal(Signal signal, int pid) {
    final long start = System.currentTimeMillis();

    if (log.isPresent()) {
      final String logLine = String.format(
        "Signaling %s (%s) to process %s",
        signal,
        signal.getCode(),
        pid
      );
      if (signal == Signal.CHECK) {
        log.get().trace(logLine);
      } else {
        log.get().info(logLine);
      }
    }

    if (signal != Signal.CHECK) {
      LOG.debug("Signaling {} ({}) to process {}", signal, signal.getCode(), pid);
    }

    try {
      final ProcessBuilder pb = new ProcessBuilder(
        "kill",
        String.format("-%s", signal.getCode()),
        Integer.toString(pid)
      );
      pb.redirectErrorStream(true);

      final Process p = pb.start();

      final int exitCode = p.waitFor();

      final String output = CharStreams.toString(
        new InputStreamReader(p.getInputStream(), Charsets.UTF_8)
      );

      Closeables.closeQuietly(p.getInputStream());

      if (log.isPresent()) {
        log
          .get()
          .trace(
            "Kill signal process for {} got exit code {} after {}",
            pid,
            exitCode,
            JavaUtils.duration(start)
          );
      }

      return new ProcessResult(exitCode, output.trim());
    } catch (InterruptedException | IOException e) {
      throw new RuntimeException(e);
    }
  }

  public int getUnixPID(Process process) {
    // older java versions have UNIXProcess, newer have ProcessImpl. Both have a pid field we can access
    Preconditions.checkArgument(
      process.getClass().getName().equals("java.lang.UNIXProcess") ||
      process.getClass().getName().equals("java.lang.ProcessImpl")
    );

    Class clazz = process.getClass();

    try {
      Field field = clazz.getDeclaredField("pid");
      field.setAccessible(true);
      Object pidObject = field.get(process);
      return (Integer) pidObject;
    } catch (
      NoSuchFieldException
      | SecurityException
      | IllegalArgumentException
      | IllegalAccessException e
    ) {
      throw new RuntimeException(e);
    }
  }

  public boolean doesProcessExist(int pid) {
    ProcessResult processResult = sendSignal(Signal.CHECK, pid);
    if (
      processResult.getExitCode() != 0 && processResult.output.contains("No such process")
    ) {
      return false;
    }
    return true;
  }
}