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

camel_processor/
load_balancer.rs
