refactor: streamline action handling and improve queue management

Made-with: Cursor

Author: smtmfft

reqactor/src/action.rs

@@ -17,8 +17,7 @@ pub enum Action {
 impl Action {
     pub fn request_key(&self) -> &RequestKey {
         match self {
-            Action::Prove { request_key, .. } => request_key,
-            Action::Cancel { request_key, .. } => request_key,
+            Action::Prove { request_key, .. } | Action::Cancel { request_key, .. } => request_key,
         }
     }
 }

reqactor/src/actor.rs

@@ -1,6 +1,5 @@
 use std::{
     collections::HashMap,
-    // ops::DerefMut,
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc,
@@ -124,43 +123,56 @@ impl Actor {
     ) -> Result<StatusWithContext, String> {
         let pool_status_opt = self.pool_get_status(&request_key).await?;
 
-        // Return successful status if the request is already succeeded
-        if matches!(
-            pool_status_opt.as_ref().map(|s| s.status()),
-            Some(Status::Success { .. })
-        ) {
+        if pool_status_opt
+            .as_ref()
+            .is_some_and(|s| matches!(s.status(), Status::Success { .. }))
+        {
             return Ok(pool_status_opt.unwrap());
         }
 
-        // Mark the request as registered in the pool
         let status = StatusWithContext::new(Status::Registered, start_time);
-        if pool_status_opt.is_none() {
-            self.pool_add_new(request_key.clone(), request_entity.clone(), status.clone())
-                .await?;
-        } else {
-            self.pool_update_status(request_key.clone(), status.clone())
-                .await?;
-        }
+        self.ensure_pool_registered(&request_key, &request_entity, &status, pool_status_opt.is_none())
+            .await?;
+
+        let queue_result = {
+            let mut queue = self.queue.lock().await;
+            if queue.contains(&request_key) {
+                Ok(())
+            } else {
+                queue.add_pending(request_key.clone(), request_entity)
+            }
+        };
 
-        // Push the request into the queue and notify to start the action
-        let mut queue = self.queue.lock().await;
-        if !queue.contains(&request_key) {
-            match queue.add_pending(request_key.clone(), request_entity) {
-                Ok(()) => {
-                    self.notify.notify_one();
-                }
-                Err(error_msg) => {
-                    // If queue is at capacity, update the status to Failed
-                    let failed_status =
-                        StatusWithContext::new(Status::Failed { error: error_msg }, start_time);
-                    self.pool_update_status(request_key.clone(), failed_status.clone())
-                        .await?;
-                    return Ok(failed_status);
-                }
+        match queue_result {
+            Ok(()) => {
+                self.notify.notify_one();
+                Ok(status)
+            }
+            Err(error_msg) => {
+                let failed_status =
+                    StatusWithContext::new(Status::Failed { error: error_msg }, start_time);
+                self.pool_update_status(request_key, failed_status.clone())
+                    .await?;
+                Ok(failed_status)
             }
         }
+    }
 
-        return Ok(status);
+    /// Add or update pool entry so the request is marked as Registered.
+    async fn ensure_pool_registered(
+        &self,
+        request_key: &RequestKey,
+        request_entity: &RequestEntity,
+        status: &StatusWithContext,
+        is_new: bool,
+    ) -> Result<(), String> {
+        if is_new {
+            self.pool_add_new(request_key.clone(), request_entity.clone(), status.clone())
+                .await
+        } else {
+            self.pool_update_status(request_key.clone(), status.clone())
+                .await
+        }
     }
 
     pub async fn pause(&self) -> Result<(), String> {

reqactor/src/backend.rs

@@ -95,49 +95,14 @@ impl Backend {
             let handle = tokio::spawn(async move {
                 let _permit = permit;
 
-                let result = match request_entity {
-                    RequestEntity::SingleProof(_)
-                    | RequestEntity::Aggregation(_)
-                    | RequestEntity::BatchProof(_)
-                    | RequestEntity::GuestInput(_)
-                    | RequestEntity::BatchGuestInput(_) => Err(
-                        "legacy single-block and batch proving are removed; use Shasta only"
-                            .to_string(),
-                    ),
-                    RequestEntity::ShastaGuestInput(entity) => {
-                        do_generate_shasta_proposal_guest_input(
-                            &mut pool_,
-                            &chain_specs,
-                            request_key_.clone(),
-                            entity,
-                        )
-                        .await
-                    }
-                    RequestEntity::ShastaProof(entity) => {
-                        do_prove_shasta_proposal(
-                            &mut pool_,
-                            &chain_specs,
-                            request_key_.clone(),
-                            entity,
-                        )
-                        .await
-                    }
-                    RequestEntity::ShastaAggregation(entity) => {
-                        do_shasta_aggregation(&mut pool_, request_key_.clone(), entity).await
-                    }
-                };
-                let status = match result {
-                    Ok(proof) => {
-                        let proof_str = format!("{}", proof);
-                        tracing::info!(
-                            "Actor Backend successfully proved {request_key_}. Proof: {proof_str}"
-                        );
-                        Status::Success { proof }
-                    }
-                    Err(e) => Status::Failed {
-                        error: e.to_string(),
-                    },
-                };
+                let result = dispatch_proof_request(
+                    &mut pool_,
+                    &chain_specs,
+                    request_key_.clone(),
+                    request_entity,
+                )
+                .await;
+                let status = result_to_status(&result, &request_key_);
                 let _ = pool_.update_status(
                     request_key_.clone(),
                     StatusWithContext::new(status, chrono::Utc::now()),
@@ -151,19 +116,62 @@ impl Backend {
             tokio::spawn(async move {
                 if let Err(e) = handle.await {
                     tracing::error!("Actor thread errored while proving {request_key}: {e:?}");
-                    let status = Status::Failed {
-                        error: e.to_string(),
-                    };
-                    let _ = pool_.update_status(request_key.clone(), status.clone().into());
+                    let status =
+                        StatusWithContext::new(Status::Failed { error: e.to_string() }, chrono::Utc::now());
+                    let _ = pool_.update_status(request_key.clone(), status);
                 }
-
-                let _res = done_tx_.send(request_key.clone()).await;
+                let _ = done_tx_.send(request_key).await;
                 notifier_.notify_one();
             });
         }
     }
 }
 
+/// Dispatches the request to the appropriate proof handler.
+async fn dispatch_proof_request(
+    pool: &mut Pool,
+    chain_specs: &SupportedChainSpecs,
+    request_key: RequestKey,
+    request_entity: RequestEntity,
+) -> Result<raiko_lib::prover::Proof, String> {
+    match request_entity {
+        RequestEntity::SingleProof(_)
+        | RequestEntity::Aggregation(_)
+        | RequestEntity::BatchProof(_)
+        | RequestEntity::GuestInput(_)
+        | RequestEntity::BatchGuestInput(_) => Err(
+            "legacy single-block and batch proving are removed; use Shasta only".to_string(),
+        ),
+        RequestEntity::ShastaGuestInput(entity) => {
+            do_generate_shasta_proposal_guest_input(pool, chain_specs, request_key, entity).await
+        }
+        RequestEntity::ShastaProof(entity) => {
+            do_prove_shasta_proposal(pool, chain_specs, request_key, entity).await
+        }
+        RequestEntity::ShastaAggregation(entity) => {
+            do_shasta_aggregation(pool, request_key, entity).await
+        }
+    }
+}
+
+/// Converts proof result to pool status.
+fn result_to_status(
+    result: &Result<raiko_lib::prover::Proof, String>,
+    request_key: &RequestKey,
+) -> Status {
+    match result {
+        Ok(proof) => {
+            tracing::info!("Actor Backend successfully proved {request_key}. Proof: {proof}");
+            Status::Success {
+                proof: proof.clone(),
+            }
+        }
+        Err(e) => Status::Failed {
+            error: e.clone(),
+        },
+    }
+}
+
 async fn do_shasta_aggregation(
     pool: &mut dyn IdWrite,
     request_key: RequestKey,

reqactor/src/queue.rs

@@ -49,39 +49,44 @@ impl Queue {
         request_key: RequestKey,
         request_entity: RequestEntity,
     ) -> Result<(), String> {
-        // Check if queue is at capacity
         if self.is_at_capacity() {
             return Err("Reached the maximum queue size, please try again later".to_string());
         }
 
-        if self.queued_keys.insert(request_key.clone()) {
-            // Check priority and add to appropriate queue using pattern matching
-            match &request_key {
-                RequestKey::Aggregation(_) => {
-                    tracing::info!("Adding aggregation request to high priority queue");
-                    self.agg_queue.push_back((request_key, request_entity));
-                }
-                RequestKey::BatchProof(_) => {
-                    tracing::info!("Adding batch proof request to medium priority queue");
-                    self.batch_queue.push_back((request_key, request_entity));
-                }
-                _ => {
-                    self.preflight_queue
-                        .push_back((request_key, request_entity));
-                }
-            }
+        if !self.queued_keys.insert(request_key.clone()) {
+            return Ok(());
         }
+
+        let queue = self.priority_queue_for(&request_key);
+        queue.push_back((request_key, request_entity));
         Ok(())
     }
 
-    /// Attempts to move a request from either the high, medium or low priority queue into the in-flight set
-    /// and starts processing it. High priority requests are processed first.
+    /// Returns the appropriate queue for the given request key (agg > batch > preflight).
+    fn priority_queue_for(
+        &mut self,
+        request_key: &RequestKey,
+    ) -> &mut VecDeque<(RequestKey, RequestEntity)> {
+        match request_key {
+            RequestKey::Aggregation(_) => {
+                tracing::info!("Adding aggregation request to high priority queue");
+                &mut self.agg_queue
+            }
+            RequestKey::BatchProof(_) => {
+                tracing::info!("Adding batch proof request to medium priority queue");
+                &mut self.batch_queue
+            }
+            _ => &mut self.preflight_queue,
+        }
+    }
+
+    /// Pops the next request (agg > batch > preflight) and marks it in-progress.
     pub fn try_next(&mut self) -> Option<(RequestKey, RequestEntity)> {
-        let (request_key, request_entity) = self.agg_queue.pop_front().or_else(|| {
-            self.batch_queue
-                .pop_front()
-                .or_else(|| self.preflight_queue.pop_front())
-        })?;
+        let (request_key, request_entity) = self
+            .agg_queue
+            .pop_front()
+            .or_else(|| self.batch_queue.pop_front())
+            .or_else(|| self.preflight_queue.pop_front())?;
 
         self.working_in_progress.insert(request_key.clone());
         Some((request_key, request_entity))

reqpool/src/config.rs

@@ -1,7 +1,7 @@
 use serde::{Deserialize, Serialize};
 
+/// Configuration for the request pool (Redis or in-memory backend).
 #[derive(Debug, Clone, Serialize, Deserialize)]
-/// The configuration for the redis-backend request pool
 pub struct RedisPoolConfig {
     /// The URL of the Redis database, e.g. "redis://localhost:6379"
     pub redis_url: String,

reqpool/src/memory_backend.rs

@@ -14,6 +14,15 @@ use lru::LruCache;
 type SingleStorage = Arc<Mutex<LruCache<Value, Value>>>;
 type GlobalStorage = Mutex<HashMap<String, SingleStorage>>;
 
+const DEFAULT_MEMORY_BACKEND_CAPACITY: usize = 512;
+
+fn parse_memory_backend_capacity() -> usize {
+    std::env::var("MEMORY_BACKEND_SIZE")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(DEFAULT_MEMORY_BACKEND_CAPACITY)
+}
+
 lazy_static! {
     // #{redis_url => single_storage}
     //
@@ -29,17 +38,13 @@ pub struct MemoryBackend {
 impl MemoryBackend {
     pub fn new(redis_url: String) -> Self {
         let mut global = GLOBAL_STORAGE.lock().unwrap();
-
-        let mem_capacity = std::env::var("MEMORY_BACKEND_SIZE")
-            .unwrap_or("2048".to_string())
-            .parse::<usize>()
-            .unwrap_or_else(|_| 2048);
+        let capacity = parse_memory_backend_capacity();
         Self {
             storage: global
                 .entry(redis_url)
                 .or_insert_with(|| {
                     Arc::new(Mutex::new(LruCache::new(
-                        NonZeroUsize::new(mem_capacity).unwrap(),
+                        NonZeroUsize::new(capacity).unwrap(),
                     )))
                 })
                 .clone(),