Limit retries on dispense (#63)

* limit retries and add better logging

* make in progress explicit

* Minimize chance of miss behaviour

---------

Co-authored-by: xgreenx <xgreenx9999@gmail.com>

src/dispense_tracker.rs

@@ -1,45 +1,30 @@
-use std::{
- cmp::Ordering,
- collections::{BinaryHeap, HashMap, HashSet},
-};
+use std::collections::BTreeMap;
+use std::collections::{HashMap, HashSet};
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use fuel_types::Address;
 
-#[derive(Debug, Eq, PartialEq)]
-pub struct Entry {
- address: Address,
- timestamp: u64,
-}
-
-impl Ord for Entry {
- fn cmp(&self, other: &Self) -> Ordering {
- self.timestamp.cmp(&other.timestamp)
- }
-}
-
-impl PartialOrd for Entry {
- fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
- Some(self.cmp(other))
- }
-}
-
 pub trait Clock: std::fmt::Debug + Send + Sync {
  fn now(&self) -> u64;
 }
 
 #[derive(Debug)]
-pub struct TokioTime {}
+pub struct StdTime {}
 
-impl Clock for TokioTime {
+impl Clock for StdTime {
  fn now(&self) -> u64 {
- tokio::time::Instant::now().elapsed().as_secs()
+ let start = SystemTime::now();
+ let since_the_epoch = start
+ .duration_since(UNIX_EPOCH)
+ .expect("Time went backwards");
+ since_the_epoch.as_secs()
  }
 }
 
 #[derive(Debug)]
 pub struct DispenseTracker {
  tracked: HashMap<Address, u64>,
- queue: BinaryHeap<Entry>,
+ queue: BTreeMap<u64, Vec<Address>>,
  in_progress: HashSet<Address>,
  clock: Box<dyn Clock>,
 }
@@ -48,9 +33,9 @@ impl Default for DispenseTracker {
  fn default() -> Self {
  Self {
  tracked: HashMap::default(),
- queue: BinaryHeap::default(),
+ queue: Default::default(),
  in_progress: HashSet::default(),
- clock: Box::new(TokioTime {}),
+ clock: Box::new(StdTime {}),
  }
  }
 }
@@ -59,7 +44,7 @@ impl DispenseTracker {
  pub fn new(clock: impl Clock + 'static) -> Self {
  Self {
  tracked: HashMap::new(),
- queue: BinaryHeap::new(),
+ queue: Default::default(),
  in_progress: HashSet::new(),
  clock: Box::new(clock),
  }
@@ -70,7 +55,7 @@ impl DispenseTracker {
 
  let timestamp = self.clock.now();
  self.tracked.insert(address, timestamp);
- self.queue.push(Entry { address, timestamp });
+ self.queue.entry(timestamp).or_default().push(address);
  }
 
  pub fn mark_in_progress(&mut self, address: Address) {
@@ -84,17 +69,24 @@ impl DispenseTracker {
  pub fn evict_expired_entries(&mut self, eviction_duration: u64) {
  let now = self.clock.now();
 
- while let Some(oldest_entry) = self.queue.peek() {
- if now - oldest_entry.timestamp > eviction_duration {
- let removed_entry = self.queue.pop().unwrap();
- self.tracked.remove(&removed_entry.address);
+ while let Some(oldest_entry) = self.queue.first_entry() {
+ if now - oldest_entry.key() > eviction_duration {
+ let (_, addresses) = oldest_entry.remove_entry();
+
+ for address in addresses {
+ self.tracked.remove(&address);
+ }
  } else {
  break;
  }
  }
  }
 
  pub fn has_tracked(&self, address: &Address) -> bool {
- self.tracked.get(address).is_some() || self.in_progress.contains(address)
+ self.tracked.get(address).is_some()
+ }
+
+ pub fn is_in_progress(&self, address: &Address) -> bool {
+ self.in_progress.contains(address)
  }
 }

src/lib.rs

@@ -42,7 +42,7 @@ mod dispense_tracker;
 mod recaptcha;
 mod routes;
 
-pub use dispense_tracker::{Clock, TokioTime};
+pub use dispense_tracker::{Clock, StdTime};
 pub use routes::THE_BIGGEST_AMOUNT;
 
 #[derive(Debug)]

src/main.rs

@@ -1,11 +1,11 @@
-use fuel_faucet::{config::Config, start_server, TokioTime};
+use fuel_faucet::{config::Config, start_server, StdTime};
 use tracing_subscriber::EnvFilter;
 
 #[tokio::main]
 async fn main() {
  let config = Config::default();
  init_logger(&config);
- let clock = TokioTime {};
+ let clock = StdTime {};
  let (_, task) = start_server(config, clock).await;
  let _ = task.await.unwrap();
 }

src/models.rs

@@ -1,3 +1,5 @@
+use std::fmt::{self, Display, Formatter};
+
 use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
 
@@ -24,3 +26,11 @@ pub struct DispenseError {
  pub status: StatusCode,
  pub error: String,
 }
+
+impl Display for DispenseError {
+ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ write!(f, "{self:?}")
+ }
+}
+
+impl std::error::Error for DispenseError {}

src/routes.rs

@@ -135,6 +135,13 @@ fn check_and_mark_dispense_limit(
  ));
  }
 
+ if tracker.is_in_progress(&address) {
+ return Err(error(
+ "Account is already in the process of receiving assets".to_string(),
+ StatusCode::TOO_MANY_REQUESTS,
+ ));
+ }
+
  tracker.mark_in_progress(address);
  Ok(())
 }
@@ -238,27 +245,37 @@ pub async fn dispense_tokens(
  }
 
  check_and_mark_dispense_limit(&dispense_tracker, address, config.dispense_limit_interval)?;
- let cleanup = || {
+
+ struct CleanUpper<Fn>(Fn)
+ where
+ Fn: FnMut();
+
+ impl<Fn> Drop for CleanUpper<Fn>
+ where
+ Fn: FnMut(),
+ {
+ fn drop(&mut self) {
+ self.0();
+ }
+ }
+
+ // We want to remove the address from `in_progress` regardless of the outcome of the transaction.
+ let _cleanup = CleanUpper(|| {
  dispense_tracker
  .lock()
  .unwrap()
  .remove_in_progress(&address);
- };
+ });
 
  let provider = wallet.provider().expect("client provider");
- let mut tx_id;
 
- loop {
+ let mut tx_id = None;
+ for _ in 0..5 {
  let mut guard = state.lock().await;
  let coin_output = if let Some(previous_coin_output) = &guard.last_output {
  *previous_coin_output
  } else {
- get_coin_output(&wallet, config.dispense_amount)
- .await
- .map_err(|e| {
- cleanup();
- e
- })?
+ get_coin_output(&wallet, config.dispense_amount).await?
  };
 
  let coin_type = CoinType::Coin(Coin {
@@ -296,55 +313,63 @@ pub async fn dispense_tokens(
  .fee_checked_from_tx(&network_config.network_info.consensus_parameters)
  .expect("Should be able to calculate fee");
 
- tx_id = script.id(network_config.network_info.consensus_parameters.chain_id);
+ let id = script.id(network_config.network_info.consensus_parameters.chain_id);
  let result = tokio::time::timeout(
  Duration::from_secs(config.timeout),
  provider.send_transaction(script),
  )
  .await
- .map(|r| {
+ .map_err(|_| {
+ error(
+ format!("Timeout while submitting transaction for address: {address:X}"),
+ StatusCode::INTERNAL_SERVER_ERROR,
+ )
+ })
+ .and_then(|r| {
  r.map_err(|e| {
  error(
- format!("Failed to submit transaction: {e}"),
+ format!(
+ "Failed to submit transaction for address: {address:X} with error: {}",
+ e
+ ),
  StatusCode::INTERNAL_SERVER_ERROR,
  )
  })
- })
- .map_err(|e| {
- error(
- format!("Timeout while submitting transaction: {e}"),
- StatusCode::INTERNAL_SERVER_ERROR,
- )
  });
 
  match result {
- Ok(Ok(_)) => {
+ Ok(_) => {
  guard.last_output = Some(CoinOutput {
- utxo_id: UtxoId::new(tx_id, 1),
+ utxo_id: UtxoId::new(id, 1),
  owner: coin_output.owner,
  amount: coin_output.amount - total_fee.min_fee() - config.dispense_amount,
  });
+ tx_id = Some(id);
  break;
  }
- _ => {
+ Err(e) => {
+ tracing::warn!("{}", e);
  guard.last_output = None;
  }
  };
  }
 
- submit_tx_with_timeout(&client, &tx_id, config.timeout)
- .await
- .map_err(|e| {
- cleanup();
- e
- })?;
+ let Some(tx_id) = tx_id else {
+ return Err(error(
+ "Failed to submit transaction".to_string(),
+ StatusCode::INTERNAL_SERVER_ERROR,
+ ));
+ };
+
+ submit_tx_with_timeout(&client, &tx_id, config.timeout).await?;
 
  info!(
  "dispensed {} tokens to {:#x}",
  config.dispense_amount, &address
  );
 
- dispense_tracker.lock().unwrap().track(address);
+ let mut tracker = dispense_tracker.lock().unwrap();
+ tracker.track(address);
 
  Ok(DispenseResponse {
  status: "Success".to_string(),