camel_processor/
load_balancer.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::task::{Context, Poll};
6
7use tower::Service;
8use tower::ServiceExt;
9
10use camel_api::{BoxProcessor, CamelError, Exchange, LoadBalanceStrategy, LoadBalancerConfig};
11
12#[derive(Clone)]
13pub struct LoadBalancerService {
14    endpoints: Vec<BoxProcessor>,
15    config: LoadBalancerConfig,
16    round_robin_index: Arc<AtomicUsize>,
17    failover_index: Arc<AtomicUsize>,
18}
19
20impl LoadBalancerService {
21    pub fn new(endpoints: Vec<BoxProcessor>, config: LoadBalancerConfig) -> Self {
22        Self {
23            endpoints,
24            config,
25            round_robin_index: Arc::new(AtomicUsize::new(0)),
26            failover_index: Arc::new(AtomicUsize::new(0)),
27        }
28    }
29}
30
31impl Service<Exchange> for LoadBalancerService {
32    type Response = Exchange;
33    type Error = CamelError;
34    type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
35
36    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
37        for endpoint in &mut self.endpoints {
38            match endpoint.poll_ready(cx) {
39                Poll::Pending => return Poll::Pending,
40                Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
41                Poll::Ready(Ok(())) => {}
42            }
43        }
44        Poll::Ready(Ok(()))
45    }
46
47    fn call(&mut self, exchange: Exchange) -> Self::Future {
48        let endpoints = self.endpoints.clone();
49        let config = self.config.clone();
50        let round_robin_index = self.round_robin_index.clone();
51        let failover_index = self.failover_index.clone();
52
53        Box::pin(async move {
54            if endpoints.is_empty() {
55                return Ok(exchange);
56            }
57
58            match &config.strategy {
59                LoadBalanceStrategy::Random => process_random(exchange, endpoints).await,
60                LoadBalanceStrategy::Weighted(weights) => {
61                    process_weighted(exchange, endpoints, weights).await
62                }
63                LoadBalanceStrategy::Failover => {
64                    process_failover(exchange, endpoints, failover_index).await
65                }
66                // RoundRobin and any future variant default to round-robin dispatch.
67                _ => process_round_robin(exchange, endpoints, round_robin_index).await,
68            }
69        })
70    }
71}
72
73async fn process_round_robin(
74    exchange: Exchange,
75    endpoints: Vec<BoxProcessor>,
76    index: Arc<AtomicUsize>,
77) -> Result<Exchange, CamelError> {
78    let len = endpoints.len();
79    let idx = index.fetch_add(1, Ordering::SeqCst) % len;
80    let mut endpoint = endpoints[idx].clone();
81    endpoint.ready().await?.call(exchange).await
82}
83
84async fn process_random(
85    exchange: Exchange,
86    endpoints: Vec<BoxProcessor>,
87) -> Result<Exchange, CamelError> {
88    let len = endpoints.len();
89    let idx = rand::random_range(0..len);
90    let mut endpoint = endpoints[idx].clone();
91    endpoint.ready().await?.call(exchange).await
92}
93
94async fn process_weighted(
95    exchange: Exchange,
96    endpoints: Vec<BoxProcessor>,
97    weights: &[(String, u32)],
98) -> Result<Exchange, CamelError> {
99    if endpoints.is_empty() || weights.is_empty() {
100        return Ok(exchange);
101    }
102
103    let numeric_weights: Vec<u64> = weights.iter().map(|(_, w)| *w as u64).collect();
104    let total: u64 = numeric_weights
105        .iter()
106        .try_fold(0u64, |acc, w| acc.checked_add(*w))
107        .ok_or_else(|| {
108            CamelError::ProcessorError("Weighted load balancer total weight overflow".to_string())
109        })?;
110
111    if total == 0 {
112        return Err(CamelError::ProcessorError(
113            "Weighted load balancer has zero total weight".to_string(),
114        ));
115    }
116
117    let mut r = rand::random::<u64>() % total;
118    let mut selected_idx = 0;
119    for (i, w) in numeric_weights.iter().enumerate() {
120        if r < *w {
121            selected_idx = i.min(endpoints.len() - 1);
122            break;
123        }
124        r -= w;
125    }
126
127    let mut endpoint = endpoints[selected_idx].clone();
128    endpoint.ready().await?.call(exchange).await
129}
130
131async fn process_failover(
132    exchange: Exchange,
133    endpoints: Vec<BoxProcessor>,
134    start_index: Arc<AtomicUsize>,
135) -> Result<Exchange, CamelError> {
136    let len = endpoints.len();
137    let start = start_index.load(Ordering::SeqCst);
138    let mut last_error = None;
139
140    for i in 0..len {
141        let idx = (start + i) % len;
142        let mut endpoint = endpoints[idx].clone();
143        match endpoint.ready().await?.call(exchange.clone()).await {
144            Ok(ex) => {
145                start_index.store((idx + 1) % len, Ordering::SeqCst);
146                return Ok(ex);
147            }
148            Err(e) => {
149                last_error = Some(e);
150            }
151        }
152    }
153
154    Err(last_error.unwrap_or_else(|| {
155        CamelError::ProcessorError("All endpoints failed in failover".to_string())
156    }))
157}
158
159// ── LoadBalanceSegment (ADR-0025 OutcomePipeline) ────────────────────────
160
161/// Outcome-aware LoadBalance segment. Holds N destinations + a strategy.
162/// On each call: strategy picks ONE destination (round-robin / failover /
163/// random / weighted), runs it. If chosen destination returns Completed,
164/// return Completed. If Stopped: return Stopped immediately (no failover —
165/// Stop is successful control flow). If Failed: strategy decides (failover
166/// retries next dest, others return Failed).
167///
168/// This differs from Multicast (which runs all branches) — LoadBalance picks
169/// exactly one. The parallel cancellation logic from T13/T15 does NOT apply.
170#[derive(Clone)]
171pub struct LoadBalanceSegment {
172    pub destinations: Vec<camel_api::OutcomeSegment>,
173    pub strategy: camel_api::LoadBalanceStrategy,
174    /// Shared round-robin index for interior mutability across cloned segments.
175    pub round_robin_index: Arc<AtomicUsize>,
176}
177
178impl camel_api::OutcomePipeline for LoadBalanceSegment {
179    fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
180        Box::new(self.clone())
181    }
182
183    fn run<'a>(
184        &'a mut self,
185        exchange: camel_api::Exchange,
186    ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
187        Box::pin(async move {
188            let len = self.destinations.len();
189            if len == 0 {
190                return camel_api::PipelineOutcome::Completed(exchange);
191            }
192
193            let start_idx = match &self.strategy {
194                camel_api::LoadBalanceStrategy::Random => rand::random_range(0..len),
195                camel_api::LoadBalanceStrategy::Weighted(weights) => pick_weighted(weights, len),
196                camel_api::LoadBalanceStrategy::Failover => 0,
197                // RoundRobin and any future variant default to round-robin indexing.
198                _ => self.round_robin_index.fetch_add(1, Ordering::SeqCst) % len,
199            };
200
201            let mut idx = start_idx;
202            let mut last_err: Option<camel_api::CamelError> = None;
203            loop {
204                if idx >= len {
205                    return camel_api::PipelineOutcome::Failed(last_err.unwrap_or_else(|| {
206                        camel_api::CamelError::ProcessorError(
207                            "load_balance: all destinations exhausted".to_string(),
208                        )
209                    }));
210                }
211                match self.destinations[idx].run(exchange.clone()).await {
212                    camel_api::PipelineOutcome::Completed(ex) => {
213                        return camel_api::PipelineOutcome::Completed(ex);
214                    }
215                    camel_api::PipelineOutcome::Stopped(ex) => {
216                        return camel_api::PipelineOutcome::Stopped(ex);
217                    }
218                    camel_api::PipelineOutcome::Failed(err) => match self.strategy {
219                        camel_api::LoadBalanceStrategy::Failover => {
220                            last_err = Some(err);
221                            idx += 1;
222                            continue;
223                        }
224                        _ => return camel_api::PipelineOutcome::Failed(err),
225                    },
226                }
227            }
228        })
229    }
230}
231
232/// Pick a destination index using weighted random selection.
233fn pick_weighted(weights: &[(String, u32)], len: usize) -> usize {
234    if weights.is_empty() || len == 0 {
235        return 0;
236    }
237    let numeric_weights: Vec<u64> = weights.iter().map(|(_, w)| *w as u64).collect();
238    let Some(total) = numeric_weights
239        .iter()
240        .try_fold(0u64, |acc, w| acc.checked_add(*w))
241    else {
242        return 0;
243    };
244    if total == 0 {
245        return 0;
246    }
247    let mut r = rand::random::<u64>() % total;
248    for (i, w) in numeric_weights.iter().enumerate() {
249        if r < *w {
250            return i.min(len - 1);
251        }
252        r -= w;
253    }
254    len - 1
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use camel_api::{BoxProcessorExt, Message};
261    use std::sync::Mutex;
262    use tower::ServiceExt;
263
264    fn counting_processor() -> (BoxProcessor, Arc<AtomicUsize>) {
265        let count = Arc::new(AtomicUsize::new(0));
266        let count_clone = count.clone();
267        let processor = BoxProcessor::from_fn(move |ex| {
268            count_clone.fetch_add(1, Ordering::SeqCst);
269            Box::pin(async move { Ok(ex) })
270        });
271        (processor, count)
272    }
273
274    #[tokio::test]
275    async fn test_round_robin_distribution() {
276        let (p1, c1) = counting_processor();
277        let (p2, c2) = counting_processor();
278        let (p3, c3) = counting_processor();
279
280        let config = LoadBalancerConfig::round_robin();
281        let mut svc = LoadBalancerService::new(vec![p1, p2, p3], config);
282
283        for _ in 0..6 {
284            let ex = Exchange::new(Message::new("test"));
285            svc.ready().await.unwrap().call(ex).await.unwrap();
286        }
287
288        assert_eq!(c1.load(Ordering::SeqCst), 2);
289        assert_eq!(c2.load(Ordering::SeqCst), 2);
290        assert_eq!(c3.load(Ordering::SeqCst), 2);
291    }
292
293    #[tokio::test]
294    async fn test_random_distribution() {
295        let (p1, c1) = counting_processor();
296        let (p2, c2) = counting_processor();
297
298        let config = LoadBalancerConfig::random();
299        let mut svc = LoadBalancerService::new(vec![p1, p2], config);
300
301        for _ in 0..100 {
302            let ex = Exchange::new(Message::new("test"));
303            svc.ready().await.unwrap().call(ex).await.unwrap();
304        }
305
306        let total = c1.load(Ordering::SeqCst) + c2.load(Ordering::SeqCst);
307        assert_eq!(total, 100);
308        assert!(c1.load(Ordering::SeqCst) > 20);
309        assert!(c2.load(Ordering::SeqCst) > 20);
310    }
311
312    #[tokio::test]
313    async fn test_failover_on_error() {
314        let failing = BoxProcessor::from_fn(|_ex| {
315            Box::pin(async { Err(CamelError::ProcessorError("fail".into())) })
316        });
317        let (success, count) = counting_processor();
318
319        let config = LoadBalancerConfig::failover();
320        let mut svc = LoadBalancerService::new(vec![failing, success], config);
321
322        let ex = Exchange::new(Message::new("test"));
323        let _result = svc.ready().await.unwrap().call(ex).await.unwrap();
324
325        assert_eq!(count.load(Ordering::SeqCst), 1);
326    }
327
328    #[tokio::test]
329    async fn test_failover_preserves_original_exchange() {
330        // Capture body seen by retry endpoint to verify it's the original
331        let seen_body: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
332        let seen_body_clone = seen_body.clone();
333
334        let failing = BoxProcessor::from_fn(|_ex| {
335            Box::pin(async { Err(CamelError::ProcessorError("fail".into())) })
336        });
337
338        let retry = BoxProcessor::from_fn(move |ex: Exchange| {
339            let seen = seen_body_clone.clone();
340            Box::pin(async move {
341                if let Some(text) = ex.input.body.as_text() {
342                    *seen.lock().unwrap() = Some(text.to_string());
343                }
344                Ok(ex)
345            })
346        });
347
348        let config = LoadBalancerConfig::failover();
349        let mut svc = LoadBalancerService::new(vec![failing, retry], config);
350
351        let ex = Exchange::new(Message::new("original body"));
352        svc.ready().await.unwrap().call(ex).await.unwrap();
353
354        assert_eq!(
355            seen_body.lock().unwrap().as_deref(),
356            Some("original body"),
357            "retry endpoint must receive the original exchange body, not a blank one"
358        );
359    }
360
361    #[tokio::test]
362    async fn test_failover_all_fail() {
363        let failing = BoxProcessor::from_fn(|_ex| {
364            Box::pin(async { Err(CamelError::ProcessorError("fail".into())) })
365        });
366
367        let config = LoadBalancerConfig::failover();
368        let mut svc = LoadBalancerService::new(vec![failing.clone(), failing], config);
369
370        let ex = Exchange::new(Message::new("test"));
371        let result = svc.ready().await.unwrap().call(ex).await;
372
373        assert!(result.is_err());
374    }
375
376    #[tokio::test]
377    async fn test_empty_endpoints() {
378        let config = LoadBalancerConfig::round_robin();
379        let mut svc = LoadBalancerService::new(vec![], config);
380
381        let ex = Exchange::new(Message::new("test"));
382        let result = svc.ready().await.unwrap().call(ex).await;
383
384        assert!(result.is_ok());
385    }
386
387    // ── LoadBalanceSegment tests (ADR-0025 OutcomePipeline parity) ───
388
389    /// OutcomePipeline body that mutates exchange body to "lb-stopped" then returns Stopped.
390    struct StoppingBody;
391    impl camel_api::OutcomePipeline for StoppingBody {
392        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
393            Box::new(StoppingBody)
394        }
395        fn run<'a>(
396            &'a mut self,
397            mut ex: Exchange,
398        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
399            Box::pin(async move {
400                ex.input.body = camel_api::Body::Text("lb-stopped".to_string());
401                camel_api::PipelineOutcome::Stopped(ex)
402            })
403        }
404    }
405
406    /// OutcomePipeline body that records invocation count via shared counter.
407    struct RecordingBody(Arc<AtomicUsize>);
408    impl camel_api::OutcomePipeline for RecordingBody {
409        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
410            Box::new(RecordingBody(Arc::clone(&self.0)))
411        }
412        fn run<'a>(
413            &'a mut self,
414            ex: Exchange,
415        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
416            let count = Arc::clone(&self.0);
417            Box::pin(async move {
418                count.fetch_add(1, Ordering::SeqCst);
419                camel_api::PipelineOutcome::Completed(ex)
420            })
421        }
422    }
423
424    /// OutcomePipeline body that always fails with ProcessorError.
425    struct FailingBody;
426    impl camel_api::OutcomePipeline for FailingBody {
427        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
428            Box::new(FailingBody)
429        }
430        fn run<'a>(
431            &'a mut self,
432            _ex: Exchange,
433        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
434            Box::pin(async {
435                camel_api::PipelineOutcome::Failed(CamelError::ProcessorError(
436                    "intentional fail".to_string(),
437                ))
438            })
439        }
440    }
441
442    /// OutcomePipeline body that mutates body to "recovered" then completes.
443    struct RecoveringBody;
444    impl camel_api::OutcomePipeline for RecoveringBody {
445        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
446            Box::new(RecoveringBody)
447        }
448        fn run<'a>(
449            &'a mut self,
450            mut ex: Exchange,
451        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
452            Box::pin(async move {
453                ex.input.body = camel_api::Body::Text("recovered".to_string());
454                camel_api::PipelineOutcome::Completed(ex)
455            })
456        }
457    }
458
459    /// Test 1: Stop inside a destination propagates immediately (no failover).
460    /// First destination mutates + Stops; second destination is NOT tried.
461    #[tokio::test]
462    async fn load_balance_child_stop_propagates() {
463        let count = Arc::new(AtomicUsize::new(0));
464        let mut seg = LoadBalanceSegment {
465            destinations: vec![
466                camel_api::OutcomeSegment::new(Box::new(StoppingBody)),
467                camel_api::OutcomeSegment::new(Box::new(RecordingBody(count.clone()))),
468            ],
469            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
470            round_robin_index: Arc::new(AtomicUsize::new(0)),
471        };
472
473        let ex = Exchange::new(Message::new("trigger"));
474        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
475
476        match result {
477            camel_api::PipelineOutcome::Stopped(ex) => {
478                assert_eq!(
479                    ex.input.body.as_text(),
480                    Some("lb-stopped"),
481                    "Stopped exchange must preserve mutation"
482                );
483            }
484            other => panic!("expected PipelineOutcome::Stopped, got {other:?}"),
485        }
486        assert_eq!(
487            count.load(Ordering::SeqCst),
488            0,
489            "second destination must NOT be tried when first is Stopped"
490        );
491    }
492
493    /// Test 2: Failover strategy retries on failure. First destination fails,
494    /// second destination succeeds.
495    #[tokio::test]
496    async fn load_balance_child_failure_retries_whole_step() {
497        let mut seg = LoadBalanceSegment {
498            destinations: vec![
499                camel_api::OutcomeSegment::new(Box::new(FailingBody)),
500                camel_api::OutcomeSegment::new(Box::new(RecoveringBody)),
501            ],
502            strategy: camel_api::LoadBalanceStrategy::Failover,
503            round_robin_index: Arc::new(AtomicUsize::new(0)),
504        };
505
506        let ex = Exchange::new(Message::new("trigger"));
507        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
508
509        match result {
510            camel_api::PipelineOutcome::Completed(ex) => {
511                assert_eq!(
512                    ex.input.body.as_text(),
513                    Some("recovered"),
514                    "failover must produce the second destination's output"
515                );
516            }
517            other => panic!("expected PipelineOutcome::Completed, got {other:?}"),
518        }
519    }
520
521    /// Test 3: Round-robin strategy distributes across destinations.
522    /// 3 sequential calls hit each destination once.
523    #[tokio::test]
524    async fn load_balance_strategy_selection_preserved() {
525        let c1 = Arc::new(AtomicUsize::new(0));
526        let c2 = Arc::new(AtomicUsize::new(0));
527        let c3 = Arc::new(AtomicUsize::new(0));
528
529        let mut seg = LoadBalanceSegment {
530            destinations: vec![
531                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c1.clone()))),
532                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c2.clone()))),
533                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c3.clone()))),
534            ],
535            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
536            round_robin_index: Arc::new(AtomicUsize::new(0)),
537        };
538
539        for _ in 0..3 {
540            let ex = Exchange::new(Message::new("test"));
541            let _result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
542        }
543
544        assert_eq!(
545            c1.load(Ordering::SeqCst),
546            1,
547            "round-robin: dest 0 call count"
548        );
549        assert_eq!(
550            c2.load(Ordering::SeqCst),
551            1,
552            "round-robin: dest 1 call count"
553        );
554        assert_eq!(
555            c3.load(Ordering::SeqCst),
556            1,
557            "round-robin: dest 2 call count"
558        );
559    }
560
561    #[tokio::test]
562    async fn test_weighted_sum_does_not_overflow_u32() {
563        let ok_processor = BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }));
564        let endpoints = vec![ok_processor.clone(), ok_processor];
565        // u32::MAX + 1 would overflow u32 sum, but u64 sum handles it.
566        let weights = vec![("a".to_string(), u32::MAX), ("b".to_string(), 1)];
567        let result =
568            process_weighted(Exchange::new(Message::new("test")), endpoints, &weights).await;
569        assert!(
570            result.is_ok(),
571            "u32::MAX + 1 must not overflow with u64 sum: {:?}",
572            result.err()
573        );
574    }
575
576    /// Test: failover exhaustion preserves the LAST destination's error,
577    /// NOT a generic "all destinations exhausted" message.
578    #[tokio::test]
579    async fn load_balance_segment_failover_exhaustion_preserves_last_error() {
580        let err1 = CamelError::ProcessorError("first-dest-failed".to_string());
581        let err2 = CamelError::ProcessorError("second-dest-failed".to_string());
582
583        struct FailWith(CamelError);
584        impl camel_api::OutcomePipeline for FailWith {
585            fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
586                Box::new(FailWith(self.0.clone()))
587            }
588            fn run<'a>(
589                &'a mut self,
590                _ex: Exchange,
591            ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
592                let e = self.0.clone();
593                Box::pin(async move { camel_api::PipelineOutcome::Failed(e) })
594            }
595        }
596
597        let mut seg = LoadBalanceSegment {
598            destinations: vec![
599                camel_api::OutcomeSegment::new(Box::new(FailWith(err1))),
600                camel_api::OutcomeSegment::new(Box::new(FailWith(err2.clone()))),
601            ],
602            strategy: camel_api::LoadBalanceStrategy::Failover,
603            round_robin_index: Arc::new(AtomicUsize::new(0)),
604        };
605
606        let ex = Exchange::new(Message::new("test"));
607        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
608
609        match result {
610            camel_api::PipelineOutcome::Failed(err) => {
611                assert_eq!(
612                    err.to_string(),
613                    err2.to_string(),
614                    "failover exhaustion must return the LAST destination error, not a generic message"
615                );
616            }
617            other => panic!(
618                "expected PipelineOutcome::Failed(last error), got {:?}",
619                other
620            ),
621        }
622    }
623}
camel_processor/load_balancer.rs

camel_processor/
load_balancer.rs
