Refactor the kraken module to automatically re-connect on errors

In order to be able to re-connect on certain errors, we model
connection errors separately from parsing errors. We also change
the API of the whole module to no longer forward all errors to
the subscribers but instead, only update the subscribers with
either a latest rate or a permanent failure in case we exhausted
all our options to re-connect the websocket.

To model all of this properly, we introduce to sub-modules so that
each submodule can have their own `Error` type.

Resolves #297.

swap/src/bin/asb.rs

@@ -92,7 +92,7 @@ async fn main() -> Result<()> {
                 bitcoin_wallet.new_address().await?
             );
 
-            let kraken_rate_updates = kraken::connect().await?;
+            let kraken_rate_updates = kraken::connect()?;
 
             let (event_loop, _) = EventLoop::new(
                 config.network.listen,

swap/src/bin/kraken_ticker.rs

@@ -3,12 +3,10 @@ use anyhow::{Context, Result};
 #[tokio::main]
 async fn main() -> Result<()> {
     tracing::subscriber::set_global_default(
-        tracing_subscriber::fmt().with_env_filter("trace").finish(),
+        tracing_subscriber::fmt().with_env_filter("debug").finish(),
     )?;
 
-    let mut ticker = swap::kraken::connect()
+    let mut ticker = swap::kraken::connect().context("Failed to connect to kraken")?;
-        .await
-        .context("Failed to connect to kraken")?;
 
     loop {
         match ticker.wait_for_update().await? {

swap/src/kraken.rs

@@ -1,95 +1,57 @@
 use crate::asb::Rate;
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
-use bitcoin::util::amount::ParseAmountError;
+use futures::{SinkExt, StreamExt, TryStreamExt};
-use futures::{SinkExt, StreamExt};
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
+use std::convert::{Infallible, TryFrom};
-use std::convert::TryFrom;
+use std::sync::Arc;
+use std::time::Duration;
 use tokio::sync::watch;
-use tokio_tungstenite::tungstenite;
 
-type RateUpdate = Result<Rate, Error>;
+/// Connect to Kraken websocket API for a constant stream of rate updates.
-
+///
-pub async fn connect() -> Result<RateUpdateStream> {
+/// If the connection fails, it will automatically be re-established.
-    let (rate_update, rate_update_receiver) = watch::channel(Err(Error::NotYetRetrieved));
+pub fn connect() -> Result<RateUpdateStream> {
-
+    let (rate_update, rate_update_receiver) = watch::channel(Err(Error::NotYetAvailable));
-    let (rate_stream, _) = tokio_tungstenite::connect_async("wss://ws.kraken.com")
+    let rate_update = Arc::new(rate_update);
-        .await
-        .context("Failed to connect to Kraken websocket API")?;
-
-    let (mut rate_stream_sink, mut rate_stream) = rate_stream.split();
 
     tokio::spawn(async move {
-        while let Some(msg) = rate_stream.next().await {
+        let result = backoff::future::retry_notify::<Infallible, _, _, _, _, _>(
-            let msg = match msg {
+            backoff::ExponentialBackoff::default(),
-                Ok(tungstenite::Message::Text(msg)) => msg,
+            || {
-                Ok(tungstenite::Message::Close(close_frame)) => {
+                let rate_update = rate_update.clone();
-                    if let Some(tungstenite::protocol::CloseFrame { code, reason }) = close_frame {
+                async move {
-                        tracing::error!(
+                    let mut stream = connection::new().await?;
-                            "Kraken rate stream was closed with code {} and reason: {}",
-                            code,
-                            reason
-                        );
-                    } else {
-                        tracing::error!("Kraken rate stream was closed without code and reason");
-                    }
-                    let _ = rate_update.send(Err(Error::ConnectionClosed));
-                    continue;
-                }
-                Ok(msg) => {
-                    tracing::trace!(
-                        "Kraken rate stream returned non text message that will be ignored: {}",
-                        msg
-                    );
-                    continue;
-                }
-                Err(e) => {
-                    tracing::error!(%e, "Error when reading from Kraken rate stream");
-                    let _ = rate_update.send(Err(e.into()));
-                    continue;
-                }
-            };
 
-            let update = match serde_json::from_str::<Event>(&msg) {
+                    while let Some(update) = stream.try_next().await.map_err(to_backoff)? {
-                Ok(Event::SystemStatus) => {
+                        let send_result = rate_update.send(Ok(update));
-                    tracing::debug!("Connected to Kraken websocket API");
+
-                    continue;
+                        if send_result.is_err() {
+                            return Err(backoff::Error::Permanent(anyhow!(
+                                "receiver disconnected"
+                            )));
                         }
-                Ok(Event::SubscriptionStatus) => {
-                    tracing::debug!("Subscribed to updates for ticker");
-                    continue;
                     }
-                Ok(Event::Heartbeat) => {
+
-                    tracing::trace!("Received heartbeat message");
+                    Err(backoff::Error::Transient(anyhow!("stream ended")))
-                    continue;
-                }
-                // if the message is not an event, it is a ticker update or an unknown event
-                Err(_) => match serde_json::from_str::<TickerUpdate>(&msg) {
-                    Ok(ticker) => ticker,
-                    Err(e) => {
-                        tracing::warn!(%e, "Failed to deserialize message '{}' as ticker update", msg);
-                        let _ = rate_update.send(Err(Error::UnknownMessage { msg }));
-                        continue;
                 }
             },
-            };
+            |error, next: Duration| {
-
+                tracing::info!(%error, "Kraken websocket connection failed, retrying in {}ms", next.as_millis());
-            let rate = match Rate::try_from(update) {
-                Ok(rate) => rate,
-                Err(e) => {
-                    let _ = rate_update.send(Err(e));
-                    continue;
             }
-            };
+        )
+        .await;
 
-            let _ = rate_update.send(Ok(rate));
+        match result {
+            Err(e) => {
+                tracing::warn!("Rate updates incurred an unrecoverable error: {:#}", e);
+
+                // in case the retries fail permanently, let the subscribers know
+                rate_update.send(Err(Error::PermanentFailure))
+            }
+            Ok(never) => match never {},
         }
     });
 
-    rate_stream_sink
-        .send(SUBSCRIBE_XMR_BTC_TICKER_PAYLOAD.into())
-        .await?;
-
     Ok(RateUpdateStream {
         inner: rate_update_receiver,
     })
@@ -112,78 +74,198 @@ impl RateUpdateStream {
     }
 }
 
-const SUBSCRIBE_XMR_BTC_TICKER_PAYLOAD: &str = r#"
+#[derive(Clone, Debug, thiserror::Error)]
-{ "event": "subscribe",
+pub enum Error {
+    #[error("Rate is not yet available")]
+    NotYetAvailable,
+    #[error("Permanently failed to retrieve rate from Kraken")]
+    PermanentFailure,
+}
+
+type RateUpdate = Result<Rate, Error>;
+
+/// Maps a [`connection::Error`] to a backoff error, effectively defining our
+/// retry strategy.
+fn to_backoff(e: connection::Error) -> backoff::Error<anyhow::Error> {
+    use backoff::Error::*;
+
+    match e {
+        // Connection closures and websocket errors will be retried
+        connection::Error::ConnectionClosed => Transient(anyhow::Error::from(e)),
+        connection::Error::WebSocket(_) => Transient(anyhow::Error::from(e)),
+
+        // Failures while parsing a message are permanent because they most likely present a
+        // programmer error
+        connection::Error::Parse(_) => Permanent(anyhow::Error::from(e)),
+    }
+}
+
+/// Kraken websocket connection module.
+///
+/// Responsible for establishing a connection to the Kraken websocket API and
+/// transforming the received websocket frames into a stream of rate updates.
+/// The connection may fail in which case it is simply terminated and the stream
+/// ends.
+mod connection {
+    use super::*;
+    use crate::kraken::wire;
+    use futures::stream::BoxStream;
+    use tokio_tungstenite::tungstenite;
+
+    pub async fn new() -> Result<BoxStream<'static, Result<Rate, Error>>> {
+        let (mut rate_stream, _) = tokio_tungstenite::connect_async("wss://ws.kraken.com")
+            .await
+            .context("Failed to connect to Kraken websocket API")?;
+
+        rate_stream
+            .send(SUBSCRIBE_XMR_BTC_TICKER_PAYLOAD.into())
+            .await?;
+
+        let stream = rate_stream.err_into().try_filter_map(parse_message).boxed();
+
+        Ok(stream)
+    }
+
+    /// Parse a websocket message into a [`Rate`].
+    ///
+    /// Messages which are not actually ticker updates are ignored and result in
+    /// `None` being returned. In the context of a [`TryStream`], these will
+    /// simply be filtered out.
+    async fn parse_message(msg: tungstenite::Message) -> Result<Option<Rate>, Error> {
+        let msg = match msg {
+            tungstenite::Message::Text(msg) => msg,
+            tungstenite::Message::Close(close_frame) => {
+                if let Some(tungstenite::protocol::CloseFrame { code, reason }) = close_frame {
+                    tracing::debug!(
+                        "Kraken rate stream was closed with code {} and reason: {}",
+                        code,
+                        reason
+                    );
+                } else {
+                    tracing::debug!("Kraken rate stream was closed without code and reason");
+                }
+
+                return Err(Error::ConnectionClosed);
+            }
+            msg => {
+                tracing::trace!(
+                    "Kraken rate stream returned non text message that will be ignored: {}",
+                    msg
+                );
+
+                return Ok(None);
+            }
+        };
+
+        let update = match serde_json::from_str::<wire::Event>(&msg) {
+            Ok(wire::Event::SystemStatus) => {
+                tracing::debug!("Connected to Kraken websocket API");
+
+                return Ok(None);
+            }
+            Ok(wire::Event::SubscriptionStatus) => {
+                tracing::debug!("Subscribed to updates for ticker");
+
+                return Ok(None);
+            }
+            Ok(wire::Event::Heartbeat) => {
+                tracing::trace!("Received heartbeat message");
+
+                return Ok(None);
+            }
+            // if the message is not an event, it is a ticker update or an unknown event
+            Err(_) => match serde_json::from_str::<wire::TickerUpdate>(&msg) {
+                Ok(ticker) => ticker,
+                Err(e) => {
+                    tracing::warn!(%e, "Failed to deserialize message '{}' as ticker update", msg);
+
+                    return Ok(None);
+                }
+            },
+        };
+
+        let update = Rate::try_from(update)?;
+
+        Ok(Some(update))
+    }
+
+    #[derive(Debug, thiserror::Error)]
+    pub enum Error {
+        #[error("The Kraken server closed the websocket connection")]
+        ConnectionClosed,
+        #[error("Failed to read message from websocket stream")]
+        WebSocket(#[from] tungstenite::Error),
+        #[error("Failed to parse rate from websocket message")]
+        Parse(#[from] wire::Error),
+    }
+
+    const SUBSCRIBE_XMR_BTC_TICKER_PAYLOAD: &str = r#"
+    { "event": "subscribe",
       "pair": [ "XMR/XBT" ],
       "subscription": {
         "name": "ticker"
       }
-}"#;
+    }"#;
-
-#[derive(Clone, Debug, thiserror::Error)]
-pub enum Error {
-    #[error("Rate has not yet been retrieved from Kraken websocket API")]
-    NotYetRetrieved,
-    #[error("The Kraken server closed the websocket connection")]
-    ConnectionClosed,
-    #[error("Websocket: {0}")]
-    WebSocket(String),
-    #[error("Received unknown message from Kraken: {msg}")]
-    UnknownMessage { msg: String },
-    #[error("Data field is missing")]
-    DataFieldMissing,
-    #[error("Ask Rate Element is of unexpected type")]
-    UnexpectedAskRateElementType,
-    #[error("Ask Rate Element is missing")]
-    MissingAskRateElementType,
-    #[error("Bitcoin amount parse error: ")]
-    BitcoinParseAmount(#[from] ParseAmountError),
 }
 
-impl From<tungstenite::Error> for Error {
+/// Kraken websocket API wire module.
-    fn from(err: tungstenite::Error) -> Self {
+///
-        Error::WebSocket(format!("{:#}", err))
+/// Responsible for parsing websocket text messages to events and rate updates.
-    }
+mod wire {
-}
+    use super::*;
+    use bitcoin::util::amount::ParseAmountError;
+    use serde_json::Value;
 
-#[derive(Debug, Serialize, Deserialize, PartialEq)]
+    #[derive(Debug, Serialize, Deserialize, PartialEq)]
-#[serde(tag = "event")]
+    #[serde(tag = "event")]
-enum Event {
+    pub enum Event {
         #[serde(rename = "systemStatus")]
         SystemStatus,
         #[serde(rename = "heartbeat")]
         Heartbeat,
         #[serde(rename = "subscriptionStatus")]
         SubscriptionStatus,
-}
+    }
 
-#[derive(Debug, Serialize, Deserialize)]
+    #[derive(Clone, Debug, thiserror::Error)]
-#[serde(transparent)]
+    pub enum Error {
-struct TickerUpdate(Vec<TickerField>);
+        #[error("Data field is missing")]
+        DataFieldMissing,
+        #[error("Ask Rate Element is of unexpected type")]
+        UnexpectedAskRateElementType,
+        #[error("Ask Rate Element is missing")]
+        MissingAskRateElementType,
+        #[error("Failed to parse Bitcoin amount")]
+        BitcoinParseAmount(#[from] ParseAmountError),
+    }
 
-#[derive(Debug, Serialize, Deserialize)]
+    #[derive(Debug, Serialize, Deserialize)]
-#[serde(untagged)]
+    #[serde(transparent)]
-enum TickerField {
+    pub struct TickerUpdate(Vec<TickerField>);
+
+    #[derive(Debug, Serialize, Deserialize)]
+    #[serde(untagged)]
+    pub enum TickerField {
         Data(TickerData),
         Metadata(Value),
-}
+    }
 
-#[derive(Debug, Serialize, Deserialize)]
+    #[derive(Debug, Serialize, Deserialize)]
-struct TickerData {
+    pub struct TickerData {
         #[serde(rename = "a")]
         ask: Vec<RateElement>,
         #[serde(rename = "b")]
         bid: Vec<RateElement>,
-}
+    }
 
-#[derive(Debug, Serialize, Deserialize)]
+    #[derive(Debug, Serialize, Deserialize)]
-#[serde(untagged)]
+    #[serde(untagged)]
-enum RateElement {
+    pub enum RateElement {
         Text(String),
         Number(u64),
-}
+    }
 
-impl TryFrom<TickerUpdate> for Rate {
+    impl TryFrom<TickerUpdate> for Rate {
         type Error = Error;
 
         fn try_from(value: TickerUpdate) -> Result<Self, Error> {
@@ -205,10 +287,10 @@ impl TryFrom<TickerUpdate> for Rate {
 
             Ok(Self { ask })
         }
-}
+    }
 
-#[cfg(test)]
+    #[cfg(test)]
-mod tests {
+    mod tests {
         use super::*;
 
         #[test]
@@ -235,4 +317,5 @@ mod tests {
 
             let _ = serde_json::from_str::<TickerUpdate>(message).unwrap();
         }
+    }
 }