瀏覽代碼

add create_persistent_stream method to p2pd

Denis Mazur 4 年之前
父節點
當前提交
c77ebd4846
共有 1 個文件被更改,包括 11 次插入0 次删除
  1. 11 0
      hivemind/p2p/p2p_daemon.py

+ 11 - 0
hivemind/p2p/p2p_daemon.py

@@ -144,6 +144,8 @@ class P2P:
 
         await self._ping_daemon_with_retries(ping_n_attempts, ping_delay)
 
+        self.persistent_streams = dict()
+
         return self
 
     async def _ping_daemon_with_retries(self, ping_n_attempts: int, ping_delay: float) -> None:
@@ -369,6 +371,15 @@ class P2P:
         finally:
             writer.close()
 
+    async def create_persistent_stream(self, peer_id: PeerID, handler_name: str):
+        if stream := self.persistent_streams.get((peer_id, handler_name)):
+            return stream
+
+        stream_info, reader, writer = await self._client.stream_open(peer_id, (handler_name,))
+        self.persistent_streams[(peer_id, handler_name)] = (stream_info, reader, writer)
+
+        return stream_info, reader, writer
+
     async def call_unary_handler(self, peer_id: PeerID, handler_name: str,
                                  request_protobuf: Any, response_proto_type: type) -> Any:
         stream_info, reader, writer = await self._client.stream_open(peer_id, (handler_name,))