camel_processor/
load_balancer.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::sync::atomic::{AtomicUsize, Ordering};
5
6// ── LoadBalanceSegment (ADR-0025 OutcomePipeline) ────────────────────────
7
8/// Outcome-aware LoadBalance segment. Holds N destinations + a strategy.
9/// On each call: strategy picks ONE destination (round-robin / failover /
10/// random / weighted), runs it. If chosen destination returns Completed,
11/// return Completed. If Stopped: return Stopped immediately (no failover —
12/// Stop is successful control flow). If Failed: strategy decides (failover
13/// retries next dest, others return Failed).
14///
15/// This differs from Multicast (which runs all branches) — LoadBalance picks
16/// exactly one. The parallel cancellation logic from T13/T15 does NOT apply.
17#[derive(Clone)]
18pub struct LoadBalanceSegment {
19    pub destinations: Vec<camel_api::OutcomeSegment>,
20    pub strategy: camel_api::LoadBalanceStrategy,
21    /// Shared round-robin index for interior mutability across cloned segments.
22    pub round_robin_index: Arc<AtomicUsize>,
23}
24
25impl camel_api::OutcomePipeline for LoadBalanceSegment {
26    fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
27        Box::new(self.clone())
28    }
29
30    fn run<'a>(
31        &'a mut self,
32        exchange: camel_api::Exchange,
33    ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
34        Box::pin(async move {
35            let len = self.destinations.len();
36            if len == 0 {
37                return camel_api::PipelineOutcome::Completed(exchange);
38            }
39
40            let start_idx = match &self.strategy {
41                camel_api::LoadBalanceStrategy::Random => rand::random_range(0..len),
42                camel_api::LoadBalanceStrategy::Weighted(weights) => pick_weighted(weights, len),
43                camel_api::LoadBalanceStrategy::Failover => 0,
44                // RoundRobin and any future variant default to round-robin indexing.
45                _ => self.round_robin_index.fetch_add(1, Ordering::SeqCst) % len,
46            };
47
48            let mut idx = start_idx;
49            let mut last_err: Option<camel_api::CamelError> = None;
50            loop {
51                if idx >= len {
52                    return camel_api::PipelineOutcome::Failed(last_err.unwrap_or_else(|| {
53                        camel_api::CamelError::ProcessorError(
54                            "load_balance: all destinations exhausted".to_string(),
55                        )
56                    }));
57                }
58                match self.destinations[idx].run(exchange.branch_clone()).await {
59                    camel_api::PipelineOutcome::Completed(ex) => {
60                        return camel_api::PipelineOutcome::Completed(ex);
61                    }
62                    camel_api::PipelineOutcome::Stopped(ex) => {
63                        return camel_api::PipelineOutcome::Stopped(ex);
64                    }
65                    camel_api::PipelineOutcome::Failed(err) => match self.strategy {
66                        camel_api::LoadBalanceStrategy::Failover => {
67                            last_err = Some(err);
68                            idx += 1;
69                            continue;
70                        }
71                        _ => return camel_api::PipelineOutcome::Failed(err),
72                    },
73                }
74            }
75        })
76    }
77}
78
79/// Pick a destination index using weighted random selection.
80fn pick_weighted(weights: &[(String, u32)], len: usize) -> usize {
81    if weights.is_empty() || len == 0 {
82        return 0;
83    }
84    let numeric_weights: Vec<u64> = weights.iter().map(|(_, w)| *w as u64).collect();
85    let Some(total) = numeric_weights
86        .iter()
87        .try_fold(0u64, |acc, w| acc.checked_add(*w))
88    else {
89        return 0;
90    };
91    if total == 0 {
92        return 0;
93    }
94    let mut r = rand::random::<u64>() % total;
95    for (i, w) in numeric_weights.iter().enumerate() {
96        if r < *w {
97            return i.min(len - 1);
98        }
99        r -= w;
100    }
101    len - 1
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use camel_api::{CamelError, Exchange, Message};
108    use std::sync::Mutex;
109
110    // ── LoadBalanceSegment tests (ADR-0025 OutcomePipeline parity) ───
111
112    /// OutcomePipeline body that mutates exchange body to "lb-stopped" then returns Stopped.
113    struct StoppingBody;
114    impl camel_api::OutcomePipeline for StoppingBody {
115        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
116            Box::new(StoppingBody)
117        }
118        fn run<'a>(
119            &'a mut self,
120            mut ex: Exchange,
121        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
122            Box::pin(async move {
123                ex.input.body = camel_api::Body::Text("lb-stopped".to_string());
124                camel_api::PipelineOutcome::Stopped(ex)
125            })
126        }
127    }
128
129    /// OutcomePipeline body that records invocation count via shared counter.
130    struct RecordingBody(Arc<AtomicUsize>);
131    impl camel_api::OutcomePipeline for RecordingBody {
132        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
133            Box::new(RecordingBody(Arc::clone(&self.0)))
134        }
135        fn run<'a>(
136            &'a mut self,
137            ex: Exchange,
138        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
139            let count = Arc::clone(&self.0);
140            Box::pin(async move {
141                count.fetch_add(1, Ordering::SeqCst);
142                camel_api::PipelineOutcome::Completed(ex)
143            })
144        }
145    }
146
147    /// OutcomePipeline body that always fails with ProcessorError.
148    struct FailingBody;
149    impl camel_api::OutcomePipeline for FailingBody {
150        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
151            Box::new(FailingBody)
152        }
153        fn run<'a>(
154            &'a mut self,
155            _ex: Exchange,
156        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
157            Box::pin(async {
158                camel_api::PipelineOutcome::Failed(CamelError::ProcessorError(
159                    "intentional fail".to_string(),
160                ))
161            })
162        }
163    }
164
165    /// OutcomePipeline body that mutates body to "recovered" then completes.
166    struct RecoveringBody;
167    impl camel_api::OutcomePipeline for RecoveringBody {
168        fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
169            Box::new(RecoveringBody)
170        }
171        fn run<'a>(
172            &'a mut self,
173            mut ex: Exchange,
174        ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
175            Box::pin(async move {
176                ex.input.body = camel_api::Body::Text("recovered".to_string());
177                camel_api::PipelineOutcome::Completed(ex)
178            })
179        }
180    }
181
182    /// Test 1: Stop inside a destination propagates immediately (no failover).
183    /// First destination mutates + Stops; second destination is NOT tried.
184    #[tokio::test]
185    async fn load_balance_child_stop_propagates() {
186        let count = Arc::new(AtomicUsize::new(0));
187        let mut seg = LoadBalanceSegment {
188            destinations: vec![
189                camel_api::OutcomeSegment::new(Box::new(StoppingBody)),
190                camel_api::OutcomeSegment::new(Box::new(RecordingBody(count.clone()))),
191            ],
192            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
193            round_robin_index: Arc::new(AtomicUsize::new(0)),
194        };
195
196        let ex = Exchange::new(Message::new("trigger"));
197        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
198
199        match result {
200            camel_api::PipelineOutcome::Stopped(ex) => {
201                assert_eq!(
202                    ex.input.body.as_text(),
203                    Some("lb-stopped"),
204                    "Stopped exchange must preserve mutation"
205                );
206            }
207            other => panic!("expected PipelineOutcome::Stopped, got {other:?}"),
208        }
209        assert_eq!(
210            count.load(Ordering::SeqCst),
211            0,
212            "second destination must NOT be tried when first is Stopped"
213        );
214    }
215
216    /// Test 2: Failover strategy retries on failure. First destination fails,
217    /// second destination succeeds.
218    #[tokio::test]
219    async fn load_balance_child_failure_retries_whole_step() {
220        let mut seg = LoadBalanceSegment {
221            destinations: vec![
222                camel_api::OutcomeSegment::new(Box::new(FailingBody)),
223                camel_api::OutcomeSegment::new(Box::new(RecoveringBody)),
224            ],
225            strategy: camel_api::LoadBalanceStrategy::Failover,
226            round_robin_index: Arc::new(AtomicUsize::new(0)),
227        };
228
229        let ex = Exchange::new(Message::new("trigger"));
230        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
231
232        match result {
233            camel_api::PipelineOutcome::Completed(ex) => {
234                assert_eq!(
235                    ex.input.body.as_text(),
236                    Some("recovered"),
237                    "failover must produce the second destination's output"
238                );
239            }
240            other => panic!("expected PipelineOutcome::Completed, got {other:?}"),
241        }
242    }
243
244    /// Test 3: Round-robin strategy distributes across destinations.
245    /// 3 sequential calls hit each destination once.
246    #[tokio::test]
247    async fn load_balance_strategy_selection_preserved() {
248        let c1 = Arc::new(AtomicUsize::new(0));
249        let c2 = Arc::new(AtomicUsize::new(0));
250        let c3 = Arc::new(AtomicUsize::new(0));
251
252        let mut seg = LoadBalanceSegment {
253            destinations: vec![
254                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c1.clone()))),
255                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c2.clone()))),
256                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c3.clone()))),
257            ],
258            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
259            round_robin_index: Arc::new(AtomicUsize::new(0)),
260        };
261
262        for _ in 0..3 {
263            let ex = Exchange::new(Message::new("test"));
264            let _result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
265        }
266
267        assert_eq!(
268            c1.load(Ordering::SeqCst),
269            1,
270            "round-robin: dest 0 call count"
271        );
272        assert_eq!(
273            c2.load(Ordering::SeqCst),
274            1,
275            "round-robin: dest 1 call count"
276        );
277        assert_eq!(
278            c3.load(Ordering::SeqCst),
279            1,
280            "round-robin: dest 2 call count"
281        );
282    }
283
284    /// Test: failover exhaustion preserves the LAST destination's error,
285    /// NOT a generic "all destinations exhausted" message.
286    #[tokio::test]
287    async fn load_balance_segment_failover_exhaustion_preserves_last_error() {
288        let err1 = CamelError::ProcessorError("first-dest-failed".to_string());
289        let err2 = CamelError::ProcessorError("second-dest-failed".to_string());
290
291        struct FailWith(CamelError);
292        impl camel_api::OutcomePipeline for FailWith {
293            fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
294                Box::new(FailWith(self.0.clone()))
295            }
296            fn run<'a>(
297                &'a mut self,
298                _ex: Exchange,
299            ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
300                let e = self.0.clone();
301                Box::pin(async move { camel_api::PipelineOutcome::Failed(e) })
302            }
303        }
304
305        let mut seg = LoadBalanceSegment {
306            destinations: vec![
307                camel_api::OutcomeSegment::new(Box::new(FailWith(err1))),
308                camel_api::OutcomeSegment::new(Box::new(FailWith(err2.clone()))),
309            ],
310            strategy: camel_api::LoadBalanceStrategy::Failover,
311            round_robin_index: Arc::new(AtomicUsize::new(0)),
312        };
313
314        let ex = Exchange::new(Message::new("test"));
315        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
316
317        match result {
318            camel_api::PipelineOutcome::Failed(err) => {
319                assert_eq!(
320                    err.to_string(),
321                    err2.to_string(),
322                    "failover exhaustion must return the LAST destination error, not a generic message"
323                );
324            }
325            other => panic!(
326                "expected PipelineOutcome::Failed(last error), got {:?}",
327                other
328            ),
329        }
330    }
331
332    // ── Ported from the removed Tower load-balancer tests ──────────────
333
334    #[tokio::test]
335    async fn load_balance_round_robin_distribution_six_calls() {
336        let c1 = Arc::new(AtomicUsize::new(0));
337        let c2 = Arc::new(AtomicUsize::new(0));
338        let c3 = Arc::new(AtomicUsize::new(0));
339
340        let mut seg = LoadBalanceSegment {
341            destinations: vec![
342                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c1.clone()))),
343                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c2.clone()))),
344                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c3.clone()))),
345            ],
346            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
347            round_robin_index: Arc::new(AtomicUsize::new(0)),
348        };
349
350        for _ in 0..6 {
351            let ex = Exchange::new(Message::new("test"));
352            let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
353            assert!(
354                matches!(result, camel_api::PipelineOutcome::Completed(_)),
355                "round-robin dispatch must complete; got {result:?}"
356            );
357        }
358
359        assert_eq!(c1.load(Ordering::SeqCst), 2);
360        assert_eq!(c2.load(Ordering::SeqCst), 2);
361        assert_eq!(c3.load(Ordering::SeqCst), 2);
362    }
363
364    #[tokio::test]
365    async fn load_balance_random_distribution() {
366        let c1 = Arc::new(AtomicUsize::new(0));
367        let c2 = Arc::new(AtomicUsize::new(0));
368
369        let mut seg = LoadBalanceSegment {
370            destinations: vec![
371                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c1.clone()))),
372                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c2.clone()))),
373            ],
374            strategy: camel_api::LoadBalanceStrategy::Random,
375            round_robin_index: Arc::new(AtomicUsize::new(0)),
376        };
377
378        for _ in 0..100 {
379            let ex = Exchange::new(Message::new("test"));
380            let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
381            assert!(
382                matches!(result, camel_api::PipelineOutcome::Completed(_)),
383                "random dispatch must complete; got {result:?}"
384            );
385        }
386
387        let total = c1.load(Ordering::SeqCst) + c2.load(Ordering::SeqCst);
388        assert_eq!(total, 100);
389        assert!(c1.load(Ordering::SeqCst) > 20);
390        assert!(c2.load(Ordering::SeqCst) > 20);
391    }
392
393    #[tokio::test]
394    async fn load_balance_failover_on_error() {
395        let count = Arc::new(AtomicUsize::new(0));
396
397        let mut seg = LoadBalanceSegment {
398            destinations: vec![
399                camel_api::OutcomeSegment::new(Box::new(FailingBody)),
400                camel_api::OutcomeSegment::new(Box::new(RecordingBody(count.clone()))),
401            ],
402            strategy: camel_api::LoadBalanceStrategy::Failover,
403            round_robin_index: Arc::new(AtomicUsize::new(0)),
404        };
405
406        let ex = Exchange::new(Message::new("test"));
407        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
408        assert!(matches!(result, camel_api::PipelineOutcome::Completed(_)));
409        assert_eq!(count.load(Ordering::SeqCst), 1);
410    }
411
412    #[tokio::test]
413    async fn load_balance_failover_preserves_original_exchange() {
414        // Failure retry must receive a branch clone of the original exchange.
415        let seen_body: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
416        let seen_body_clone = Arc::clone(&seen_body);
417
418        struct CaptureBody(Arc<Mutex<Option<String>>>);
419        impl camel_api::OutcomePipeline for CaptureBody {
420            fn clone_box(&self) -> Box<dyn camel_api::OutcomePipeline> {
421                Box::new(CaptureBody(Arc::clone(&self.0)))
422            }
423            fn run<'a>(
424                &'a mut self,
425                ex: Exchange,
426            ) -> Pin<Box<dyn Future<Output = camel_api::PipelineOutcome> + Send + 'a>> {
427                let seen = Arc::clone(&self.0);
428                Box::pin(async move {
429                    if let Some(text) = ex.input.body.as_text() {
430                        *seen.lock().unwrap() = Some(text.to_string());
431                    }
432                    camel_api::PipelineOutcome::Completed(ex)
433                })
434            }
435        }
436
437        let mut seg = LoadBalanceSegment {
438            destinations: vec![
439                camel_api::OutcomeSegment::new(Box::new(FailingBody)),
440                camel_api::OutcomeSegment::new(Box::new(CaptureBody(seen_body_clone))),
441            ],
442            strategy: camel_api::LoadBalanceStrategy::Failover,
443            round_robin_index: Arc::new(AtomicUsize::new(0)),
444        };
445
446        let ex = Exchange::new(Message::new("original body"));
447        let _result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
448
449        assert_eq!(
450            seen_body.lock().unwrap().as_deref(),
451            Some("original body"),
452            "retry destination must receive the original exchange body, not a blank one"
453        );
454    }
455
456    #[tokio::test]
457    async fn load_balance_empty_destinations_completes() {
458        let mut seg = LoadBalanceSegment {
459            destinations: vec![],
460            strategy: camel_api::LoadBalanceStrategy::RoundRobin,
461            round_robin_index: Arc::new(AtomicUsize::new(0)),
462        };
463
464        let ex = Exchange::new(Message::new("test"));
465        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
466
467        match result {
468            camel_api::PipelineOutcome::Completed(ex) => {
469                assert_eq!(ex.input.body.as_text(), Some("test"));
470            }
471            other => panic!("expected Completed(original), got {other:?}"),
472        }
473    }
474
475    #[tokio::test]
476    async fn load_balance_weighted_sum_does_not_overflow_u32() {
477        // u32::MAX + 1 overflows a u32 sum; the segment's u64 selection must
478        // not panic and must still dispatch to a destination.
479        let c1 = Arc::new(AtomicUsize::new(0));
480        let c2 = Arc::new(AtomicUsize::new(0));
481
482        let mut seg = LoadBalanceSegment {
483            destinations: vec![
484                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c1.clone()))),
485                camel_api::OutcomeSegment::new(Box::new(RecordingBody(c2.clone()))),
486            ],
487            strategy: camel_api::LoadBalanceStrategy::Weighted(vec![
488                ("a".to_string(), u32::MAX),
489                ("b".to_string(), 1),
490            ]),
491            round_robin_index: Arc::new(AtomicUsize::new(0)),
492        };
493
494        let ex = Exchange::new(Message::new("test"));
495        let result = camel_api::OutcomePipeline::run(&mut seg, ex).await;
496        assert!(
497            matches!(result, camel_api::PipelineOutcome::Completed(_)),
498            "u32::MAX + 1 must not overflow with u64 sum: {result:?}"
499        );
500        assert_eq!(c1.load(Ordering::SeqCst) + c2.load(Ordering::SeqCst), 1);
501    }
502}
camel_processor/load_balancer.rs

camel_processor/
load_balancer.rs
