Jelajahi Sumber

Implement DHT.p2p property

Aleksandr Borzunov 4 tahun lalu
induk
melakukan
e907eb25a2
1 mengubah file dengan 19 tambahan dan 1 penghapusan
  1. 19 1
      hivemind/dht/__init__.py

+ 19 - 1
hivemind/dht/__init__.py

@@ -23,7 +23,8 @@ from typing import Awaitable, Callable, Iterable, List, Optional, Sequence, Type
 
 from multiaddr import Multiaddr
 
-from hivemind.dht.node import DHTID, DHTNode
+from hivemind.dht.node import DHTNode
+from hivemind.p2p import P2P
 from hivemind.dht.routing import DHTKey, DHTValue, Subkey
 from hivemind.dht.validation import CompositeValidator, RecordValidatorBase
 from hivemind.utils import DHTExpiration, MPFuture, ValueWithExpiration, await_cancelled, get_logger, switch_to_uvloop
@@ -85,6 +86,7 @@ class DHT(mp.Process):
         self.shutdown_timeout = shutdown_timeout
         self.ready = mp.Event()
         self.daemon = daemon
+        self._p2p_replica = None
         if start:
             self.run_in_background(await_ready=True)
 
@@ -266,6 +268,22 @@ class DHT(mp.Process):
     async def _get_visible_maddrs(self, node: DHTNode, latest: bool = False) -> List[Multiaddr]:
         return await node.get_visible_maddrs(latest=latest)
 
+    @property
+    def p2p(self) -> P2P:
+        """
+        Get a replica of a P2P instance used in the DHT process internally.
+        """
+
+        if self._p2p_replica is not None:
+            return self._p2p_replica
+
+        daemon_listen_maddr = self.run_coroutine(DHT._get_p2p_daemon_listen_maddr)
+        self._p2p_replica = P2P.replicate(daemon_listen_maddr)
+        return self._p2p_replica
+
+    async def _get_p2p_daemon_listen_maddr(self, node: DHTNode) -> Multiaddr:
+        return node.p2p.daemon_listen_maddr
+
     def __del__(self):
         if self._parent_pid == os.getpid() and self.is_alive():
             self.shutdown()