added PingService

libp2p/host/ping.py

@@ -1,6 +1,7 @@
 import logging
-
+import time
 import trio
+import secrets
 
 from libp2p.network.stream.exceptions import (
     StreamClosed,
@@ -15,6 +16,11 @@
     TProtocol,
 )
 
+from libp2p.host.host_interface import (
+    IHost,
+)
+from typing import List
+
 ID = TProtocol("/ipfs/ping/1.0.0")
 PING_LENGTH = 32
 RESP_TIMEOUT = 60
@@ -65,7 +71,43 @@ async def handle_ping(stream: INetStream) -> None:
         try:
             should_continue = await _handle_ping(stream, peer_id)
             if not should_continue:
+                stream.close()
                 return
         except Exception:
             await stream.reset()
             return
+
+
+async def _ping(stream: INetStream) -> int:
+    """ 
+    Helper function to perform a single ping operation on a given stream,
+    returns integer value rtt - which denotes round trip time for a ping request in ms
+    """
+    ping_bytes = secrets.token_bytes(PING_LENGTH)
+    before = time.time()
+    await stream.write(ping_bytes)
+    pong_bytes = await stream.read(PING_LENGTH)
+    rtt = int((time.time() - before) * (10 ** 6)) 
+    if ping_bytes != pong_bytes:
+        logger.debug("invalid pong response")
+        raise
+    return rtt
+
+class PingService:
+    """ PingService executes pings and returns RTT in miliseconds. """
+
+    def __init__(self, host: IHost):
+        self._host = host
+
+    async def ping(self, peer_id: PeerID, ping_amt: int = 1) -> List[int]:
+        stream = await self._host.new_stream(peer_id, [ID])
+
+        try:
+            rtts = [
+                await _ping(stream) for _ in range(ping_amt)
+            ]
+            await stream.close()
+            return rtts
+        except Exception:
+            await stream.close()
+            raise