com.hubspot.singularity.runner.base.shared.SafeProcessManager 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.Joiner;
import com.google.common.base.Preconditions;
import com.hubspot.mesos.JavaUtils;
import java.io.IOException;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.slf4j.Logger;

public abstract class SafeProcessManager {
  private final Logger log;
  private final Lock processLock;

  // unnecessary code here : we shouldn't really subclass this when all we need is the following
  // (1) artifact manager needs to kill to ensure it doesn't start something after killing.
  // (2) process lock is unnecessary since we don't need to run multiple processes except in artifact manager
  // (3) we could probably hand off these process variables to make them final in some other object.

  private volatile Optional currentProcessCmd;
  private volatile Optional currentProcess;
  private volatile Optional currentProcessPid;
  private volatile Optional currentProcessStart;

  private final AtomicBoolean killed;
  private final ProcessUtils processUtils;

  public SafeProcessManager(Logger log) {
    this.log = log;
    this.processUtils = new ProcessUtils(log);

    this.currentProcessCmd = Optional.empty();
    this.currentProcess = Optional.empty();
    this.currentProcessPid = Optional.empty();

    this.processLock = new ReentrantLock();

    this.killed = new AtomicBoolean(false);
  }

  public Logger getLog() {
    return log;
  }

  public boolean wasKilled() {
    return killed.get();
  }

  public void markKilled() {
    this.processLock.lock();

    try {
      killed.set(true);
    } finally {
      this.processLock.unlock();
    }
  }

  private void lockInterruptibly() {
    try {
      this.processLock.lockInterruptibly();
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }

  public Process startProcess(ProcessBuilder builder) {
    final String cmd = builder.command().get(0);

    log.debug("Starting process {}", Joiner.on(" ").join(builder.command()));

    processLock.lock();

    Process process = null;

    try {
      Preconditions.checkState(!killed.get(), "Can not start new process, killed is set");
      Preconditions.checkState(
        !currentProcess.isPresent(),
        "Can not start new process, already had process"
      );

      currentProcessStart = Optional.of(System.currentTimeMillis());

      process = builder.start();

      currentProcessPid = Optional.of(processUtils.getUnixPID(process));
      currentProcess = Optional.of(process);
      currentProcessCmd = Optional.of(cmd);

      log.debug("Started process {}", getCurrentProcessToString());
    } catch (IOException e) {
      throw new RuntimeException(e);
    } finally {
      processLock.unlock();
    }

    return process;
  }

  private void resetCurrentVariables() {
    currentProcess = Optional.empty();
    currentProcessPid = Optional.empty();
    currentProcessCmd = Optional.empty();
    currentProcessStart = Optional.empty();
  }

  protected void processFinished(Optional exitCode) {
    lockInterruptibly();

    try {
      if (currentProcessCmd.isPresent() && currentProcessStart.isPresent()) {
        if (exitCode.isPresent()) {
          log.debug(
            "Process {} exited with {} after {}",
            currentProcessCmd.get(),
            exitCode.get(),
            JavaUtils.duration(currentProcessStart.get())
          );
        } else {
          log.debug(
            "Process {} abandoned after {}",
            currentProcessCmd.get(),
            JavaUtils.duration(currentProcessStart.get())
          );
        }
      } else {
        log.warn("Process finished called on an empty process manager");
      }

      resetCurrentVariables();
    } finally {
      processLock.unlock();
    }
  }

  public Optional getCurrentPid() {
    return currentProcessPid;
  }

  public String getCurrentProcessToString() {
    return String.format(
      "%s - (pid: %s)",
      currentProcessCmd.orElse(""),
      currentProcessPid.orElse(0)
    );
  }

  public void signalTermToProcessIfActive() {
    this.processLock.lock();

    try {
      if (currentProcessPid.isPresent()) {
        processUtils.sendSignal(Signal.SIGTERM, currentProcessPid.get());
      }
    } finally {
      this.processLock.unlock();
    }
  }

  public void signalKillToProcessIfActive() {
    this.processLock.lock();

    try {
      if (currentProcess.isPresent()) {
        processUtils.sendSignal(Signal.SIGKILL, currentProcessPid.get());
      }
    } finally {
      this.processLock.unlock();
    }
  }
}