1use crate::batch_metadata::BatchMetadataBuilder;
2use crate::prelude::*;
3
4#[derive(Clone)]
5pub struct BatchOp<T> {
6 ops: Vec<Arc<dyn Op<T>>>,
7 continue_on_error: bool,
8}
9
10impl<T> std::fmt::Debug for BatchOp<T> {
11 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12 f.debug_struct("BatchOp")
13 .field("ops_count", &self.ops.len())
14 .field("continue_on_error", &self.continue_on_error)
15 .finish()
16 }
17}
18
19impl<T> BatchOp<T>
20where
21 T: Send + Sync + 'static,
22{
23 pub fn new(ops: Vec<Arc<dyn Op<T>>>) -> Self {
24 Self {
25 ops,
26 continue_on_error: false,
27 }
28 }
29
30 pub fn with_continue_on_error(mut self, continue_on_error: bool) -> Self {
31 self.continue_on_error = continue_on_error;
32 self
33 }
34
35 pub fn add_op(&mut self, op: Arc<dyn Op<T>>) {
36 self.ops.push(op);
37 }
38
39 pub fn len(&self) -> usize {
40 self.ops.len()
41 }
42
43 pub fn is_empty(&self) -> bool {
44 self.ops.is_empty()
45 }
46
47 async fn rollback_succeeded_ops(
48 &self,
49 succeeded_ops: &[Arc<dyn Op<T>>],
50 dry: &mut DryContext,
51 wet: &mut WetContext,
52 ) {
53 for op in succeeded_ops.iter().rev() {
55 if let Err(rollback_error) = op.rollback(dry, wet).await {
56 error!(
57 "Failed to rollback op {}: {}",
58 op.metadata().name,
59 rollback_error
60 );
61 } else {
62 debug!("Successfully rolled back op {}", op.metadata().name);
63 }
64 }
65 }
66}
67
68#[async_trait]
69impl<T> Op<Vec<T>> for BatchOp<T>
70where
71 T: Send + Sync + 'static,
72{
73 async fn perform(&self, dry: &mut DryContext, wet: &mut WetContext) -> OpResult<Vec<T>> {
74 let mut results = Vec::with_capacity(self.ops.len());
75 let mut errors = Vec::new();
76 let mut succeeded_ops = Vec::new(); for (index, op) in self.ops.iter().enumerate() {
79 if dry.is_aborted() {
81 self.rollback_succeeded_ops(&succeeded_ops, dry, wet).await;
83 let reason = dry
84 .abort_reason()
85 .cloned()
86 .unwrap_or_else(|| "Batch operation aborted".to_string());
87 return Err(OpError::Aborted(reason));
88 }
89
90 match op.perform(dry, wet).await {
91 Ok(result) => {
92 results.push(result);
93 succeeded_ops.push(op.clone());
94 }
95 Err(OpError::Aborted(reason)) => {
96 self.rollback_succeeded_ops(&succeeded_ops, dry, wet).await;
98 return Err(OpError::Aborted(reason));
99 }
100 Err(error) => {
101 if self.continue_on_error {
102 errors.push((index, error));
103 } else {
104 self.rollback_succeeded_ops(&succeeded_ops, dry, wet).await;
106 let chain =
107 format!("Op {}-{} failed: {}", index, op.metadata().name, error);
108 return Err(match error.failure_code() {
112 Some(code) => OpError::WrappedClassified {
113 chain,
114 code: code.to_string(),
115 class: error.attribution_class(),
116 reason: error.failure_reason(),
117 arg_urn: error.failure_arg_urn().map(str::to_string),
118 },
119 None => OpError::BatchFailed(chain),
120 });
121 }
122 }
123 }
124 }
125
126 if !errors.is_empty() && !self.continue_on_error {
127 self.rollback_succeeded_ops(&succeeded_ops, dry, wet).await;
128 return Err(OpError::BatchFailed(format!(
129 "Batch op had {} errors",
130 errors.len()
131 )));
132 }
133
134 Ok(results)
135 }
136
137 fn metadata(&self) -> OpMetadata {
138 BatchMetadataBuilder::new(&self.ops).build()
140 }
141}
142
143#[cfg(test)]
144mod tests {
145 use super::*;
146 use serde_json::json;
147
148 struct TestOp {
149 value: i32,
150 should_fail: bool,
151 }
152
153 #[async_trait]
154 impl Op<i32> for TestOp {
155 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
156 if self.should_fail {
157 Err(OpError::ExecutionFailed("Test failure".to_string()))
158 } else {
159 Ok(self.value)
160 }
161 }
162
163 fn metadata(&self) -> OpMetadata {
164 OpMetadata::builder("TestOp").build()
165 }
166 }
167
168 #[tokio::test]
170 async fn test0049_batch_op_success() {
171 let ops = vec![
172 Arc::new(TestOp {
173 value: 1,
174 should_fail: false,
175 }) as Arc<dyn Op<i32>>,
176 Arc::new(TestOp {
177 value: 2,
178 should_fail: false,
179 }) as Arc<dyn Op<i32>>,
180 ];
181
182 let batch = BatchOp::new(ops);
183 let mut dry = DryContext::new();
184 let mut wet = WetContext::new();
185
186 let results = batch.perform(&mut dry, &mut wet).await.unwrap();
187 assert_eq!(results, vec![1, 2]);
188 }
189
190 #[tokio::test]
192 async fn test0050_batch_op_failure() {
193 let ops = vec![
194 Arc::new(TestOp {
195 value: 1,
196 should_fail: false,
197 }) as Arc<dyn Op<i32>>,
198 Arc::new(TestOp {
199 value: 2,
200 should_fail: true,
201 }) as Arc<dyn Op<i32>>,
202 ];
203
204 let batch = BatchOp::new(ops);
205 let mut dry = DryContext::new();
206 let mut wet = WetContext::new();
207
208 let result = batch.perform(&mut dry, &mut wet).await;
209 assert!(result.is_err());
210 }
211
212 #[tokio::test]
214 async fn test0051_batch_op_returns_all_results() {
215 let ops = vec![
216 Arc::new(TestOp {
217 value: 1,
218 should_fail: false,
219 }) as Arc<dyn Op<i32>>,
220 Arc::new(TestOp {
221 value: 2,
222 should_fail: false,
223 }) as Arc<dyn Op<i32>>,
224 ];
225
226 let batch = BatchOp::new(ops);
227 let mut dry = DryContext::new();
228 let mut wet = WetContext::new();
229
230 let results = batch.perform(&mut dry, &mut wet).await.unwrap();
231 assert_eq!(results.len(), 2);
232 assert!(results.contains(&1));
233 assert!(results.contains(&2));
234 }
235
236 #[tokio::test]
238 async fn test0052_batch_metadata_data_flow() {
239 struct ProducerOp;
241 struct ConsumerOp;
242
243 #[async_trait]
244 impl Op<()> for ProducerOp {
245 async fn perform(&self, dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
246 let initial_value = dry.get_required::<String>("initial_value")?;
247 dry.insert("produced_value", format!("processed_{}", initial_value));
248 Ok(())
249 }
250
251 fn metadata(&self) -> OpMetadata {
252 OpMetadata::builder("ProducerOp")
253 .input_schema(json!({
254 "type": "object",
255 "properties": {
256 "initial_value": { "type": "string" }
257 },
258 "required": ["initial_value"]
259 }))
260 .output_schema(json!({
261 "type": "object",
262 "properties": {
263 "produced_value": { "type": "string" }
264 }
265 }))
266 .build()
267 }
268 }
269
270 #[async_trait]
271 impl Op<()> for ConsumerOp {
272 async fn perform(&self, dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
273 let produced = dry.get_required::<String>("produced_value")?;
274 let extra = dry.get_required::<i32>("extra_param")?;
275 dry.insert("final_result", format!("{}_extra_{}", produced, extra));
276 Ok(())
277 }
278
279 fn metadata(&self) -> OpMetadata {
280 OpMetadata::builder("ConsumerOp")
281 .input_schema(json!({
282 "type": "object",
283 "properties": {
284 "produced_value": { "type": "string" },
285 "extra_param": { "type": "integer" }
286 },
287 "required": ["produced_value", "extra_param"]
288 }))
289 .output_schema(json!({
290 "type": "object",
291 "properties": {
292 "final_result": { "type": "string" }
293 }
294 }))
295 .build()
296 }
297 }
298
299 let ops: Vec<Arc<dyn Op<()>>> = vec![Arc::new(ProducerOp), Arc::new(ConsumerOp)];
300
301 let batch = BatchOp::new(ops);
302 let metadata = batch.metadata();
303
304 if let Some(input_schema) = metadata.input_schema {
307 let required = input_schema
308 .get("required")
309 .and_then(|r| r.as_array())
310 .unwrap();
311
312 let required_fields: Vec<&str> = required.iter().filter_map(|v| v.as_str()).collect();
313
314 assert_eq!(required_fields.len(), 2);
315 assert!(required_fields.contains(&"initial_value"));
316 assert!(required_fields.contains(&"extra_param"));
317 assert!(!required_fields.contains(&"produced_value")); }
319 }
320
321 #[tokio::test]
323 async fn test0053_batch_reference_schema_merging() {
324 struct ServiceAOp;
325 struct ServiceBOp;
326
327 #[async_trait]
328 impl Op<()> for ServiceAOp {
329 async fn perform(&self, _dry: &mut DryContext, wet: &mut WetContext) -> OpResult<()> {
330 let _service = wet.get_required::<String>("service_a")?;
331 Ok(())
332 }
333
334 fn metadata(&self) -> OpMetadata {
335 OpMetadata::builder("ServiceAOp")
336 .reference_schema(json!({
337 "type": "object",
338 "properties": {
339 "service_a": { "type": "ServiceA" },
340 "shared_service": { "type": "SharedService" }
341 },
342 "required": ["service_a", "shared_service"]
343 }))
344 .build()
345 }
346 }
347
348 #[async_trait]
349 impl Op<()> for ServiceBOp {
350 async fn perform(&self, _dry: &mut DryContext, wet: &mut WetContext) -> OpResult<()> {
351 let _service = wet.get_required::<String>("service_b")?;
352 Ok(())
353 }
354
355 fn metadata(&self) -> OpMetadata {
356 OpMetadata::builder("ServiceBOp")
357 .reference_schema(json!({
358 "type": "object",
359 "properties": {
360 "service_b": { "type": "ServiceB" },
361 "shared_service": { "type": "SharedService" }
362 },
363 "required": ["service_b", "shared_service"]
364 }))
365 .build()
366 }
367 }
368
369 let ops: Vec<Arc<dyn Op<()>>> = vec![Arc::new(ServiceAOp), Arc::new(ServiceBOp)];
370
371 let batch = BatchOp::new(ops);
372 let metadata = batch.metadata();
373
374 if let Some(ref_schema) = metadata.reference_schema {
376 let required = ref_schema
377 .get("required")
378 .and_then(|r| r.as_array())
379 .unwrap();
380
381 let required_services: Vec<&str> = required.iter().filter_map(|v| v.as_str()).collect();
382
383 assert_eq!(required_services.len(), 3);
384 assert!(required_services.contains(&"service_a"));
385 assert!(required_services.contains(&"service_b"));
386 assert!(required_services.contains(&"shared_service")); }
388 }
389
390 #[tokio::test]
392 async fn test0054_batch_rollback_on_failure() {
393 use std::sync::{Arc, Mutex};
394
395 struct RollbackTrackingOp {
396 id: u32,
397 should_fail: bool,
398 performed: Arc<Mutex<bool>>,
399 rolled_back: Arc<Mutex<bool>>,
400 }
401
402 #[async_trait]
403 impl Op<u32> for RollbackTrackingOp {
404 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<u32> {
405 *self.performed.lock().unwrap() = true;
406 if self.should_fail {
407 Err(OpError::ExecutionFailed(format!("Op {} failed", self.id)))
408 } else {
409 Ok(self.id)
410 }
411 }
412
413 async fn rollback(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
414 *self.rolled_back.lock().unwrap() = true;
415 Ok(())
416 }
417
418 fn metadata(&self) -> OpMetadata {
419 OpMetadata::builder(&format!("RollbackTrackingOp{}", self.id)).build()
420 }
421 }
422
423 let op1_performed = Arc::new(Mutex::new(false));
425 let op1_rolled_back = Arc::new(Mutex::new(false));
426 let op2_performed = Arc::new(Mutex::new(false));
427 let op2_rolled_back = Arc::new(Mutex::new(false));
428 let op3_performed = Arc::new(Mutex::new(false));
429 let op3_rolled_back = Arc::new(Mutex::new(false));
430
431 let ops = vec![
432 Arc::new(RollbackTrackingOp {
433 id: 1,
434 should_fail: false,
435 performed: op1_performed.clone(),
436 rolled_back: op1_rolled_back.clone(),
437 }) as Arc<dyn Op<u32>>,
438 Arc::new(RollbackTrackingOp {
439 id: 2,
440 should_fail: false,
441 performed: op2_performed.clone(),
442 rolled_back: op2_rolled_back.clone(),
443 }) as Arc<dyn Op<u32>>,
444 Arc::new(RollbackTrackingOp {
445 id: 3,
446 should_fail: true, performed: op3_performed.clone(),
448 rolled_back: op3_rolled_back.clone(),
449 }) as Arc<dyn Op<u32>>,
450 ];
451
452 let batch = BatchOp::new(ops);
453 let mut dry = DryContext::new();
454 let mut wet = WetContext::new();
455
456 let result = batch.perform(&mut dry, &mut wet).await;
458 assert!(result.is_err());
459
460 assert!(
462 *op1_performed.lock().unwrap(),
463 "Op1 should have been performed"
464 );
465 assert!(
466 *op2_performed.lock().unwrap(),
467 "Op2 should have been performed"
468 );
469 assert!(
470 *op3_performed.lock().unwrap(),
471 "Op3 should have been performed (and failed)"
472 );
473
474 assert!(
476 *op1_rolled_back.lock().unwrap(),
477 "Op1 should have been rolled back"
478 );
479 assert!(
480 *op2_rolled_back.lock().unwrap(),
481 "Op2 should have been rolled back"
482 );
483 assert!(
484 !*op3_rolled_back.lock().unwrap(),
485 "Op3 should NOT have been rolled back (it failed)"
486 );
487 }
488
489 #[tokio::test]
491 async fn test0055_batch_rollback_order() {
492 use std::sync::{Arc, Mutex};
493
494 struct OrderTrackingOp {
495 id: u32,
496 rollback_order: Arc<Mutex<Vec<u32>>>,
497 }
498
499 #[async_trait]
500 impl Op<u32> for OrderTrackingOp {
501 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<u32> {
502 Ok(self.id)
503 }
504
505 async fn rollback(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
506 self.rollback_order.lock().unwrap().push(self.id);
507 Ok(())
508 }
509
510 fn metadata(&self) -> OpMetadata {
511 OpMetadata::builder(&format!("OrderTrackingOp{}", self.id)).build()
512 }
513 }
514
515 struct FailingOp;
516
517 #[async_trait]
518 impl Op<u32> for FailingOp {
519 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<u32> {
520 Err(OpError::ExecutionFailed("Intentional failure".to_string()))
521 }
522
523 fn metadata(&self) -> OpMetadata {
524 OpMetadata::builder("FailingOp").build()
525 }
526 }
527
528 let rollback_order = Arc::new(Mutex::new(Vec::new()));
529
530 let ops = vec![
531 Arc::new(OrderTrackingOp {
532 id: 1,
533 rollback_order: rollback_order.clone(),
534 }) as Arc<dyn Op<u32>>,
535 Arc::new(OrderTrackingOp {
536 id: 2,
537 rollback_order: rollback_order.clone(),
538 }) as Arc<dyn Op<u32>>,
539 Arc::new(OrderTrackingOp {
540 id: 3,
541 rollback_order: rollback_order.clone(),
542 }) as Arc<dyn Op<u32>>,
543 Arc::new(FailingOp) as Arc<dyn Op<u32>>, ];
545
546 let batch = BatchOp::new(ops);
547 let mut dry = DryContext::new();
548 let mut wet = WetContext::new();
549
550 let result = batch.perform(&mut dry, &mut wet).await;
552 assert!(result.is_err());
553
554 let order = rollback_order.lock().unwrap();
556 assert_eq!(
557 *order,
558 vec![3, 2, 1],
559 "Rollback should happen in reverse order"
560 );
561 }
562
563 #[tokio::test]
565 async fn test0056_batch_rollback_on_failure_partial() {
566 use std::sync::{Arc, Mutex};
567
568 struct RollbackTrackingOp {
569 id: u32,
570 should_fail: bool,
571 performed: Arc<Mutex<bool>>,
572 rolled_back: Arc<Mutex<bool>>,
573 }
574
575 #[async_trait]
576 impl Op<u32> for RollbackTrackingOp {
577 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<u32> {
578 *self.performed.lock().unwrap() = true;
579 if self.should_fail {
580 Err(OpError::ExecutionFailed(format!("Op {} failed", self.id)))
581 } else {
582 Ok(self.id)
583 }
584 }
585
586 async fn rollback(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
587 *self.rolled_back.lock().unwrap() = true;
588 Ok(())
589 }
590
591 fn metadata(&self) -> OpMetadata {
592 OpMetadata::builder(&format!("RollbackTrackingOp{}", self.id)).build()
593 }
594 }
595
596 let op1_performed = Arc::new(Mutex::new(false));
598 let op1_rolled_back = Arc::new(Mutex::new(false));
599 let op2_performed = Arc::new(Mutex::new(false));
600 let op2_rolled_back = Arc::new(Mutex::new(false));
601
602 let ops = vec![
603 Arc::new(RollbackTrackingOp {
604 id: 1,
605 should_fail: false,
606 performed: op1_performed.clone(),
607 rolled_back: op1_rolled_back.clone(),
608 }) as Arc<dyn Op<u32>>,
609 Arc::new(RollbackTrackingOp {
610 id: 2,
611 should_fail: true, performed: op2_performed.clone(),
613 rolled_back: op2_rolled_back.clone(),
614 }) as Arc<dyn Op<u32>>,
615 ];
616
617 let batch = BatchOp::new(ops);
618 let mut dry = DryContext::new();
619 let mut wet = WetContext::new();
620
621 let result = batch.perform(&mut dry, &mut wet).await;
623 assert!(result.is_err());
624
625 assert!(
627 *op1_performed.lock().unwrap(),
628 "Op1 should have been performed"
629 );
630 assert!(
631 *op2_performed.lock().unwrap(),
632 "Op2 should have been performed (and failed)"
633 );
634
635 assert!(
637 *op1_rolled_back.lock().unwrap(),
638 "Op1 should have been rolled back"
639 );
640 assert!(
641 !*op2_rolled_back.lock().unwrap(),
642 "Op2 should NOT have been rolled back (it failed)"
643 );
644 }
645
646 #[test]
648 fn test0093_batch_len_and_is_empty() {
649 let empty: BatchOp<i32> = BatchOp::new(vec![]);
650 assert_eq!(empty.len(), 0);
651 assert!(empty.is_empty());
652
653 let nonempty = BatchOp::new(vec![Arc::new(TestOp {
654 value: 1,
655 should_fail: false,
656 }) as Arc<dyn Op<i32>>]);
657 assert_eq!(nonempty.len(), 1);
658 assert!(!nonempty.is_empty());
659 }
660
661 #[tokio::test]
663 async fn test0094_batch_add_op() {
664 let mut batch = BatchOp::new(vec![Arc::new(TestOp {
665 value: 10,
666 should_fail: false,
667 }) as Arc<dyn Op<i32>>]);
668 batch.add_op(Arc::new(TestOp {
669 value: 20,
670 should_fail: false,
671 }));
672
673 let mut dry = DryContext::new();
674 let mut wet = WetContext::new();
675 let results = batch.perform(&mut dry, &mut wet).await.unwrap();
676 assert_eq!(results, vec![10, 20]);
677 }
678
679 #[tokio::test]
681 async fn test0095_batch_continue_on_error() {
682 let ops = vec![
683 Arc::new(TestOp {
684 value: 1,
685 should_fail: false,
686 }) as Arc<dyn Op<i32>>,
687 Arc::new(TestOp {
688 value: 2,
689 should_fail: true,
690 }) as Arc<dyn Op<i32>>,
691 Arc::new(TestOp {
692 value: 3,
693 should_fail: false,
694 }) as Arc<dyn Op<i32>>,
695 ];
696 let batch = BatchOp::new(ops).with_continue_on_error(true);
697 let mut dry = DryContext::new();
698 let mut wet = WetContext::new();
699 let results = batch.perform(&mut dry, &mut wet).await.unwrap();
700 assert_eq!(results, vec![1, 3]);
702 }
703
704 #[tokio::test]
706 async fn test0096_empty_batch_returns_empty() {
707 let batch: BatchOp<i32> = BatchOp::new(vec![]);
708 let mut dry = DryContext::new();
709 let mut wet = WetContext::new();
710 let results = batch.perform(&mut dry, &mut wet).await.unwrap();
711 assert!(results.is_empty());
712 }
713
714 #[tokio::test]
716 async fn test0097_nested_batch_rollback() {
717 use std::sync::{Arc as StdArc, Mutex};
718
719 let rollback_log: StdArc<Mutex<Vec<&'static str>>> = StdArc::new(Mutex::new(vec![]));
720
721 struct TrackingOp {
722 name: &'static str,
723 should_fail: bool,
724 log: StdArc<Mutex<Vec<&'static str>>>,
725 }
726
727 #[async_trait]
728 impl Op<i32> for TrackingOp {
729 async fn perform(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<i32> {
730 if self.should_fail {
731 Err(OpError::ExecutionFailed(format!("{} failed", self.name)))
732 } else {
733 Ok(0)
734 }
735 }
736 async fn rollback(&self, _dry: &mut DryContext, _wet: &mut WetContext) -> OpResult<()> {
737 self.log.lock().unwrap().push(self.name);
738 Ok(())
739 }
740 fn metadata(&self) -> OpMetadata {
741 OpMetadata::builder(self.name).build()
742 }
743 }
744
745 let log = rollback_log.clone();
746 let inner_ops: Vec<Arc<dyn Op<i32>>> = vec![
747 Arc::new(TrackingOp {
748 name: "inner_a",
749 should_fail: false,
750 log: log.clone(),
751 }),
752 Arc::new(TrackingOp {
753 name: "inner_b",
754 should_fail: false,
755 log: log.clone(),
756 }),
757 ];
758 let inner_batch = Arc::new(BatchOp::new(inner_ops));
759
760 struct FailingOp;
762 #[async_trait]
763 impl Op<Vec<i32>> for FailingOp {
764 async fn perform(
765 &self,
766 _dry: &mut DryContext,
767 _wet: &mut WetContext,
768 ) -> OpResult<Vec<i32>> {
769 Err(OpError::ExecutionFailed("outer fail".to_string()))
770 }
771 fn metadata(&self) -> OpMetadata {
772 OpMetadata::builder("FailingOp").build()
773 }
774 }
775
776 let outer_ops: Vec<Arc<dyn Op<Vec<i32>>>> = vec![inner_batch, Arc::new(FailingOp)];
777 let outer_batch = BatchOp::new(outer_ops);
778 let mut dry = DryContext::new();
779 let mut wet = WetContext::new();
780 let result = outer_batch.perform(&mut dry, &mut wet).await;
781 assert!(result.is_err());
782 match result.unwrap_err() {
785 OpError::BatchFailed(_) => {}
786 e => panic!("Expected BatchFailed, got {:?}", e),
787 }
788 }
789}