Mevshield fast runtime

node/Cargo.toml

@@ -137,6 +137,7 @@ subtensor-custom-rpc-runtime-api = { workspace = true, features = ["std"] }
 pallet-subtensor-swap-rpc = { workspace = true, features = ["std"] }
 pallet-subtensor-swap-runtime-api = { workspace = true, features = ["std"] }
 subtensor-macros.workspace = true
+runtime-common.workspace = true
 
 [build-dependencies]
 substrate-build-script-utils.workspace = true

node/src/mev_shield/author.rs

@@ -12,6 +12,7 @@ use sp_runtime::{AccountId32, KeyTypeId};
 use std::sync::{Arc, Mutex};
 use subtensor_macros::freeze_struct;
 use tokio::time::sleep;
+use runtime_common::prod_or_fast;
 
 /// Parameters controlling time windows inside the slot.
 #[freeze_struct("5c7ce101b36950de")]
@@ -130,7 +131,7 @@ pub fn aead_decrypt(
 const AURA_KEY_TYPE: KeyTypeId = KeyTypeId(*b"aura");
 
 /// Start background tasks:
-///  - per-slot ML‑KEM key rotation
+///  - ML‑KEM key rotation (every prod_or_fast!(1, 4) blocks)
 ///  - at ~announce_at_ms announce the next key bytes on chain,
 pub fn spawn_author_tasks<B, C, Pool>(
     task_spawner: &sc_service::SpawnTaskHandle,
@@ -204,6 +205,12 @@ where
 
             let mut import_stream = client_clone.import_notification_stream();
 
+            // 🧮 Number of blocks for which a key is reused:
+            //  - prod: rotate every block (1)
+            //  - fast-runtime / fast-blocks: rotate every 4 blocks
+            let blocks_per_key: u32 = prod_or_fast!(1u32, 4u32);
+            let mut blocks_since_rotation: u32 = 0;
+
             while let Some(notif) = import_stream.next().await {
                 // Only act on blocks that this node authored.
                 if notif.origin != BlockOrigin::Own {
@@ -223,15 +230,16 @@ where
 
                 log::debug!(
                     target: "mev-shield",
-                    "Slot start (local author): (pk sizes: curr={curr_pk_len}B, next={next_pk_len}B)",
+                    "Slot start (local author): (pk sizes: curr={curr_pk_len}B, next={next_pk_len}B) \
+                     blocks_since_rotation={blocks_since_rotation} blocks_per_key={blocks_per_key}",
                 );
 
                 // Wait until the announce window in this slot.
                 if announce_at_ms > 0 {
                     sleep(std::time::Duration::from_millis(announce_at_ms)).await;
                 }
 
-                // Read the next key we intend to use for the following block.
+                // Read the next key we intend to use for the *next epoch*.
                 let next_pk = match ctx_clone.keys.lock() {
                     Ok(k) => k.next_pk.clone(),
                     Err(e) => {
@@ -282,21 +290,34 @@ where
                     sleep(std::time::Duration::from_millis(tail_ms)).await;
                 }
 
-                // Roll keys for the next block.
-                match ctx_clone.keys.lock() {
-                    Ok(mut k) => {
-                        k.roll_for_next_slot();
-                        log::debug!(
-                            target: "mev-shield",
-                            "Rolled ML-KEM key at slot boundary",
-                        );
-                    }
-                    Err(e) => {
-                        log::debug!(
-                            target: "mev-shield",
-                            "spawn_author_tasks: failed to lock ShieldKeys for roll_for_next_slot: {e:?}",
-                        );
+                // Rotate keys only every `blocks_per_key` blocks.
+                blocks_since_rotation = blocks_since_rotation.saturating_add(1);
+
+                if blocks_since_rotation >= blocks_per_key {
+                    match ctx_clone.keys.lock() {
+                        Ok(mut k) => {
+                            k.roll_for_next_slot();
+                            log::debug!(
+                                target: "mev-shield",
+                                "Rolled ML-KEM key at key-epoch boundary (blocks_per_key={blocks_per_key})",
+                            );
+                        }
+                        Err(e) => {
+                            log::debug!(
+                                target: "mev-shield",
+                                "spawn_author_tasks: failed to lock ShieldKeys for roll_for_next_slot: {e:?}",
+                            );
+                        }
                     }
+
+                    // Reset the epoch counter after rotation.
+                    blocks_since_rotation = 0;
+                } else {
+                    log::debug!(
+                        target: "mev-shield",
+                        "Within key epoch; not rotating ML-KEM key this block \
+                         (blocks_since_rotation={blocks_since_rotation}, blocks_per_key={blocks_per_key})",
+                    );
                 }
             }
         },

runtime/src/lib.rs

@@ -43,7 +43,7 @@ use pallet_subtensor::{CommitmentsInterface, ProxyInterface};
 use pallet_subtensor_proxy as pallet_proxy;
 use pallet_subtensor_swap_runtime_api::SimSwapResult;
 use pallet_subtensor_utility as pallet_utility;
-use runtime_common::prod_or_fast;
+pub use runtime_common::prod_or_fast;
 use sp_api::impl_runtime_apis;
 use sp_consensus_aura::sr25519::AuthorityId as AuraId;
 use sp_consensus_babe::BabeConfiguration;