simplify the ws socket handling

Author: bplatak

apps/pyth-lazer-agent/config/config.toml

@@ -1,4 +1,5 @@
-relayer_urls = ["ws://relayer-0.pyth-lazer.dourolabs.app/v1/transaction", "ws://relayer-0.pyth-lazer.dourolabs.app/v1/transaction"]
-publish_keypair_path = "/path/to/solana/id.json"
+relayer_urls = ["ws://localhost:10001"]
+publish_keypair_path = "/Users/bartplatak/workspace/pyth-crosschain/apps/pyth-lazer-agent/config/test_keypair.json"
 listen_address = "0.0.0.0:8910"
 publish_interval_duration = "25ms"
+authorization_token="token1"

apps/pyth-lazer-agent/src/config.rs

@@ -17,6 +17,7 @@ pub struct Config {
     pub publish_keypair_path: PathBuf,
     #[serde(with = "humantime_serde", default = "default_publish_interval")]
     pub publish_interval_duration: Duration,
+    pub history_service_url: Option<Url>
 }
 
 fn default_publish_interval() -> Duration {

apps/pyth-lazer-agent/src/http_server.rs

@@ -1,29 +1,23 @@
-use crate::jrpc_handle::jrpc_handler_inner;
-use crate::publisher_handle::publisher_inner_handler;
-use crate::websocket_utils::{handle_websocket_error, send_text};
+use crate::jrpc_handle::{JrpcConnectionContext, handle_jrpc};
+use crate::publisher_handle::handle_publisher;
 use crate::{
-    config::Config, lazer_publisher::LazerPublisher, publisher_handle::PublisherConnectionContext,
+    config::Config, http_server, lazer_publisher::LazerPublisher,
+    publisher_handle::PublisherConnectionContext,
 };
-use anyhow::{Context, Result, bail};
-use futures_util::io::{BufReader, BufWriter};
+use anyhow::{Context, Result};
 use hyper::body::Incoming;
 use hyper::{Response, StatusCode, body::Bytes, server::conn::http1, service::service_fn};
 use hyper_util::rt::TokioIo;
-use pyth_lazer_protocol::publisher::{ServerResponse, UpdateDeserializationErrorResponse};
 use soketto::{
     BoxedError,
     handshake::http::{Server, is_upgrade_request},
 };
 use std::fmt::Debug;
-use std::pin::Pin;
 use std::{io, net::SocketAddr};
 use tokio::net::{TcpListener, TcpStream};
-use tokio::{pin, select};
-use tokio_util::compat::TokioAsyncReadCompatExt;
-use tracing::{debug, error, info, instrument, warn};
+use tracing::{debug, info, instrument, warn};
 
 type FullBody = http_body_util::Full<Bytes>;
-pub type InnerHandlerResult = Pin<Box<dyn Future<Output = Result<Option<String>>> + Send>>;
 
 #[derive(Debug, Copy, Clone)]
 pub enum PublisherRequest {
@@ -53,8 +47,10 @@ pub async fn run(config: Config, lazer_publisher: LazerPublisher) -> Result<()>
         let stream_addr = listener.accept().await;
         let lazer_publisher_clone = lazer_publisher.clone();
         let config = config.clone();
-        tokio::spawn(async {
-            if let Err(err) = try_handle_connection(config, stream_addr, lazer_publisher_clone).await {
+        tokio::spawn(async move {
+            if let Err(err) =
+                try_handle_connection(config, stream_addr, lazer_publisher_clone).await
+            {
                 warn!("error while handling connection: {err:?}");
             }
         });
@@ -74,7 +70,12 @@ async fn try_handle_connection(
             TokioIo::new(stream),
             service_fn(move |r| {
                 let request = RelayerRequest(r);
-                request_handler(config.clone(), request, remote_addr, lazer_publisher.clone())
+                request_handler(
+                    config.clone(),
+                    request,
+                    remote_addr,
+                    lazer_publisher.clone(),
+                )
             }),
         )
         .with_upgrades()
@@ -136,27 +137,23 @@ async fn request_handler(
                         request_type: publisher_request_type,
                         _remote_addr: remote_addr,
                     };
-
-                    tokio::spawn(handle_ws(
-                        config,
+                    tokio::spawn(handle_publisher(
                         server,
                         request.0,
-                        lazer_publisher,
                         publisher_connection_context,
-                        publisher_inner_handler,
+                        lazer_publisher,
                     ));
                     Ok(response.map(|()| FullBody::default()))
                 }
                 Request::JrpcV1 => {
-                    tokio::spawn(handle_ws(
-                        config,
+                    let publisher_connection_context = JrpcConnectionContext {};
+                    tokio::spawn(handle_jrpc(
+                        config.clone(),
                         server,
                         request.0,
+                        publisher_connection_context,
                         lazer_publisher,
-                        (),
-                        jrpc_handler_inner,
                     ));
-
                     Ok(response.map(|()| FullBody::default()))
                 }
             }
@@ -170,88 +167,3 @@ async fn request_handler(
         }
     }
 }
-
-#[instrument(
-    skip(server, request, lazer_publisher),
-    fields(component = "publisher_ws")
-)]
-async fn handle_ws<T: Debug + Copy>(
-    config: Config,
-    server: Server,
-    request: http::Request<Incoming>,
-    lazer_publisher: LazerPublisher,
-    context: T,
-    inner_handler: fn(Config, Vec<u8>, LazerPublisher, T) -> InnerHandlerResult,
-) {
-    if let Err(err) = try_handle_ws(config, server, request, lazer_publisher, context, inner_handler).await
-    {
-        handle_websocket_error(err);
-    }
-}
-
-#[instrument(
-    skip(server, request, lazer_publisher),
-    fields(component = "publisher_ws")
-)]
-async fn try_handle_ws<T: Debug + Copy>(
-    config: Config,
-    server: Server,
-    request: http::Request<Incoming>,
-    lazer_publisher: LazerPublisher,
-    context: T,
-    inner_handler: fn(Config, Vec<u8>, LazerPublisher, T) -> InnerHandlerResult,
-) -> Result<()> {
-    let stream = hyper::upgrade::on(request).await?;
-    let io = TokioIo::new(stream);
-    let stream = BufReader::new(BufWriter::new(io.compat()));
-    let (mut ws_sender, mut ws_receiver) = server.into_builder(stream).finish();
-
-    let mut receive_buf = Vec::new();
-
-    let mut error_count = 0u32;
-    const MAX_ERROR_LOG: u32 = 10u32;
-    const MAX_ERROR_DISCONNECT: u32 = 100u32;
-
-    loop {
-        receive_buf.clear();
-        {
-            // soketto is not cancel-safe, so we need to store the future and poll it
-            // in the inner loop.
-            let receive = async { ws_receiver.receive(&mut receive_buf).await };
-            pin!(receive);
-            loop {
-                select! {
-                    _result = &mut receive => {
-                        break
-                    }
-                }
-            }
-        }
-
-        match inner_handler(config.clone(), receive_buf.clone(), lazer_publisher.clone(), context).await {
-            Ok(response) => {
-                if let Some(response) = response {
-                    send_text(&mut ws_sender, &response).await?;
-                }
-            }
-            Err(err) => {
-                error_count += 1;
-                if error_count <= MAX_ERROR_LOG {
-                    warn!("Error decoding message error: {err}");
-                }
-                if error_count >= MAX_ERROR_DISCONNECT {
-                    error!("Error threshold reached; disconnecting");
-                    bail!("Error threshold reached");
-                }
-                let error_json = &serde_json::to_string::<ServerResponse>(
-                    &UpdateDeserializationErrorResponse {
-                        error: format!("failed to parse a binary message: {err}"),
-                    }
-                    .into(),
-                )?;
-                send_text(&mut ws_sender, error_json).await?;
-                continue;
-            }
-        }
-    }
-}