1use std::{
38 collections::HashMap,
39 fmt,
40 future::Future,
41 pin::Pin,
42 sync::Arc,
43 time::{Duration, Instant},
44};
45
46use tokio::task::JoinSet;
47
48pub type TaskFactory = Arc<
57 dyn Fn(
58 Arc<HashMap<String, String>>,
59 ) -> Pin<Box<dyn Future<Output = Result<i64, String>> + Send>>
60 + Send
61 + Sync,
62>;
63
64#[derive(Clone)]
69pub struct FanTask {
70 pub name: String,
72 factory: TaskFactory,
74 pub timeout: Option<Duration>,
76 pub weight: u32,
78}
79
80impl fmt::Debug for FanTask {
81 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82 f.debug_struct("FanTask")
83 .field("name", &self.name)
84 .field("timeout", &self.timeout)
85 .field("weight", &self.weight)
86 .finish()
87 }
88}
89
90impl FanTask {
91 pub fn new<N, F, Fut>(name: N, factory: F) -> Self
96 where
97 N: Into<String>,
98 F: Fn(Arc<HashMap<String, String>>) -> Fut + Send + Sync + 'static,
99 Fut: Future<Output = Result<i64, String>> + Send + 'static,
100 {
101 Self {
102 name: name.into(),
103 factory: Arc::new(move |ctx| Box::pin(factory(ctx))),
104 timeout: None,
105 weight: 1,
106 }
107 }
108
109 #[must_use]
111 pub fn with_timeout(mut self, timeout: Duration) -> Self {
112 self.timeout = Some(timeout);
113 self
114 }
115
116 #[must_use]
119 pub fn with_weight(mut self, weight: u32) -> Self {
120 self.weight = weight;
121 self
122 }
123
124 fn invoke(
126 &self,
127 ctx: Arc<HashMap<String, String>>,
128 ) -> Pin<Box<dyn Future<Output = Result<i64, String>> + Send>> {
129 (self.factory)(ctx)
130 }
131}
132
133#[derive(Debug, Default, Clone)]
142pub struct FanOut {
143 pub context: HashMap<String, String>,
145}
146
147impl FanOut {
148 #[must_use]
150 pub fn new() -> Self {
151 Self::default()
152 }
153
154 pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
156 self.context.insert(key.into(), value.into());
157 self
158 }
159
160 #[must_use]
162 pub fn with(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
163 self.context.insert(key.into(), value.into());
164 self
165 }
166}
167
168#[derive(Debug, Clone, Copy, PartialEq, Eq)]
170pub enum AggregationStrategy {
171 CollectAll,
173 FirstSuccess,
175 Sum,
177 Max,
179 Min,
181 WeightedSum,
183}
184
185impl Default for AggregationStrategy {
186 fn default() -> Self {
187 Self::CollectAll
188 }
189}
190
191#[derive(Debug, Clone)]
193pub struct FanInResult {
194 pub task_results: HashMap<String, Result<i64, String>>,
196 pub aggregated: Option<i64>,
201 pub elapsed: Duration,
203 pub success_count: usize,
205 pub failure_count: usize,
207}
208
209impl FanInResult {
210 #[must_use]
212 pub fn all_succeeded(&self) -> bool {
213 self.failure_count == 0
214 }
215
216 #[must_use]
218 pub fn any_succeeded(&self) -> bool {
219 self.success_count > 0
220 }
221}
222
223#[derive(Debug)]
232pub struct FanPattern {
233 pub name: String,
235 pub tasks: Vec<FanTask>,
237 pub fan_out: FanOut,
239 pub strategy: AggregationStrategy,
241 pub default_timeout: Option<Duration>,
243 pub fail_fast: bool,
246}
247
248impl FanPattern {
249 #[must_use]
253 pub fn new(name: impl Into<String>, tasks: Vec<FanTask>) -> Self {
254 Self {
255 name: name.into(),
256 tasks,
257 fan_out: FanOut::new(),
258 strategy: AggregationStrategy::CollectAll,
259 default_timeout: None,
260 fail_fast: false,
261 }
262 }
263
264 #[must_use]
266 pub fn with_strategy(mut self, strategy: AggregationStrategy) -> Self {
267 self.strategy = strategy;
268 self
269 }
270
271 #[must_use]
273 pub fn with_default_timeout(mut self, timeout: Duration) -> Self {
274 self.default_timeout = Some(timeout);
275 self
276 }
277
278 #[must_use]
280 pub fn with_fail_fast(mut self, fail_fast: bool) -> Self {
281 self.fail_fast = fail_fast;
282 self
283 }
284
285 #[must_use]
287 pub fn with_context(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
288 self.fan_out.insert(key, value);
289 self
290 }
291
292 #[must_use]
294 pub fn task_count(&self) -> usize {
295 self.tasks.len()
296 }
297}
298
299#[derive(Debug, thiserror::Error)]
305pub enum FanError {
306 #[error("All {count} fan tasks failed")]
308 AllFailed {
309 count: usize,
311 errors: HashMap<String, String>,
313 },
314 #[error("{failure_count} of {total} fan tasks failed")]
316 SomeFailed {
317 failure_count: usize,
319 total: usize,
321 errors: HashMap<String, String>,
323 },
324 #[error("Fan task '{name}' timed out after {timeout_ms}ms")]
326 TaskTimeout {
327 name: String,
329 timeout_ms: u64,
331 },
332 #[error("Fan task '{name}' panicked: {reason}")]
334 JoinError {
335 name: String,
337 reason: String,
339 },
340}
341
342#[derive(Debug, Default, Clone)]
351pub struct FanExecutor {
352 pub global_timeout: Option<Duration>,
355}
356
357impl FanExecutor {
358 #[must_use]
360 pub fn new() -> Self {
361 Self::default()
362 }
363
364 #[must_use]
366 pub fn with_global_timeout(mut self, timeout: Duration) -> Self {
367 self.global_timeout = Some(timeout);
368 self
369 }
370
371 pub async fn execute(&self, pattern: FanPattern) -> Result<FanInResult, FanError> {
380 let start = Instant::now();
381 let ctx = Arc::new(pattern.fan_out.context.clone());
382 let strategy = pattern.strategy;
383
384 let mut join_set: JoinSet<(String, u32, Result<i64, String>)> = JoinSet::new();
386
387 for task in pattern.tasks {
388 let task_name = task.name.clone();
389 let task_weight = task.weight;
390 let task_ctx = ctx.clone();
391 let effective_timeout = task
392 .timeout
393 .or(pattern.default_timeout)
394 .or(self.global_timeout);
395 let fut = task.invoke(task_ctx);
396
397 join_set.spawn(async move {
398 match effective_timeout {
399 Some(timeout) => match tokio::time::timeout(timeout, fut).await {
400 Ok(result) => (task_name, task_weight, result),
401 Err(_elapsed) => (
402 task_name.clone(),
403 task_weight,
404 Err(format!("task '{}' timed out", task_name)),
405 ),
406 },
407 None => {
408 let result = fut.await;
409 (task_name, task_weight, result)
410 }
411 }
412 });
413 }
414
415 let mut task_results: HashMap<String, Result<i64, String>> = HashMap::new();
417 let mut weights: HashMap<String, u32> = HashMap::new();
418 let mut first_success: Option<(String, i64)> = None;
419 let mut errors: HashMap<String, String> = HashMap::new();
420
421 while let Some(join_result) = join_set.join_next().await {
422 match join_result {
423 Ok((name, weight, outcome)) => {
424 weights.insert(name.clone(), weight);
425 match &outcome {
426 Ok(val) => {
427 if first_success.is_none() {
428 first_success = Some((name.clone(), *val));
429 }
430 }
431 Err(e) => {
432 errors.insert(name.clone(), e.clone());
433 }
434 }
435 task_results.insert(name, outcome);
436 }
437 Err(join_err) => {
438 let name = format!("<unknown-task-{}>", errors.len());
440 let reason = join_err.to_string();
441 errors.insert(name.clone(), reason.clone());
442 task_results.insert(name, Err(format!("join error: {reason}")));
443 }
444 }
445 }
446
447 let elapsed = start.elapsed();
448 let total = task_results.len();
449 let failure_count = errors.len();
450 let success_count = total - failure_count;
451
452 let aggregated = Self::aggregate(&task_results, &weights, strategy, &first_success);
454
455 match strategy {
457 AggregationStrategy::CollectAll => {
458 if failure_count > 0 {
459 return Err(FanError::SomeFailed {
460 failure_count,
461 total,
462 errors,
463 });
464 }
465 }
466 AggregationStrategy::FirstSuccess => {
467 if first_success.is_none() {
468 return Err(FanError::AllFailed {
469 count: total,
470 errors,
471 });
472 }
473 }
474 AggregationStrategy::Sum
475 | AggregationStrategy::Max
476 | AggregationStrategy::Min
477 | AggregationStrategy::WeightedSum => {
478 if success_count == 0 {
479 return Err(FanError::AllFailed {
480 count: total,
481 errors,
482 });
483 }
484 }
485 }
486
487 Ok(FanInResult {
488 task_results,
489 aggregated,
490 elapsed,
491 success_count,
492 failure_count,
493 })
494 }
495
496 fn aggregate(
497 task_results: &HashMap<String, Result<i64, String>>,
498 weights: &HashMap<String, u32>,
499 strategy: AggregationStrategy,
500 first_success: &Option<(String, i64)>,
501 ) -> Option<i64> {
502 let successes: Vec<(i64, u32)> = task_results
503 .iter()
504 .filter_map(|(name, r)| {
505 r.as_ref().ok().map(|&v| {
506 let w = weights.get(name).copied().unwrap_or(1);
507 (v, w)
508 })
509 })
510 .collect();
511
512 match strategy {
513 AggregationStrategy::CollectAll => {
514 if task_results.values().all(|r| r.is_ok()) {
516 Some(successes.iter().map(|(v, _)| *v).sum())
517 } else {
518 None
519 }
520 }
521 AggregationStrategy::FirstSuccess => first_success.as_ref().map(|(_, v)| *v),
522 AggregationStrategy::Sum => {
523 if successes.is_empty() {
524 None
525 } else {
526 Some(successes.iter().map(|(v, _)| *v).sum())
527 }
528 }
529 AggregationStrategy::Max => successes.iter().map(|(v, _)| *v).max(),
530 AggregationStrategy::Min => successes.iter().map(|(v, _)| *v).min(),
531 AggregationStrategy::WeightedSum => {
532 if successes.is_empty() {
533 None
534 } else {
535 Some(
536 successes
537 .iter()
538 .map(|(v, w)| v.saturating_mul(i64::from(*w)))
539 .fold(0i64, i64::saturating_add),
540 )
541 }
542 }
543 }
544 }
545}
546
547#[cfg(test)]
552mod tests {
553 use super::*;
554
555 fn make_ok_task(name: &str, value: i64) -> FanTask {
556 FanTask::new(name.to_string(), move |_ctx| async move {
557 Ok::<i64, String>(value)
558 })
559 }
560
561 fn make_fail_task(name: &str, reason: &str) -> FanTask {
562 let reason = reason.to_string();
563 FanTask::new(name.to_string(), move |_ctx| {
564 let r = reason.clone();
565 async move { Err::<i64, String>(r) }
566 })
567 }
568
569 #[test]
574 fn fan_out_insert_and_with() {
575 let mut fo = FanOut::new();
576 fo.insert("k1", "v1");
577 let fo = fo.with("k2".to_string(), "v2".to_string());
578 assert_eq!(fo.context.get("k1").map(String::as_str), Some("v1"));
579 assert_eq!(fo.context.get("k2").map(String::as_str), Some("v2"));
580 }
581
582 #[test]
587 fn fan_task_weight_default_is_one() {
588 let t = make_ok_task("t", 42);
589 assert_eq!(t.weight, 1);
590 }
591
592 #[test]
593 fn fan_task_with_weight() {
594 let t = make_ok_task("t", 0).with_weight(5);
595 assert_eq!(t.weight, 5);
596 }
597
598 #[test]
599 fn fan_task_with_timeout() {
600 let t = make_ok_task("t", 0).with_timeout(Duration::from_millis(100));
601 assert!(t.timeout.is_some());
602 }
603
604 #[test]
609 fn fan_pattern_task_count() {
610 let p = FanPattern::new("p", vec![make_ok_task("a", 1), make_ok_task("b", 2)]);
611 assert_eq!(p.task_count(), 2);
612 }
613
614 #[test]
615 fn fan_pattern_defaults() {
616 let p = FanPattern::new("p", vec![]);
617 assert_eq!(p.strategy, AggregationStrategy::CollectAll);
618 assert!(!p.fail_fast);
619 assert!(p.default_timeout.is_none());
620 }
621
622 #[tokio::test]
627 async fn execute_collect_all_success() {
628 let tasks = vec![
629 make_ok_task("a", 10),
630 make_ok_task("b", 20),
631 make_ok_task("c", 30),
632 ];
633 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::CollectAll);
634 let result = FanExecutor::new()
635 .execute(pattern)
636 .await
637 .expect("all tasks succeed");
638
639 assert_eq!(result.success_count, 3);
640 assert_eq!(result.failure_count, 0);
641 assert!(result.all_succeeded());
642 assert_eq!(result.aggregated, Some(60));
644 }
645
646 #[tokio::test]
647 async fn execute_sum_strategy() {
648 let tasks = vec![
649 make_ok_task("a", 5),
650 make_ok_task("b", 15),
651 make_fail_task("c", "fail"),
652 ];
653 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::Sum);
654 let result = FanExecutor::new()
655 .execute(pattern)
656 .await
657 .expect("sum succeeds when at least one succeeds");
658
659 assert_eq!(result.success_count, 2);
660 assert_eq!(result.failure_count, 1);
661 assert_eq!(result.aggregated, Some(20));
662 }
663
664 #[tokio::test]
665 async fn execute_max_strategy() {
666 let tasks = vec![
667 make_ok_task("a", 3),
668 make_ok_task("b", 100),
669 make_ok_task("c", 7),
670 ];
671 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::Max);
672 let result = FanExecutor::new()
673 .execute(pattern)
674 .await
675 .expect("max succeeds");
676 assert_eq!(result.aggregated, Some(100));
677 }
678
679 #[tokio::test]
680 async fn execute_min_strategy() {
681 let tasks = vec![
682 make_ok_task("a", 50),
683 make_ok_task("b", 3),
684 make_ok_task("c", 25),
685 ];
686 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::Min);
687 let result = FanExecutor::new()
688 .execute(pattern)
689 .await
690 .expect("min succeeds");
691 assert_eq!(result.aggregated, Some(3));
692 }
693
694 #[tokio::test]
695 async fn execute_weighted_sum_strategy() {
696 let tasks = vec![
697 make_ok_task("a", 10).with_weight(3), make_ok_task("b", 5).with_weight(2), ];
700 let pattern =
701 FanPattern::new("test", tasks).with_strategy(AggregationStrategy::WeightedSum);
702 let result = FanExecutor::new()
703 .execute(pattern)
704 .await
705 .expect("weighted sum ok");
706 assert_eq!(result.aggregated, Some(40));
707 }
708
709 #[tokio::test]
710 async fn execute_first_success_strategy() {
711 let tasks = vec![
713 make_fail_task("a", "oops"),
714 make_ok_task("b", 77),
715 make_ok_task("c", 99),
716 ];
717 let pattern =
718 FanPattern::new("test", tasks).with_strategy(AggregationStrategy::FirstSuccess);
719 let result = FanExecutor::new()
720 .execute(pattern)
721 .await
722 .expect("first_success finds at least one success");
723 assert!(result.any_succeeded());
724 assert!(result.aggregated.is_some());
725 }
726
727 #[tokio::test]
728 async fn execute_collect_all_fails_on_any_failure() {
729 let tasks = vec![
730 make_ok_task("a", 1),
731 make_fail_task("b", "boom"),
732 make_ok_task("c", 3),
733 ];
734 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::CollectAll);
735 let err = FanExecutor::new()
736 .execute(pattern)
737 .await
738 .expect_err("should fail");
739
740 assert!(matches!(
741 err,
742 FanError::SomeFailed {
743 failure_count: 1,
744 total: 3,
745 ..
746 }
747 ));
748 }
749
750 #[tokio::test]
751 async fn execute_all_fail_returns_all_failed() {
752 let tasks = vec![make_fail_task("a", "err-a"), make_fail_task("b", "err-b")];
753 let pattern = FanPattern::new("test", tasks).with_strategy(AggregationStrategy::Sum);
754 let err = FanExecutor::new()
755 .execute(pattern)
756 .await
757 .expect_err("all fail");
758 assert!(matches!(err, FanError::AllFailed { count: 2, .. }));
759 }
760
761 #[tokio::test]
762 async fn execute_with_context_passed_to_tasks() {
763 let tasks = vec![FanTask::new(
764 "ctx-reader",
765 |ctx: Arc<HashMap<String, String>>| async move {
766 let val: i64 = ctx
767 .get("multiplier")
768 .and_then(|v| v.parse().ok())
769 .unwrap_or(1);
770 Ok::<i64, String>(val * 10)
771 },
772 )];
773 let pattern = FanPattern::new("ctx-test", tasks)
774 .with_context("multiplier", "7")
775 .with_strategy(AggregationStrategy::CollectAll);
776
777 let result = FanExecutor::new().execute(pattern).await.expect("ok");
778 let output = result.task_results["ctx-reader"].as_ref().expect("ok");
779 assert_eq!(*output, 70);
780 }
781
782 #[tokio::test]
783 async fn execute_empty_tasks_collect_all() {
784 let pattern =
785 FanPattern::new("empty", vec![]).with_strategy(AggregationStrategy::CollectAll);
786 let result = FanExecutor::new()
787 .execute(pattern)
788 .await
789 .expect("empty is fine");
790 assert_eq!(result.success_count, 0);
791 assert_eq!(result.failure_count, 0);
792 }
793}