asimov_protocol/
ping_protocol.rs

1// This is free and unencumbered software released into the public domain.
2
3//! The internode ping protocol.
4
5use alloc::{boxed::Box, sync::Arc};
6use core::{error::Error, result::Result};
7use iroh::{
8    Endpoint, EndpointAddr,
9    endpoint::Connection,
10    protocol::{AcceptError, ProtocolHandler},
11};
12use iroh_metrics::{Counter, MetricsGroup};
13use tokio::time::{Duration, Instant};
14
15/// The ALPN string for the ping protocol.
16pub const PING_ALPN: &[u8] = b"asimov/ping";
17
18/// Metrics for the ping protocol.
19#[derive(Debug, Default, MetricsGroup)]
20#[metrics(name = "ping")]
21pub struct PingMetrics {
22    /// The count of valid ping messages sent.
23    pub pings_sent: Counter,
24
25    /// The count of valid ping messages received.
26    pub pings_recv: Counter,
27}
28
29/// The ping protocol for use with `iroh::protocol::Router`.
30#[derive(Debug, Clone)]
31pub struct PingProtocol {
32    /// Shared state for use across incoming connections.
33    metrics: Arc<PingMetrics>,
34}
35
36impl Default for PingProtocol {
37    fn default() -> Self {
38        Self::new()
39    }
40}
41
42impl PingProtocol {
43    /// Creates a new ping protocol state.
44    pub fn new() -> Self {
45        Self {
46            metrics: Arc::new(PingMetrics::default()),
47        }
48    }
49
50    /// Returns a handle to the ping metrics.
51    pub fn metrics(&self) -> &Arc<PingMetrics> {
52        &self.metrics
53    }
54
55    /// Sends a ping on the provided endpoint to a given node address.
56    pub async fn ping(
57        &self,
58        endpoint: &Endpoint,
59        addr: EndpointAddr,
60    ) -> Result<Duration, Box<dyn Error>> {
61        // Open a connection to the accepting node:
62        let conn = endpoint.connect(addr, PING_ALPN).await?;
63
64        // Open a bidirectional QUIC stream:
65        let (mut send, mut recv) = conn.open_bi().await?;
66
67        // Begin measuring elapsed time:
68        let start = Instant::now();
69
70        // Send the PING request and finish the send stream:
71        send.write_all(b"PING").await?;
72        send.finish()?;
73
74        // Read the PONG response:
75        let response = recv.read_to_end(4).await?;
76        assert_eq!(&response, b"PONG");
77
78        // Measure the duration of this interaction:
79        let duration = start.elapsed();
80
81        // Update the metrics counters:
82        self.metrics.pings_sent.inc();
83
84        // Close the connection explicitly and gracefully:
85        conn.close(0u32.into(), b"BYE!");
86
87        Ok(duration)
88    }
89}
90
91impl ProtocolHandler for PingProtocol {
92    /// Each incoming connection for our ALPN results in a call to `accept`.
93    ///
94    /// The returned future runs on a newly spawned Tokio task, so it can run
95    /// indefinitely as long as the connection remains open.
96    async fn accept(&self, connection: Connection) -> n0_error::Result<(), AcceptError> {
97        let node_id = connection.remote_id();
98        #[cfg(feature = "std")]
99        std::eprintln!("Accepted a connection from node {node_id}"); // DEBUG
100
101        // Expect the connecting peer to open a bidirectional QUIC stream:
102        let (mut send, mut recv) = connection.accept_bi().await?;
103
104        // Read the PING request:
105        let req = recv.read_to_end(4).await.map_err(AcceptError::from_err)?;
106        assert_eq!(&req, b"PING");
107
108        // Update the metrics counters:
109        self.metrics.pings_recv.inc();
110
111        // Send the PONG response and finish the send stream:
112        send.write_all(b"PONG")
113            .await
114            .map_err(AcceptError::from_err)?;
115        send.finish()?;
116
117        // Wait for the remote end to explicitly and gracefully close the
118        // connection after receiving our response:
119        connection.closed().await;
120
121        Ok(())
122    }
123}
asimov_protocol/ping_protocol.rs

asimov_protocol/
ping_protocol.rs
