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