1use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::{Arc, Condvar, Mutex};
19use std::thread;
20use std::time::{Duration, Instant};
21
22use crate::error::{CoreError, CoreResult, ErrorContext, ErrorLocation};
23
24fn lock_err(ctx: &'static str, e: impl std::fmt::Display) -> CoreError {
27 CoreError::MutexError(
28 ErrorContext::new(format!("{ctx}: mutex poisoned: {e}"))
29 .with_location(ErrorLocation::new(file!(), line!())),
30 )
31}
32
33fn wait_err(ctx: &'static str, e: impl std::fmt::Display) -> CoreError {
34 CoreError::MutexError(
35 ErrorContext::new(format!("{ctx}: condvar wait poisoned: {e}"))
36 .with_location(ErrorLocation::new(file!(), line!())),
37 )
38}
39
40pub struct Semaphore {
62 inner: Mutex<usize>,
63 condvar: Condvar,
64}
65
66impl Semaphore {
67 pub fn new(permits: usize) -> Self {
69 Self {
70 inner: Mutex::new(permits),
71 condvar: Condvar::new(),
72 }
73 }
74
75 pub fn acquire(&self) -> CoreResult<()> {
77 let mut g = self
78 .inner
79 .lock()
80 .map_err(|e| lock_err("Semaphore::acquire", e))?;
81 loop {
82 if *g > 0 {
83 *g -= 1;
84 return Ok(());
85 }
86 g = self
87 .condvar
88 .wait(g)
89 .map_err(|e| wait_err("Semaphore::acquire", e))?;
90 }
91 }
92
93 pub fn acquire_timeout(&self, timeout: Duration) -> CoreResult<bool> {
97 let deadline = Instant::now() + timeout;
98 let mut g = self
99 .inner
100 .lock()
101 .map_err(|e| lock_err("Semaphore::acquire_timeout", e))?;
102 loop {
103 if *g > 0 {
104 *g -= 1;
105 return Ok(true);
106 }
107 let remaining = deadline.saturating_duration_since(Instant::now());
108 if remaining.is_zero() {
109 return Ok(false);
110 }
111 let (ng, _) = self
112 .condvar
113 .wait_timeout(g, remaining)
114 .map_err(|e| wait_err("Semaphore::acquire_timeout", e))?;
115 g = ng;
116 }
117 }
118
119 pub fn try_acquire(&self) -> CoreResult<bool> {
121 let mut g = self
122 .inner
123 .lock()
124 .map_err(|e| lock_err("Semaphore::try_acquire", e))?;
125 if *g > 0 {
126 *g -= 1;
127 Ok(true)
128 } else {
129 Ok(false)
130 }
131 }
132
133 pub fn release(&self) {
135 if let Ok(mut g) = self.inner.lock() {
136 *g += 1;
137 self.condvar.notify_one();
138 }
139 }
140
141 pub fn release_n(&self, n: usize) {
143 if let Ok(mut g) = self.inner.lock() {
144 *g += n;
145 if n == 1 {
146 self.condvar.notify_one();
147 } else {
148 self.condvar.notify_all();
149 }
150 }
151 }
152
153 pub fn available(&self) -> usize {
155 self.inner.lock().map(|g| *g).unwrap_or(0)
156 }
157}
158
159pub struct SemaphoreGuard<'a> {
163 sem: &'a Semaphore,
164}
165
166impl<'a> SemaphoreGuard<'a> {
167 pub fn acquire(sem: &'a Semaphore) -> CoreResult<Self> {
169 sem.acquire()?;
170 Ok(Self { sem })
171 }
172}
173
174impl Drop for SemaphoreGuard<'_> {
175 fn drop(&mut self) {
176 self.sem.release();
177 }
178}
179
180pub struct TokenBucketRateLimiter {
199 inner: Mutex<TokenBucketState>,
200 condvar: Condvar,
201 rate: f64,
203 capacity: f64,
205}
206
207struct TokenBucketState {
208 tokens: f64,
209 last_refill: Instant,
210}
211
212impl TokenBucketRateLimiter {
213 pub fn new(capacity: f64, rate: f64) -> Self {
218 let capacity = capacity.max(1.0);
219 let rate = rate.max(f64::MIN_POSITIVE);
220 Self {
221 inner: Mutex::new(TokenBucketState {
222 tokens: capacity,
223 last_refill: Instant::now(),
224 }),
225 condvar: Condvar::new(),
226 rate,
227 capacity,
228 }
229 }
230
231 fn refill(state: &mut TokenBucketState, rate: f64, capacity: f64) {
233 let now = Instant::now();
234 let elapsed = now.duration_since(state.last_refill).as_secs_f64();
235 state.tokens = (state.tokens + elapsed * rate).min(capacity);
236 state.last_refill = now;
237 }
238
239 pub fn acquire(&self) -> CoreResult<()> {
241 self.acquire_n(1.0)
242 }
243
244 pub fn acquire_n(&self, n: f64) -> CoreResult<()> {
246 let n = n.max(0.0);
247 loop {
248 let wait_duration = {
249 let mut g = self
250 .inner
251 .lock()
252 .map_err(|e| lock_err("TokenBucketRateLimiter::acquire_n", e))?;
253 Self::refill(&mut g, self.rate, self.capacity);
254 if g.tokens >= n {
255 g.tokens -= n;
256 return Ok(());
257 }
258 let deficit = n - g.tokens;
260 Duration::from_secs_f64(deficit / self.rate)
261 };
262 thread::sleep(wait_duration.min(Duration::from_millis(50)));
263 }
264 }
265
266 pub fn try_acquire(&self) -> CoreResult<bool> {
270 let mut g = self
271 .inner
272 .lock()
273 .map_err(|e| lock_err("TokenBucketRateLimiter::try_acquire", e))?;
274 Self::refill(&mut g, self.rate, self.capacity);
275 if g.tokens >= 1.0 {
276 g.tokens -= 1.0;
277 Ok(true)
278 } else {
279 Ok(false)
280 }
281 }
282
283 pub fn available_tokens(&self) -> f64 {
285 self.inner.lock().map(|g| g.tokens).unwrap_or(0.0)
286 }
287}
288
289#[derive(Debug, Clone, Copy, PartialEq, Eq)]
293pub enum BackoffStrategy {
294 Constant,
296 Exponential,
298 ExponentialWithJitter,
300}
301
302#[derive(Debug, Clone)]
319pub struct RetryPolicy {
320 pub max_attempts: u32,
322 pub initial_delay: Duration,
324 pub max_delay: Duration,
326 pub strategy: BackoffStrategy,
328}
329
330impl RetryPolicy {
331 pub fn new(max_attempts: u32, initial_delay: Duration, strategy: BackoffStrategy) -> Self {
333 Self {
334 max_attempts: max_attempts.max(1),
335 initial_delay,
336 max_delay: Duration::from_secs(60),
337 strategy,
338 }
339 }
340
341 pub fn with_max_delay(mut self, max_delay: Duration) -> Self {
343 self.max_delay = max_delay;
344 self
345 }
346
347 pub fn retry<T, E, F>(&self, mut f: F) -> Result<T, E>
351 where
352 F: FnMut() -> Result<T, E>,
353 {
354 let mut last_err = None;
355 for attempt in 0..self.max_attempts {
356 match f() {
357 Ok(v) => return Ok(v),
358 Err(e) => {
359 last_err = Some(e);
360 if attempt + 1 < self.max_attempts {
361 let delay = self.compute_delay(attempt);
362 thread::sleep(delay);
363 }
364 }
365 }
366 }
367 Err(last_err.expect("retry: loop did not execute"))
369 }
370
371 pub fn retry_with_timeout<T, E, F>(&self, total_timeout: Duration, mut f: F) -> CoreResult<T>
373 where
374 F: FnMut() -> Result<T, E>,
375 E: std::fmt::Display,
376 {
377 let deadline = Instant::now() + total_timeout;
378 let mut last_msg = String::new();
379 for attempt in 0..self.max_attempts {
380 if Instant::now() >= deadline {
381 return Err(CoreError::TimeoutError(ErrorContext::new(format!(
382 "RetryPolicy: total timeout exceeded after {attempt} attempts. Last error: {last_msg}"
383 ))));
384 }
385 match f() {
386 Ok(v) => return Ok(v),
387 Err(e) => {
388 last_msg = e.to_string();
389 if attempt + 1 < self.max_attempts {
390 let delay = self
391 .compute_delay(attempt)
392 .min(deadline.saturating_duration_since(Instant::now()));
393 if !delay.is_zero() {
394 thread::sleep(delay);
395 }
396 }
397 }
398 }
399 }
400 Err(CoreError::ComputationError(ErrorContext::new(format!(
401 "RetryPolicy: all {max} attempts failed. Last error: {last_msg}",
402 max = self.max_attempts,
403 ))))
404 }
405
406 fn compute_delay(&self, attempt: u32) -> Duration {
407 let base = match self.strategy {
408 BackoffStrategy::Constant => self.initial_delay,
409 BackoffStrategy::Exponential | BackoffStrategy::ExponentialWithJitter => {
410 let factor = 1u64.checked_shl(attempt.min(30)).unwrap_or(u64::MAX);
411 self.initial_delay.saturating_mul(factor as u32)
412 }
413 };
414
415 let delay = base.min(self.max_delay);
416
417 if self.strategy == BackoffStrategy::ExponentialWithJitter {
418 let seed = Instant::now().elapsed().subsec_nanos() as u64;
420 let pseudo_rand = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
421 let jitter_pct = (pseudo_rand % 50) as f64 / 100.0 - 0.25; let jitter_ns = (delay.as_nanos() as f64 * jitter_pct) as i64;
423 let ns = delay.as_nanos() as i64 + jitter_ns;
424 Duration::from_nanos(ns.max(0) as u64)
425 } else {
426 delay
427 }
428 }
429}
430
431pub struct FutureExecutor {
451 tx: Arc<Mutex<std::collections::VecDeque<(u64, Box<dyn FnOnce() + Send + 'static>)>>>,
452 cond: Arc<Condvar>,
453 stop: Arc<std::sync::atomic::AtomicBool>,
454 handles: Vec<thread::JoinHandle<()>>,
455 next_id: AtomicU64,
456}
457
458impl FutureExecutor {
459 pub fn new(n_workers: usize) -> Self {
461 let n = n_workers.max(1);
462 let tx: Arc<Mutex<std::collections::VecDeque<(u64, Box<dyn FnOnce() + Send + 'static>)>>> =
463 Arc::new(Mutex::new(std::collections::VecDeque::new()));
464 let cond = Arc::new(Condvar::new());
465 let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
466 let mut handles = Vec::with_capacity(n);
467
468 for _ in 0..n {
469 let tx2 = Arc::clone(&tx);
470 let cond2 = Arc::clone(&cond);
471 let stop2 = Arc::clone(&stop);
472 let handle = thread::Builder::new()
473 .name("future-exec-worker".into())
474 .spawn(move || loop {
475 let task = {
476 let mut g = tx2.lock().expect("executor queue lock");
477 loop {
478 if let Some(t) = g.pop_front() {
479 break Some(t);
480 }
481 if stop2.load(Ordering::Relaxed) {
482 break None;
483 }
484 g = cond2.wait(g).expect("executor condvar wait");
485 }
486 };
487 match task {
488 Some((_, f)) => f(),
489 None => break,
490 }
491 })
492 .expect("spawn executor worker");
493 handles.push(handle);
494 }
495
496 Self {
497 tx,
498 cond,
499 stop,
500 handles,
501 next_id: AtomicU64::new(0),
502 }
503 }
504
505 pub fn spawn<T, F>(&self, f: F) -> CoreResult<JoinFuture<T>>
507 where
508 T: Send + 'static,
509 F: FnOnce() -> T + Send + 'static,
510 {
511 let result: Arc<Mutex<Option<T>>> = Arc::new(Mutex::new(None));
512 let cond = Arc::new(Condvar::new());
513 let result2 = Arc::clone(&result);
514 let cond2 = Arc::clone(&cond);
515
516 let task = Box::new(move || {
517 let v = f();
518 if let Ok(mut g) = result2.lock() {
519 *g = Some(v);
520 cond2.notify_one();
521 }
522 });
523
524 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
525 let mut q = self
526 .tx
527 .lock()
528 .map_err(|e| lock_err("FutureExecutor::spawn", e))?;
529 q.push_back((id, task));
530 self.cond.notify_one();
531
532 Ok(JoinFuture { result, cond })
533 }
534
535 pub fn shutdown(self) -> CoreResult<()> {
537 {
538 let _guard = self
541 .tx
542 .lock()
543 .map_err(|e| lock_err("FutureExecutor::shutdown", e))?;
544 self.stop.store(true, Ordering::SeqCst);
545 }
546 self.cond.notify_all();
547 for h in self.handles {
548 h.join().map_err(|_| {
549 CoreError::SchedulerError(
550 ErrorContext::new("FutureExecutor: worker panicked on shutdown")
551 .with_location(ErrorLocation::new(file!(), line!())),
552 )
553 })?;
554 }
555 Ok(())
556 }
557}
558
559pub struct JoinFuture<T> {
562 result: Arc<Mutex<Option<T>>>,
563 cond: Arc<Condvar>,
564}
565
566impl<T> JoinFuture<T> {
567 pub fn join(self) -> CoreResult<T> {
569 let mut g = self
570 .result
571 .lock()
572 .map_err(|e| lock_err("JoinFuture::join", e))?;
573 loop {
574 if g.is_some() {
575 return g.take().ok_or_else(|| {
576 CoreError::MutexError(ErrorContext::new("JoinFuture: value already taken"))
577 });
578 }
579 g = self
580 .cond
581 .wait(g)
582 .map_err(|e| wait_err("JoinFuture::join", e))?;
583 }
584 }
585
586 pub fn join_timeout(self, timeout: Duration) -> CoreResult<Option<T>> {
588 let deadline = Instant::now() + timeout;
589 let mut g = self
590 .result
591 .lock()
592 .map_err(|e| lock_err("JoinFuture::join_timeout", e))?;
593 loop {
594 if g.is_some() {
595 return Ok(g.take());
596 }
597 let remaining = deadline.saturating_duration_since(Instant::now());
598 if remaining.is_zero() {
599 return Ok(None);
600 }
601 let (ng, _) = self
602 .cond
603 .wait_timeout(g, remaining)
604 .map_err(|e| wait_err("JoinFuture::join_timeout", e))?;
605 g = ng;
606 }
607 }
608
609 pub fn try_join(&self) -> CoreResult<Option<T>> {
611 let mut g = self
612 .result
613 .lock()
614 .map_err(|e| lock_err("JoinFuture::try_join", e))?;
615 Ok(g.take())
616 }
617}
618
619#[cfg(test)]
622mod tests {
623 use super::*;
624 use std::sync::atomic::{AtomicU32, Ordering as AO};
625
626 #[test]
629 fn semaphore_basic_acquire_release() {
630 let sem = Semaphore::new(3);
631 sem.acquire().expect("acq 1");
632 sem.acquire().expect("acq 2");
633 sem.acquire().expect("acq 3");
634 assert_eq!(sem.available(), 0);
635 sem.release();
636 assert_eq!(sem.available(), 1);
637 }
638
639 #[test]
640 fn semaphore_try_acquire() {
641 let sem = Semaphore::new(1);
642 assert!(sem.try_acquire().expect("try 1"));
643 assert!(!sem.try_acquire().expect("try 2 (empty)"));
644 sem.release();
645 assert!(sem.try_acquire().expect("try 3 after release"));
646 }
647
648 #[test]
649 fn semaphore_concurrent_access() {
650 let sem = Arc::new(Semaphore::new(2));
651 let counter = Arc::new(AtomicU32::new(0));
652 let mut handles = Vec::new();
653
654 for _ in 0..8 {
655 let s = Arc::clone(&sem);
656 let c = Arc::clone(&counter);
657 handles.push(thread::spawn(move || {
658 s.acquire().expect("acquire");
659 c.fetch_add(1, AO::Relaxed);
660 thread::sleep(Duration::from_millis(5));
661 s.release();
662 }));
663 }
664 for h in handles {
665 h.join().expect("thread");
666 }
667 assert_eq!(counter.load(AO::Relaxed), 8);
668 }
669
670 #[test]
671 fn semaphore_guard() {
672 let sem = Semaphore::new(1);
673 {
674 let _guard = SemaphoreGuard::acquire(&sem).expect("guard");
675 assert_eq!(sem.available(), 0);
676 }
677 assert_eq!(sem.available(), 1);
678 }
679
680 #[test]
681 fn semaphore_acquire_timeout_succeeds() {
682 let sem = Semaphore::new(1);
683 let ok = sem
684 .acquire_timeout(Duration::from_millis(100))
685 .expect("timeout acq");
686 assert!(ok);
687 }
688
689 #[test]
690 fn semaphore_acquire_timeout_expires() {
691 let sem = Semaphore::new(0); let ok = sem
693 .acquire_timeout(Duration::from_millis(20))
694 .expect("timeout acq");
695 assert!(!ok);
696 }
697
698 #[test]
701 fn rate_limiter_basic() {
702 let rl = TokenBucketRateLimiter::new(10.0, 1000.0);
704 for _ in 0..5 {
705 rl.acquire().expect("acquire token");
706 }
707 }
708
709 #[test]
710 fn rate_limiter_try_acquire() {
711 let rl = TokenBucketRateLimiter::new(2.0, 100.0);
712 assert!(rl.try_acquire().expect("t1"));
713 assert!(rl.try_acquire().expect("t2"));
714 assert!(!rl.try_acquire().expect("t3 empty"));
715 }
716
717 #[test]
720 fn retry_succeeds_on_nth_attempt() {
721 let counter = std::sync::atomic::AtomicU32::new(0);
722 let policy = RetryPolicy::new(5, Duration::from_millis(1), BackoffStrategy::Constant);
723 let result: Result<u32, &str> = policy.retry(|| {
724 let n = counter.fetch_add(1, AO::Relaxed);
725 if n < 3 {
726 Err("not yet")
727 } else {
728 Ok(n)
729 }
730 });
731 assert!(result.is_ok());
732 }
733
734 #[test]
735 fn retry_exhausts_all_attempts() {
736 let policy = RetryPolicy::new(3, Duration::from_millis(1), BackoffStrategy::Constant);
737 let result: Result<u32, &str> = policy.retry(|| Err("always fail"));
738 assert!(result.is_err());
739 }
740
741 #[test]
742 fn retry_exponential_backoff() {
743 let policy = RetryPolicy::new(4, Duration::from_millis(1), BackoffStrategy::Exponential);
744 let counter = std::sync::atomic::AtomicU32::new(0);
745 let _: Result<u32, &str> = policy.retry(|| {
746 counter.fetch_add(1, AO::Relaxed);
747 Err("fail")
748 });
749 assert_eq!(counter.load(AO::Relaxed), 4);
750 }
751
752 #[test]
755 fn future_executor_basic() {
756 let exec = FutureExecutor::new(4);
757 let h1 = exec.spawn(|| 2u64 + 2).expect("spawn h1");
758 let h2 = exec.spawn(|| 10u64 * 10).expect("spawn h2");
759 assert_eq!(h1.join().expect("join h1"), 4);
760 assert_eq!(h2.join().expect("join h2"), 100);
761 exec.shutdown().expect("shutdown");
762 }
763
764 #[test]
765 fn future_executor_many_tasks() {
766 let exec = FutureExecutor::new(4);
767 let handles: Vec<_> = (0u64..50)
768 .map(|i| exec.spawn(move || i * i).expect("spawn"))
769 .collect();
770 for (i, h) in handles.into_iter().enumerate() {
771 assert_eq!(h.join().expect("join"), (i as u64) * (i as u64));
772 }
773 exec.shutdown().expect("shutdown");
774 }
775
776 #[test]
777 fn future_join_timeout_succeeds() {
778 let exec = FutureExecutor::new(2);
779 let h = exec.spawn(|| 42u64).expect("spawn");
780 let v = h
781 .join_timeout(Duration::from_secs(5))
782 .expect("timeout join");
783 assert_eq!(v, Some(42));
784 exec.shutdown().expect("shutdown");
785 }
786}