apollo_storage: implement storage reader server with axum

crates/apollo_storage/Cargo.toml

@@ -14,6 +14,7 @@ apollo_config.workspace = true
 apollo_metrics.workspace = true
 apollo_proc_macros.workspace = true
 async-trait.workspace = true
+axum.workspace = true
 byteorder.workspace = true
 cairo-lang-casm = { workspace = true, features = ["parity-scale-codec"] }
 cairo-lang-starknet-classes.workspace = true

crates/apollo_storage/src/storage_reader_server.rs

@@ -1,9 +1,16 @@
+use std::io;
 use std::marker::PhantomData;
 use std::net::SocketAddr;
 
 use async_trait::async_trait;
+use axum::extract::State;
+use axum::http::StatusCode;
+use axum::response::{IntoResponse, Response};
+use axum::routing::post;
+use axum::{Json, Router};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
+use tracing::{error, info};
 
 use crate::{StorageError, StorageReader};
 
@@ -13,16 +20,14 @@ use crate::{StorageError, StorageReader};
 pub struct ServerConfig {
     /// The socket address to bind the server to.
     socket: SocketAddr,
-    /// Maximum number of concurrent requests the server can handle.
-    max_concurrency: usize,
     /// Whether the server is enabled.
     enable: bool,
 }
 
 impl ServerConfig {
     /// Creates a new server configuration.
-    pub fn new(socket: SocketAddr, max_concurrency: usize, enable: bool) -> Self {
-        Self { socket, max_concurrency, enable }
+    pub fn new(socket: SocketAddr, enable: bool) -> Self {
+        Self { socket, enable }
     }
 }
 
@@ -42,19 +47,39 @@ pub trait StorageReaderServerHandler<Request, Response> {
 /// A server for handling remote storage reader queries via a configurable request handler.
 pub struct StorageReaderServer<RequestHandler, Request, Response>
 where
-    RequestHandler: StorageReaderServerHandler<Request, Response>,
+    RequestHandler: StorageReaderServerHandler<Request, Response> + Clone,
     Request: Serialize + DeserializeOwned + Send + 'static,
     Response: Serialize + DeserializeOwned + Send + 'static,
+{
+    app_state: AppState<RequestHandler, Request, Response>,
+    config: ServerConfig,
+}
+
+struct AppState<RequestHandler, Request, Response>
+where
+    RequestHandler: StorageReaderServerHandler<Request, Response> + Clone,
 {
     storage_reader: StorageReader,
     request_handler: RequestHandler,
-    config: ServerConfig,
     _req_res: PhantomData<(Request, Response)>,
 }
 
+impl<RequestHandler, Request, Response> Clone for AppState<RequestHandler, Request, Response>
+where
+    RequestHandler: StorageReaderServerHandler<Request, Response> + Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            storage_reader: self.storage_reader.clone(),
+            request_handler: self.request_handler.clone(),
+            _req_res: PhantomData,
+        }
+    }
+}
+
 impl<RequestHandler, Request, Response> StorageReaderServer<RequestHandler, Request, Response>
 where
-    RequestHandler: StorageReaderServerHandler<Request, Response>,
+    RequestHandler: StorageReaderServerHandler<Request, Response> + Clone,
     Request: Serialize + DeserializeOwned + Send + 'static,
     Response: Serialize + DeserializeOwned + Send + 'static,
 {
@@ -64,11 +89,78 @@ where
         request_handler: RequestHandler,
         config: ServerConfig,
     ) -> Self {
-        Self { storage_reader, request_handler, config, _req_res: PhantomData }
+        let app_state = AppState { storage_reader, request_handler, _req_res: PhantomData };
+        Self { app_state, config }
+    }
+
+    /// Creates the axum router with configured routes and state.
+    pub fn app(&self) -> Router
+    where
+        RequestHandler: Send + Sync + 'static,
+        Request: Send + Sync + 'static,
+        Response: Send + Sync + 'static,
+    {
+        Router::new()
+            .route(
+                "/storage/query",
+                post(handle_request_endpoint::<RequestHandler, Request, Response>),
+            )
+            .with_state(self.app_state.clone())
     }
 
-    /// Starts the server to handle incoming requests.
-    pub fn start(&mut self) {
-        unimplemented!()
+    /// Runs the server to handle incoming requests.
+    pub async fn run(self) -> Result<(), StorageError>
+    where
+        RequestHandler: Send + Sync + 'static,
+        Request: Send + Sync + 'static,
+        Response: Send + Sync + 'static,
+    {
+        if !self.config.enable {
+            info!("Storage reader server is disabled, not starting");
+            return Ok(());
+        }
+        info!("Starting storage reader server on {}", self.config.socket);
+        let app = self.app();
+        info!("Storage reader server listening on {}", self.config.socket);
+
+        // Start the server
+        axum::Server::bind(&self.config.socket).serve(app.into_make_service()).await.map_err(|e| {
+            error!("Storage reader server error: {}", e);
+            StorageError::IOError(io::Error::other(e))
+        })
+    }
+}
+
+/// Axum handler for storage query requests.
+async fn handle_request_endpoint<RequestHandler, Request, Response>(
+    State(app_state): State<AppState<RequestHandler, Request, Response>>,
+    Json(request): Json<Request>,
+) -> Result<Json<Response>, StorageServerError>
+where
+    RequestHandler: StorageReaderServerHandler<Request, Response> + Clone,
+    Request: Send + Sync + 'static,
+    Response: Send + Sync + 'static,
+{
+    let response =
+        app_state.request_handler.handle_request(&app_state.storage_reader, request).await?;
+
+    Ok(Json(response))
+}
+
+/// Error type for HTTP responses.
+#[derive(Debug)]
+struct StorageServerError(StorageError);
+
+impl From<StorageError> for StorageServerError {
+    fn from(error: StorageError) -> Self {
+        StorageServerError(error)
+    }
+}
+
+impl IntoResponse for StorageServerError {
+    fn into_response(self) -> Response {
+        let error_message = format!("Storage error: {}", self.0);
+        error!("{}", error_message);
+        (StatusCode::INTERNAL_SERVER_ERROR, error_message).into_response()
     }
 }