You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@aurora.apache.org by sa...@apache.org on 2018/04/20 20:26:33 UTC

aurora git commit: Add --pid-file flag to `aurora task ssh` to write the PID of the underlying SSH command to a specified file.

Repository: aurora
Updated Branches:
  refs/heads/master ad0bc5f9d -> 38aae0fdc


Add --pid-file flag to `aurora task ssh` to write the PID of the underlying SSH command to a specified file.

My team has some scripts to start devel shards which create tunnels:

```
aurora task ssh -L 8002:http --ssh-options "-f -N" "$DC/$USER/devel/proxyapp/0"
aurora task ssh -L 9002:health --ssh-options "-f -N" "$DC/$USER/devel/proxyapp/0"
```

We use fixed local port numbers because that way we can run dependent services locally that look for locally-running copies of the
same service on a fixed port, but then those requests get tunnelled through to the devel shard.

When the devel shard is restarted, however, the tunnel is still running so the subsequent call to create a new tunnel fails because
it can't bind to the fixed port.

If we save the SSH process PID to a file, we can then kill existing tunnel to the old instance before starting up the new tunnel to the
new instance.

Testing Done:
```
$ ./pants test src/test/python/apache/aurora/client::
```

And when applying the same patch to our local repo at Twitter:

```
$ ./pants run twitter/src/main/python/twitter/aurora/client/cli_internal:aurora_internal -- task ssh -L 8005:http --ssh-options "-n -N" --pid-file /tmp/p "smf1/sbrenn/devel/proxyapp/0" &
$ ps -p `cat /tmp/p`
  PID TTY           TIME CMD
34729 ttys000    0:00.05 ssh -t -n -N -L 8005:smf1-aki-27-sr1.prod.twitter.com:31794 sbrenn@smf1-aki-27-sr1.prod.twitter.com cd /var/lib/mesos/slaves/*/frameworks/*/exec
```

Reviewed at https://reviews.apache.org/r/66697/


Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/38aae0fd
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/38aae0fd
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/38aae0fd

Branch: refs/heads/master
Commit: 38aae0fdc88450069c4d0a7a250759a8f057b6fd
Parents: ad0bc5f
Author: Sameer Brenn <sb...@twitter.com>
Authored: Fri Apr 20 13:26:08 2018 -0700
Committer: Santhosh Kumar <ss...@twitter.com>
Committed: Fri Apr 20 13:26:08 2018 -0700

----------------------------------------------------------------------
 .../python/apache/aurora/client/cli/task.py     | 11 +++-
 .../apache/aurora/client/cli/test_task.py       | 56 ++++++++++++++++++--
 2 files changed, 62 insertions(+), 5 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/38aae0fd/src/main/python/apache/aurora/client/cli/task.py
----------------------------------------------------------------------
diff --git a/src/main/python/apache/aurora/client/cli/task.py b/src/main/python/apache/aurora/client/cli/task.py
index 652a545..44dc908 100644
--- a/src/main/python/apache/aurora/client/cli/task.py
+++ b/src/main/python/apache/aurora/client/cli/task.py
@@ -109,6 +109,9 @@ class SshCommand(Verb):
         CommandOption('--command', '-c', dest='command', type=str, default=None,
             metavar="unix_command_line",
             help="Command to execute through the ssh connection."),
+        CommandOption('--pid-file', '-p', dest='pid_file', type=str, default=None,
+            metavar="pid_file",
+            help="File in which to store the PID of the resulting ssh call")
     ]
 
   def execute(self, context):
@@ -154,7 +157,13 @@ class SshCommand(Verb):
           '-L', '%d:%s:%d' % (port, slave_host, assigned.assignedPorts[name])]
 
     ssh_command += ['%s@%s' % (context.options.ssh_user or role, slave_host), command]
-    return subprocess.call(ssh_command)
+    process = subprocess.Popen(ssh_command)
+
+    if context.options.pid_file:
+      with open(context.options.pid_file, "w") as f:
+        f.write(str(process.pid))
+
+    return process.wait()
 
 
 class ScpCommand(Verb):

http://git-wip-us.apache.org/repos/asf/aurora/blob/38aae0fd/src/test/python/apache/aurora/client/cli/test_task.py
----------------------------------------------------------------------
diff --git a/src/test/python/apache/aurora/client/cli/test_task.py b/src/test/python/apache/aurora/client/cli/test_task.py
index a543d4a..effcaf6 100644
--- a/src/test/python/apache/aurora/client/cli/test_task.py
+++ b/src/test/python/apache/aurora/client/cli/test_task.py
@@ -13,6 +13,7 @@
 #
 
 import contextlib
+import tempfile
 
 import pytest
 from mock import Mock, patch
@@ -110,6 +111,14 @@ class TestRunCommand(AuroraClientCommandTest):
 
 
 class TestSshCommand(AuroraClientCommandTest):
+  MOCKED_PID = 12312
+
+  @classmethod
+  def create_mock_process(cls):
+    process = Mock()
+    process.pid = cls.MOCKED_PID
+    process.wait.return_value = 0
+    return process
 
   @classmethod
   def create_status_response(cls):
@@ -128,6 +137,45 @@ class TestSshCommand(AuroraClientCommandTest):
   def create_failed_status_response(cls):
     return cls.create_blank_response(ResponseCode.INVALID_REQUEST, 'No tasks found for query')
 
+  def test_successful_ssh_with_pidfile(self):
+    """Test the ssh command with a saved PID file."""
+    with tempfile.NamedTemporaryFile() as pid_file:
+      (mock_api, mock_scheduler_proxy) = self.create_mock_api()
+      mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response()
+      sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'}
+      with contextlib.nested(
+          patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
+          patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
+              return_value=sandbox_args),
+          patch('subprocess.Popen', return_value=self.create_mock_process())) as (
+              mock_scheduler_proxy_class,
+              mock_runner_args_patch,
+              mock_subprocess):
+        cmd = AuroraCommandLine()
+        cmd.execute([
+          'task',
+          'ssh',
+          '--pid-file={}'.format(pid_file.name),
+          '--ssh-options=-v',
+          'west/bozo/test/hello/1',
+          '--command=ls'
+        ])
+
+        # The status command sends a getTasksStatus query to the scheduler,
+        # and then prints the result.
+        mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery(
+            jobKeys=[JobKey(role='bozo', environment='test', name='hello')],
+            instanceIds=set([1]),
+            statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING,
+                ScheduleStatus.PREEMPTING, ScheduleStatus.PARTITIONED, ScheduleStatus.DRAINING])),
+            retry=True)
+        mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost',
+            'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/'
+            'slaverun/sandbox;ls'])
+
+        pid = pid_file.read()
+        assert(pid == str(self.MOCKED_PID))
+
   def test_successful_ssh(self):
     """Test the ssh command."""
     (mock_api, mock_scheduler_proxy) = self.create_mock_api()
@@ -137,7 +185,7 @@ class TestSshCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
             return_value=sandbox_args),
-        patch('subprocess.call', return_value=0)) as (
+        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
             mock_scheduler_proxy_class,
             mock_runner_args_patch,
             mock_subprocess):
@@ -165,7 +213,7 @@ class TestSshCommand(AuroraClientCommandTest):
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
         patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args',
             return_value=sandbox_args),
-        patch('subprocess.call', return_value=0)) as (
+        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
             mock_scheduler_proxy_class,
             mock_runner_args_patch,
             mock_subprocess):
@@ -190,7 +238,7 @@ class TestSshCommand(AuroraClientCommandTest):
     mock_scheduler_proxy.getTasksStatus.return_value = self.create_nojob_status_response()
     with contextlib.nested(
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
-        patch('subprocess.call', return_value=0)) as (
+        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
             mock_scheduler_proxy_class,
             mock_subprocess):
       cmd = AuroraCommandLine()
@@ -204,7 +252,7 @@ class TestSshCommand(AuroraClientCommandTest):
     mock_scheduler_proxy.getTasksStatus.return_value = self.create_nojob_status_response()
     with contextlib.nested(
         patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy),
-        patch('subprocess.call', return_value=0)) as (
+        patch('subprocess.Popen', return_value=self.create_mock_process())) as (
             mock_scheduler_proxy_class,
             mock_subprocess):
       cmd = AuroraCommandLine()