backend: fix ws stream shutdown (#19)

CHANGELOG.md

@@ -25,8 +25,9 @@ The minor version will be incremented upon a breaking change and the patch versi
 - geyser: wait all transactions before process block ([#10](https://github.com/solana-stream-solutions/solfees/pull/10))
 - frontend: init ([#8](https://github.com/solana-stream-solutions/solfees/pull/8))
 - geyser: use process_compute_budget_instructions ([#15](https://github.com/solana-stream-solutions/solfees/pull/15))
-- api: improve parallelism ([#16](https://github.com/solana-stream-solutions/solfees/pull/16))
+- backend: improve parallelism ([#16](https://github.com/solana-stream-solutions/solfees/pull/16))
 - geyser: do not stream outdated data ([#17](https://github.com/solana-stream-solutions/solfees/pull/17))
-- api: add metrics of used resources ([#18](https://github.com/solana-stream-solutions/solfees/pull/18))
+- backend: add metrics of used resources ([#18](https://github.com/solana-stream-solutions/solfees/pull/18))
+- backend: fix ws stream shutdown ([#19](https://github.com/solana-stream-solutions/solfees/pull/19))
 
 ### Breaking

solfees-be/src/bin/solfees-ws-client.rs

@@ -68,12 +68,14 @@ async fn main() -> anyhow::Result<()> {
                 Some(Ok(Message::Ping(_))) => continue,
                 Some(Ok(Message::Pong(_))) => continue,
                 Some(Ok(Message::Frame(_))) => continue,
-                Some(Ok(Message::Close(_))) => return Ok(()),
+                Some(Ok(Message::Close(_))) => anyhow::bail!("close message received"),
                 Some(Err(error)) => anyhow::bail!(error),
-                None => anyhow::bail!("stream is closed"),
+                None => anyhow::bail!("stream finished"),
             };
             info!("new message: {text}");
         }
+        #[allow(unreachable_code)]
+        Ok::<(), anyhow::Error>(())
     };
 
     tokio::try_join!(req_to_ws, ws_to_stdout).map(|_| ())

solfees-be/src/rpc_server.rs

@@ -22,7 +22,7 @@ use {
         net::TcpListener,
         sync::{broadcast, Notify},
     },
-    tracing::{error, info},
+    tracing::{debug, error, info},
 };
 
 pub async fn run_admin(addr: SocketAddr, shutdown: Arc<Notify>) -> anyhow::Result<()> {
@@ -73,6 +73,9 @@ pub async fn run_solfees(
     let listener = TcpListener::bind(addr).await?;
     info!(%addr, "Start Solfees RPC server");
 
+    let (ws_tx, _ws_rx) = broadcast::channel(1);
+    let ws_tx = Arc::new(ws_tx);
+
     let http = ServerBuilder::new(TokioExecutor::new());
     let graceful = GracefulShutdown::new();
     loop {
@@ -83,13 +86,13 @@ pub async fn run_solfees(
 
         let solana_rpc = solana_rpc.clone();
         let config_metrics = Arc::clone(&config_metrics);
-        let (shutdown_tx, shutdown_rx) = broadcast::channel(1);
+        let ws_tx = Arc::clone(&ws_tx);
         let connection = http.serve_connection_with_upgrades(
             TokioIo::new(Box::pin(stream)),
             service_fn(move |mut req: Request<BodyIncoming>| {
                 let solana_rpc = solana_rpc.clone();
                 let config_metrics = Arc::clone(&config_metrics);
-                let shutdown_rx = shutdown_rx.resubscribe();
+                let ws_tx = Arc::clone(&ws_tx);
                 async move {
                     #[derive(Debug, Clone, Copy, PartialEq, Eq)]
                     enum ReqType {
@@ -129,7 +132,6 @@ pub async fn run_solfees(
                                         client_id,
                                         solana_rpc_mode,
                                         body.aggregate(),
-                                        shutdown_rx,
                                     )
                                 })
                                 .await
@@ -166,6 +168,7 @@ pub async fn run_solfees(
                                         client_id,
                                         solana_rpc_mode,
                                         websocket,
+                                        ws_tx.subscribe(),
                                     ));
                                     let (parts, body) = response.into_parts();
                                     Ok(Response::from_parts(parts, body.boxed()))
@@ -179,15 +182,16 @@ pub async fn run_solfees(
                 }
             }),
         );
-        let fut = graceful.watch(connection.into_owned());
-
+        let connection = graceful.watch(connection.into_owned());
         tokio::spawn(async move {
-            let _ = fut.await;
-            drop(shutdown_tx);
+            if let Err(error) = connection.await {
+                debug!(error, "connection error");
+            }
         });
     }
 
     drop(listener);
+    drop(ws_tx);
     graceful.shutdown().await;
 
     Ok::<(), anyhow::Error>(())

solfees-be/src/rpc_solana.rs

@@ -192,7 +192,6 @@ impl SolanaRpc {
         client_id: ClientId,
         mode: SolanaRpcMode,
         body: impl Buf,
-        mut shutdown_rx: broadcast::Receiver<()>,
     ) -> anyhow::Result<(RpcRequestsStats, Vec<u8>)> {
         let timer = client_id.start_timer_cpu();
         let mut stats = RpcRequestsStats::default();
@@ -534,14 +533,6 @@ impl SolanaRpc {
             let response_rx = join_all(rxs);
 
             tokio::select! {
-                value = shutdown_rx.recv() => match value {
-                    Ok(()) => unreachable!(),
-                    Err(broadcast::error::RecvError::Closed) => {
-                        shutdown.store(true, Ordering::Relaxed);
-                        anyhow::bail!("connection closed");
-                    },
-                    Err(broadcast::error::RecvError::Lagged(_)) => unreachable!(),
-                },
                 () = sleep(self.request_timeout) => {
                     shutdown.store(true, Ordering::Relaxed);
                     anyhow::bail!("request timeout");
@@ -590,6 +581,7 @@ impl SolanaRpc {
         client_id: ClientId,
         mode: SolanaRpcMode,
         websocket: HyperWebsocket,
+        mut shutdown_rx: broadcast::Receiver<()>,
     ) {
         let ws_frontend = match mode {
             SolanaRpcMode::Solfees => false,
@@ -627,6 +619,8 @@ impl SolanaRpc {
             };
 
             tokio::select! {
+                _ = shutdown_rx.recv() => break Some(None),
+
                 flush_result = websocket_tx_flush => match flush_result {
                     Ok(()) => {
                         flush_required = false;