1use futures::{StreamExt, pin_mut};
2use std::future::Future;
3use std::pin::Pin;
4use std::task::{Context, Poll};
5use tokio_util::sync::CancellationToken;
6use tower::Service;
7
8use camel_api::{
9 AggregationStrategy, Body, BoxProcessor, CamelError, Exchange, StreamingSplitExpression, Value,
10};
11
12pub const CAMEL_SPLIT_INDEX: &str = "CamelSplitIndex";
13pub const CAMEL_SPLIT_SIZE: &str = "CamelSplitSize";
14pub const CAMEL_SPLIT_COMPLETE: &str = "CamelSplitComplete";
15
16#[derive(Clone)]
17pub struct StreamingSplitterService {
18 expression: StreamingSplitExpression,
19 sub_pipeline: BoxProcessor,
20 aggregation: AggregationStrategy,
21 stop_on_exception: bool,
22 cancel_token: CancellationToken,
23}
24
25impl StreamingSplitterService {
26 pub fn new(
27 expression: StreamingSplitExpression,
28 sub_pipeline: BoxProcessor,
29 aggregation: AggregationStrategy,
30 stop_on_exception: bool,
31 ) -> Self {
32 Self {
33 expression,
34 sub_pipeline,
35 aggregation,
36 stop_on_exception,
37 cancel_token: CancellationToken::new(),
38 }
39 }
40
41 pub fn cancel(&self) {
42 self.cancel_token.cancel();
43 }
44
45 pub fn is_cancelled(&self) -> bool {
46 self.cancel_token.is_cancelled()
47 }
48}
49
50impl Service<Exchange> for StreamingSplitterService {
51 type Response = Exchange;
52 type Error = CamelError;
53 type Future = Pin<Box<dyn Future<Output = Result<Exchange, CamelError>> + Send>>;
54
55 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
56 self.sub_pipeline.poll_ready(cx)
57 }
58
59 fn call(&mut self, exchange: Exchange) -> Self::Future {
60 let mut original = exchange.clone();
61 if matches!(original.input.body, Body::Stream(_)) {
62 original.input.body = Body::Empty;
63 }
64 let expression = self.expression.clone();
65 let sub_pipeline = self.sub_pipeline.clone();
66 let aggregation = self.aggregation.clone();
67 let stop_on_exception = self.stop_on_exception;
68 let cancel_token = self.cancel_token.clone();
69
70 Box::pin(async move {
71 let stream = expression(exchange);
72 pin_mut!(stream);
73
74 let mut acc: Option<Exchange> = None;
75 let mut acc_bodies: Vec<Value> = Vec::new();
76 let mut index: u64 = 0;
77
78 let mut current = stream.next().await;
95
96 while let Some(fragment_result) = current.take() {
97 if cancel_token.is_cancelled() {
98 return Err(CamelError::ProcessorError(
99 "StreamingSplitter cancelled".to_string(),
100 ));
101 }
102
103 let fragment = fragment_result?;
104
105 let next = stream.next().await;
107 let is_last = next.is_none();
108
109 let mut fragment = fragment;
110 fragment.set_property(CAMEL_SPLIT_INDEX, Value::from(index));
111 fragment.set_property(CAMEL_SPLIT_COMPLETE, Value::Bool(is_last));
112 if is_last {
115 fragment.set_property(CAMEL_SPLIT_SIZE, Value::from(index + 1));
116 }
117
118 let mut pipeline = sub_pipeline.clone();
119 let ready = tower::ServiceExt::ready(&mut pipeline).await;
120 let result = match ready {
121 Ok(svc) => svc.call(fragment).await,
122 Err(e) => Err(e),
123 };
124
125 match result {
126 Ok(processed) => {
127 match &aggregation {
128 AggregationStrategy::CollectAll => {
129 let v = match &processed.input.body {
130 Body::Text(s) => Value::String(s.clone()),
131 Body::Json(v) => v.clone(),
132 Body::Xml(s) => Value::String(s.clone()),
133 Body::Bytes(b) => {
134 Value::String(String::from_utf8_lossy(b).into_owned())
135 }
136 Body::Stream(_) => {
137 return Err(CamelError::TypeConversionFailed(
138 "StreamingSplitter CollectAll cannot aggregate Body::Stream — use 'stream_cache' or 'convert_body_to' before this step".to_string(),
139 ));
140 }
141 _ => Value::Null,
143 };
144 acc_bodies.push(v);
145 }
146 AggregationStrategy::Custom(fold_fn) => {
147 acc = Some(match acc {
148 Some(prev) => fold_fn(prev, processed),
149 None => processed,
150 });
151 }
152 _ => {
153 acc = Some(processed);
154 }
155 }
156 index += 1;
157 }
158 Err(e) => {
159 if stop_on_exception {
160 return Err(e);
161 }
162 index += 1;
163 }
164 }
165
166 current = next;
167 }
168
169 match &aggregation {
170 AggregationStrategy::Original => Ok(original),
171 AggregationStrategy::CollectAll => {
172 let mut out = original;
173 out.input.body = Body::Json(Value::Array(acc_bodies));
174 Ok(out)
175 }
176 _ => Ok(acc.unwrap_or(original)),
179 }
180 })
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187 use bytes::Bytes;
188 use camel_api::{BoxProcessorExt, Message, StreamBody, StreamMetadata};
189 use futures::stream;
190 use std::sync::Arc;
191 use tokio::sync::Mutex;
192 use tower::ServiceExt;
193
194 use crate::stream_codec::{StreamSplitInput, resolve_format, resolve_incremental_codec};
195
196 fn passthrough_pipeline() -> BoxProcessor {
197 BoxProcessor::from_fn(|ex| Box::pin(async move { Ok(ex) }))
198 }
199
200 fn uppercase_pipeline() -> BoxProcessor {
201 BoxProcessor::from_fn(|mut ex: Exchange| {
202 Box::pin(async move {
203 if let Body::Text(s) = &ex.input.body {
204 ex.input.body = Body::Text(s.to_uppercase());
205 }
206 Ok(ex)
207 })
208 })
209 }
210
211 fn make_exchange(text: &str) -> Exchange {
212 Exchange::new(Message::new(text))
213 }
214
215 fn test_expression(fragments: Vec<Exchange>) -> StreamingSplitExpression {
216 Arc::new(move |_| {
217 let frags = fragments.clone();
218 Box::pin(stream::iter(frags.into_iter().map(Ok)))
219 })
220 }
221
222 fn error_expression() -> StreamingSplitExpression {
223 Arc::new(|_| {
224 Box::pin(stream::iter(vec![Err(CamelError::ProcessorError(
225 "stream error".to_string(),
226 ))]))
227 })
228 }
229
230 fn ndjson_stream_expression(config: camel_api::StreamSplitConfig) -> StreamingSplitExpression {
234 Arc::new(move |exchange: Exchange| {
235 let config = config.clone();
236 let (stream_body, parent) = match &exchange.input.body {
237 Body::Stream(sb) => (sb.clone(), {
238 let mut p = exchange.clone();
239 p.input.body = Body::Empty;
240 p
241 }),
242 _ => {
243 return Box::pin(futures::stream::once(async {
244 Err(CamelError::ProcessorError(
245 "streaming split requires Body::Stream".into(),
246 ))
247 }));
248 }
249 };
250
251 let stream = match stream_body.stream.try_lock() {
252 Ok(mut guard) => match guard.take() {
253 Some(s) => s,
254 None => {
255 return Box::pin(futures::stream::once(async {
256 Err(CamelError::ProcessorError(
257 "stream body already consumed".into(),
258 ))
259 }));
260 }
261 },
262 Err(_) => {
263 return Box::pin(futures::stream::once(async {
264 Err(CamelError::ProcessorError("stream body locked".into()))
265 }));
266 }
267 };
268
269 let input = StreamSplitInput {
270 parent,
271 stream,
272 metadata: stream_body.metadata,
273 };
274
275 match resolve_format(&config.format, &input.metadata) {
276 Ok(f) => {
277 let codec = resolve_incremental_codec(&f);
278 let codec = match codec {
279 Ok(c) => c,
280 Err(e) => return Box::pin(futures::stream::once(async { Err(e) })),
281 };
282 codec.split(input, config)
283 }
284 Err(e) => Box::pin(futures::stream::once(async { Err(e) })),
285 }
286 })
287 }
288
289 #[tokio::test]
294 async fn test_ndjson_body_stream_streaming_split() {
295 let ndjson_lines: Vec<Result<Bytes, CamelError>> = vec![
298 Ok(Bytes::from("{\"id\":1,\"name\":\"a\"}\n")),
299 Ok(Bytes::from("{\"id\":2,\"name\":\"b\"}\n")),
300 Ok(Bytes::from("{\"id\":3,\"name\":\"c\"}\n")),
301 ];
302 let byte_stream = futures::stream::iter(ndjson_lines);
303
304 let stream_body = StreamBody {
305 stream: Arc::new(Mutex::new(Some(Box::pin(byte_stream)))),
306 metadata: StreamMetadata {
307 content_type: Some("application/x-ndjson".into()),
308 size_hint: None,
309 origin: Some("test://ndjson".into()),
310 },
311 };
312
313 let ex = Exchange::new(Message::new(Body::Stream(stream_body)));
314
315 let split_config = camel_api::StreamSplitConfig {
317 format: camel_api::StreamSplitFormat::Ndjson,
318 ..Default::default()
319 };
320
321 #[allow(clippy::type_complexity)]
323 let fragments: Arc<
324 Mutex<Vec<(Option<serde_json::Value>, Option<Value>, Option<Value>)>>,
325 > = Arc::new(Mutex::new(Vec::new()));
326 let fragments_clone = Arc::clone(&fragments);
327 let recorder = BoxProcessor::from_fn(move |ex: Exchange| {
328 let frags = Arc::clone(&fragments_clone);
329 Box::pin(async move {
330 let body_json = match &ex.input.body {
331 Body::Json(v) => Some(v.clone()),
332 _ => None,
333 };
334 let split_index = ex.property(CAMEL_SPLIT_INDEX).cloned();
335 let split_complete = ex.property(CAMEL_SPLIT_COMPLETE).cloned();
336 let mut guard = frags.lock().await;
337 guard.push((body_json, split_index, split_complete));
338 Ok(ex)
339 })
340 });
341
342 let expression = ndjson_stream_expression(split_config);
343
344 let mut splitter = StreamingSplitterService::new(
346 expression,
347 recorder,
348 AggregationStrategy::CollectAll,
349 true, );
351
352 let result = splitter
353 .ready()
354 .await
355 .expect("splitter ready")
356 .call(ex)
357 .await
358 .expect("splitter call");
359
360 let guard = fragments.lock().await;
362
363 assert_eq!(guard.len(), 3, "expected 3 NDJSON fragments");
365
366 for (i, (body_json, _idx, _complete)) in guard.iter().enumerate() {
368 assert!(
369 body_json.is_some(),
370 "fragment {i}: expected Body::Json body, got non-Json"
371 );
372 }
373
374 for (i, (_body, idx, _complete)) in guard.iter().enumerate() {
376 assert_eq!(
377 *idx,
378 Some(Value::Number(serde_json::Number::from(i as u64))),
379 "fragment {i}: CamelSplitIndex mismatch"
380 );
381 }
382
383 assert_eq!(
385 guard[0].2,
386 Some(Value::Bool(false)),
387 "first fragment: CamelSplitComplete should be false"
388 );
389 assert_eq!(
390 guard[1].2,
391 Some(Value::Bool(false)),
392 "second fragment: CamelSplitComplete should be false"
393 );
394 assert_eq!(
395 guard[2].2,
396 Some(Value::Bool(true)),
397 "last fragment: CamelSplitComplete should be true"
398 );
399
400 match &result.input.body {
402 Body::Json(v) => {
403 let arr = v.as_array().expect("CollectAll result should be array");
404 assert_eq!(arr.len(), 3);
405 assert_eq!(arr[0], serde_json::json!({"id":1,"name":"a"}));
406 assert_eq!(arr[1], serde_json::json!({"id":2,"name":"b"}));
407 assert_eq!(arr[2], serde_json::json!({"id":3,"name":"c"}));
408 }
409 other => panic!("expected Body::Json from CollectAll, got {other:?}"),
410 }
411
412 assert!(
415 matches!(result.input.body, Body::Json(_)),
416 "aggregate body should be Json, not Stream"
417 );
418 }
419
420 #[tokio::test]
425 async fn test_ndjson_body_stream_empty_stream() {
426 let byte_stream = futures::stream::iter(Vec::<Result<Bytes, CamelError>>::new());
429
430 let stream_body = StreamBody {
431 stream: Arc::new(Mutex::new(Some(Box::pin(byte_stream)))),
432 metadata: StreamMetadata {
433 content_type: Some("application/x-ndjson".into()),
434 size_hint: None,
435 origin: None,
436 },
437 };
438
439 let mut ex = Exchange::new(Message::new(Body::Stream(stream_body)));
440 ex.set_property("trace_id", Value::String("empty-test".into()));
441
442 let split_config = camel_api::StreamSplitConfig {
443 format: camel_api::StreamSplitFormat::Ndjson,
444 ..Default::default()
445 };
446
447 let expression = ndjson_stream_expression(split_config);
448
449 let mut splitter = StreamingSplitterService::new(
451 expression,
452 passthrough_pipeline(),
453 AggregationStrategy::CollectAll,
454 true,
455 );
456
457 let result = splitter
458 .ready()
459 .await
460 .expect("splitter ready")
461 .call(ex)
462 .await
463 .expect("splitter call");
464
465 match &result.input.body {
468 Body::Json(v) => {
469 let arr = v.as_array().expect("CollectAll result should be array");
470 assert!(
471 arr.is_empty(),
472 "empty stream should produce empty array, got {arr:?}"
473 );
474 }
475 other => {
476 panic!("expected Body::Json([]) from CollectAll on empty stream, got {other:?}")
477 }
478 }
479
480 assert_eq!(
482 result.property("trace_id"),
483 Some(&Value::String("empty-test".into()))
484 );
485 }
486
487 #[tokio::test]
488 async fn test_streaming_sequential_last_wins() {
489 let expr = test_expression(vec![
490 make_exchange("a"),
491 make_exchange("b"),
492 make_exchange("c"),
493 ]);
494 let mut svc = StreamingSplitterService::new(
495 expr,
496 uppercase_pipeline(),
497 AggregationStrategy::LastWins,
498 true,
499 );
500
501 let result = svc
502 .ready()
503 .await
504 .unwrap()
505 .call(make_exchange("original"))
506 .await
507 .unwrap();
508 assert_eq!(result.input.body.as_text(), Some("C"));
509 }
510
511 #[tokio::test]
512 async fn test_streaming_sequential_original() {
513 let expr = test_expression(vec![make_exchange("a"), make_exchange("b")]);
514 let mut svc = StreamingSplitterService::new(
515 expr,
516 uppercase_pipeline(),
517 AggregationStrategy::Original,
518 true,
519 );
520
521 let result = svc
522 .ready()
523 .await
524 .unwrap()
525 .call(make_exchange("original"))
526 .await
527 .unwrap();
528 assert_eq!(result.input.body.as_text(), Some("original"));
529 }
530
531 #[tokio::test]
532 async fn test_streaming_stop_on_exception() {
533 let expr = test_expression(vec![make_exchange("a"), make_exchange("b")]);
534 let fail_pipeline = BoxProcessor::from_fn(|_| {
535 Box::pin(async { Err(CamelError::ProcessorError("boom".into())) })
536 });
537 let mut svc =
538 StreamingSplitterService::new(expr, fail_pipeline, AggregationStrategy::LastWins, true);
539
540 let result = svc
541 .ready()
542 .await
543 .unwrap()
544 .call(make_exchange("original"))
545 .await;
546 assert!(result.is_err());
547 }
548
549 #[tokio::test]
550 async fn test_streaming_empty_stream() {
551 let expr: StreamingSplitExpression = Arc::new(|_| Box::pin(futures::stream::empty()));
552 let mut svc = StreamingSplitterService::new(
553 expr,
554 passthrough_pipeline(),
555 AggregationStrategy::LastWins,
556 true,
557 );
558
559 let mut ex = make_exchange("original");
560 ex.set_property("marker", Value::Bool(true));
561 let result = svc.ready().await.unwrap().call(ex).await.unwrap();
562 assert_eq!(result.input.body.as_text(), Some("original"));
563 assert_eq!(result.property("marker"), Some(&Value::Bool(true)));
564 }
565
566 #[tokio::test]
567 async fn test_streaming_error_in_expression() {
568 let mut svc = StreamingSplitterService::new(
569 error_expression(),
570 passthrough_pipeline(),
571 AggregationStrategy::LastWins,
572 true,
573 );
574
575 let result = svc
576 .ready()
577 .await
578 .unwrap()
579 .call(make_exchange("original"))
580 .await;
581 assert!(result.is_err());
582 }
583
584 #[tokio::test]
585 async fn test_streaming_cancellation() {
586 let expr = test_expression(vec![make_exchange("a"), make_exchange("b")]);
587 let slow_pipeline = BoxProcessor::from_fn(|ex| {
588 Box::pin(async move {
589 tokio::time::sleep(std::time::Duration::from_secs(60)).await;
590 Ok(ex)
591 })
592 });
593 let svc =
594 StreamingSplitterService::new(expr, slow_pipeline, AggregationStrategy::LastWins, true);
595 svc.cancel();
596
597 let mut svc_clone = svc.clone();
598 let result = svc_clone
599 .ready()
600 .await
601 .unwrap()
602 .call(make_exchange("original"))
603 .await;
604 assert!(result.is_err());
605 }
606
607 #[tokio::test]
608 async fn test_streaming_sequential_collect_all() {
609 let expr = test_expression(vec![
610 make_exchange("a"),
611 make_exchange("b"),
612 make_exchange("c"),
613 ]);
614 let mut svc = StreamingSplitterService::new(
615 expr,
616 uppercase_pipeline(),
617 AggregationStrategy::CollectAll,
618 true,
619 );
620
621 let result = svc
622 .ready()
623 .await
624 .unwrap()
625 .call(make_exchange("original"))
626 .await
627 .unwrap();
628 let expected = serde_json::json!(["A", "B", "C"]);
629 match &result.input.body {
630 Body::Json(v) => assert_eq!(*v, expected),
631 other => panic!("expected JSON body, got {other:?}"),
632 }
633 }
634
635 #[tokio::test]
636 async fn test_streaming_sequential_custom_aggregation() {
637 let joiner: Arc<dyn Fn(Exchange, Exchange) -> Exchange + Send + Sync> =
638 Arc::new(|mut acc: Exchange, next: Exchange| {
639 let acc_text = acc.input.body.as_text().unwrap_or("").to_string();
640 let next_text = next.input.body.as_text().unwrap_or("").to_string();
641 acc.input.body = Body::Text(format!("{acc_text}+{next_text}"));
642 acc
643 });
644
645 let expr = test_expression(vec![
646 make_exchange("a"),
647 make_exchange("b"),
648 make_exchange("c"),
649 ]);
650 let mut svc = StreamingSplitterService::new(
651 expr,
652 uppercase_pipeline(),
653 AggregationStrategy::Custom(joiner),
654 true,
655 );
656
657 let result = svc
658 .ready()
659 .await
660 .unwrap()
661 .call(make_exchange("original"))
662 .await
663 .unwrap();
664 assert_eq!(result.input.body.as_text(), Some("A+B+C"));
665 }
666
667 #[tokio::test]
668 async fn test_streaming_error_continue_on_exception() {
669 let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
670 let count_clone = call_count.clone();
671 let fail_on_first = BoxProcessor::from_fn(move |ex: Exchange| {
672 let count = count_clone.clone();
673 Box::pin(async move {
674 let n = count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
675 if n == 0 {
676 Err(CamelError::ProcessorError("first fails".into()))
677 } else {
678 Ok(ex)
679 }
680 })
681 });
682
683 let expr = test_expression(vec![make_exchange("a"), make_exchange("b")]);
684 let mut svc = StreamingSplitterService::new(
685 expr,
686 fail_on_first,
687 AggregationStrategy::LastWins,
688 false,
689 );
690
691 let result = svc
692 .ready()
693 .await
694 .unwrap()
695 .call(make_exchange("original"))
696 .await
697 .unwrap();
698 assert_eq!(result.input.body.as_text(), Some("b"));
699 assert_eq!(call_count.load(std::sync::atomic::Ordering::SeqCst), 2);
700 }
701
702 #[tokio::test]
703 async fn test_streaming_metadata_lookahead() {
704 let recorder = BoxProcessor::from_fn(|ex: Exchange| {
705 Box::pin(async move {
706 let idx = ex.property(CAMEL_SPLIT_INDEX).cloned();
707 let complete = ex.property(CAMEL_SPLIT_COMPLETE).cloned();
708 let body = serde_json::json!({
709 "index": idx,
710 "complete": complete,
711 });
712 let mut out = ex;
713 out.input.body = Body::Json(body);
714 Ok(out)
715 })
716 });
717
718 let expr = test_expression(vec![
719 make_exchange("x"),
720 make_exchange("y"),
721 make_exchange("z"),
722 ]);
723 let mut svc =
724 StreamingSplitterService::new(expr, recorder, AggregationStrategy::CollectAll, true);
725
726 let result = svc
727 .ready()
728 .await
729 .unwrap()
730 .call(make_exchange("original"))
731 .await
732 .unwrap();
733 let expected = serde_json::json!([
734 {"index": 0, "complete": false},
735 {"index": 1, "complete": false},
736 {"index": 2, "complete": true},
737 ]);
738 match &result.input.body {
739 Body::Json(v) => assert_eq!(*v, expected),
740 other => panic!("expected JSON body, got {other:?}"),
741 }
742 }
743
744 #[tokio::test]
745 async fn test_streaming_split_sanitizes_stream_body_in_original() {
746 let chunks = vec![Ok(Bytes::from("line1\n"))];
747 let stream = futures::stream::iter(chunks);
748 let sb = StreamBody {
749 stream: Arc::new(Mutex::new(Some(Box::pin(stream)))),
750 metadata: Default::default(),
751 };
752 let ex = Exchange::new(Message::new(Body::Stream(sb)));
753
754 let expression =
755 test_expression(vec![Exchange::new(Message::new(Body::Text("frag".into())))]);
756 let sub_pipeline = passthrough_pipeline();
757 let mut splitter = StreamingSplitterService::new(
758 expression,
759 sub_pipeline,
760 AggregationStrategy::Original,
761 true,
762 );
763
764 let result = splitter
765 .ready()
766 .await
767 .expect("ready")
768 .call(ex)
769 .await
770 .expect("call");
771 assert!(
772 matches!(result.input.body, Body::Empty),
773 "original body should be sanitized to Empty"
774 );
775 }
776
777 #[tokio::test]
784 async fn test_streaming_split_size_set_only_on_last_fragment() {
785 let recorder = BoxProcessor::from_fn(|ex: Exchange| {
786 Box::pin(async move {
787 let idx = ex.property(CAMEL_SPLIT_INDEX).cloned();
788 let size = ex.property(CAMEL_SPLIT_SIZE).cloned();
789 let complete = ex.property(CAMEL_SPLIT_COMPLETE).cloned();
790 let body = serde_json::json!({
791 "index": idx,
792 "size": size,
793 "complete": complete,
794 });
795 let mut out = ex;
796 out.input.body = Body::Json(body);
797 Ok(out)
798 })
799 });
800
801 let expr = test_expression(vec![
802 make_exchange("x"),
803 make_exchange("y"),
804 make_exchange("z"),
805 ]);
806 let mut svc =
807 StreamingSplitterService::new(expr, recorder, AggregationStrategy::CollectAll, true);
808
809 let result = svc
810 .ready()
811 .await
812 .unwrap()
813 .call(make_exchange("original"))
814 .await
815 .unwrap();
816
817 let expected = serde_json::json!([
819 {"index": 0, "size": serde_json::Value::Null, "complete": false},
820 {"index": 1, "size": serde_json::Value::Null, "complete": false},
821 {"index": 2, "size": 3, "complete": true},
822 ]);
823 match &result.input.body {
824 Body::Json(v) => assert_eq!(*v, expected),
825 other => panic!("expected JSON body, got {other:?}"),
826 }
827 }
828}