From 13cdd135284d1571240013d0697210c03a965cd7 Mon Sep 17 00:00:00 2001
From: Alexander Borzunov <borzunov.alexander@gmail.com>
Date: Tue, 11 Oct 2022 19:10:41 +0400
Subject: [PATCH] Follow up #501 and #511 with minor fixes (#513)

- In `hivemind.Server`, use the graceful shutdown for `ConnectionHandler`
- In `hivemind.P2P`, if we are the first peer, skip checking if the provided identity is free
---
 hivemind/moe/server/server.py | 14 +++++---------
 hivemind/p2p/p2p_daemon.py    |  5 +++--
 tests/test_start_server.py    |  3 ---
 3 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/hivemind/moe/server/server.py b/hivemind/moe/server/server.py
index f4d7d7a77..6f47d1e40 100644
--- a/hivemind/moe/server/server.py
+++ b/hivemind/moe/server/server.py
@@ -247,10 +247,8 @@ def run(self):
         if self.checkpoint_saver is not None:
             self.checkpoint_saver.start()
 
-        for process in self.conn_handlers:
-            if not process.is_alive():
-                process.start()
-            process.ready.result()
+        for handler in self.conn_handlers:
+            handler.run_in_background()
 
         try:
             self.runtime.run()
@@ -287,9 +285,8 @@ def shutdown(self):
         """
         self.ready.clear()
 
-        for process in self.conn_handlers:
-            process.terminate()
-            process.join()
+        for handler in self.conn_handlers:
+            handler.shutdown()
         logger.debug("Connection handlers terminated")
 
         if self.module_backends:
@@ -301,11 +298,10 @@ def shutdown(self):
             self.checkpoint_saver.join()
 
         self.dht.shutdown()
-        self.dht.join()
 
         logger.debug(f"Shutting down runtime")
-
         self.runtime.shutdown()
+
         logger.info("Server shutdown succesfully")
 
 
diff --git a/hivemind/p2p/p2p_daemon.py b/hivemind/p2p/p2p_daemon.py
index 366d9c39c..2b046ee96 100644
--- a/hivemind/p2p/p2p_daemon.py
+++ b/hivemind/p2p/p2p_daemon.py
@@ -131,7 +131,8 @@ async def create(
         :param quic: Deprecated, has no effect since libp2p 0.17.0
         :param use_relay_hop: Deprecated, has no effect since libp2p 0.17.0
         :param use_relay_discovery: Deprecated, has no effect since libp2p 0.17.0
-        :param check_if_identity_free: If enabled (default) and ``identity_path`` is provided,
+        :param check_if_identity_free: If enabled (default), ``identity_path`` is provided,
+                                       and we are connecting to an existing swarm,
                                        ensure that this identity is not used by other peers already.
                                        This slows down ``P2P.create()`` but protects from unintuitive libp2p errors
                                        appearing in case of the identity collision.
@@ -176,7 +177,7 @@ async def create(
 
         if identity_path is not None:
             if os.path.isfile(identity_path):
-                if check_if_identity_free:
+                if check_if_identity_free and need_bootstrap:
                     logger.info(f"Checking that identity from `{identity_path}` is not used by other peers")
                     if await cls.is_identity_taken(
                         identity_path,
diff --git a/tests/test_start_server.py b/tests/test_start_server.py
index 2a31d7c3b..b84dd5407 100644
--- a/tests/test_start_server.py
+++ b/tests/test_start_server.py
@@ -52,9 +52,6 @@ def test_cli_run_server_identity_path():
             encoding="utf-8",
         )
 
-        line = server_2_proc.stderr.readline()
-        assert re.search(r"Checking that identity.+is not used by other peers", line) is not None
-
         line = server_2_proc.stderr.readline()
         addrs_pattern_result = re.search(pattern, line)
         assert addrs_pattern_result is not None, line