feat: get last max fee limit message

Author: MarcosNicolau

crates/batcher/src/lib.rs

@@ -33,9 +33,9 @@ use aligned_sdk::common::constants::{
     RESPOND_TO_TASK_FEE_LIMIT_PERCENTAGE_MULTIPLIER,
 };
 use aligned_sdk::common::types::{
-    ClientMessage, GetNonceResponseMessage, NoncedVerificationData, ProofInvalidReason,
-    ProvingSystemId, SubmitProofMessage, SubmitProofResponseMessage, VerificationCommitmentBatch,
-    VerificationData, VerificationDataCommitment,
+    ClientMessage, GetLastMaxFeeResponseMessage, GetNonceResponseMessage, NoncedVerificationData,
+    ProofInvalidReason, ProvingSystemId, SubmitProofMessage, SubmitProofResponseMessage,
+    VerificationCommitmentBatch, VerificationData, VerificationDataCommitment,
 };
 
 use aws_sdk_s3::client::Client as S3Client;
@@ -916,6 +916,11 @@ impl Batcher {
                     .handle_submit_proof_msg(msg, ws_conn_sink)
                     .await
             }
+            ClientMessage::GetLastMaxFee(address) => {
+                self.clone()
+                    .handle_get_last_max_fee(address, ws_conn_sink)
+                    .await
+            }
         }
     }
 
@@ -1004,6 +1009,83 @@ impl Batcher {
         Ok(())
     }
 
+    async fn handle_get_last_max_fee(
+        self: Arc<Self>,
+        mut address: Address,
+        ws_conn_sink: WsMessageSink,
+    ) -> Result<(), Error> {
+        // If the address is not paying, we will return the last max fee of the aligned_payment_address
+        if !self.has_to_pay(&address) {
+            info!("Handling nonpaying message");
+            let Some(non_paying_config) = self.non_paying_config.as_ref() else {
+                warn!(
+                    "There isn't a non-paying configuration loaded. This message will be ignored"
+                );
+                send_message(
+                    ws_conn_sink.clone(),
+                    GetLastMaxFeeResponseMessage::InvalidRequest(
+                        "There isn't a non-paying configuration loaded.".to_string(),
+                    ),
+                )
+                .await;
+                return Ok(());
+            };
+            let replacement_addr = non_paying_config.replacement.address();
+            address = replacement_addr;
+        }
+
+        let user_states_guard = match timeout(MESSAGE_HANDLER_LOCK_TIMEOUT, self.user_states.read())
+            .await
+        {
+            Ok(guard) => guard,
+            Err(_) => {
+                warn!("User states read lock acquisition timed out in handle_get_last_max_fee_for_address_msg");
+                self.metrics.inc_message_handler_user_states_lock_timeouts();
+                send_message(ws_conn_sink, GetLastMaxFeeResponseMessage::ServerBusy).await;
+                return Ok(());
+            }
+        };
+
+        let Some(usr_ref) = user_states_guard.get(&address).cloned() else {
+            drop(user_states_guard);
+            // if the user isn't in the cache, return zero
+            send_message(
+                ws_conn_sink.clone(),
+                GetLastMaxFeeResponseMessage::LastMaxFee(U256::MAX),
+            )
+            .await;
+            return Ok(());
+        };
+
+        let Some(usr_lock) = self
+            .try_user_lock_with_timeout(address, usr_ref.lock())
+            .await
+        else {
+            send_message(
+                ws_conn_sink.clone(),
+                GetLastMaxFeeResponseMessage::ServerBusy,
+            )
+            .await;
+            return Ok(());
+        };
+
+        let proofs_in_queue = usr_lock.proofs_in_batch;
+        let max_fee = if proofs_in_queue > 0 {
+            usr_lock.last_max_fee_limit
+        } else {
+            U256::MAX
+        };
+        drop(usr_lock);
+
+        send_message(
+            ws_conn_sink.clone(),
+            GetLastMaxFeeResponseMessage::LastMaxFee(max_fee),
+        )
+        .await;
+
+        Ok(())
+    }
+
     /// Returns the Aligned-funded address that will be used to pay for proofs when users don't need to pay themselves.
     /// This function assumes that the non-paying configuration is set.
     fn aligned_payment_address(&self) -> Address {

crates/sdk/src/common/types.rs

@@ -295,6 +295,7 @@ pub enum ClientMessage {
     // Needs to be wrapped in box as the message is 3x bigger than the others
     // see https://rust-lang.github.io/rust-clippy/master/index.html#large_enum_variant
     SubmitProof(Box<SubmitProofMessage>),
+    GetLastMaxFee(Address),
 }
 
 impl Display for ClientMessage {
@@ -303,6 +304,7 @@ impl Display for ClientMessage {
         match self {
             ClientMessage::GetNonceForAddress(_) => write!(f, "GetNonceForAddress"),
             ClientMessage::SubmitProof(_) => write!(f, "SubmitProof"),
+            ClientMessage::GetLastMaxFee(_) => write!(f, "GetLastMaxFee"),
         }
     }
 }
@@ -464,6 +466,13 @@ pub enum GetNonceResponseMessage {
     ServerBusy,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum GetLastMaxFeeResponseMessage {
+    LastMaxFee(U256),
+    InvalidRequest(String),
+    ServerBusy,
+}
+
 #[derive(Debug, Clone)]
 pub enum Network {
     Devnet,