Wire Kafka consumer and ResourceAnnotator into runner

niran · niran · commit c381fa97b570 · 2025-12-22T23:06:14.000-06:00
Completes the metering pipeline by:
- Creating MeteringRuntime to hold cache, estimator, and channels
- Spawning ResourceAnnotator task to correlate transactions with flashblocks
- Spawning KafkaBundleConsumer task to consume AcceptedBundle events
- Adding CompositeFlashblocksReceiver that forwards to both FlashblocksState
  and the metering pipeline via FlashblockInclusion events
- Adding flashblock_inclusion_from_flashblock helper to extract tx hashes
- Supporting Kafka properties file loading

The metering cache is now populated when:
1. Kafka consumer receives AcceptedBundle events (transaction data)
2. FlashblocksSubscriber receives flashblocks (inclusion position)
3. ResourceAnnotator correlates both to update the cache
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/runner/Cargo.toml b/crates/runner/Cargo.toml
@@ -27,6 +27,9 @@ reth-optimism-chainspec.workspace = true
 # alloy
 alloy-primitives.workspace = true
 
+# flashblocks
+base-flashtypes.workspace = true
+
 # misc
 eyre.workspace = true
 futures-util.workspace = true
@@ -35,3 +38,5 @@ tracing.workspace = true
 url.workspace = true
 parking_lot.workspace = true
 derive_more = { workspace = true, features = ["debug"] }
+rdkafka.workspace = true
+tokio.workspace = true
diff --git a/crates/runner/src/extensions/rpc.rs b/crates/runner/src/extensions/rpc.rs
@@ -2,22 +2,115 @@
 
 use std::sync::Arc;
 
-use alloy_primitives::U256;
-use base_reth_flashblocks::{FlashblocksState, FlashblocksSubscriber};
+use alloy_primitives::{B256, U256, keccak256};
+use base_flashtypes::Flashblock;
+use base_reth_flashblocks::{FlashblocksReceiver, FlashblocksState, FlashblocksSubscriber};
 use base_reth_rpc::{
-    EthApiExt, EthApiOverrideServer, EthPubSub, EthPubSubApiServer, MeteringApiImpl,
-    MeteringApiServer, MeteringCache, PriorityFeeEstimator, ResourceLimits,
+    EthApiExt, EthApiOverrideServer, EthPubSub, EthPubSubApiServer, FlashblockInclusion,
+    KafkaBundleConsumer, KafkaBundleConsumerConfig, MeteredTransaction, MeteringApiImpl,
+    MeteringApiServer, MeteringCache, PriorityFeeEstimator, ResourceAnnotator, ResourceLimits,
     TransactionStatusApiImpl, TransactionStatusApiServer,
 };
 use parking_lot::RwLock;
-use tracing::info;
+use rdkafka::ClientConfig;
+use tokio::sync::mpsc;
+use tracing::{error, info, warn};
 use url::Url;
 
 use crate::{
     BaseNodeConfig, FlashblocksConfig, MeteringConfig,
     extensions::{BaseNodeExtension, ConfigurableBaseNodeExtension, FlashblocksCell, OpBuilder},
 };
 
+/// Runtime state for the metering pipeline.
+#[derive(Clone)]
+struct MeteringRuntime {
+    /// Shared cache for metered transactions.
+    cache: Arc<RwLock<MeteringCache>>,
+    /// Priority fee estimator.
+    estimator: Arc<PriorityFeeEstimator>,
+    /// Sender for metered transactions from Kafka.
+    tx_sender: mpsc::UnboundedSender<MeteredTransaction>,
+    /// Sender for flashblock inclusions.
+    flashblock_sender: mpsc::UnboundedSender<FlashblockInclusion>,
+}
+
+/// Composite receiver that forwards flashblocks to both FlashblocksState and the metering pipeline.
+struct CompositeFlashblocksReceiver<Client> {
+    state: Arc<FlashblocksState<Client>>,
+    /// Optional channel for the metering pipeline; flashblocks RPC still needs the stream even
+    /// when metering is disabled, so we only forward inclusions if a sender is provided.
+    metering_sender: Option<mpsc::UnboundedSender<FlashblockInclusion>>,
+}
+
+impl<Client> CompositeFlashblocksReceiver<Client> {
+    fn new(
+        state: Arc<FlashblocksState<Client>>,
+        metering_sender: Option<mpsc::UnboundedSender<FlashblockInclusion>>,
+    ) -> Self {
+        Self { state, metering_sender }
+    }
+}
+
+impl<Client> FlashblocksReceiver for CompositeFlashblocksReceiver<Client>
+where
+    FlashblocksState<Client>: FlashblocksReceiver,
+{
+    fn on_flashblock_received(&self, flashblock: Flashblock) {
+        // Forward to the state first
+        self.state.on_flashblock_received(flashblock.clone());
+
+        // Then forward to metering if enabled
+        let Some(sender) = &self.metering_sender else {
+            return;
+        };
+        let Some(inclusion) = flashblock_inclusion_from_flashblock(&flashblock) else {
+            return;
+        };
+
+        if sender.send(inclusion).is_err() {
+            warn!(
+                target: "metering::flashblocks",
+                "Failed to forward flashblock inclusion to metering"
+            );
+        }
+    }
+}
+
+/// Converts a flashblock to a FlashblockInclusion for the metering pipeline.
+fn flashblock_inclusion_from_flashblock(flashblock: &Flashblock) -> Option<FlashblockInclusion> {
+    if flashblock.diff.transactions.is_empty() {
+        return None;
+    }
+
+    let ordered_tx_hashes: Vec<B256> =
+        flashblock.diff.transactions.iter().map(|tx_bytes| keccak256(tx_bytes)).collect();
+
+    Some(FlashblockInclusion {
+        block_number: flashblock.metadata.block_number,
+        flashblock_index: flashblock.index,
+        ordered_tx_hashes,
+    })
+}
+
+/// Loads Kafka configuration from a properties file.
+fn load_kafka_config_from_file(
+    path: &str,
+) -> Result<Vec<(String, String)>, Box<dyn std::error::Error + Send + Sync>> {
+    let content = std::fs::read_to_string(path)?;
+    let mut props = Vec::new();
+    for line in content.lines() {
+        let line = line.trim();
+        if line.is_empty() || line.starts_with('#') {
+            continue;
+        }
+        if let Some((key, value)) = line.split_once('=') {
+            props.push((key.trim().to_string(), value.trim().to_string()));
+        }
+    }
+    Ok(props)
+}
+
 /// Helper struct that wires the custom RPC modules into the node builder.
 #[derive(Debug, Clone)]
 pub struct BaseRpcExtension {
@@ -52,38 +145,102 @@ impl BaseNodeExtension for BaseRpcExtension {
         let sequencer_rpc = self.sequencer_rpc.clone();
 
         builder.extend_rpc_modules(move |ctx| {
-            if metering.enabled {
-                info!(message = "Starting Metering RPC");
-
-                // Create priority fee estimator if configured
-                let estimator = if metering.kafka.is_some() {
-                    info!(message = "Enabling priority fee estimation");
-                    let cache = Arc::new(RwLock::new(MeteringCache::new(metering.cache_size)));
-                    let limits = ResourceLimits {
-                        gas_used: Some(metering.resource_limits.gas_limit),
-                        execution_time_us: Some(metering.resource_limits.execution_time_us as u128),
-                        state_root_time_us: metering
-                            .resource_limits
-                            .state_root_time_us
-                            .map(|v| v as u128),
-                        data_availability_bytes: Some(metering.resource_limits.da_bytes),
-                    };
-                    let default_fee = U256::from(metering.uncongested_priority_fee);
-                    let estimator = Arc::new(PriorityFeeEstimator::new(
-                        cache,
-                        metering.priority_fee_percentile,
-                        limits,
-                        default_fee,
-                        None, // Dynamic DA config not wired yet
-                    ));
-                    Some(estimator)
-                } else {
-                    None
+            // Set up metering runtime if enabled with Kafka
+            let metering_runtime = if metering.enabled && metering.kafka.is_some() {
+                info!(message = "Starting Metering RPC with priority fee estimation");
+
+                let cache = Arc::new(RwLock::new(MeteringCache::new(metering.cache_size)));
+                let limits = ResourceLimits {
+                    gas_used: Some(metering.resource_limits.gas_limit),
+                    execution_time_us: Some(metering.resource_limits.execution_time_us as u128),
+                    state_root_time_us: metering
+                        .resource_limits
+                        .state_root_time_us
+                        .map(|v| v as u128),
+                    data_availability_bytes: Some(metering.resource_limits.da_bytes),
                 };
+                let default_fee = U256::from(metering.uncongested_priority_fee);
+                let estimator = Arc::new(PriorityFeeEstimator::new(
+                    cache.clone(),
+                    metering.priority_fee_percentile,
+                    limits,
+                    default_fee,
+                    None, // Dynamic DA config not wired yet
+                ));
+
+                // Create channels for the annotator
+                let (tx_sender, tx_receiver) = mpsc::unbounded_channel::<MeteredTransaction>();
+                let (flashblock_sender, flashblock_receiver) =
+                    mpsc::unbounded_channel::<FlashblockInclusion>();
+
+                // Spawn the resource annotator
+                let annotator_cache = cache.clone();
+                tokio::spawn(async move {
+                    ResourceAnnotator::new(annotator_cache, tx_receiver, flashblock_receiver)
+                        .run()
+                        .await;
+                });
+
+                Some(MeteringRuntime { cache, estimator, tx_sender, flashblock_sender })
+            } else {
+                None
+            };
 
-                let metering_api = estimator.map_or_else(
+            // Spawn Kafka consumer if configured
+            if let (Some(runtime), Some(kafka_cfg)) = (&metering_runtime, &metering.kafka) {
+                info!(message = "Starting Kafka consumer for metering");
+
+                let mut client_config = ClientConfig::new();
+                client_config.set("bootstrap.servers", &kafka_cfg.brokers);
+                client_config.set("group.id", &kafka_cfg.group_id);
+                client_config.set("enable.partition.eof", "false");
+                client_config.set("session.timeout.ms", "6000");
+                client_config.set("enable.auto.commit", "true");
+                client_config.set("auto.offset.reset", "earliest");
+
+                if let Some(path) = kafka_cfg.properties_file.as_ref() {
+                    match load_kafka_config_from_file(path) {
+                        Ok(props) => {
+                            for (key, value) in props {
+                                client_config.set(key, value);
+                            }
+                        }
+                        Err(err) => {
+                            warn!(
+                                message = "Failed to load Kafka properties file",
+                                file = %path,
+                                %err
+                            );
+                        }
+                    }
+                }
+
+                let tx_sender = runtime.tx_sender.clone();
+                let topic = kafka_cfg.topic.clone();
+                tokio::spawn(async move {
+                    let config = KafkaBundleConsumerConfig { client_config, topic };
+
+                    match KafkaBundleConsumer::new(config, tx_sender) {
+                        Ok(consumer) => consumer.run().await,
+                        Err(err) => error!(
+                            target: "metering::kafka",
+                            %err,
+                            "Failed to initialize Kafka consumer"
+                        ),
+                    }
+                });
+            }
+
+            // Register metering RPC
+            if metering.enabled {
+                let metering_api = metering_runtime.as_ref().map_or_else(
                     || MeteringApiImpl::new(ctx.provider().clone()),
-                    |est| MeteringApiImpl::with_estimator(ctx.provider().clone(), est),
+                    |rt| {
+                        MeteringApiImpl::with_estimator(
+                            ctx.provider().clone(),
+                            rt.estimator.clone(),
+                        )
+                    },
                 );
                 ctx.modules.merge_configured(metering_api.into_rpc())?;
             }
@@ -107,7 +264,13 @@ impl BaseNodeExtension for BaseRpcExtension {
                     .clone();
                 fb.start();
 
-                let mut flashblocks_client = FlashblocksSubscriber::new(fb.clone(), ws_url);
+                // Create composite receiver that forwards to both flashblocks state and metering
+                let metering_sender =
+                    metering_runtime.as_ref().map(|rt| rt.flashblock_sender.clone());
+                let receiver =
+                    Arc::new(CompositeFlashblocksReceiver::new(fb.clone(), metering_sender));
+
+                let mut flashblocks_client = FlashblocksSubscriber::new(receiver, ws_url);
                 flashblocks_client.start();
 
                 let api_ext = EthApiExt::new(
