1use crate::LimitStrategy;
5use std::sync::Arc;
6use tokio::sync::Semaphore;
7use std::time::Instant;
8use crate::error::FlowError; pub struct FlowGuard<S: LimitStrategy> {
11 strategy: Arc<S>,
12 semaphore: Arc<Semaphore>,
13}
14
15impl<S: LimitStrategy + 'static> FlowGuard<S> {
16 pub fn new(strategy: S) -> Self {
17 let initial_limit = strategy.current_limit();
18 Self {
19 strategy: Arc::new(strategy),
20 semaphore: Arc::new(Semaphore::new(initial_limit)),
22 }
23 }
24
25 pub async fn run<F, T, E>(&self, f: F) -> Result<T, FlowError<E>>
26 where
27 F: std::future::Future<Output = Result<T, E>>,
28 {
29 let _permit = self.semaphore.acquire().await
31 .map_err(|_| FlowError::Dropped)?;
32
33 let start = Instant::now();
34
35 let result = f.await;
37
38 let duration = start.elapsed();
39
40 match &result {
42 Ok(_) => self.strategy.on_success(duration),
43 Err(_) => self.strategy.on_error(),
44 }
45
46 result.map_err(FlowError::AppError)
48 }
49}