Skip to main content

oxigdal_streaming/pipeline/
executor.rs

1//! Pipeline executor for running streaming pipelines.
2
3use super::builder::PipelineConfig;
4use super::stage::{PipelineStage, StageResult};
5use crate::error::{Result, StreamingError};
6use crate::pipeline::zerocopy::ZeroCopyBuffer;
7use bytes::Bytes;
8use std::sync::Arc;
9use std::time::Instant;
10use tracing::{debug, info};
11
12/// Statistics about a completed pipeline execution.
13#[derive(Debug, Clone)]
14pub struct ExecutionStats {
15    /// Total number of items processed
16    pub items_processed: usize,
17
18    /// Total bytes processed
19    pub bytes_processed: usize,
20
21    /// Total execution time in milliseconds
22    pub total_time_ms: u64,
23
24    /// Average time per item in milliseconds
25    pub avg_time_per_item_ms: f64,
26}
27
28/// Executes a pipeline of stages on a stream of data.
29pub struct PipelineExecutor {
30    config: PipelineConfig,
31    stages: Vec<Arc<dyn PipelineStage>>,
32}
33
34impl PipelineExecutor {
35    /// Create a new pipeline executor.
36    pub async fn new(config: PipelineConfig, stages: Vec<Arc<dyn PipelineStage>>) -> Result<Self> {
37        if stages.is_empty() {
38            return Err(StreamingError::ConfigError(
39                "Pipeline must have at least one stage".to_string(),
40            ));
41        }
42
43        // Initialize all stages
44        for stage in &stages {
45            stage.initialize().await?;
46        }
47
48        info!("Created pipeline executor with {} stages", stages.len());
49
50        Ok(Self { config, stages })
51    }
52
53    /// Process a single item through all pipeline stages.
54    pub async fn process(&self, data: Bytes) -> Result<StageResult> {
55        let start = Instant::now();
56        let mut current = ZeroCopyBuffer::new(data);
57
58        for stage in &self.stages {
59            debug!("Executing stage: {}", stage.name());
60            let result = stage.process(current).await?;
61            current = result.data;
62        }
63
64        let elapsed_ms = start.elapsed().as_millis() as u64;
65        Ok(StageResult::new(current, elapsed_ms))
66    }
67
68    /// Process multiple items and return aggregated statistics.
69    pub async fn process_batch(&self, items: Vec<Bytes>) -> Result<ExecutionStats> {
70        let start = Instant::now();
71        let mut total_bytes = 0usize;
72        let item_count = items.len();
73
74        for item in items {
75            let byte_len = item.len();
76            self.process(item).await?;
77            total_bytes += byte_len;
78        }
79
80        let total_ms = start.elapsed().as_millis() as u64;
81        let avg_ms = if item_count > 0 {
82            total_ms as f64 / item_count as f64
83        } else {
84            0.0
85        };
86
87        let stats = ExecutionStats {
88            items_processed: item_count,
89            bytes_processed: total_bytes,
90            total_time_ms: total_ms,
91            avg_time_per_item_ms: avg_ms,
92        };
93
94        info!(
95            "Processed {} items ({} bytes) in {}ms",
96            stats.items_processed, stats.bytes_processed, stats.total_time_ms
97        );
98
99        Ok(stats)
100    }
101
102    /// Finalize and close the pipeline.
103    pub async fn finalize(self) -> Result<()> {
104        for stage in &self.stages {
105            stage.finalize().await?;
106        }
107        info!("Pipeline finalized");
108        Ok(())
109    }
110
111    /// Get the pipeline configuration.
112    pub fn config(&self) -> &PipelineConfig {
113        &self.config
114    }
115
116    /// Get the number of stages.
117    pub fn stage_count(&self) -> usize {
118        self.stages.len()
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use crate::pipeline::stage::TransformStage;
126    use bytes::Bytes;
127
128    #[tokio::test]
129    async fn test_pipeline_executor() {
130        let config = PipelineConfig::default();
131        // TransformStage: Fn(&[u8]) -> Result<Vec<u8>>
132        let stage = Arc::new(TransformStage::new(
133            "identity".to_string(),
134            |data: &[u8]| Ok(data.to_vec()),
135        ));
136
137        let executor = PipelineExecutor::new(config, vec![stage])
138            .await
139            .expect("executor should be created");
140
141        let data = Bytes::from(vec![1u8, 2, 3, 4, 5]);
142        let result = executor
143            .process(data)
144            .await
145            .expect("process should succeed");
146
147        assert_eq!(result.bytes_processed, 5);
148    }
149}