Add optional shadow builder to args

Author: avalonche

crates/rollup-boost/src/cli.rs

@@ -12,13 +12,16 @@ use crate::{
     get_version, init_metrics,
     probe::ProbeLayer,
 };
-use crate::{FlashblocksService, RpcClient};
+use crate::{FlashblocksService, RpcClient, ShadowBuilderArgs};
 
 #[derive(Clone, Debug, clap::Args)]
 pub struct RollupBoostLibArgs {
     #[clap(flatten)]
     pub builder: BuilderArgs,
 
+    #[clap(flatten)]
+    pub shadow_builder: Option<ShadowBuilderArgs>,
+
     #[clap(flatten)]
     pub l2_client: L2ClientArgs,
 
@@ -117,6 +120,12 @@ impl RollupBoostServiceArgs {
         let builder_client_args: ClientArgs = self.lib.builder.clone().into();
         let builder_http_client = builder_client_args.new_http_client(PayloadSource::Builder)?;
 
+        let shadow_builder_client_args: Option<ClientArgs> =
+            self.lib.shadow_builder.clone().map(Into::into);
+        let shadow_builder_http_client = shadow_builder_client_args
+            .map(|client| client.new_http_client(PayloadSource::Builder))
+            .transpose()?;
+
         let (probe_layer, probes) = ProbeLayer::new();
 
         let (health_handle, rpc_module) = if self.lib.flashblocks.flashblocks {
@@ -148,8 +157,9 @@ impl RollupBoostServiceArgs {
             tower::ServiceBuilder::new()
                 .layer(probe_layer)
                 .layer(ProxyLayer::new(
-                    l2_http_client.clone(),
-                    builder_http_client.clone(),
+                    l2_http_client,
+                    builder_http_client,
+                    shadow_builder_http_client,
                 ));
 
         let server = Server::builder()

crates/rollup-boost/src/client/rpc.rs

@@ -467,7 +467,11 @@ macro_rules! define_client_args {
     };
 }
 
-define_client_args!((BuilderArgs, builder), (L2ClientArgs, l2));
+define_client_args!(
+    (BuilderArgs, builder),
+    (L2ClientArgs, l2),
+    (ShadowBuilderArgs, shadow_builder)
+);
 
 #[cfg(test)]
 pub mod tests {

crates/rollup-boost/src/proxy.rs

@@ -6,7 +6,7 @@ use jsonrpsee::server::HttpBody;
 use std::task::{Context, Poll};
 use std::{future::Future, pin::Pin};
 use tower::{Layer, Service};
-use tracing::info;
+use tracing::{debug, info};
 
 const ENGINE_METHOD: &str = "engine_";
 
@@ -24,13 +24,19 @@ const FORWARD_REQUESTS: [&str; 6] = [
 pub struct ProxyLayer {
     l2_client: HttpClient,
     builder_client: HttpClient,
+    shadow_builder_client: Option<HttpClient>,
 }
 
 impl ProxyLayer {
-    pub fn new(l2_client: HttpClient, builder_client: HttpClient) -> Self {
+    pub fn new(
+        l2_client: HttpClient,
+        builder_client: HttpClient,
+        shadow_builder_client: Option<HttpClient>,
+    ) -> Self {
         ProxyLayer {
             l2_client,
             builder_client,
+            shadow_builder_client,
         }
     }
 }
@@ -43,6 +49,7 @@ impl<S> Layer<S> for ProxyLayer {
             inner,
             l2_client: self.l2_client.clone(),
             builder_client: self.builder_client.clone(),
+            shadow_builder_client: self.shadow_builder_client.clone(),
         }
     }
 }
@@ -52,6 +59,7 @@ pub struct ProxyService<S> {
     inner: S,
     l2_client: HttpClient,
     builder_client: HttpClient,
+    shadow_builder_client: Option<HttpClient>,
 }
 
 // Consider using `RpcServiceT` when https://github.com/paritytech/jsonrpsee/pull/1521 is merged
@@ -108,13 +116,26 @@ where
                 let method_clone = method.clone();
                 let buffered_clone = buffered.clone();
                 let mut builder_client = service.builder_client.clone();
-
+                debug!(target: "proxy::call", message = "forwarding request to builder", ?method);
                 // Fire and forget the builder request
                 tokio::spawn(async move {
                     let _ = builder_client.forward(buffered_clone, method_clone).await;
                 });
             }
 
+            // If the shadow builder is enabled, forward the request to the shadow builder
+            if let Some(shadow_builder_client) = service.shadow_builder_client.clone() {
+                let method_clone = method.clone();
+                let buffered_clone = buffered.clone();
+                let mut shadow_builder_client = shadow_builder_client.clone();
+                debug!(target: "proxy::call", message = "proxying request to shadow builder", ?method);
+                tokio::spawn(async move {
+                    let _ = shadow_builder_client
+                        .forward(buffered_clone, method_clone)
+                        .await;
+                });
+            }
+
             // Return the response from the L2 client
             service
                 .l2_client
@@ -196,6 +217,7 @@ mod tests {
                     timeout: 1,
                 }
                 .new_http_client(PayloadSource::Builder)?,
+                None,
             ));
 
             let temp_listener = TcpListener::bind("127.0.0.1:0").await?;
@@ -512,6 +534,7 @@ mod tests {
             }
             .new_http_client(PayloadSource::Builder)
             .unwrap(),
+            None,
         );
 
         // Create a layered server
@@ -825,6 +848,7 @@ mod tests {
                 timeout: 200,
             }
             .new_http_client(PayloadSource::Builder)?,
+            None,
         );
 
         // Start proxy server
@@ -926,6 +950,7 @@ mod tests {
                 timeout: 200,
             }
             .new_http_client(PayloadSource::Builder)?,
+            None,
         );
 
         // Start proxy on dynamic port

crates/rollup-boost/src/server.rs

@@ -909,7 +909,7 @@ pub mod tests {
 
             let http_middleware = tower::ServiceBuilder::new()
                 .layer(probe_layer)
-                .layer(ProxyLayer::new(l2_http_client, builder_http_client));
+                .layer(ProxyLayer::new(l2_http_client, builder_http_client, None));
 
             let server = Server::builder()
                 .set_http_middleware(http_middleware)