add PID replacement

HubSpot · Apr 28, 2015 · 301d517 · 301d517
1 parent c9aa233
commit 301d517
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 2 deletions.
diff --git a/...ava/com/hubspot/singularity/executor/SingularityExecutorMesosFrameworkMessageHandler.java b/...ava/com/hubspot/singularity/executor/SingularityExecutorMesosFrameworkMessageHandler.java
@@ -11,10 +11,12 @@
 import com.google.common.base.Optional;
 import com.google.inject.Inject;
 import com.hubspot.singularity.SingularityTaskShellCommandRequest;
+import com.hubspot.singularity.SingularityTaskShellCommandUpdate.UpdateType;
 import com.hubspot.singularity.executor.config.SingularityExecutorConfiguration;
 import com.hubspot.singularity.executor.shells.SingularityExecutorShellCommandRunner;
 import com.hubspot.singularity.executor.shells.SingularityExecutorShellCommandUpdater;
 import com.hubspot.singularity.executor.task.SingularityExecutorTask;
+import com.hubspot.singularity.executor.task.SingularityExecutorTaskProcessCallable;
 
 public class SingularityExecutorMesosFrameworkMessageHandler {
 
@@ -44,8 +46,15 @@ public void handleMessage(byte[] data) {
 
       SingularityExecutorShellCommandUpdater updater = new SingularityExecutorShellCommandUpdater(objectMapper, shellRequest, matchingTask.get());
 
+      Optional<SingularityExecutorTaskProcessCallable> taskProcess = monitor.getTaskProcess(shellRequest.getTaskId().getId());
+
+      if (!taskProcess.isPresent()) {
+        updater.sendUpdate(UpdateType.INVALID, Optional.of("No task process found"));
+        return;
+      }
+
       SingularityExecutorShellCommandRunner shellRunner = new SingularityExecutorShellCommandRunner(shellRequest, executorConfiguration, matchingTask.get(),
-          monitor.createExecutorService(shellRequest.getTaskId().getId()), updater);
+          taskProcess.get(), monitor.createExecutorService(shellRequest.getTaskId().getId()), updater);
 
       shellRunner.start();
 

diff --git a/...tyExecutor/src/main/java/com/hubspot/singularity/executor/SingularityExecutorMonitor.java b/...tyExecutor/src/main/java/com/hubspot/singularity/executor/SingularityExecutorMonitor.java
@@ -246,6 +246,10 @@ public Collection<SingularityExecutorTaskProcessCallable> getRunningTasks() {
     return processRunningTasks.values();
   }
 
+  public Optional<SingularityExecutorTaskProcessCallable> getTaskProcess(String taskId) {
+    return Optional.fromNullable(processRunningTasks.get(taskId));
+  }
+
   public Optional<SingularityExecutorTask> getTask(String taskId) {
     return Optional.fromNullable(tasks.get(taskId));
   }

diff --git a/...c/main/java/com/hubspot/singularity/executor/config/SingularityExecutorConfiguration.java b/...c/main/java/com/hubspot/singularity/executor/config/SingularityExecutorConfiguration.java
@@ -214,10 +214,21 @@ public class SingularityExecutorConfiguration extends BaseRunnerConfiguration {
   @JsonProperty
   public List<SingularityExecutorShellCommandDescriptor> shellCommands = Collections.emptyList();
 
+  @JsonProperty
+  private String pidCommandPlaceholder = "{PID}";
+
   public SingularityExecutorConfiguration() {
     super(Optional.of("singularity-executor.log"));
   }
 
+  public String getPidCommandPlaceholder() {
+    return pidCommandPlaceholder;
+  }
+
+  public void setPidCommandPlaceholder(String pidCommandPlaceholder) {
+    this.pidCommandPlaceholder = pidCommandPlaceholder;
+  }
+
   public List<String> getLogrotateAdditionalFiles() {
     return logrotateAdditionalFiles;
   }

diff --git a/...n/java/com/hubspot/singularity/executor/shells/SingularityExecutorShellCommandRunner.java b/...n/java/com/hubspot/singularity/executor/shells/SingularityExecutorShellCommandRunner.java
@@ -15,11 +15,13 @@
 import com.hubspot.singularity.SingularityTaskShellCommandUpdate.UpdateType;
 import com.hubspot.singularity.executor.config.SingularityExecutorConfiguration;
 import com.hubspot.singularity.executor.task.SingularityExecutorTask;
+import com.hubspot.singularity.executor.task.SingularityExecutorTaskProcessCallable;
 
 public class SingularityExecutorShellCommandRunner {
 
   private final SingularityTaskShellCommandRequest shellRequest;
   private final SingularityExecutorTask task;
+  private final SingularityExecutorTaskProcessCallable taskProcess;
   private final ListeningExecutorService shellCommandExecutorService;
   private final SingularityExecutorShellCommandUpdater shellCommandUpdater;
   private final SingularityExecutorConfiguration executorConfiguration;
@@ -34,10 +36,11 @@ public InvalidShellCommandException(String message) {
   }
 
   public SingularityExecutorShellCommandRunner(SingularityTaskShellCommandRequest shellRequest, SingularityExecutorConfiguration executorConfiguration, SingularityExecutorTask task,
-      ListeningExecutorService shellCommandExecutorService, SingularityExecutorShellCommandUpdater shellCommandUpdater) {
+      SingularityExecutorTaskProcessCallable taskProcess, ListeningExecutorService shellCommandExecutorService, SingularityExecutorShellCommandUpdater shellCommandUpdater) {
     this.shellRequest = shellRequest;
     this.executorConfiguration = executorConfiguration;
     this.task = task;
+    this.taskProcess = taskProcess;
     this.shellCommandUpdater = shellCommandUpdater;
     this.shellCommandExecutorService = shellCommandExecutorService;
   }
@@ -107,6 +110,15 @@ public boolean apply(SingularityExecutorShellCommandDescriptor input) {
 
     List<String> command = new ArrayList<>(shellCommandDescriptor.getCommand());
 
+    for (int i = 0; i < command.size(); i++) {
+      if (command.get(i).equals(executorConfiguration.getPidCommandPlaceholder())) {
+        if (!taskProcess.getCurrentPid().isPresent()) {
+          throw new InvalidShellCommandException("No PID found");
+        }
+        command.set(i, Integer.toString(taskProcess.getCurrentPid().get()));
+      }
+    }
+
     for (SingularityExecutorShellCommandOptionDescriptor option : shellCommandDescriptor.getOptions()) {
       if (shellRequest.getShellCommand().getOptions().contains(option.getName())) {
         command.add(option.getFlag());