diff --git a/irpc-iroh/examples/bidi.rs b/irpc-iroh/examples/bidi.rs new file mode 100644 index 0000000..87c18df --- /dev/null +++ b/irpc-iroh/examples/bidi.rs @@ -0,0 +1,317 @@ +//! Bidirectional RPC over a single iroh connection. +//! +//! In the usual irpc setup one peer is the client (it opens streams and sends requests) and the +//! other is the server (it accepts streams and answers them). A QUIC connection is symmetric, +//! though: either peer can open a bidirectional stream at any time. This example uses that to let +//! *both* peers act as client and server at once, over one shared connection. +//! +//! Two independent protocols run in opposite directions: +//! +//! - [`proto::ClientToServer`]: the connecting peer asks the listening peer to `Get`/`Set` entries +//! in a key-value store. +//! - [`proto::ServerToClient`]: the listening peer periodically `Ping`s each connected peer to +//! check liveness. +//! +//! The trick is that each peer holds *both* an [`irpc::Client`] (to open streams and send its own +//! requests) *and* a [`irpc_iroh::read_request`] loop (to accept and answer the other peer's +//! requests). Both are built from the same [`iroh::endpoint::Connection`]. Because `accept_bi` +//! only ever yields streams opened by the *remote* side, the two directions never collide: the +//! listener's accept loop only sees `ClientToServer` streams, the connector's accept loop only +//! sees `ServerToClient` streams. +//! +//! Run the listener, note its endpoint id, then drive it from a second process: +//! +//! ```text +//! cargo run --example bidi -- listen +//! cargo run --example bidi -- connect set foo bar +//! cargo run --example bidi -- connect get foo +//! ``` + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt::init(); + cli::run().await +} + +/// The protocol definitions and the peer implementations for both roles. +mod proto { + use std::{ + collections::{BTreeMap, HashMap}, + sync::{Arc, Mutex}, + time::Duration, + }; + + use anyhow::Result; + use iroh::{ + Endpoint, EndpointId, + endpoint::{Connection, presets}, + protocol::{AcceptError, ProtocolHandler, Router}, + }; + use irpc::{Client, WithChannels, channel::oneshot, rpc_requests}; + use irpc_iroh::{IrohRemoteConnection, read_request}; + use serde::{Deserialize, Serialize}; + use tokio::time::Instant; + + /// ALPN identifying this example protocol on the iroh endpoint. + /// + /// Both peers must agree on it: the listener registers it with the [`Router`], the connector + /// passes it to `connect`. The trailing `/1` leaves room to bump the version later. + const ALPN: &[u8] = b"iroh-irpc/example-bidi/1"; + + /// Requests the connecting peer sends to the listening peer. + /// + /// These operate on the listener's key-value store. Each variant carries a `tx` oneshot channel + /// on which the handler sends its single response. + #[rpc_requests(message = ClientToServerMsg)] + #[derive(Debug, Serialize, Deserialize)] + pub enum ClientToServer { + /// Looks up a key, responding with the current value or `None` if absent. + #[rpc(tx=oneshot::Sender>)] + #[wrap(GetRequest, derive(Clone))] + Get(String), + + /// Stores a value under a key, responding with the previous value if there was one. + #[rpc(tx=oneshot::Sender>)] + #[wrap(SetRequest)] + Set { + /// The key to store under. + key: String, + /// The value to store. + value: String, + }, + } + + /// Requests the listening peer sends back to a connected peer. + /// + /// This is the reverse direction: the listener drives it from its own [`irpc::Client`] built on + /// top of the accepted connection. Here it is just a liveness check. + #[rpc_requests(message = ServerToClientMsg)] + #[derive(Debug, Serialize, Deserialize)] + pub enum ServerToClient { + /// Asks the peer to reply so the listener can confirm the connection is still alive. + #[rpc(tx=oneshot::Sender<()>)] + #[wrap(PingRequest)] + Ping {}, + } + + /// Spawn the listening peer. + /// + /// Binds an endpoint, registers the [`Server`] as the [`ProtocolHandler`] for [`ALPN`], and + /// spawns the background [`Server::ping_loop`] that pings connected peers. + pub async fn listen() -> Result { + let endpoint = Endpoint::bind(presets::N0).await?; + let server = Server::default(); + + // The router drives `Server::accept` for every incoming connection on this ALPN. + let router = Router::builder(endpoint.clone()) + .accept(ALPN, server.clone()) + .spawn(); + // Printed so the connecting side can be pointed at this peer. + println!("endpoint id: {}", router.endpoint().id()); + + // Spawn a loop that pings all clients once per second. + // The task is terminted once the endpoint closes through `Endpoint::closed`. + tokio::spawn(async move { endpoint.closed().run_until(server.ping_loop()).await }); + + Ok(router) + } + + /// The listening state: a shared key-value store plus a registry of reverse clients. + #[derive(Debug, Clone, Default)] + pub struct Server { + /// The key-value store that [`ClientToServer`] requests read and write. + state: Arc>>, + /// A reverse [`irpc::Client`] per connected peer, used to send [`ServerToClient`] pings. + clients: Arc>>>, + } + + impl ProtocolHandler for Server { + /// Serves one accepted connection for its whole lifetime. + /// + /// This wires up both roles on the listener side: it first builds a reverse client over the + /// same connection so the ping loop can call *back* to this peer, then runs the accept loop + /// that answers this peer's requests. + async fn accept(&self, conn: Connection) -> Result<(), AcceptError> { + // Build a client on top of the incoming connection and register it under the peer's id. + // The ping loop opens `ServerToClient` streams through this, in the reverse direction + // to the `ClientToServer` streams we accept below. + let client = Client::boxed(IrohRemoteConnection::new(conn.clone())); + self.clients + .lock() + .unwrap() + .insert(conn.remote_id(), client); + // Accept and handle this peer's requests one at a time until it closes the connection. + // `read_request` returns `None` on a clean shutdown, ending the loop. + while let Some(msg) = read_request::(&conn).await? { + self.handle_message(msg).await; + } + // The request loop ended; wait for the connection to fully close before returning. The + // reverse client stays registered until a ping fails and the ping loop evicts it. + conn.closed().await; + Ok(()) + } + } + + impl Server { + /// Handles a single decoded [`ClientToServer`] request against the store. + /// + /// Each arm destructures [`WithChannels`] to split the request payload (`inner`) from its + /// response channel (`tx`), then sends the reply. `tx.send(..).await.ok()` ignores a send + /// error, which happens only if the requesting peer has already disconnected. + async fn handle_message(&self, msg: ClientToServerMsg) { + match msg { + ClientToServerMsg::Get(msg) => { + let WithChannels { inner, tx, .. } = msg; + println!("handle request: {inner:?}"); + let GetRequest(key) = inner; + let value = self.state.lock().unwrap().get(&key).cloned(); + tx.send(value).await.ok(); + } + ClientToServerMsg::Set(msg) => { + let WithChannels { inner, tx, .. } = msg; + println!("handle request: {inner:?}"); + let SetRequest { key, value } = inner; + let prev_value = self.state.lock().unwrap().insert(key, value); + tx.send(prev_value).await.ok(); + } + } + } + + /// Pings every connected peer once a second, evicting any whose ping fails. + /// + /// This is the listener acting as a client: it calls out over each reverse client stored in + /// [`Server::clients`]. Each ping runs in its own task so one slow or dead peer does not + /// hold up the others. A failed ping means the connection is gone, so its entry is removed. + async fn ping_loop(self) { + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + // Hold the lock only for a synchronous loop: each iteration clones out what a + // task needs and spawns it. The actual ping awaits happen inside those tasks, which + // own their clones, so the std mutex is never held across an await. + for (&remote_id, client) in self.clients.lock().unwrap().iter() { + let client = client.clone(); + let clients = self.clients.clone(); + tokio::spawn(async move { + // Send a `PingRequest` in to the client and await its reply. + let short_id = remote_id.fmt_short(); + let now = Instant::now(); + match client.rpc(PingRequest {}).await { + Ok(()) => println!("ping {short_id}: OK ({:?})", now.elapsed()), + Err(err) => { + // The ping failed. Remove the peer from our map of client connections. + println!("ping {short_id}: FAIL {err:#} ({:?})", now.elapsed()); + clients.lock().unwrap().remove(&remote_id); + } + } + }); + } + } + } + } + + /// Connects to a listening peer and sets up both directions. + /// + /// Returns the endpoint (kept alive by the caller) and an [`irpc::Client`] for sending + /// [`ClientToServer`] requests. It also spawns a background accept loop that answers the + /// other peers's [`ServerToClient`] requests. + pub async fn connect(endpoint_id: EndpointId) -> Result<(Endpoint, Client)> { + println!("connecting to {endpoint_id}"); + let endpoint = Endpoint::bind(presets::N0).await?; + let conn = endpoint.connect(endpoint_id, ALPN).await?; + // The outgoing client with which we send `ClientToServerRequests` to the peer. + let client = Client::boxed(IrohRemoteConnection::new(conn.clone())); + // Spawn a task that reads `ServerToClient` requests from the peer, on the same connection we + // use for outgoing requests. The task terminates once the connection closes. + tokio::spawn(async move { + if let Err(err) = pong_loop(conn).await { + println!("pong loop failed: {err:#}"); + } + }); + Ok((endpoint, client)) + } + + /// Accepts RPC requests from the server and answers them. + async fn pong_loop(conn: Connection) -> Result<()> { + // Wait for the next server-initiated request. `Ok(None)` is a clean close; an error + // means the connection dropped. Either way, stop serving. + while let Some(msg) = read_request::(&conn).await? { + match msg { + ServerToClientMsg::Ping(msg) => { + // Answer the ping by sending a unit on its response channel. + let WithChannels { tx, .. } = msg; + println!("Received ping from server, sending pong"); + tx.send(()).await?; + } + } + } + println!("ping loop closed: remote closed the connection"); + anyhow::Ok(()) + } +} + +/// Command-line interface: selects the listener or connector role and runs it. +mod cli { + use anyhow::Result; + use clap::Parser; + use iroh::EndpointId; + + use crate::proto::{GetRequest, SetRequest, connect, listen}; + + /// Runs the example as either a listening or a connecting peer. + #[derive(Debug, Parser)] + enum Cli { + /// Run the listening peer and print its endpoint id. + Listen, + /// Connect to a listening peer and issue one request. + Connect { + endpoint_id: EndpointId, + #[clap(subcommand)] + command: Command, + }, + } + + /// The single request a connector issues after connecting. + #[derive(Debug, Parser)] + enum Command { + Get { key: String }, + Set { key: String, value: String }, + } + + /// Parses arguments and runs the chosen role to completion. + pub async fn run() -> Result<()> { + match Cli::parse() { + Cli::Listen => { + let router = listen().await?; + println!("waiting for ctrl-c"); + tokio::signal::ctrl_c().await.ok(); + router.shutdown().await?; + } + Cli::Connect { + endpoint_id, + command, + } => { + let (endpoint, client) = connect(endpoint_id).await?; + // A single request/response RPC. `client.rpc` opens a bi stream, sends the request, + // and awaits the one response the listener sends on the reply channel. + match command { + Command::Get { key } => { + println!("get '{key}'"); + let value = client.rpc(GetRequest(key)).await?; + println!("{value:?}"); + } + Command::Set { key, value } => { + println!("set '{key}' to '{value}'"); + let value = client.rpc(SetRequest { key, value }).await?; + println!("OK (previous: {value:?})"); + } + } + // Stay connected so the spawned accept loop keeps answering the listener's pings; + // this is what makes the ping output visible on both sides until ctrl-c. + println!("waiting for ctrl-c"); + tokio::signal::ctrl_c().await.ok(); + endpoint.close().await; + } + } + Ok(()) + } +}