feat(pyth-lazer) Implement JRPC endpoint for the lazer agent

Author: bplatak

apps/pyth-lazer-agent/Cargo.lock

@@ -4,8 +4,8 @@ version = "0.1.2"
 edition = "2024"
 
 [dependencies]
-pyth-lazer-publisher-sdk = "0.1.5"
-pyth-lazer-protocol = "0.7.2"
+pyth-lazer-publisher-sdk = { path = "../../lazer/publisher_sdk/rust" }
+pyth-lazer-protocol = {path = "../../lazer/sdk/rust/protocol"}
 
 anyhow = "1.0.98"
 backoff = "0.4.0"
@@ -20,7 +20,7 @@ futures-util = "0.3.31"
 http = "1.3.1"
 http-body-util = "0.1.3"
 humantime-serde = "1.1.1"
-hyper = { version = "1.6.0", features = ["http1", "server"] }
+hyper = { version = "1.6.0", features = ["http1", "server", "client"] }
 hyper-util = { version = "0.1.10", features = ["tokio"] }
 protobuf = "3.7.2"
 serde = { version = "1.0.219", features = ["derive"] }
@@ -33,6 +33,7 @@ tokio-util = { version = "0.7.14", features = ["compat"] }
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter", "json"] }
 url = { version = "2.5.4", features = ["serde"] }
+reqwest = "0.12.22"
 
 [dev-dependencies]
 tempfile = "3.20.0"

apps/pyth-lazer-agent/Cargo.toml

@@ -1,32 +1,46 @@
-use anyhow::{Context, Result};
+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::{
+    config::Config, lazer_publisher::LazerPublisher, publisher_handle::PublisherConnectionContext,
+};
+use anyhow::{Context, Result, bail};
+use futures_util::io::{BufReader, BufWriter};
+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 tracing::{debug, info, instrument, warn};
-
-use crate::{
-    config::Config,
-    lazer_publisher::LazerPublisher,
-    publisher_handle::{PublisherConnectionContext, handle_publisher},
-};
+use tokio::{pin, select};
+use tokio_util::compat::TokioAsyncReadCompatExt;
+use tracing::{debug, error, info, instrument, warn};
 
 type FullBody = http_body_util::Full<Bytes>;
+pub type InnerHandlerResult = Pin<Box<dyn Future<Output = Result<Option<String>>> + Send>>;
 
-#[derive(Debug)]
-pub enum Request {
+#[derive(Debug, Copy, Clone)]
+pub enum PublisherRequest {
     PublisherV1,
     PublisherV2,
 }
 
-pub struct RelayerRequest(pub http::Request<hyper::body::Incoming>);
+pub enum Request {
+    PublisherRequest(PublisherRequest),
+    JrpcV1,
+}
 
-const PUBLISHER_WS_URI: &str = "/v1/publisher";
+pub struct RelayerRequest(pub http::Request<Incoming>);
+
+const PUBLISHER_WS_URI_V1: &str = "/v1/publisher";
 const PUBLISHER_WS_URI_V2: &str = "/v2/publisher";
+const JRPC_WS_URI_V1: &str = "/v1/jprc";
 
 const READINESS_PROBE_PATH: &str = "/ready";
 const LIVENESS_PROBE_PATH: &str = "/live";
@@ -38,15 +52,17 @@ pub async fn run(config: Config, lazer_publisher: LazerPublisher) -> Result<()>
     loop {
         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(stream_addr, lazer_publisher_clone).await {
+            if let Err(err) = try_handle_connection(config, stream_addr, lazer_publisher_clone).await {
                 warn!("error while handling connection: {err:?}");
             }
         });
     }
 }
 
 async fn try_handle_connection(
+    config: Config,
     stream_addr: io::Result<(TcpStream, SocketAddr)>,
     lazer_publisher: LazerPublisher,
 ) -> Result<()> {
@@ -58,7 +74,7 @@ async fn try_handle_connection(
             TokioIo::new(stream),
             service_fn(move |r| {
                 let request = RelayerRequest(r);
-                request_handler(request, remote_addr, lazer_publisher.clone())
+                request_handler(config.clone(), request, remote_addr, lazer_publisher.clone())
             }),
         )
         .with_upgrades()
@@ -68,15 +84,17 @@ async fn try_handle_connection(
 
 #[instrument(skip_all, fields(component = "http_server", remote_addr = remote_addr.to_string()))]
 async fn request_handler(
+    config: Config,
     request: RelayerRequest,
     remote_addr: SocketAddr,
     lazer_publisher: LazerPublisher,
 ) -> Result<Response<FullBody>, BoxedError> {
     let path = request.0.uri().path();
 
     let request_type = match path {
-        PUBLISHER_WS_URI => Request::PublisherV1,
-        PUBLISHER_WS_URI_V2 => Request::PublisherV2,
+        PUBLISHER_WS_URI_V1 => Request::PublisherRequest(PublisherRequest::PublisherV1),
+        PUBLISHER_WS_URI_V2 => Request::PublisherRequest(PublisherRequest::PublisherV2),
+        JRPC_WS_URI_V1 => Request::JrpcV1,
         LIVENESS_PROBE_PATH => {
             let response = Response::builder().status(StatusCode::OK);
             return Ok(response.body(FullBody::default())?);
@@ -113,17 +131,32 @@ async fn request_handler(
         Ok(response) => {
             info!("accepted connection from publisher");
             match request_type {
-                Request::PublisherV1 | Request::PublisherV2 => {
+                Request::PublisherRequest(publisher_request_type) => {
                     let publisher_connection_context = PublisherConnectionContext {
-                        request_type,
+                        request_type: publisher_request_type,
                         _remote_addr: remote_addr,
                     };
-                    tokio::spawn(handle_publisher(
+
+                    tokio::spawn(handle_ws(
+                        config,
                         server,
                         request.0,
+                        lazer_publisher,
                         publisher_connection_context,
+                        publisher_inner_handler,
+                    ));
+                    Ok(response.map(|()| FullBody::default()))
+                }
+                Request::JrpcV1 => {
+                    tokio::spawn(handle_ws(
+                        config,
+                        server,
+                        request.0,
                         lazer_publisher,
+                        (),
+                        jrpc_handler_inner,
                     ));
+
                     Ok(response.map(|()| FullBody::default()))
                 }
             }
@@ -137,3 +170,88 @@ 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;
+            }
+        }
+    }
+}

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

@@ -0,0 +1,59 @@
+use anyhow::{bail, Error};
+use futures_util::future::err;
+use reqwest::Client;
+use pyth_lazer_protocol::jrpc::{GetMetadataParams, JrpcResponse, JsonRpcVersion, JrpcParams, PythLazerAgentJrpcV1};
+use crate::config::Config;
+use crate::http_server::InnerHandlerResult;
+use crate::lazer_publisher::LazerPublisher;
+
+pub fn jrpc_handler_inner(config: Config, receive_buf: Vec<u8>, lazer_publisher: LazerPublisher, _context: ()) -> InnerHandlerResult {
+    Box::pin(async move {
+        let (data, _) = bincode::serde::decode_from_slice::<PythLazerAgentJrpcV1, _>(
+            &receive_buf,
+            bincode::config::legacy(),
+        )?;
+
+        match data.params {
+            JrpcParams::SendUpdates(update_params) => {
+                lazer_publisher
+                    .push_feed_update(update_params.into())
+                    .await?;
+
+                return Ok(Some(serde_json::to_string::<JrpcResponse<()>>(&JrpcResponse {
+                    jsonrpc: JsonRpcVersion::V2,
+                    result: (),
+                    id: data.id,
+                })?))
+            }
+            JrpcParams::GetMetadata(params) => {
+                match get_metadata(config, params).await {
+                    Ok(result) => {
+                        return Ok(Some(serde_json::to_string::<JrpcResponse<()>>(&JrpcResponse {
+                            jsonrpc: JsonRpcVersion::V2,
+                            result: result,
+                            id: data.id,
+                        })?))
+                    }
+                    Err(_) => {}
+                }
+            }
+        }
+
+        Ok(None)
+    })
+}
+
+async fn get_metadata(config: Config, params: GetMetadataParams) -> Result<Option<String>, anyhow::Error> {
+    let client = Client::new();
+
+    let resp = client
+        .post("https://httpbin.org/post")
+        .json(&payload)
+        .send()
+        .await?
+        .json::<HttpBinPostResponse>()
+        .await?;
+
+    let
+    Ok(None)
+}

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

@@ -24,7 +24,7 @@ use tokio::{
 };
 use tracing::error;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct LazerPublisher {
     sender: Sender<FeedUpdate>,
     pub(crate) is_ready: Arc<AtomicBool>,

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

@@ -12,6 +12,7 @@ mod lazer_publisher;
 mod publisher_handle;
 mod relayer_session;
 mod websocket_utils;
+mod jrpc_handle;
 
 #[derive(Parser)]
 #[command(version)]