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 [email protected] 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 <[email protected]> Authored: Fri Apr 20 13:26:08 2018 -0700 Committer: Santhosh Kumar <[email protected]> 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()
