Implement connection in yet another way

Cargo.lock

@@ -192,6 +192,7 @@ dependencies = [
  "edit",
  "futures",
  "parking_lot",
+ "rand",
  "rusqlite",
  "serde",
  "serde_json",

cove-tui/Cargo.toml

@@ -12,6 +12,7 @@ directories = "4.0.1"
 edit = "0.1.4"
 futures = "0.3.21"
 parking_lot = "0.12.1"
+rand = "0.8.5"
 rusqlite = { version = "0.27.0", features = ["chrono"] }
 serde = { version = "1.0.137", features = ["derive"] }
 serde_json = "1.0.81"

cove-tui/src/euph.rs

@@ -1,5 +1,6 @@
 mod api;
-mod room;
+mod conn;
+// mod room;
 
 pub use api::{Message, SessionView, Snowflake, Time, UserId};
-pub use room::Room;
+// pub use room::Room;

cove-tui/src/euph/api/packet.rs

@@ -184,6 +184,14 @@ impl fmt::Display for PacketType {
     }
 }
 
+#[derive(Debug, thiserror::Error)]
+pub enum PacketError {
+    #[error("throttled: {0}")]
+    Throttled(String),
+    #[error("error: {0}")]
+    Error(String),
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Packet {
     pub id: Option<String>,
@@ -195,6 +203,21 @@ pub struct Packet {
     pub throttled_reason: Option<String>,
 }
 
+impl Packet {
+    pub fn data(self) -> Result<Value, PacketError> {
+        if self.throttled {
+            let reason = self
+                .throttled_reason
+                .unwrap_or_else(|| "no reason given".to_string());
+            Err(PacketError::Throttled(reason))
+        } else if let Some(error) = self.error {
+            Err(PacketError::Error(error))
+        } else {
+            Ok(self.data.unwrap_or_default())
+        }
+    }
+}
+
 pub trait HasPacketType {
     fn packet_type() -> PacketType;
 }
@@ -234,14 +257,10 @@ pub enum DecodeError {
         expected: PacketType,
         actual: PacketType,
     },
-    #[error("throttled: {0}")]
-    Throttled(String),
-    #[error("error: {0}")]
-    Error(String),
-    #[error("no data")]
-    NoData,
     #[error("{0}")]
     SerdeJson(#[from] serde_json::Error),
+    #[error("{0}")]
+    Packet(#[from] PacketError),
 }
 
 pub trait FromPacket: Sized {
@@ -255,15 +274,8 @@ impl<T: HasPacketType + DeserializeOwned> FromPacket for T {
                 expected: Self::packet_type(),
                 actual: packet.r#type,
             })
-        } else if packet.throttled {
-            let reason = packet
-                .throttled_reason
-                .unwrap_or_else(|| "no reason given".to_string());
-            Err(DecodeError::Throttled(reason))
-        } else if let Some(error) = packet.error {
-            Err(DecodeError::Error(error))
         } else {
-            let data = packet.data.unwrap_or_default();
+            let data = packet.data()?;
             Ok(serde_json::from_value(data)?)
         }
     }

cove-tui/src/euph/api/types.rs

@@ -10,7 +10,6 @@ use std::fmt;
 
 use chrono::{DateTime, Utc};
 use serde::{de, ser, Deserialize, Serialize};
-use serde_json::Value;
 
 /// Describes an account and its preferred name.
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -146,7 +145,7 @@ impl<'de> Deserialize<'de> for Snowflake {
 
 /// Time is specified as a signed 64-bit integer, giving the number of seconds
 /// since the Unix Epoch.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Time(#[serde(with = "chrono::serde::ts_seconds")] pub DateTime<Utc>);
 
 /// Identifies a user.
@@ -156,5 +155,5 @@ pub struct Time(#[serde(with = "chrono::serde::ts_seconds")] pub DateTime<Utc>);
 ///
 /// It is possible for this value to have no prefix and colon, and there is no
 /// fixed format for the unique value.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub struct UserId(pub String);

cove-tui/src/euph/conn.rs

@@ -0,0 +1,404 @@
+//! Connection state modeling.
+
+use std::collections::HashMap;
+use std::convert::Infallible;
+use std::time::Duration;
+
+use anyhow::bail;
+use chrono::Utc;
+use futures::channel::oneshot;
+use futures::stream::{SplitSink, SplitStream};
+use futures::{SinkExt, StreamExt};
+use rand::Rng;
+use tokio::net::TcpStream;
+use tokio::sync::mpsc;
+use tokio::{select, task, time};
+use tokio_tungstenite::{tungstenite, MaybeTlsStream, WebSocketStream};
+
+use crate::replies::{self, Replies};
+
+use super::api::{
+    BounceEvent, FromPacket, HelloEvent, JoinEvent, NetworkEvent, NickEvent, NickReply, Packet,
+    PacketType, PartEvent, PersonalAccountView, Ping, PingEvent, PingReply, SendEvent,
+    SnapshotEvent, ToPacket,
+};
+use super::{SessionView, Time, UserId};
+
+pub type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("connection closed")]
+    ConnectionClosed,
+    #[error("packet timed out")]
+    TimedOut,
+}
+
+#[derive(Debug)]
+enum Event {
+    Message(tungstenite::Message),
+    Send(Packet, oneshot::Sender<Result<Packet, Error>>),
+    Status(oneshot::Sender<Status>),
+    DoPings,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct Joining {
+    hello: Option<HelloEvent>,
+    snapshot: Option<SnapshotEvent>,
+    bounce: Option<BounceEvent>,
+}
+
+impl Joining {
+    fn on_packet(&mut self, packet: Packet) -> anyhow::Result<()> {
+        match packet.r#type {
+            PacketType::BounceEvent => self.bounce = Some(BounceEvent::from_packet(packet)?),
+            PacketType::HelloEvent => self.hello = Some(HelloEvent::from_packet(packet)?),
+            PacketType::SnapshotEvent => self.snapshot = Some(SnapshotEvent::from_packet(packet)?),
+            _ => {}
+        }
+        Ok(())
+    }
+
+    fn joined(&self) -> Option<Joined> {
+        if let (Some(hello), Some(snapshot)) = (&self.hello, &self.snapshot) {
+            let listing = snapshot
+                .listing
+                .iter()
+                .cloned()
+                .map(|s| (s.id.clone(), s))
+                .collect::<HashMap<_, _>>();
+            Some(Joined {
+                session: hello.session.clone(),
+                account: hello.account.clone(),
+                listing,
+            })
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Joined {
+    session: SessionView,
+    account: Option<PersonalAccountView>,
+    listing: HashMap<UserId, SessionView>,
+}
+
+impl Joined {
+    fn on_packet(&mut self, packet: Packet) -> anyhow::Result<()> {
+        match packet.r#type {
+            PacketType::JoinEvent => {
+                let packet = JoinEvent::from_packet(packet)?;
+                self.listing.insert(packet.0.id.clone(), packet.0);
+            }
+            PacketType::SendEvent => {
+                let packet = SendEvent::from_packet(packet)?;
+                self.listing
+                    .insert(packet.0.sender.id.clone(), packet.0.sender);
+            }
+            PacketType::PartEvent => {
+                let packet = PartEvent::from_packet(packet)?;
+                self.listing.remove(&packet.0.id);
+            }
+            PacketType::NetworkEvent => {
+                let p = NetworkEvent::from_packet(packet)?;
+                if p.r#type == "partition" {
+                    self.listing.retain(|_, s| {
+                        !(s.server_id == p.server_id && s.server_era == p.server_era)
+                    });
+                }
+            }
+            PacketType::NickEvent => {
+                let packet = NickEvent::from_packet(packet)?;
+                if let Some(session) = self.listing.get_mut(&packet.id) {
+                    session.name = packet.to;
+                }
+            }
+            PacketType::NickReply => {
+                // Since this is a reply, it may contain errors, for example if
+                // the user specified an invalid nick. We can't just die if that
+                // happens, so we ignore the error case.
+                if let Ok(packet) = NickReply::from_packet(packet) {
+                    assert_eq!(self.session.id, packet.id);
+                    self.session.name = packet.to;
+                }
+            }
+            // The who reply is broken and can't be trusted right now, so we'll
+            // not even look at it.
+            _ => {}
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum Status {
+    Joining(Joining),
+    Joined(Joined),
+}
+
+struct State {
+    ws_tx: SplitSink<WsStream, tungstenite::Message>,
+    last_id: usize,
+    replies: Replies<String, Packet>,
+
+    packet_tx: mpsc::UnboundedSender<Packet>,
+
+    last_ws_ping: Option<Vec<u8>>,
+    last_ws_pong: Option<Vec<u8>>,
+    last_euph_ping: Option<Time>,
+    last_euph_pong: Option<Time>,
+
+    status: Status,
+}
+
+impl State {
+    async fn run(
+        ws: WsStream,
+        tx_canary: oneshot::Receiver<Infallible>,
+        rx_canary: oneshot::Receiver<Infallible>,
+        event_tx: mpsc::UnboundedSender<Event>,
+        mut event_rx: mpsc::UnboundedReceiver<Event>,
+        packet_tx: mpsc::UnboundedSender<Packet>,
+    ) {
+        let (ws_tx, mut ws_rx) = ws.split();
+        let state = Self {
+            ws_tx,
+            last_id: 0,
+            replies: Replies::new(Duration::from_secs(10)), // TODO Make configurable
+            packet_tx,
+            last_ws_ping: None,
+            last_ws_pong: None,
+            last_euph_ping: None,
+            last_euph_pong: None,
+            status: Status::Joining(Joining::default()),
+        };
+
+        select! {
+            _ = Self::listen(&mut ws_rx, &event_tx) => (),
+            _ = Self::send_ping_events(&event_tx) => (),
+            _ = state.handle_events(&event_tx, &mut event_rx) => (),
+        }
+    }
+
+    async fn listen(
+        ws_rx: &mut SplitStream<WsStream>,
+        event_tx: &mpsc::UnboundedSender<Event>,
+    ) -> anyhow::Result<()> {
+        while let Some(msg) = ws_rx.next().await {
+            event_tx.send(Event::Message(msg?))?;
+        }
+        Ok(())
+    }
+
+    async fn send_ping_events(event_tx: &mpsc::UnboundedSender<Event>) -> anyhow::Result<()> {
+        loop {
+            event_tx.send(Event::DoPings)?;
+            time::sleep(Duration::from_secs(10)).await; // TODO Make configurable
+        }
+    }
+
+    async fn handle_events(
+        mut self,
+        event_tx: &mpsc::UnboundedSender<Event>,
+        event_rx: &mut mpsc::UnboundedReceiver<Event>,
+    ) -> anyhow::Result<()> {
+        while let Some(ev) = event_rx.recv().await {
+            match ev {
+                Event::Message(msg) => self.on_msg(msg, event_tx)?,
+                Event::Send(packet, reply_tx) => self.on_send(packet, reply_tx).await?,
+                Event::Status(reply_tx) => self.on_status(reply_tx),
+                Event::DoPings => self.do_pings(event_tx).await?,
+            }
+        }
+        Ok(())
+    }
+
+    fn on_msg(
+        &mut self,
+        msg: tungstenite::Message,
+        event_tx: &mpsc::UnboundedSender<Event>,
+    ) -> anyhow::Result<()> {
+        match msg {
+            tungstenite::Message::Text(t) => self.on_packet(serde_json::from_str(&t)?, event_tx)?,
+            tungstenite::Message::Binary(_) => bail!("unexpected binary message"),
+            tungstenite::Message::Ping(_) => {}
+            tungstenite::Message::Pong(p) => self.last_ws_pong = Some(p),
+            tungstenite::Message::Close(_) => {}
+            tungstenite::Message::Frame(_) => {}
+        }
+        Ok(())
+    }
+
+    fn on_packet(
+        &mut self,
+        packet: Packet,
+        event_tx: &mpsc::UnboundedSender<Event>,
+    ) -> anyhow::Result<()> {
+        if packet.r#type == PacketType::PingReply {
+            let packet = PingReply::from_packet(packet.clone())?;
+            self.last_euph_pong = packet.time;
+        } else if packet.r#type == PacketType::PingEvent {
+            let time = Some(PingEvent::from_packet(packet.clone())?.time);
+            Self::send_unconditionally(event_tx, PingReply { time }, packet.id.clone())?;
+        }
+
+        if let Some(id) = &packet.id {
+            self.replies.complete(id, packet.clone());
+        }
+
+        self.packet_tx.send(packet.clone())?;
+
+        // TODO Handle disconnect event?
+
+        match &mut self.status {
+            Status::Joining(joining) => {
+                joining.on_packet(packet)?;
+                if let Some(joined) = joining.joined() {
+                    self.status = Status::Joined(joined);
+                }
+            }
+            Status::Joined(joined) => joined.on_packet(packet)?,
+        }
+
+        Ok(())
+    }
+
+    async fn on_send(
+        &mut self,
+        mut packet: Packet,
+        reply_tx: oneshot::Sender<Result<Packet, Error>>,
+    ) -> anyhow::Result<()> {
+        let id = if let Some(id) = packet.id.clone() {
+            id
+        } else {
+            // Overkill of universe-heat-death-like proportions
+            self.last_id = self.last_id.wrapping_add(1);
+            format!("{}", self.last_id)
+        };
+        packet.id = Some(id.clone());
+
+        let pending_reply = self.replies.wait_for(id);
+
+        let msg = tungstenite::Message::Text(serde_json::to_string(&packet)?);
+        self.ws_tx.send(msg).await?;
+
+        let reply = match pending_reply.get().await {
+            Ok(reply) => Ok(reply),
+            Err(replies::Error::TimedOut) => Err(Error::TimedOut),
+            // We could also send an Error::ConnectionClosed here, but that
+            // happens automatically in the send function once we drop reply_tx.
+            Err(replies::Error::Canceled) => return Ok(()),
+        };
+        let _ = reply_tx.send(reply);
+
+        Ok(())
+    }
+
+    fn on_status(&mut self, reply_tx: oneshot::Sender<Status>) {
+        let _ = reply_tx.send(self.status.clone());
+    }
+
+    async fn do_pings(&mut self, event_tx: &mpsc::UnboundedSender<Event>) -> anyhow::Result<()> {
+        // Check old ws ping
+        if self.last_ws_ping.is_some() && self.last_ws_ping != self.last_ws_pong {
+            bail!("server missed ws ping")
+        }
+
+        // Send new ws ping
+        let mut ws_payload = [0_u8; 8];
+        rand::thread_rng().fill(&mut ws_payload);
+        self.ws_tx
+            .send(tungstenite::Message::Ping(ws_payload.to_vec()))
+            .await?;
+
+        // Check old euph ping
+        if self.last_euph_ping.is_some() && self.last_euph_ping != self.last_euph_pong {
+            bail!("server missed euph ping")
+        }
+
+        // Send new euph ping
+        let euph_payload = Time(Utc::now());
+        Self::send_unconditionally(event_tx, Ping { time: euph_payload }, None)?;
+
+        Ok(())
+    }
+
+    fn send_unconditionally<T: ToPacket>(
+        event_tx: &mpsc::UnboundedSender<Event>,
+        packet: T,
+        id: Option<String>,
+    ) -> anyhow::Result<()> {
+        let (tx, _) = oneshot::channel();
+        event_tx.send(Event::Send(packet.to_packet(id), tx))?;
+        Ok(())
+    }
+}
+
+pub struct ConnTx {
+    canary: oneshot::Sender<Infallible>,
+    event_tx: mpsc::UnboundedSender<Event>,
+}
+
+impl ConnTx {
+    pub async fn send<T: ToPacket>(&self, packet: T) -> Result<Packet, Error> {
+        let (tx, rx) = oneshot::channel();
+        let event = Event::Send(packet.to_packet(None), tx);
+        self.event_tx
+            .send(event)
+            .map_err(|_| Error::ConnectionClosed)?;
+        match rx.await {
+            Ok(result) => result,
+            Err(_) => Err(Error::ConnectionClosed),
+        }
+    }
+
+    pub async fn status(&self) -> Result<Status, Error> {
+        let (tx, rx) = oneshot::channel();
+        self.event_tx
+            .send(Event::Status(tx))
+            .map_err(|_| Error::ConnectionClosed)?;
+        rx.await.map_err(|_| Error::ConnectionClosed)
+    }
+}
+
+pub struct ConnRx {
+    canary: oneshot::Sender<Infallible>,
+    packet_rx: mpsc::UnboundedReceiver<Packet>,
+}
+
+impl ConnRx {
+    pub async fn recv(&mut self) -> Result<Packet, Error> {
+        self.packet_rx.recv().await.ok_or(Error::ConnectionClosed)
+    }
+}
+
+// TODO Combine ConnTx and ConnRx and implement Stream + Sink?
+
+pub fn wrap(ws: WsStream) -> (ConnTx, ConnRx) {
+    let (tx_canary_tx, tx_canary_rx) = oneshot::channel();
+    let (rx_canary_tx, rx_canary_rx) = oneshot::channel();
+    let (event_tx, event_rx) = mpsc::unbounded_channel();
+    let (packet_tx, packet_rx) = mpsc::unbounded_channel();
+
+    task::spawn(State::run(
+        ws,
+        tx_canary_rx,
+        rx_canary_rx,
+        event_tx.clone(),
+        event_rx,
+        packet_tx,
+    ));
+
+    let tx = ConnTx {
+        canary: tx_canary_tx,
+        event_tx,
+    };
+    let rx = ConnRx {
+        canary: rx_canary_tx,
+        packet_rx,
+    };
+    (tx, rx)
+}

cove-tui/src/euph/room.rs

@@ -3,93 +3,84 @@ use std::time::Duration;
 
 use futures::stream::{SplitSink, SplitStream};
 use futures::StreamExt;
-use tokio::net::TcpStream;
 use tokio::sync::{mpsc, oneshot};
 use tokio::{select, task, time};
-use tokio_tungstenite::{tungstenite, MaybeTlsStream, WebSocketStream};
+use tokio_tungstenite::tungstenite;
 
-type WsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
+use super::conn::{State, Status, WsStream};
 
 #[derive(Debug)]
 enum Event {
     Connected(SplitSink<WsStream, tungstenite::Message>),
     Disconnected,
-    Message(tungstenite::Message),
+    WsMessage(tungstenite::Message),
-    Ping,
+    DoPings,
+    GetStatus(oneshot::Sender<Option<Status>>),
 }
 
-#[derive(Debug)]
+async fn run(
-struct Connected {
+    canary: oneshot::Receiver<Infallible>,
-    tx: SplitSink<WsStream, tungstenite::Message>,
+    tx: mpsc::UnboundedSender<Event>,
+    rx: mpsc::UnboundedReceiver<Event>,
+    url: String,
+) {
+    let state = State::default();
+    select! {
+        _ = canary => (),
+        _ = respond_to_events(state, rx) => (),
+        _ = maintain_connection(tx, url) => (),
+    }
 }
 
-#[derive(Debug)]
+async fn respond_to_events(
-enum State {
+    mut state: State,
-    Connecting,
+    mut rx: mpsc::UnboundedReceiver<Event>,
-    Connected(Connected),
+) -> anyhow::Result<()> {
+    while let Some(event) = rx.recv().await {
+        match event {
+            Event::Connected(tx) => state.on_connected(tx),
+            Event::Disconnected => state.on_disconnected(),
+            Event::WsMessage(msg) => state.on_ws_message(msg)?,
+            Event::DoPings => state.on_do_pings()?,
+            Event::GetStatus(tx) => {
+                let _ = tx.send(state.status());
+            }
+        }
+    }
+    Ok(())
 }
 
-impl State {
+async fn maintain_connection(tx: mpsc::UnboundedSender<Event>, url: String) -> anyhow::Result<()> {
-    async fn run(
+    loop {
-        self,
+        // TODO Cookies
-        canary: oneshot::Receiver<Infallible>,
+        let (ws, _) = tokio_tungstenite::connect_async(&url).await?;
-        tx: mpsc::UnboundedSender<Event>,
+        let (ws_tx, ws_rx) = ws.split();
-        rx: mpsc::UnboundedReceiver<Event>,
+        tx.send(Event::Connected(ws_tx))?;
-        url: String,
-    ) {
         select! {
-            _ = canary => (),
+            _ = receive_messages(&tx, ws_rx) => (),
-            _ = Self::maintain_connection(tx, url) => (),
+            _ = prompt_pings(&tx) => ()
-            _ = self.respond_to_events(rx) => (),
         }
+        tx.send(Event::Disconnected)?;
+        // TODO Make reconnect delay configurable
+        time::sleep(Duration::from_secs(5)).await;
     }
+}
 
-    async fn maintain_connection(
+async fn receive_messages(
-        tx: mpsc::UnboundedSender<Event>,
+    tx: &mpsc::UnboundedSender<Event>,
-        url: String,
+    mut rx: SplitStream<WsStream>,
-    ) -> anyhow::Result<()> {
+) -> anyhow::Result<()> {
-        loop {
+    while let Some(msg) = rx.next().await {
-            // TODO Cookies
+        tx.send(Event::WsMessage(msg?))?;
-            let (ws, _) = tokio_tungstenite::connect_async(&url).await?;
-            let (ws_tx, ws_rx) = ws.split();
-            tx.send(Event::Connected(ws_tx))?;
-            select! {
-                _ = Self::receive_messages(&tx, ws_rx) => (),
-                _ = Self::prompt_pings(&tx) => ()
-            }
-            tx.send(Event::Disconnected)?;
-            // TODO Make reconnect delay configurable
-            time::sleep(Duration::from_secs(5)).await;
-        }
     }
+    Ok(())
+}
 
-    async fn receive_messages(
+async fn prompt_pings(tx: &mpsc::UnboundedSender<Event>) -> anyhow::Result<()> {
-        tx: &mpsc::UnboundedSender<Event>,
+    loop {
-        mut rx: SplitStream<WsStream>,
+        // TODO Make ping delay configurable
-    ) -> anyhow::Result<()> {
+        time::sleep(Duration::from_secs(10)).await;
-        while let Some(msg) = rx.next().await {
+        tx.send(Event::DoPings)?;
-            tx.send(Event::Message(msg?))?;
-        }
-        Ok(())
-    }
-
-    async fn prompt_pings(tx: &mpsc::UnboundedSender<Event>) -> anyhow::Result<()> {
-        loop {
-            // TODO Make ping delay configurable
-            time::sleep(Duration::from_secs(10)).await;
-            tx.send(Event::Ping)?;
-        }
-    }
-
-    async fn respond_to_events(mut self, mut rx: mpsc::UnboundedReceiver<Event>) {
-        while let Some(event) = rx.recv().await {
-            match event {
-                Event::Connected(tx) => self = State::Connected(Connected { tx }),
-                Event::Disconnected => self = State::Connecting,
-                Event::Message(_) => todo!(),
-                Event::Ping => todo!(),
-            }
-        }
     }
 }
 
@@ -103,11 +94,17 @@ impl Room {
         let (event_tx, event_rx) = mpsc::unbounded_channel();
         let (canary_tx, canary_rx) = oneshot::channel();
 
-        task::spawn(State::Connecting.run(canary_rx, event_tx.clone(), event_rx, url));
+        task::spawn(run(canary_rx, event_tx.clone(), event_rx, url));
 
         Self {
             canary: canary_tx,
             tx: event_tx,
         }
     }
+
+    pub async fn status(&self) -> anyhow::Result<Option<Status>> {
+        let (tx, rx) = oneshot::channel();
+        self.tx.send(Event::GetStatus(tx))?;
+        Ok(rx.await?)
+    }
 }

cove-tui/src/main.rs

@@ -3,6 +3,7 @@
 mod chat;
 mod euph;
 mod log;
+mod replies;
 mod store;
 mod ui;
 mod vault;

cove-tui/src/replies.rs

@@ -0,0 +1,70 @@
+use std::collections::HashMap;
+use std::hash::Hash;
+use std::result;
+use std::time::Duration;
+
+use tokio::sync::oneshot::{self, Receiver, Sender};
+use tokio::time;
+
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("timed out")]
+    TimedOut,
+    #[error("canceled")]
+    Canceled,
+}
+
+pub type Result<T> = result::Result<T, Error>;
+
+pub struct PendingReply<R> {
+    timeout: Duration,
+    result: Receiver<R>,
+}
+
+impl<R> PendingReply<R> {
+    pub async fn get(self) -> Result<R> {
+        let result = time::timeout(self.timeout, self.result).await;
+        match result {
+            Err(_) => Err(Error::TimedOut),
+            Ok(Err(_)) => Err(Error::Canceled),
+            Ok(Ok(value)) => Ok(value),
+        }
+    }
+}
+
+pub struct Replies<I, R> {
+    timeout: Duration,
+    pending: HashMap<I, Sender<R>>,
+}
+
+impl<I: Eq + Hash, R> Replies<I, R> {
+    pub fn new(timeout: Duration) -> Self {
+        Self {
+            timeout,
+            pending: HashMap::new(),
+        }
+    }
+
+    pub fn wait_for(&mut self, id: I) -> PendingReply<R> {
+        let (tx, rx) = oneshot::channel();
+        self.pending.insert(id, tx);
+        PendingReply {
+            timeout: self.timeout,
+            result: rx,
+        }
+    }
+
+    pub fn complete(&mut self, id: &I, result: R) {
+        if let Some(tx) = self.pending.remove(id) {
+            let _ = tx.send(result);
+        }
+    }
+
+    pub fn cancel(&mut self, id: &I) {
+        self.pending.remove(id);
+    }
+
+    pub fn purge(&mut self) {
+        self.pending.retain(|_, tx| !tx.is_closed());
+    }
+}