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