Repository: jcrist/skein
Pull request #126 (2019-01-02 17:27:47)

Add ``skein driver stop --force``
Previously the global driver would only be stopped if skein could verify
that the corresponding PID was actually a skein driver. In the case of
credential mismatch this could result in an inability to stop a running
driver using the skein cli. This adds a `--force` option to always kill
the process, even if we're not sure if it's a driver.

Addresses comment here: https://github.com/jcrist/skein/issues/124#issuecomment-450892520


Patch:
diff --git a/skein/cli.py b/skein/cli.py
index 836e992..f2497aa 100644
--- a/skein/cli.py
+++ b/skein/cli.py
@@ -6,7 +6,7 @@ import sys
 import traceback
 
 from . import __version__
-from .core import Client, ApplicationClient, properties
+from .core import Client, ApplicationClient, properties, _read_driver
 from .compatibility import read_stdin_bytes, write_stdout_bytes
 from .exceptions import context, SkeinError, DriverNotRunningError
 from .model import ApplicationSpec, ContainerState, ApplicationState, Security
@@ -173,16 +173,31 @@ def driver_start(keytab=None, principal=None, log=False,
 
 @driver_and_daemon('address', 'The address of the running driver')
 def driver_address():
-    try:
-        client = Client.from_global_driver()
-        print(client.address)
-    except DriverNotRunningError:
+    address, _ = _read_driver()
+    if address is None:
         fail("No skein driver is running")
+    else:
+        print(address)
+
+
+@subcommand(driver.subs,
+            'pid',
+            'The pid of the running driver')
+def driver_pid():
+    _, pid = _read_driver()
+    if pid is None:
+        fail("No skein driver is running")
+    else:
+        print(pid)
 
 
-@driver_and_daemon('stop', 'Stop the skein driver')
-def driver_stop():
-    Client.stop_global_driver()
+@driver_and_daemon('stop', 'Stop the skein driver',
+                   arg('--force', '-f', action='store_true',
+                       help=("Stop the process associated with the driver PID, "
+                             "even if unable to verify it corresponds to a "
+                             "skein driver")))
+def driver_stop(force=False):
+    Client.stop_global_driver(force=force)
 
 
 @driver_and_daemon('restart', 'Restart the skein driver',
diff --git a/skein/core.py b/skein/core.py
index 97b95a9..b128fa5 100644
--- a/skein/core.py
+++ b/skein/core.py
@@ -1,5 +1,6 @@
 from __future__ import print_function, division, absolute_import
 
+import errno
 import json
 import os
 import select
@@ -407,24 +408,37 @@ class Client(_ClientBase):
         return address
 
     @staticmethod
-    def stop_global_driver():
+    def stop_global_driver(force=False):
         """Stops the global driver if running.
 
-        No-op if no global driver is running."""
+        No-op if no global driver is running.
+
+        Parameters
+        ----------
+        force : bool, optional
+            By default skein will check that the process associated with the
+            driver PID is actually a skein driver. Setting ``force`` to
+            ``True`` will kill the process in all cases.
+        """
         address, pid = _read_driver()
         if address is None:
             return
 
-        try:
+        if not force:
+            # Attempt to connect first, errors on failure
             Client(address=address)
-        except DriverNotRunningError:
-            pass
-        else:
+        try:
             os.kill(pid, signal.SIGTERM)
+        except OSError as exc:
+            # If we're forcing a kill, ignore EPERM as well, as we're not sure
+            # if the process is a driver.
+            ignore = (errno.ESRCH, errno.EPERM) if force else (errno.ESRCH,)
+            if exc.errno not in ignore:  # pragma: no cover
+                raise
 
         try:
             os.remove(os.path.join(properties.config_dir, 'driver'))
-        except OSError:
+        except OSError:  # pragma: no cover
             pass
 
     # TODO: deprecated, remove after next release cycle


Test patch:
diff --git a/skein/test/conftest.py b/skein/test/conftest.py
index 3e9addc..ce4cf98 100644
--- a/skein/test/conftest.py
+++ b/skein/test/conftest.py
@@ -10,6 +10,14 @@ import pytest
 import skein
 
 
+def pid_exists(pid):
+    try:
+        os.kill(pid, 0)
+    except OSError:
+        return False
+    return True
+
+
 @contextmanager
 def set_skein_config(tmpdir):
     tmpdir = str(tmpdir)
diff --git a/skein/test/test_cli.py b/skein/test/test_cli.py
index ec2d078..14a26eb 100644
--- a/skein/test/test_cli.py
+++ b/skein/test/test_cli.py
@@ -2,8 +2,10 @@ from __future__ import absolute_import, print_function, division
 
 import io
 import os
+import socket
+import subprocess
 import sys
-from contextlib import contextmanager
+from contextlib import contextmanager, closing
 
 import pytest
 
@@ -11,9 +13,11 @@ import skein
 from skein.compatibility import PY2
 from skein.exceptions import context
 from skein.cli import main
+from skein.core import _write_driver
 from skein.test.conftest import (run_application, sleep_until_killed,
                                  check_is_shutdown, wait_for_containers,
-                                 set_skein_config, ensure_shutdown)
+                                 set_skein_config, ensure_shutdown,
+                                 pid_exists)
 
 
 bad_spec_yaml = """
@@ -66,6 +70,7 @@ def global_client(kinit, tmpdir_factory):
                           'driver stop',
                           'driver restart',
                           'driver address',
+                          'driver pid',
                           'application',
                           'application submit',
                           'application status',
@@ -180,6 +185,12 @@ def test_cli_driver(capsys, skein_config):
         assert not err
         assert out2 == out
 
+        # Get pid
+        run_command('driver pid')
+        out2, err = capsys.readouterr()
+        assert not err
+        int(out2)  # smoketest is integer
+
         # Restart driver
         run_command('driver restart')
         out2, err = capsys.readouterr()
@@ -199,6 +210,49 @@ def test_cli_driver(capsys, skein_config):
         assert not err
 
 
+def test_cli_driver_force_stop(tmpdir, capsys):
+    with set_skein_config(str(tmpdir)):
+        run_command('config gencerts')
+        driver_file = os.path.join(skein.properties.config_dir, 'driver')
+
+        proc = subprocess.Popen(
+            [sys.executable, '-c', '"import time;time.sleep(10)"']
+        )
+        sock = socket.socket()
+        sock.bind(('', 0))
+        address = 'localhost:%d' % sock.getsockname()[1]
+        with closing(sock):
+            # PID is not a skein driver
+            _write_driver(address, proc.pid)
+            assert os.path.exists(driver_file)
+
+            run_command('driver stop', error=True)
+            out, err = capsys.readouterr()
+            assert not out
+            assert err
+            assert os.path.exists(driver_file)
+
+            run_command('driver stop --force')
+            out, err = capsys.readouterr()
+            assert not out
+            assert not err
+            assert not os.path.exists(driver_file)
+            assert proc.wait() is not None
+
+            # Find a PID that doesn't exist
+            pid = 1234
+            while pid_exists(pid):
+                pid += 1
+            _write_driver(address, pid)
+            assert os.path.exists(driver_file)
+
+            run_command('driver stop --force')
+            out, err = capsys.readouterr()
+            assert not out
+            assert not err
+            assert not os.path.exists(driver_file)
+
+
 def test_cli_application_submit_errors(tmpdir, capsys, global_client):
     spec_path = os.path.join(str(tmpdir), 'spec.yaml')
 
diff --git a/skein/test/test_core.py b/skein/test/test_core.py
index f4d2ce5..9155c5f 100644
--- a/skein/test/test_core.py
+++ b/skein/test/test_core.py
@@ -11,7 +11,8 @@ import skein
 from skein.core import Properties
 from skein.exceptions import FileNotFoundError, FileExistsError
 from skein.test.conftest import (run_application, wait_for_containers,
-                                 wait_for_completion, get_logs, KEYTAB_PATH)
+                                 wait_for_completion, get_logs, KEYTAB_PATH,
+                                 pid_exists)
 
 
 def test_properties():
@@ -104,14 +105,6 @@ def test_security_auto_inits(skein_config):
     assert sec == sec2
 
 
-def pid_exists(pid):
-    try:
-        os.kill(pid, 0)
-    except OSError:
-        return False
-    return True
-
-
 def test_client(security, kinit, tmpdir):
     logpath = str(tmpdir.join("log.txt"))