feat(runner): Generic Extensions (#268)

refcell · haardikk21 · commit 352cc438015d · 2025-12-11T21:38:15.000+05:30
* chore(runner): node handle

* feat(runner): make extensions generic
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bin/node/Cargo.toml b/bin/node/Cargo.toml
@@ -24,7 +24,7 @@ reth-cli-util.workspace = true
 
 # misc
 clap.workspace = true
-alloy-rlp = "0.3.12"
+once_cell.workspace = true
 
 [features]
 default = []
diff --git a/bin/node/src/cli.rs b/bin/node/src/cli.rs
@@ -1,6 +1,9 @@
 //! Contains the CLI arguments
 
-use base_reth_runner::{BaseNodeConfig, FlashblocksConfig, TracingConfig};
+use std::sync::Arc;
+
+use base_reth_runner::{BaseNodeConfig, FlashblocksCell, FlashblocksConfig, TracingConfig};
+use once_cell::sync::OnceCell;
 use reth_optimism_node::args::RollupArgs;
 
 /// CLI Arguments
@@ -49,6 +52,7 @@ impl Args {
 
 impl From<Args> for BaseNodeConfig {
     fn from(args: Args) -> Self {
+        let flashblocks_cell: FlashblocksCell = Arc::new(OnceCell::new());
         let flashblocks = args.websocket_url.map(|websocket_url| FlashblocksConfig {
             websocket_url,
             max_pending_blocks_depth: args.max_pending_blocks_depth,
@@ -62,6 +66,7 @@ impl From<Args> for BaseNodeConfig {
                 logs_enabled: args.enable_transaction_tracing_logs,
             },
             metering_enabled: args.enable_metering,
+            flashblocks_cell,
         }
     }
 }
diff --git a/bin/node/src/main.rs b/bin/node/src/main.rs
@@ -5,7 +5,9 @@
 
 pub mod cli;
 
-use base_reth_runner::BaseNodeRunner;
+use base_reth_runner::{
+    BaseNodeRunner, BaseRpcExtension, FlashblocksCanonExtension, TransactionTracingExtension,
+};
 
 #[global_allocator]
 static ALLOC: reth_cli_util::allocator::Allocator = reth_cli_util::allocator::new_allocator();
@@ -21,9 +23,12 @@ fn main() {
 
     // Step 3: Hand the parsed CLI to the node runner so it can build and launch the Base node.
     cli.run(|builder, args| async move {
-        let runner = BaseNodeRunner::new(args);
-        let handle = runner.build(builder).await?;
-        runner.run(handle).await
+        let mut runner = BaseNodeRunner::new(args);
+        runner.install_ext::<FlashblocksCanonExtension>()?;
+        runner.install_ext::<TransactionTracingExtension>()?;
+        runner.install_ext::<BaseRpcExtension>()?;
+        let handle = runner.run(builder);
+        handle.await
     })
     .unwrap();
 }
diff --git a/crates/runner/Cargo.toml b/crates/runner/Cargo.toml
@@ -30,3 +30,4 @@ futures-util.workspace = true
 once_cell.workspace = true
 tracing.workspace = true
 url.workspace = true
+derive_more = { workspace = true, features = ["debug"] }
diff --git a/crates/runner/src/config.rs b/crates/runner/src/config.rs
@@ -2,6 +2,8 @@
 
 use reth_optimism_node::args::RollupArgs;
 
+use crate::extensions::FlashblocksCell;
+
 /// Captures the pieces of CLI configuration that the node logic cares about.
 #[derive(Debug, Clone)]
 pub struct BaseNodeConfig {
@@ -13,6 +15,8 @@ pub struct BaseNodeConfig {
     pub tracing: TracingConfig,
     /// Indicates whether the metering RPC surface should be installed.
     pub metering_enabled: bool,
+    /// Shared Flashblocks state cache.
+    pub flashblocks_cell: FlashblocksCell,
 }
 
 impl BaseNodeConfig {
diff --git a/crates/runner/src/extensions/canon.rs b/crates/runner/src/extensions/canon.rs
@@ -8,8 +8,8 @@ use futures_util::TryStreamExt;
 use reth_exex::ExExEvent;
 
 use crate::{
-    FlashblocksConfig,
-    extensions::{FlashblocksCell, OpBuilder},
+    BaseNodeConfig, FlashblocksConfig,
+    extensions::{BaseNodeExtension, ConfigurableBaseNodeExtension, FlashblocksCell, OpBuilder},
 };
 
 /// Helper struct that wires the Flashblocks canon ExEx into the node builder.
@@ -23,12 +23,14 @@ pub struct FlashblocksCanonExtension {
 
 impl FlashblocksCanonExtension {
     /// Create a new Flashblocks canon extension helper.
-    pub const fn new(cell: FlashblocksCell, config: Option<FlashblocksConfig>) -> Self {
-        Self { cell, config }
+    pub fn new(config: &BaseNodeConfig) -> Self {
+        Self { cell: config.flashblocks_cell.clone(), config: config.flashblocks.clone() }
     }
+}
 
+impl BaseNodeExtension for FlashblocksCanonExtension {
     /// Applies the extension to the supplied builder.
-    pub fn apply(&self, builder: OpBuilder) -> OpBuilder {
+    fn apply(&self, builder: OpBuilder) -> OpBuilder {
         let flashblocks = self.config.clone();
         let flashblocks_enabled = flashblocks.is_some();
         let flashblocks_cell = self.cell.clone();
@@ -64,3 +66,9 @@ impl FlashblocksCanonExtension {
         })
     }
 }
+
+impl ConfigurableBaseNodeExtension for FlashblocksCanonExtension {
+    fn build(config: &BaseNodeConfig) -> eyre::Result<Self> {
+        Ok(Self::new(config))
+    }
+}
diff --git a/crates/runner/src/extensions/extension.rs b/crates/runner/src/extensions/extension.rs
@@ -0,0 +1,19 @@
+//! Traits describing configurable node builder extensions.
+
+use std::fmt::Debug;
+
+use eyre::Result;
+
+use crate::{BaseNodeConfig, extensions::OpBuilder};
+
+/// A node builder extension that can apply additional wiring to the builder.
+pub trait BaseNodeExtension: Send + Sync + Debug {
+    /// Applies the extension to the supplied builder.
+    fn apply(&self, builder: OpBuilder) -> OpBuilder;
+}
+
+/// An extension that can be constructed from [`BaseNodeConfig`].
+pub trait ConfigurableBaseNodeExtension: BaseNodeExtension + Sized + 'static {
+    /// Builds the extension from the node config.
+    fn build(config: &BaseNodeConfig) -> Result<Self>;
+}
diff --git a/crates/runner/src/extensions/mod.rs b/crates/runner/src/extensions/mod.rs
@@ -3,6 +3,9 @@
 //! Builder extensions for the node nicely modularizes parts
 //! of the node building process.
 
+mod extension;
+pub use extension::{BaseNodeExtension, ConfigurableBaseNodeExtension};
+
 mod canon;
 pub use canon::FlashblocksCanonExtension;
 
diff --git a/crates/runner/src/extensions/rpc.rs b/crates/runner/src/extensions/rpc.rs
@@ -11,8 +11,8 @@ use tracing::info;
 use url::Url;
 
 use crate::{
-    FlashblocksConfig,
-    extensions::{FlashblocksCell, OpBuilder},
+    BaseNodeConfig, FlashblocksConfig,
+    extensions::{BaseNodeExtension, ConfigurableBaseNodeExtension, FlashblocksCell, OpBuilder},
 };
 
 /// Helper struct that wires the custom RPC modules into the node builder.
@@ -30,17 +30,19 @@ pub struct BaseRpcExtension {
 
 impl BaseRpcExtension {
     /// Creates a new RPC extension helper.
-    pub const fn new(
-        flashblocks_cell: FlashblocksCell,
-        flashblocks: Option<FlashblocksConfig>,
-        metering_enabled: bool,
-        sequencer_rpc: Option<String>,
-    ) -> Self {
-        Self { flashblocks_cell, flashblocks, metering_enabled, sequencer_rpc }
+    pub fn new(config: &BaseNodeConfig) -> Self {
+        Self {
+            flashblocks_cell: config.flashblocks_cell.clone(),
+            flashblocks: config.flashblocks.clone(),
+            metering_enabled: config.metering_enabled,
+            sequencer_rpc: config.rollup_args.sequencer.clone(),
+        }
     }
+}
 
+impl BaseNodeExtension for BaseRpcExtension {
     /// Applies the extension to the supplied builder.
-    pub fn apply(&self, builder: OpBuilder) -> OpBuilder {
+    fn apply(&self, builder: OpBuilder) -> OpBuilder {
         let flashblocks_cell = self.flashblocks_cell.clone();
         let flashblocks = self.flashblocks.clone();
         let metering_enabled = self.metering_enabled;
@@ -93,3 +95,9 @@ impl BaseRpcExtension {
         })
     }
 }
+
+impl ConfigurableBaseNodeExtension for BaseRpcExtension {
+    fn build(config: &BaseNodeConfig) -> eyre::Result<Self> {
+        Ok(Self::new(config))
+    }
+}
diff --git a/crates/runner/src/extensions/tracing.rs b/crates/runner/src/extensions/tracing.rs
@@ -3,7 +3,10 @@
 
 use base_tracex::tracex_exex;
 
-use crate::{TracingConfig, extensions::OpBuilder};
+use crate::{
+    BaseNodeConfig, TracingConfig,
+    extensions::{BaseNodeExtension, ConfigurableBaseNodeExtension, OpBuilder},
+};
 
 /// Helper struct that wires the transaction tracing ExEx into the node builder.
 #[derive(Debug, Clone, Copy)]
@@ -14,15 +17,23 @@ pub struct TransactionTracingExtension {
 
 impl TransactionTracingExtension {
     /// Creates a new transaction tracing extension helper.
-    pub const fn new(config: TracingConfig) -> Self {
-        Self { config }
+    pub const fn new(config: &BaseNodeConfig) -> Self {
+        Self { config: config.tracing }
     }
+}
 
+impl BaseNodeExtension for TransactionTracingExtension {
     /// Applies the extension to the supplied builder.
-    pub fn apply(&self, builder: OpBuilder) -> OpBuilder {
+    fn apply(&self, builder: OpBuilder) -> OpBuilder {
         let tracing = self.config;
         builder.install_exex_if(tracing.enabled, "tracex", move |ctx| async move {
             Ok(tracex_exex(ctx, tracing.logs_enabled))
         })
     }
 }
+
+impl ConfigurableBaseNodeExtension for TransactionTracingExtension {
+    fn build(config: &BaseNodeConfig) -> eyre::Result<Self> {
+        Ok(Self::new(config))
+    }
+}
diff --git a/crates/runner/src/handle.rs b/crates/runner/src/handle.rs
@@ -0,0 +1,66 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use derive_more::Debug;
+use eyre::Result;
+use futures_util::{FutureExt, future::BoxFuture};
+use reth::builder::NodeHandleFor;
+use reth_optimism_node::OpNode;
+
+/// Handle to a launched Base node.
+///
+/// This wraps the underlying [`NodeHandleFor`] so callers can await it directly to wait for node
+/// shutdown.
+#[must_use = "Dropping the handle will stop the node immediately"]
+#[derive(Debug, Default)]
+pub struct BaseNodeHandle {
+    #[debug(skip)]
+    build_fut: Option<BoxFuture<'static, Result<NodeHandleFor<OpNode>>>>,
+    #[debug(skip)]
+    handle: Option<Box<NodeHandleFor<OpNode>>>,
+}
+
+impl BaseNodeHandle {
+    pub(crate) fn new(
+        fut: impl Future<Output = Result<NodeHandleFor<OpNode>>> + Send + 'static,
+    ) -> Self {
+        Self { build_fut: Some(fut.boxed()), handle: None }
+    }
+}
+
+impl Future for BaseNodeHandle {
+    type Output = Result<()>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if let Some(build_fut) = self.build_fut.as_mut() {
+            match build_fut.as_mut().poll(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(Ok(handle)) => {
+                    self.handle = Some(Box::new(handle));
+                    self.build_fut = None;
+                }
+                Poll::Ready(Err(err)) => {
+                    self.build_fut = None;
+                    return Poll::Ready(Err(err));
+                }
+            }
+        }
+
+        if let Some(handle) = self.handle.as_mut() {
+            // SAFETY: The handle is stored on the heap so its address is stable; we never move the
+            // inner value after taking this reference.
+            let node_exit_future = unsafe { Pin::new_unchecked(&mut handle.node_exit_future) };
+            let poll = node_exit_future.poll(cx);
+            if poll.is_ready() {
+                self.handle = None;
+            }
+            return poll;
+        }
+
+        // If neither future is present we completed successfully and have nothing left to drive.
+        Poll::Ready(Ok(()))
+    }
+}
diff --git a/crates/runner/src/lib.rs b/crates/runner/src/lib.rs
@@ -6,6 +6,9 @@
 mod context;
 pub use context::BaseNodeBuilder;
 
+mod handle;
+pub use handle::BaseNodeHandle;
+
 mod runner;
 pub use runner::BaseNodeRunner;
 
@@ -14,6 +17,6 @@ pub use config::{BaseNodeConfig, FlashblocksConfig, TracingConfig};
 
 mod extensions;
 pub use extensions::{
-    BaseRpcExtension, FlashblocksCanonExtension, FlashblocksCell, OpBuilder, OpProvider,
-    TransactionTracingExtension,
+    BaseNodeExtension, BaseRpcExtension, ConfigurableBaseNodeExtension, FlashblocksCanonExtension,
+    FlashblocksCell, OpBuilder, OpProvider, TransactionTracingExtension,
 };
diff --git a/crates/runner/src/runner.rs b/crates/runner/src/runner.rs
