use new rblib checkpoint context to track flashblocks

Author: akundaz

src/bundle.rs

@@ -9,7 +9,7 @@
 //! - The `eth_sendBundle` input parameters and their validation.
 
 use {
-	crate::platform::Flashblocks,
+	crate::{platform::Flashblocks, state::FlashblockNumber},
 	core::convert::Infallible,
 	rblib::{
 		alloy::{
@@ -126,7 +126,7 @@ impl Bundle<Flashblocks> for FlashblocksBundle {
 	fn is_eligible(
 		&self,
 		block: &BlockContext<Flashblocks>,
-		_ctx: &(),
+		_ctx: &FlashblockNumber,
 	) -> Eligibility {
 		if self.txs.is_empty() {
 			// empty bundles are never eligible

src/lib.rs

@@ -1,5 +1,6 @@
 mod bundle;
 mod platform;
 mod primitives;
+mod state;
 
-pub use {bundle::*, platform::*, primitives::*};
+pub use {bundle::*, platform::*, primitives::*, state::*};

src/limits.rs

@@ -5,7 +5,10 @@
 //! flashblock partitioning logic.
 
 use {
-	crate::{Flashblocks, state::FlashblockNumber},
+	crate::{
+		Flashblocks,
+		state::{FlashblockNumber, TargetFlashblocks},
+	},
 	core::time::Duration,
 	rblib::{alloy::consensus::BlockHeader, prelude::*},
 	std::sync::{Arc, Mutex},
@@ -33,7 +36,7 @@ pub struct FlashblockState {
 	current_block: Option<u64>,
 	/// Current flashblock number. Used to check if we're on the first
 	/// flashblock or to adjust the target number of flashblocks for a block.
-	flashblock_number: Arc<FlashblockNumber>,
+	target_flashblocks: Arc<TargetFlashblocks>,
 	/// Duration for the first flashblock, which may be shortened to absorb
 	/// timing variance.
 	first_flashblock_interval: Duration,
@@ -43,20 +46,16 @@ pub struct FlashblockState {
 }
 
 impl FlashblockState {
-	fn current_gas_limit(&self) -> u64 {
+	fn current_gas_limit(&self, flashblock_number: &FlashblockNumber) -> u64 {
 		self
 			.gas_per_flashblock
-			.saturating_mul(self.flashblock_number.current())
+			.saturating_mul(flashblock_number.current())
 	}
 }
 
 impl FlashblockLimits {
-	pub fn new(
-		interval: Duration,
-		flashblock_number: Arc<FlashblockNumber>,
-	) -> Self {
+	pub fn new(interval: Duration) -> Self {
 		let state = FlashblockState {
-			flashblock_number,
 			..Default::default()
 		};
 		FlashblockLimits {
@@ -97,10 +96,7 @@ impl FlashblockLimits {
 				.unwrap_or(enclosing.gas_limit);
 			state.current_block = Some(payload.block().number());
 			state.first_flashblock_interval = first_flashblock_interval;
-			state.flashblock_number.reset_current_flashblock();
-			state
-				.flashblock_number
-				.set_target_flashblocks(target_flashblock);
+			state.target_flashblocks.set(target_flashblock);
 		}
 	}
 
@@ -111,18 +107,19 @@ impl FlashblockLimits {
 	pub fn get_limits(
 		&self,
 		enclosing: &Limits<Flashblocks>,
+		flashblock_number: &FlashblockNumber,
 	) -> Limits<Flashblocks> {
 		let state = self.state.lock().expect("mutex is not poisoned");
 		// If flashblock number == 1, we're building the first flashblock
-		let deadline = if state.flashblock_number.current() == 1 {
+		let deadline = if flashblock_number.current() == 1 {
 			state.first_flashblock_interval
 		} else {
 			self.interval
 		};
 
 		enclosing
 			.with_deadline(deadline)
-			.with_gas_limit(state.current_gas_limit())
+			.with_gas_limit(state.current_gas_limit(flashblock_number))
 	}
 
 	/// Calculates the number of flashblocks and first flashblock interval for
@@ -153,27 +150,30 @@ impl ScopedLimits<Flashblocks> for FlashblockLimits {
 		payload: &Checkpoint<Flashblocks>,
 		enclosing: &Limits<Flashblocks>,
 	) -> Limits<Flashblocks> {
+		let flashblock_number = payload.context();
 		// Check the state and reset if we started building next block
 		self.update_state(payload, enclosing);
 
-		let limits = self.get_limits(enclosing);
+		let limits = self.get_limits(enclosing, flashblock_number);
 
 		let state = self.state.lock().expect("mutex is not poisoned");
-		if state.flashblock_number.in_bounds() {
+		let flashblock_number = payload.context();
+		if flashblock_number.current() <= state.target_flashblocks.get() {
 			let gas_used = payload.cumulative_gas_used();
 			let remaining_gas = enclosing.gas_limit.saturating_sub(gas_used);
 			tracing::info!(
 				"Creating flashblocks limits: {}, payload txs: {}, gas used: {} \
 				 ({}%), gas_remaining: {} ({}%), next_block_gas_limit: {} ({}%), gas \
 				 per block: {} ({}%), remaining_time: {}ms, gas_limit: {}",
-				state.flashblock_number,
+				flashblock_number,
 				payload.history().transactions().count(),
 				gas_used,
 				(gas_used * 100 / enclosing.gas_limit),
 				remaining_gas,
 				(remaining_gas * 100 / enclosing.gas_limit),
-				state.current_gas_limit(),
-				(state.current_gas_limit() * 100 / enclosing.gas_limit),
+				state.current_gas_limit(flashblock_number),
+				(state.current_gas_limit(flashblock_number) * 100
+					/ enclosing.gas_limit),
 				state.gas_per_flashblock,
 				(state.gas_per_flashblock * 100 / enclosing.gas_limit),
 				limits.deadline.expect("deadline is set").as_millis(),

src/main.rs

@@ -6,7 +6,7 @@ use {
 		publish::{PublishFlashblock, WebSocketSink},
 		rpc::TransactionStatusRpc,
 		signer::BuilderSigner,
-		state::FlashblockNumber,
+		state::TargetFlashblocks,
 		stop::BreakAfterMaxFlashblocks,
 	},
 	platform::Flashblocks,
@@ -104,9 +104,7 @@ fn build_pipeline(
 		.clone()
 		.unwrap_or(BuilderSigner::random());
 
-	// Multiple steps need to access flashblock number state, so we need to
-	// initialize it outside
-	let flashblock_number = Arc::new(FlashblockNumber::new());
+	let target_flashblocks = Arc::new(TargetFlashblocks::new());
 
 	let pipeline = Pipeline::<Flashblocks>::named("top")
 		.with_step(OptimismPrologue)
@@ -143,17 +141,13 @@ fn build_pipeline(
 								)
 								.with_epilogue(PublishFlashblock::new(
 									ws.clone(),
-									flashblock_number.clone(),
 									cli_args.flashblocks_args.calculate_state_root,
 								))
-								.with_limits(FlashblockLimits::new(
-									interval,
-									flashblock_number.clone(),
-								)),
+								.with_limits(FlashblockLimits::new(interval)),
 						)
 						.with_step(BreakAfterDeadline),
 				)
-				.with_step(BreakAfterMaxFlashblocks::new(flashblock_number)),
+				.with_step(BreakAfterMaxFlashblocks::new(target_flashblocks)),
 		)
 		.with_limits(Scaled::default().deadline(total_building_time));
 

src/platform.rs

@@ -1,5 +1,5 @@
 use {
-	crate::bundle::FlashblocksBundle,
+	crate::{bundle::FlashblocksBundle, state::FlashblockNumber},
 	rblib::{prelude::*, reth::providers::StateProvider},
 	serde::{Deserialize, Serialize},
 	std::sync::Arc,
@@ -19,7 +19,7 @@ pub struct Flashblocks;
 
 impl Platform for Flashblocks {
 	type Bundle = FlashblocksBundle;
-	type CheckpointContext = ();
+	type CheckpointContext = FlashblockNumber;
 	type DefaultLimits = types::DefaultLimits<Optimism>;
 	type EvmConfig = types::EvmConfig<Optimism>;
 	type ExtraLimits = types::ExtraLimits<Optimism>;

src/publish.rs

@@ -10,7 +10,7 @@
 //!   building jobs.
 
 use {
-	crate::{Flashblocks, primitives::*, state::FlashblockNumber},
+	crate::{Flashblocks, primitives::*},
 	atomic_time::AtomicOptionInstant,
 	core::{net::SocketAddr, sync::atomic::Ordering},
 	futures::{SinkExt, StreamExt},
@@ -50,10 +50,6 @@ pub struct PublishFlashblock {
 	/// subscribers.
 	sink: Arc<WebSocketSink>,
 
-	/// Reference to the current flashblock number within the payload job.
-	/// Once we build a flashblock we increment this.
-	flashblock_number: Arc<FlashblockNumber>,
-
 	/// Set once at the begining of the payload job, captures immutable
 	/// information about the payload that is being built. This info is derived
 	/// from the payload attributes parameter on the FCU from the EL node.
@@ -70,13 +66,11 @@ pub struct PublishFlashblock {
 impl PublishFlashblock {
 	pub fn new(
 		sink: Arc<WebSocketSink>,
-		flashblock_number: Arc<FlashblockNumber>,
 		// TODO: Will be implemented later
 		_calculate_state_root: bool,
 	) -> Self {
 		Self {
 			sink,
-			flashblock_number,
 			block_base: RwLock::new(None),
 			metrics: Metrics::default(),
 			times: Times::default(),
@@ -120,10 +114,10 @@ impl Step<Flashblocks> for PublishFlashblock {
 				.expect("withdrawals_root is present"),
 		};
 
+		let flashblock_number = payload.context();
+
 		// Get 0-index to use in flashblock
-		let index = self.flashblock_number.index();
-		// Increment flashblock number since we've built the flashblock
-		self.flashblock_number.advance();
+		let index = flashblock_number.index();
 
 		// Push the contents of the payload
 		if let Err(e) = self.sink.publish(&FlashblocksPayloadV1 {
@@ -152,10 +146,13 @@ impl Step<Flashblocks> for PublishFlashblock {
 		self.times.on_published_block(&self.metrics);
 		self.capture_payload_metrics(&this_block_span);
 
+		// Increment flashblock number since we've built the flashblock
+		let next_flashblock_number = flashblock_number.advance();
+
 		// Place a barrier after each published flashblock to freeze the contents
 		// of the payload up to this point, since this becomes a publicly committed
 		// state.
-		ControlFlow::Ok(payload.barrier())
+		ControlFlow::Ok(payload.barrier_with_context(next_flashblock_number))
 	}
 
 	/// Before the payload job starts prepare the contents of the
@@ -200,9 +197,6 @@ impl Step<Flashblocks> for PublishFlashblock {
 	) -> Result<(), PayloadBuilderError> {
 		self.times.on_job_ended(&self.metrics);
 
-		// Reset current flashblock number since we're done with the whole block
-		let count = self.flashblock_number.reset_current_flashblock();
-		self.metrics.blocks_per_payload_job.record(count as f64);
 		*self.block_base.write() = None;
 
 		Ok(())
@@ -232,7 +226,10 @@ impl PublishFlashblock {
 		payload: &Checkpoint<Flashblocks>,
 	) -> Span<Flashblocks> {
 		// If we haven't published flashblock return whole history
-		payload.history()
+		let previous_flashblock_number = payload.context().clone();
+		payload
+			.history_since_last_context(&previous_flashblock_number)
+			.unwrap_or(payload.history())
 	}
 
 	/// Called for each flashblock to capture metrics about the produced
@@ -273,9 +270,6 @@ struct Metrics {
 	/// Histogram of the number of bundles per flashblock.
 	pub bundles_per_block: Histogram,
 
-	/// Histogram of flashblocks per job.
-	pub blocks_per_payload_job: Histogram,
-
 	/// The time interval flashblocks within one block.
 	pub intra_block_interval: Histogram,
 

src/state.rs

@@ -1,68 +1,56 @@
-use std::{
-	fmt::Display,
-	sync::atomic::{AtomicU64, Ordering},
+use {
+	crate::Flashblocks,
+	rblib::prelude::CheckpointContext,
+	std::{
+		fmt::Display,
+		sync::atomic::{AtomicU64, Ordering},
+	},
 };
 
-#[derive(Debug)]
-pub struct FlashblockNumber {
-	/// Current flashblock number (1-indexed).
-	current_flashblock: AtomicU64,
-	/// Number of flashblocks we're targeting to build for this block.
-	target_flashblocks: AtomicU64,
-}
-
-impl FlashblockNumber {
-	pub fn new() -> Self {
-		Self {
-			current_flashblock: AtomicU64::new(1),
-			target_flashblocks: AtomicU64::new(0),
-		}
-	}
+/// Current flashblock number (1-indexed).
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct FlashblockNumber(u64);
 
-	pub fn current(&self) -> u64 {
-		self.current_flashblock.load(Ordering::Relaxed)
-	}
+/// Number of flashblocks we're targeting to build for this block.
+#[derive(Debug, Default)]
+pub struct TargetFlashblocks(AtomicU64);
 
-	pub fn max(&self) -> u64 {
-		self.target_flashblocks.load(Ordering::Relaxed)
-	}
+// TODO fix the default impl
 
+impl FlashblockNumber {
 	/// Returns current flashblock in 0-index format
 	pub fn index(&self) -> u64 {
-		self
-			.current_flashblock
-			.load(Ordering::Relaxed)
-			.saturating_sub(1)
+		self.0 - 1
 	}
 
-	pub fn advance(&self) -> u64 {
-		self.current_flashblock.fetch_add(1, Ordering::Relaxed)
+	pub fn current(&self) -> u64 {
+		self.0
 	}
 
-	pub fn in_bounds(&self) -> bool {
-		self.current_flashblock.load(Ordering::Relaxed)
-			<= self.target_flashblocks.load(Ordering::Relaxed)
+	#[must_use]
+	pub fn advance(&self) -> Self {
+		Self(self.0 + 1)
 	}
+}
 
-	pub fn set_target_flashblocks(&self, target_flashblock: u64) {
-		self
-			.target_flashblocks
-			.store(target_flashblock, Ordering::Relaxed);
+impl Display for FlashblockNumber {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.current())
 	}
+}
+
+impl CheckpointContext<Flashblocks> for FlashblockNumber {}
 
-	pub fn reset_current_flashblock(&self) -> u64 {
-		self.current_flashblock.swap(1, Ordering::Relaxed)
+impl TargetFlashblocks {
+	pub fn new() -> Self {
+		Self(AtomicU64::default())
 	}
-}
 
-impl Default for FlashblockNumber {
-	fn default() -> Self {
-		Self::new()
+	pub fn get(&self) -> u64 {
+		self.0.load(Ordering::Relaxed)
 	}
-}
 
-impl Display for FlashblockNumber {
-	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		write!(f, "{}/{}", self.current(), self.max())
+	pub fn set(&self, val: u64) {
+		self.0.store(val, Ordering::Relaxed);
 	}
 }