Keep long-running tasks alive

coordinator/src/message.rs

@@ -2,6 +2,7 @@ use crate::db::user;
 use crate::notifications::FcmToken;
 use crate::notifications::Notification;
 use crate::notifications::NotificationKind;
+use anyhow::Context;
 use anyhow::Result;
 use bitcoin::secp256k1::PublicKey;
 use commons::Message;
@@ -16,8 +17,10 @@ use std::sync::Arc;
 use tokio::sync::broadcast;
 use tokio::sync::broadcast::error::RecvError;
 use tokio::sync::mpsc;
+use tokio::sync::mpsc::Sender;
+use tokio::task::spawn_blocking;
 
-/// This value is arbitrarily set to 100 and defines the message accepted in the message
+/// This value is arbitrarily set to 100 and defines theff message accepted in the message
 /// channel buffer.
 const NOTIFICATION_BUFFER_SIZE: usize = 100;
 
@@ -33,14 +36,14 @@ pub enum OrderbookMessage {
 #[derive(Clone)]
 pub struct NewUserMessage {
  pub new_user: PublicKey,
- pub sender: mpsc::Sender<Message>,
+ pub sender: Sender<Message>,
 }
 
 pub fn spawn_delivering_messages_to_authenticated_users(
  pool: Pool<ConnectionManager<PgConnection>>,
- notification_sender: mpsc::Sender<Notification>,
+ notification_sender: Sender<Notification>,
  tx_user_feed: broadcast::Sender<NewUserMessage>,
-) -> (RemoteHandle<Result<()>>, mpsc::Sender<OrderbookMessage>) {
+) -> (RemoteHandle<()>, Sender<OrderbookMessage>) {
  let (sender, mut receiver) = mpsc::channel::<OrderbookMessage>(NOTIFICATION_BUFFER_SIZE);
 
  let authenticated_users = Arc::new(RwLock::new(HashMap::new()));
@@ -57,12 +60,12 @@ pub fn spawn_delivering_messages_to_authenticated_users(
  .insert(new_user_msg.new_user, new_user_msg.sender);
  }
  Err(RecvError::Closed) => {
- tracing::error!("New user message sender died! Channel closed.");
+ tracing::error!("New user message sender died! Channel closed");
  break;
  }
- Err(RecvError::Lagged(skip)) => tracing::warn!(%skip,
- "Lagging behind on new user message."
- ),
+ Err(RecvError::Lagged(skip)) => {
+ tracing::warn!(%skip, "Lagging behind on new user message")
+ }
  }
  }
  }
@@ -71,49 +74,19 @@ pub fn spawn_delivering_messages_to_authenticated_users(
  let (fut, remote_handle) = {
  async move {
  while let Some(notification) = receiver.recv().await {
- let mut conn = pool.get()?;
- match notification {
- OrderbookMessage::TraderMessage { trader_id, message , notification} => {
- tracing::info!(%trader_id, "Sending trader message: {message:?}");
-
- let trader = authenticated_users.read().get(&trader_id).cloned();
-
- match trader {
- Some(sender) => {
- if let Err(e) = sender.send(message).await {
- tracing::warn!(%trader_id, "Connection lost to trader {e:#}");
- } else {
- tracing::trace!(%trader_id, "Skipping optional push notifications as the user was successfully notified via the websocket.");
- continue;
- }
- }
- None => tracing::warn!(%trader_id, "Trader is not connected."),
- };
-
- if let (Some(notification_kind),Some(user)) = (notification, user::by_id(&mut conn, trader_id.to_string())?) {
- tracing::debug!(%trader_id, "Sending push notification to user");
-
- match FcmToken::new(user.fcm_token) {
- Ok(fcm_token) => {
- if let Err(e) = notification_sender
- .send(Notification {
- user_fcm_token: fcm_token,
- notification_kind,
- })
- .await {
- tracing::error!(%trader_id, "Failed to send push notification. Error: {e:#}");
- }
- }
- Err(error) => {
- tracing::error!(%trader_id, "Could not send notification to user. Error: {error:#}");
- }
- }
- }
- }
+ if let Err(e) = process_orderbook_message(
+ pool.clone(),
+ &authenticated_users,
+ &notification_sender,
+ notification,
+ )
+ .await
+ {
+ tracing::error!("Failed to process orderbook message: {e:#}");
  }
  }
 
- Ok(())
+ tracing::error!("Channel closed");
  }
  .remote_handle()
  };
@@ -122,3 +95,63 @@ pub fn spawn_delivering_messages_to_authenticated_users(
 
  (remote_handle, sender)
 }
+
+async fn process_orderbook_message(
+ pool: Pool<ConnectionManager<PgConnection>>,
+ authenticated_users: &RwLock<HashMap<PublicKey, Sender<Message>>>,
+ notification_sender: &Sender<Notification>,
+ notification: OrderbookMessage,
+) -> Result<()> {
+ let mut conn = spawn_blocking(move || pool.get())
+ .await
+ .expect("task to complete")?;
+
+ match notification {
+ OrderbookMessage::TraderMessage {
+ trader_id,
+ message,
+ notification,
+ } => {
+ tracing::info!(%trader_id, ?message, "Sending trader message");
+
+ let trader = authenticated_users.read().get(&trader_id).cloned();
+
+ match trader {
+ Some(sender) => {
+ if let Err(e) = sender.send(message).await {
+ tracing::warn!(%trader_id, "Connection lost to trader: {e:#}");
+ } else {
+ tracing::trace!(
+ %trader_id,
+ "Skipping optional push notifications as the user was successfully \
+ notified via the websocket"
+ );
+ return Ok(());
+ }
+ }
+ None => tracing::warn!(%trader_id, "Trader is not connected"),
+ };
+
+ let user = user::by_id(&mut conn, trader_id.to_string())
+ .context("Failed to get user by ID")?;
+
+ if let (Some(notification_kind), Some(user)) = (notification, user) {
+ tracing::debug!(%trader_id, "Sending push notification to user");
+
+ let fcm_token = FcmToken::new(user.fcm_token)?;
+
+ notification_sender
+ .send(Notification {
+ user_fcm_token: fcm_token,
+ notification_kind,
+ })
+ .await
+ .with_context(|| {
+ format!("Failed to send push notification to trader {trader_id}")
+ })?;
+ }
+ }
+ }
+
+ Ok(())
+}

coordinator/src/node/rollover.rs

@@ -31,6 +31,7 @@ use time::OffsetDateTime;
 use tokio::sync::broadcast;
 use tokio::sync::broadcast::error::RecvError;
 use tokio::sync::mpsc;
+use tokio::task::spawn_blocking;
 use trade::ContractSymbol;
 
 #[derive(Debug, Clone)]
@@ -51,20 +52,20 @@ pub fn monitor(
  notifier: mpsc::Sender<OrderbookMessage>,
  network: Network,
  node: Node,
-) -> RemoteHandle<Result<()>> {
+) -> RemoteHandle<()> {
  let mut user_feed = tx_user_feed.subscribe();
  let (fut, remote_handle) = async move {
  loop {
  match user_feed.recv().await {
  Ok(new_user_msg) => {
  tokio::spawn({
- let mut conn = pool.get()?;
  let notifier = notifier.clone();
  let node = node.clone();
+ let pool = pool.clone();
  async move {
  if let Err(e) = node
  .check_if_eligible_for_rollover(
- &mut conn,
+ pool,
  notifier,
  new_user_msg.new_user,
  network,
@@ -87,7 +88,6 @@ pub fn monitor(
  ),
  }
  }
- Ok(())
  }
  .remote_handle();
 
@@ -156,14 +156,18 @@ impl Rollover {
 impl Node {
  async fn check_if_eligible_for_rollover(
  &self,
- conn: &mut PgConnection,
+ pool: Pool<ConnectionManager<PgConnection>>,
  notifier: mpsc::Sender<OrderbookMessage>,
  trader_id: PublicKey,
  network: Network,
  ) -> Result<()> {
+ let mut conn = spawn_blocking(move || pool.get())
+ .await
+ .expect("task to complete")?;
+
  tracing::debug!(%trader_id, "Checking if the users positions is eligible for rollover");
  if let Some(position) = positions::Position::get_position_by_trader(
- conn,
+ &mut conn,
  trader_id,
  vec![PositionState::Open, PositionState::Rollover],
  )? {
@@ -172,7 +176,9 @@ impl Node {
  .get_signed_channel_by_trader_id(position.trader)?;
 
  let (retry_rollover, contract_id) = match position.position_state {
- PositionState::Rollover => self.rollback_channel_if_needed(conn, signed_channel)?,
+ PositionState::Rollover => {
+ self.rollback_channel_if_needed(&mut conn, signed_channel)?
+ }
  PositionState::Open => (false, signed_channel.get_contract_id()),
  _ => bail!("Unexpected position state {:?}", position.position_state),
  };

coordinator/src/orderbook/async_match.rs

@@ -22,25 +22,37 @@ use time::OffsetDateTime;
 use tokio::sync::broadcast;
 use tokio::sync::broadcast::error::RecvError;
 use tokio::sync::mpsc;
+use tokio::task::spawn_blocking;
 
 pub fn monitor(
  pool: Pool<ConnectionManager<PgConnection>>,
  tx_user_feed: broadcast::Sender<NewUserMessage>,
  notifier: mpsc::Sender<OrderbookMessage>,
  network: Network,
  oracle_pk: XOnlyPublicKey,
-) -> RemoteHandle<Result<()>> {
+) -> RemoteHandle<()> {
  let mut user_feed = tx_user_feed.subscribe();
  let (fut, remote_handle) = async move {
  loop {
  match user_feed.recv().await {
  Ok(new_user_msg) => {
  tokio::spawn({
- let mut conn = pool.get()?;
  let notifier = notifier.clone();
+ let pool = pool.clone();
  async move {
- tracing::debug!(trader_id=%new_user_msg.new_user, "Checking if the user needs to be notified about pending matches");
- if let Err(e) = process_pending_match(&mut conn, notifier, new_user_msg.new_user, network, oracle_pk).await {
+ tracing::debug!(
+ trader_id=%new_user_msg.new_user,
+ "Checking if the user needs to be notified about pending matches"
+ );
+ if let Err(e) = process_pending_match(
+ pool,
+ notifier,
+ new_user_msg.new_user,
+ network,
+ oracle_pk,
+ )
+ .await
+ {
  tracing::error!("Failed to process pending match. Error: {e:#}");
  }
  }
@@ -50,13 +62,16 @@ pub fn monitor(
  tracing::error!("New user message sender died! Channel closed.");
  break;
  }
- Err(RecvError::Lagged(skip)) => tracing::warn!(%skip,
+ Err(RecvError::Lagged(skip)) => {
+ tracing::warn!(
+ %skip,
  "Lagging behind on new user message."
- ),
+ )
+ }
  }
  }
-  Ok(())
- }.remote_handle();
+ }
+ .remote_handle();
 
  tokio::spawn(fut);
 
@@ -65,16 +80,22 @@ pub fn monitor(
 
 /// Checks if there are any pending matches
 async fn process_pending_match(
- conn: &mut PgConnection,
+ pool: Pool<ConnectionManager<PgConnection>>,
  notifier: mpsc::Sender<OrderbookMessage>,
  trader_id: PublicKey,
  network: Network,
  oracle_pk: XOnlyPublicKey,
 ) -> Result<()> {
- if let Some(order) = orders::get_by_trader_id_and_state(conn, trader_id, OrderState::Matched)? {
+ let mut conn = spawn_blocking(move || pool.get())
+ .await
+ .expect("task to complete")?;
+
+ if let Some(order) =
+ orders::get_by_trader_id_and_state(&mut conn, trader_id, OrderState::Matched)?
+ {
  tracing::debug!(%trader_id, order_id=%order.id, "Notifying trader about pending match");
 
- let matches = matches::get_matches_by_order_id(conn, order.id)?;
+ let matches = matches::get_matches_by_order_id(&mut conn, order.id)?;
  let filled_with = get_filled_with_from_matches(matches, network, oracle_pk)?;
 
  let message = match order.order_reason {

coordinator/src/orderbook/collaborative_revert.rs

@@ -16,20 +16,21 @@ use rust_decimal::Decimal;
 use tokio::sync::broadcast;
 use tokio::sync::broadcast::error::RecvError;
 use tokio::sync::mpsc;
+use tokio::task::spawn_blocking;
 
 pub fn monitor(
  pool: Pool<ConnectionManager<PgConnection>>,
  tx_user_feed: broadcast::Sender<NewUserMessage>,
  notifier: mpsc::Sender<OrderbookMessage>,
-) -> RemoteHandle<Result<()>> {
+) -> RemoteHandle<()> {
  let mut user_feed = tx_user_feed.subscribe();
  let (fut, remote_handle) = async move {
  loop {
  match user_feed.recv().await {
  Ok(new_user_msg) => {
  tokio::spawn({
- let mut conn = pool.get()?;
  let notifier = notifier.clone();
+ let pool = pool.clone();
  async move {
  tracing::debug!(
  trader_id=%new_user_msg.new_user,
@@ -38,7 +39,7 @@ pub fn monitor(
  );
 
  if let Err(e) = process_pending_collaborative_revert(
- &mut conn,
+ pool,
  notifier,
  new_user_msg.new_user,
  )
@@ -60,7 +61,6 @@ pub fn monitor(
  ),
  }
  }
- Ok(())
  }
  .remote_handle();
 
@@ -71,11 +71,15 @@ pub fn monitor(
 
 /// Checks if there are any pending collaborative reverts
 async fn process_pending_collaborative_revert(
- conn: &mut PgConnection,
+ pool: Pool<ConnectionManager<PgConnection>>,
  notifier: mpsc::Sender<OrderbookMessage>,
  trader_id: PublicKey,
 ) -> Result<()> {
- match collaborative_reverts::by_trader_pubkey(trader_id.to_string().as_str(), conn)? {
+ let mut conn = spawn_blocking(move || pool.get())
+ .await
+ .expect("task to complete")?;
+
+ match collaborative_reverts::by_trader_pubkey(trader_id.to_string().as_str(), &mut conn)? {
  None => {
  // nothing to revert
  }