feat: add p2p daemon (#164)

* Add p2p daemon

* Test p2p daemon exits correctly

* Impose restriction on elapsed time

Co-authored-by: Ilya Kobelev <[email protected]>

hivemind/__init__.py

@@ -1,5 +1,6 @@
 from hivemind.client import *
 from hivemind.dht import *
+from hivemind.p2p import *
 from hivemind.server import *
 from hivemind.utils import *
 from hivemind.optim import *

hivemind/p2p/__init__.py

@@ -0,0 +1 @@
+from hivemind.p2p.p2p_daemon import P2P

hivemind/p2p/p2p_daemon.py

@@ -0,0 +1,45 @@
+import subprocess
+import typing as tp
+
+
+class P2P(object):
+    """
+    Forks a child process and executes p2pd command with given arguments.
+    Sends SIGKILL to the child in destructor and on exit from contextmanager.
+    """
+
+    LIBP2P_CMD = 'p2pd'
+
+    def __init__(self, *args, **kwargs):
+        self._child = subprocess.Popen(args=self._make_process_args(args, kwargs))
+        try:
+            stdout, stderr = self._child.communicate(timeout=0.2)
+        except subprocess.TimeoutExpired:
+            pass
+        else:
+            raise RuntimeError(f'p2p daemon exited with stderr: {stderr}')
+
+    def __enter__(self):
+        return self._child
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._kill_child()
+
+    def __del__(self):
+        self._kill_child()
+
+    def _kill_child(self):
+        if self._child.poll() is None:
+            self._child.kill()
+            self._child.wait()
+
+    def _make_process_args(self, args: tp.Tuple[tp.Any],
+                           kwargs: tp.Dict[str, tp.Any]) -> tp.List[str]:
+        proc_args = [self.LIBP2P_CMD]
+        proc_args.extend(
+            str(entry) for entry in args
+        )
+        proc_args.extend(
+            f'-{key}={str(value)}' for key, value in kwargs.items()
+        )
+        return proc_args

tests/test_p2p_daemon.py

@@ -0,0 +1,55 @@
+import subprocess
+from time import perf_counter
+
+import pytest
+
+import hivemind.p2p
+from hivemind.p2p import P2P
+
+RUNNING = 'running'
+NOT_RUNNING = 'not running'
+CHECK_PID_CMD = '''
+if ps -p {0} > /dev/null;
+then
+    echo "{1}"
+else
+    echo "{2}"
+fi
+'''
+
+
+def is_process_running(pid: int) -> bool:
+    cmd = CHECK_PID_CMD.format(pid, RUNNING, NOT_RUNNING)
+    return subprocess.check_output(cmd, shell=True).decode('utf-8').strip() == RUNNING
+
+
+@pytest.fixture()
+def mock_p2p_class():
+    P2P.LIBP2P_CMD = "sleep"
+
+
+def test_daemon_killed_on_del(mock_p2p_class):
+    start = perf_counter()
+    p2p_daemon = P2P('10s')
+
+    child_pid = p2p_daemon._child.pid
+    assert is_process_running(child_pid)
+
+    del p2p_daemon
+    assert not is_process_running(child_pid)
+    assert perf_counter() - start < 1
+
+
+def test_daemon_killed_on_exit(mock_p2p_class):
+    start = perf_counter()
+    with P2P('10s') as daemon:
+        child_pid = daemon.pid
+        assert is_process_running(child_pid)
+
+    assert not is_process_running(child_pid)
+    assert perf_counter() - start < 1
+
+
+def test_daemon_raises_on_faulty_args():
+    with pytest.raises(RuntimeError):
+        P2P(faulty='argument')