feat: make addons stateful (paradigmxyz#11204)

Co-authored-by: Arsenii Kulikov <klkvrr@gmail.com>
testmachine-ai · Oct 14, 2024 · b87abfc · b87abfc
1 parent 69db5b5
commit b87abfc
Show file tree

Hide file tree

Showing 19 changed files with 75 additions and 25 deletions.
diff --git a/bin/reth/src/main.rs b/bin/reth/src/main.rs
@@ -60,7 +60,7 @@ fn main() {
                     let handle = builder
                         .with_types_and_provider::<EthereumNode, BlockchainProvider2<_>>()
                         .with_components(EthereumNode::components())
-                        .with_add_ons::<EthereumAddOns>()
+                        .with_add_ons(EthereumAddOns::default())
                         .launch_with_fn(|builder| {
                             let launcher = EngineNodeLauncher::new(
                                 builder.task_executor().clone(),

diff --git a/crates/ethereum/node/src/node.rs b/crates/ethereum/node/src/node.rs
@@ -77,7 +77,8 @@ impl NodeTypesWithEngine for EthereumNode {
 }
 
 /// Add-ons w.r.t. l1 ethereum.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
+#[non_exhaustive]
 pub struct EthereumAddOns;
 
 impl<N: FullNodeComponents> NodeAddOns<N> for EthereumAddOns {
@@ -104,6 +105,10 @@ where
     fn components_builder(&self) -> Self::ComponentsBuilder {
         Self::components()
     }
+
+    fn add_ons(&self) -> Self::AddOns {
+        EthereumAddOns::default()
+    }
 }
 
 /// A regular ethereum evm and executor builder.

diff --git a/crates/ethereum/node/tests/it/builder.rs b/crates/ethereum/node/tests/it/builder.rs
@@ -22,7 +22,7 @@ fn test_basic_setup() {
         .with_database(db)
         .with_types::<EthereumNode>()
         .with_components(EthereumNode::components())
-        .with_add_ons::<EthereumAddOns>()
+        .with_add_ons(EthereumAddOns::default())
         .on_component_initialized(move |ctx| {
             let _provider = ctx.provider();
             println!("{msg}");
@@ -54,7 +54,7 @@ async fn test_eth_launcher() {
                 NodeTypesWithDBAdapter<EthereumNode, Arc<TempDatabase<DatabaseEnv>>>,
             >>()
             .with_components(EthereumNode::components())
-            .with_add_ons::<EthereumAddOns>()
+            .with_add_ons(EthereumAddOns::default())
             .launch_with_fn(|builder| {
                 let launcher = EngineNodeLauncher::new(
                     tasks.executor(),

diff --git a/crates/ethereum/node/tests/it/exex.rs b/crates/ethereum/node/tests/it/exex.rs
@@ -33,7 +33,7 @@ fn basic_exex() {
         .with_database(db)
         .with_types::<EthereumNode>()
         .with_components(EthereumNode::components())
-        .with_add_ons::<EthereumAddOns>()
+        .with_add_ons(EthereumAddOns::default())
         .install_exex("dummy", move |ctx| future::ok(DummyExEx { _ctx: ctx }))
         .check_launch();
 }
diff --git a/crates/exex/test-utils/src/lib.rs b/crates/exex/test-utils/src/lib.rs
@@ -154,6 +154,10 @@ where
             .consensus(TestConsensusBuilder::default())
             .engine_validator(EthereumEngineValidatorBuilder::default())
     }
+
+    fn add_ons(&self) -> Self::AddOns {
+        EthereumAddOns::default()
+    }
 }
 
 /// A shared [`TempDatabase`] used for testing

diff --git a/crates/node/builder/src/builder/add_ons.rs b/crates/node/builder/src/builder/add_ons.rs
@@ -14,6 +14,8 @@ pub struct AddOns<Node: FullNodeComponents, AddOns: NodeAddOns<Node>> {
     pub exexs: Vec<(String, Box<dyn BoxedLaunchExEx<Node>>)>,
     /// Additional RPC add-ons.
     pub rpc: RpcAddOns<Node, AddOns::EthApi>,
+    /// Additional captured addons.
+    pub addons: AddOns,
 }
 
 /// Captures node specific addons that can be installed on top of the type configured node and are

diff --git a/crates/node/builder/src/builder/mod.rs b/crates/node/builder/src/builder/mod.rs
@@ -243,7 +243,7 @@ where
     where
         N: Node<RethFullAdapter<DB, N>, ChainSpec = ChainSpec>,
     {
-        self.with_types().with_components(node.components_builder()).with_add_ons::<N::AddOns>()
+        self.with_types().with_components(node.components_builder()).with_add_ons(node.add_ons())
     }
 }
 
@@ -311,7 +311,7 @@ where
     where
         N: Node<RethFullAdapter<DB, N>, ChainSpec = ChainSpec>,
     {
-        self.with_types().with_components(node.components_builder()).with_add_ons::<N::AddOns>()
+        self.with_types().with_components(node.components_builder()).with_add_ons(node.add_ons())
     }
 
     /// Launches a preconfigured [Node]
@@ -375,12 +375,15 @@ where
 {
     /// Advances the state of the node builder to the next state where all customizable
     /// [`NodeAddOns`] types are configured.
-    pub fn with_add_ons<AO>(self) -> WithLaunchContext<NodeBuilderWithComponents<T, CB, AO>>
+    pub fn with_add_ons<AO>(
+        self,
+        add_ons: AO,
+    ) -> WithLaunchContext<NodeBuilderWithComponents<T, CB, AO>>
     where
         AO: NodeAddOns<NodeAdapter<T, CB::Components>>,
     {
         WithLaunchContext {
-            builder: self.builder.with_add_ons::<AO>(),
+            builder: self.builder.with_add_ons(add_ons),
             task_executor: self.task_executor,
         }
     }

diff --git a/crates/node/builder/src/builder/states.rs b/crates/node/builder/src/builder/states.rs
@@ -58,6 +58,7 @@ impl<T: FullNodeTypes> NodeBuilderWithTypes<T> {
                 hooks: NodeHooks::default(),
                 rpc: RpcAddOns { hooks: RpcHooks::default() },
                 exexs: Vec::new(),
+                addons: (),
             },
         }
     }
@@ -168,7 +169,7 @@ where
 {
     /// Advances the state of the node builder to the next state where all customizable
     /// [`NodeAddOns`] types are configured.
-    pub fn with_add_ons<AO>(self) -> NodeBuilderWithComponents<T, CB, AO>
+    pub fn with_add_ons<AO>(self, addons: AO) -> NodeBuilderWithComponents<T, CB, AO>
     where
         AO: NodeAddOns<NodeAdapter<T, CB::Components>>,
     {
@@ -182,6 +183,7 @@ where
                 hooks: NodeHooks::default(),
                 rpc: RpcAddOns { hooks: RpcHooks::default() },
                 exexs: Vec::new(),
+                addons,
             },
         }
     }

diff --git a/crates/node/builder/src/launch/engine.rs b/crates/node/builder/src/launch/engine.rs
@@ -91,7 +91,7 @@ where
         let NodeBuilderWithComponents {
             adapter: NodeTypesAdapter { database },
             components_builder,
-            add_ons: AddOns { hooks, rpc, exexs: installed_exex },
+            add_ons: AddOns { hooks, rpc, exexs: installed_exex, .. },
             config,
         } = target;
         let NodeHooks { on_component_initialized, on_node_started, .. } = hooks;

diff --git a/crates/node/builder/src/launch/mod.rs b/crates/node/builder/src/launch/mod.rs
@@ -126,7 +126,7 @@ where
         let NodeBuilderWithComponents {
             adapter: NodeTypesAdapter { database },
             components_builder,
-            add_ons: AddOns { hooks, rpc, exexs: installed_exex },
+            add_ons: AddOns { hooks, rpc, exexs: installed_exex, .. },
             config,
         } = target;
         let NodeHooks { on_component_initialized, on_node_started, .. } = hooks;

diff --git a/crates/node/builder/src/node.rs b/crates/node/builder/src/node.rs
@@ -34,21 +34,29 @@ pub trait Node<N: FullNodeTypes>: NodeTypesWithEngine + Clone {
 
     /// Returns a [`NodeComponentsBuilder`] for the node.
     fn components_builder(&self) -> Self::ComponentsBuilder;
+
+    /// Returns the node add-ons.
+    fn add_ons(&self) -> Self::AddOns;
 }
 
 /// A [`Node`] type builder
 #[derive(Clone, Default, Debug)]
-pub struct AnyNode<N = (), C = (), AO = ()>(PhantomData<(N, AO)>, C);
+pub struct AnyNode<N = (), C = (), AO = ()>(PhantomData<N>, C, AO);
 
-impl<N, C> AnyNode<N, C> {
+impl<N, C, AO> AnyNode<N, C, AO> {
     /// Configures the types of the node.
-    pub fn types<T>(self) -> AnyNode<T, C> {
-        AnyNode::<T, C>(PhantomData::<(T, ())>, self.1)
+    pub fn types<T>(self) -> AnyNode<T, C, AO> {
+        AnyNode(PhantomData, self.1, self.2)
     }
 
     /// Sets the node components builder.
-    pub const fn components_builder<T>(&self, value: T) -> AnyNode<N, T> {
-        AnyNode::<N, T>(PhantomData::<(N, ())>, value)
+    pub fn components_builder<T>(self, value: T) -> AnyNode<N, T, AO> {
+        AnyNode(PhantomData, value, self.2)
+    }
+
+    /// Sets the node add-ons.
+    pub fn add_ons<T>(self, value: T) -> AnyNode<N, C, T> {
+        AnyNode(PhantomData, self.1, value)
     }
 }
 
@@ -84,6 +92,10 @@ where
     fn components_builder(&self) -> Self::ComponentsBuilder {
         self.1.clone()
     }
+
+    fn add_ons(&self) -> Self::AddOns {
+        self.2.clone()
+    }
 }
 
 /// The launched node with all components including RPC handlers.

diff --git a/crates/optimism/bin/src/main.rs b/crates/optimism/bin/src/main.rs
@@ -35,7 +35,7 @@ fn main() {
                     let handle = builder
                         .with_types_and_provider::<OptimismNode, BlockchainProvider2<_>>()
                         .with_components(OptimismNode::components(rollup_args))
-                        .with_add_ons::<OptimismAddOns>()
+                        .with_add_ons(OptimismAddOns::new(sequencer_http_arg.clone()))
                         .extend_rpc_modules(move |ctx| {
                             // register sequencer tx forwarder
                             if let Some(sequencer_http) = sequencer_http_arg {

diff --git a/crates/optimism/node/src/node.rs b/crates/optimism/node/src/node.rs
@@ -102,6 +102,10 @@ where
         let Self { args } = self;
         Self::components(args.clone())
     }
+
+    fn add_ons(&self) -> Self::AddOns {
+        OptimismAddOns::new(self.args.sequencer_http.clone())
+    }
 }
 
 impl NodeTypes for OptimismNode {
@@ -115,7 +119,21 @@ impl NodeTypesWithEngine for OptimismNode {
 
 /// Add-ons w.r.t. optimism.
 #[derive(Debug, Clone)]
-pub struct OptimismAddOns;
+pub struct OptimismAddOns {
+    sequencer_http: Option<String>,
+}
+
+impl OptimismAddOns {
+    /// Create a new instance with the given `sequencer_http` URL.
+    pub const fn new(sequencer_http: Option<String>) -> Self {
+        Self { sequencer_http }
+    }
+
+    /// Returns the sequencer HTTP URL.
+    pub fn sequencer_http(&self) -> Option<&str> {
+        self.sequencer_http.as_deref()
+    }
+}
 
 impl<N: FullNodeComponents> NodeAddOns<N> for OptimismAddOns {
     type EthApi = OpEthApi<N>;

diff --git a/crates/optimism/node/tests/it/builder.rs b/crates/optimism/node/tests/it/builder.rs
@@ -15,7 +15,7 @@ fn test_basic_setup() {
         .with_database(db)
         .with_types::<OptimismNode>()
         .with_components(OptimismNode::components(Default::default()))
-        .with_add_ons::<OptimismAddOns>()
+        .with_add_ons(OptimismAddOns::new(None))
         .on_component_initialized(move |ctx| {
             let _provider = ctx.provider();
             Ok(())

diff --git a/examples/custom-engine-types/src/main.rs b/examples/custom-engine-types/src/main.rs
@@ -253,6 +253,10 @@ where
             .consensus(EthereumConsensusBuilder::default())
             .engine_validator(CustomEngineValidatorBuilder::default())
     }
+
+    fn add_ons(&self) -> Self::AddOns {
+        EthereumAddOns::default()
+    }
 }
 
 /// A custom payload service builder that supports the custom engine types

diff --git a/examples/custom-evm/src/main.rs b/examples/custom-evm/src/main.rs
@@ -226,7 +226,7 @@ async fn main() -> eyre::Result<()> {
                 .executor(MyExecutorBuilder::default())
                 .payload(MyPayloadBuilder::default()),
         )
-        .with_add_ons::<EthereumAddOns>()
+        .with_add_ons(EthereumAddOns::default())
         .launch()
         .await
         .unwrap();

diff --git a/examples/custom-node-components/src/main.rs b/examples/custom-node-components/src/main.rs
@@ -25,7 +25,7 @@ fn main() {
                 // Configure the components of the node
                 // use default ethereum components but use our custom pool
                 .with_components(EthereumNode::components().pool(CustomPoolBuilder::default()))
-                .with_add_ons::<EthereumAddOns>()
+                .with_add_ons(EthereumAddOns::default())
                 .launch()
                 .await?;
 

diff --git a/examples/custom-payload-builder/src/main.rs b/examples/custom-payload-builder/src/main.rs
@@ -81,7 +81,7 @@ fn main() {
                 .with_components(
                     EthereumNode::components().payload(CustomPayloadBuilder::default()),
                 )
-                .with_add_ons::<EthereumAddOns>()
+                .with_add_ons(EthereumAddOns::default())
                 .launch()
                 .await?;
 

diff --git a/examples/stateful-precompile/src/main.rs b/examples/stateful-precompile/src/main.rs
@@ -263,7 +263,7 @@ async fn main() -> eyre::Result<()> {
         .with_types::<EthereumNode>()
         // use default ethereum components but with our executor
         .with_components(EthereumNode::components().executor(MyExecutorBuilder::default()))
-        .with_add_ons::<EthereumAddOns>()
+        .with_add_ons(EthereumAddOns::default())
         .launch()
         .await
         .unwrap();