use crate::aggregation::{Aggregator, CompositeAggregator, AggregationResult};
use crate::core::{EventTimeExtractor, KeyExtractor, ProcessorEvent};
use crate::error::{ProcessorError, WindowError};
use crate::watermark::{BoundedOutOfOrdernessWatermark, WatermarkGenerator};
use crate::window::{
Window, WindowAssigner, WindowManager, WindowTrigger, WindowMetadata,
TumblingWindowAssigner, SlidingWindowAssigner, SessionWindowAssigner,
OnWatermarkTrigger, ProcessingTimeTrigger, CountTrigger,
};
use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use futures::Stream;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock};
use tokio_stream::wrappers::ReceiverStream;
use tracing::{debug, error, info, trace, warn};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WindowResult {
pub window: Window,
pub key: String,
pub results: AggregationResult,
pub event_count: u64,
pub processed_at: DateTime<Utc>,
pub watermark: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
pub struct StreamProcessorConfig {
pub allow_late_events: bool,
pub late_event_threshold: Duration,
pub watermark_interval: Duration,
pub window_retention: Duration,
pub max_windows_per_key: usize,
pub result_buffer_size: usize,
}
impl Default for StreamProcessorConfig {
fn default() -> Self {
Self {
allow_late_events: true,
late_event_threshold: Duration::minutes(5),
watermark_interval: Duration::seconds(1),
window_retention: Duration::hours(1),
max_windows_per_key: 1000,
result_buffer_size: 1000,
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ProcessorStats {
pub events_processed: u64,
pub events_dropped: u64,
pub windows_created: u64,
pub windows_fired: u64,
pub active_windows: usize,
pub current_watermark: Option<DateTime<Utc>>,
pub latency_p50_ms: f64,
pub latency_p95_ms: f64,
pub latency_p99_ms: f64,
pub last_update: DateTime<Utc>,
}
struct KeyState<A: Aggregator> {
aggregators: HashMap<String, A>,
last_event_time: Option<DateTime<Utc>>,
}
impl<A: Aggregator> KeyState<A> {
fn new() -> Self {
Self {
aggregators: HashMap::new(),
last_event_time: None,
}
}
}
pub struct StreamProcessor<T, A, W, Tr>
where
A: Aggregator<Input = f64, Output = f64> + Clone + Send + Sync + 'static,
W: WindowAssigner + Clone + Send + Sync + 'static,
Tr: WindowTrigger + Clone + Send + Sync + 'static,
{
config: StreamProcessorConfig,
window_manager: Arc<RwLock<WindowManager<W, Tr>>>,
watermark_generator: Arc<RwLock<dyn WatermarkGenerator>>,
aggregators: Arc<DashMap<String, KeyState<A>>>,
result_tx: mpsc::Sender<WindowResult>,
result_rx: Option<mpsc::Receiver<WindowResult>>,
stats: Arc<RwLock<ProcessorStats>>,
aggregator_factory: Arc<dyn Fn(&str) -> A + Send + Sync>,
_phantom: PhantomData<T>,
}
impl<T, A, W, Tr> StreamProcessor<T, A, W, Tr>
where
A: Aggregator<Input = f64, Output = f64> + Clone + Send + Sync + 'static,
W: WindowAssigner + Clone + Send + Sync + 'static,
Tr: WindowTrigger + Clone + Send + Sync + 'static,
{
pub fn new<F>(
config: StreamProcessorConfig,
assigner: W,
trigger: Tr,
aggregator_factory: F,
) -> Self
where
F: Fn(&str) -> A + Send + Sync + 'static,
{
let window_manager = WindowManager::with_late_events(
assigner,
trigger,
config.allow_late_events,
config.late_event_threshold,
);
let (result_tx, result_rx) = mpsc::channel(config.result_buffer_size);
let watermark_threshold = config.late_event_threshold.to_std().unwrap_or(std::time::Duration::from_secs(300));
Self {
config,
window_manager: Arc::new(RwLock::new(window_manager)),
watermark_generator: Arc::new(RwLock::new(
BoundedOutOfOrdernessWatermark::new(
watermark_threshold,
None,
)
)),
aggregators: Arc::new(DashMap::new()),
result_tx,
result_rx: Some(result_rx),
stats: Arc::new(RwLock::new(ProcessorStats::default())),
aggregator_factory: Arc::new(aggregator_factory),
_phantom: PhantomData,
}
}
pub async fn process_event(
&self,
key: String,
event_time: DateTime<Utc>,
value: f64,
) -> Result<(), ProcessorError> {
let start_time = Utc::now();
{
let mut stats = self.stats.write().await;
stats.events_processed += 1;
stats.last_update = Utc::now();
}
let windows = {
let mut manager = self.window_manager.write().await;
match manager.assign_event(event_time) {
Ok(windows) => windows,
Err(e) => {
warn!(
error = %e,
event_time = %event_time,
"Failed to assign event to windows"
);
let mut stats = self.stats.write().await;
stats.events_dropped += 1;
return Err(e);
}
}
};
if windows.is_empty() {
let mut stats = self.stats.write().await;
stats.events_dropped += 1;
return Ok(());
}
trace!(
key = %key,
event_time = %event_time,
window_count = windows.len(),
"Processing event"
);
let mut key_state = self.aggregators
.entry(key.clone())
.or_insert_with(KeyState::new);
for window in &windows {
let aggregator = key_state.aggregators
.entry(window.id.clone())
.or_insert_with(|| (self.aggregator_factory)(&key));
if let Err(e) = aggregator.update(value) {
error!(
window_id = %window.id,
error = %e,
"Failed to update aggregator"
);
return Err(ProcessorError::Aggregation(
crate::error::AggregationError::CorruptedState {
aggregation_type: "unknown".to_string(),
details: e.to_string(),
}
));
}
}
key_state.last_event_time = Some(event_time);
drop(key_state);
let watermark = {
let mut wm_gen = self.watermark_generator.write().await;
wm_gen.on_event(event_time.timestamp_millis(), 0);
wm_gen.current_watermark()
};
let ready_windows = {
let mut manager = self.window_manager.write().await;
manager.advance_watermark(watermark.to_datetime())?
};
for window_id in ready_windows {
self.fire_window(&key, &window_id, watermark.to_datetime()).await?;
}
let latency = (Utc::now() - start_time).num_milliseconds() as f64;
self.update_latency_stats(latency).await;
Ok(())
}
async fn fire_window(
&self,
key: &str,
window_id: &str,
watermark: DateTime<Utc>,
) -> Result<(), ProcessorError> {
let metadata = {
let manager = self.window_manager.read().await;
manager.get_window(window_id)
};
let metadata = match metadata {
Some(m) => m,
None => {
warn!(window_id = %window_id, "Window not found");
return Ok(());
}
};
let results = {
let mut key_state = match self.aggregators.get_mut(key) {
Some(state) => state,
None => {
warn!(key = %key, window_id = %window_id, "Key state not found");
return Ok(());
}
};
let aggregator = match key_state.aggregators.remove(window_id) {
Some(agg) => agg,
None => {
warn!(key = %key, window_id = %window_id, "Aggregator not found");
return Ok(());
}
};
if let Some(composite) = (&aggregator as &dyn std::any::Any).downcast_ref::<CompositeAggregator>() {
composite.finalize().map_err(|e| {
ProcessorError::Aggregation(crate::error::AggregationError::CorruptedState {
aggregation_type: "composite".to_string(),
details: e.to_string(),
})
})?
} else {
let value = aggregator.finalize().map_err(|e| {
ProcessorError::Aggregation(crate::error::AggregationError::CorruptedState {
aggregation_type: "simple".to_string(),
details: e.to_string(),
})
})?;
let mut result = AggregationResult::new();
result.set_value("value", value);
result.set_count("count", aggregator.count());
result
}
};
let window_result = WindowResult {
window: metadata.window.clone(),
key: key.to_string(),
results,
event_count: metadata.stats.event_count,
processed_at: Utc::now(),
watermark: Some(watermark),
};
debug!(
window_id = %window_id,
key = %key,
event_count = metadata.stats.event_count,
"Window fired"
);
if let Err(e) = self.result_tx.send(window_result).await {
error!(error = %e, "Failed to send window result");
}
{
let mut stats = self.stats.write().await;
stats.windows_fired += 1;
stats.current_watermark = Some(watermark);
}
Ok(())
}
async fn update_latency_stats(&self, latency_ms: f64) {
let mut stats = self.stats.write().await;
if stats.latency_p50_ms == 0.0 {
stats.latency_p50_ms = latency_ms;
stats.latency_p95_ms = latency_ms;
stats.latency_p99_ms = latency_ms;
} else {
stats.latency_p50_ms = 0.9 * stats.latency_p50_ms + 0.1 * latency_ms;
stats.latency_p95_ms = 0.95 * stats.latency_p95_ms + 0.05 * latency_ms.max(stats.latency_p95_ms);
stats.latency_p99_ms = 0.99 * stats.latency_p99_ms + 0.01 * latency_ms.max(stats.latency_p99_ms);
}
}
pub fn results(&mut self) -> Option<impl Stream<Item = WindowResult>> {
self.result_rx.take().map(ReceiverStream::new)
}
pub async fn stats(&self) -> ProcessorStats {
let mut stats = self.stats.read().await.clone();
let manager = self.window_manager.read().await;
stats.active_windows = manager.active_window_count();
stats.current_watermark = manager.current_watermark();
stats
}
pub async fn cleanup(&self) -> Result<usize, ProcessorError> {
let mut manager = self.window_manager.write().await;
let removed = manager.cleanup_old_windows(self.config.window_retention);
if removed > 0 {
info!(removed, "Cleaned up old windows");
}
Ok(removed)
}
pub async fn reset(&self) {
let mut manager = self.window_manager.write().await;
manager.reset();
self.aggregators.clear();
let mut stats = self.stats.write().await;
*stats = ProcessorStats::default();
info!("Stream processor reset");
}
}
pub struct StreamProcessorBuilder<T, A> {
config: StreamProcessorConfig,
_phantom: PhantomData<(T, A)>,
}
impl<T, A> Default for StreamProcessorBuilder<T, A>
where
A: Aggregator<Input = f64, Output = f64> + Clone + Send + Sync + 'static,
{
fn default() -> Self {
Self::new()
}
}
impl<T, A> StreamProcessorBuilder<T, A>
where
A: Aggregator<Input = f64, Output = f64> + Clone + Send + Sync + 'static,
{
pub fn new() -> Self {
Self {
config: StreamProcessorConfig::default(),
_phantom: PhantomData,
}
}
pub fn with_config(mut self, config: StreamProcessorConfig) -> Self {
self.config = config;
self
}
pub fn allow_late_events(mut self, allow: bool, threshold: Duration) -> Self {
self.config.allow_late_events = allow;
self.config.late_event_threshold = threshold;
self
}
pub fn watermark_interval(mut self, interval: Duration) -> Self {
self.config.watermark_interval = interval;
self
}
pub fn build_tumbling<F>(
self,
window_size: Duration,
aggregator_factory: F,
) -> StreamProcessor<T, A, TumblingWindowAssigner, OnWatermarkTrigger>
where
F: Fn(&str) -> A + Send + Sync + 'static,
{
let assigner = TumblingWindowAssigner::new(window_size);
let trigger = OnWatermarkTrigger::new();
StreamProcessor::new(self.config, assigner, trigger, aggregator_factory)
}
pub fn build_sliding<F>(
self,
window_size: Duration,
slide: Duration,
aggregator_factory: F,
) -> StreamProcessor<T, A, SlidingWindowAssigner, OnWatermarkTrigger>
where
F: Fn(&str) -> A + Send + Sync + 'static,
{
let assigner = SlidingWindowAssigner::new(window_size, slide);
let trigger = OnWatermarkTrigger::new();
StreamProcessor::new(self.config, assigner, trigger, aggregator_factory)
}
pub fn build_session<F>(
self,
gap: Duration,
aggregator_factory: F,
) -> StreamProcessor<T, A, SessionWindowAssigner, OnWatermarkTrigger>
where
F: Fn(&str) -> A + Send + Sync + 'static,
{
let assigner = SessionWindowAssigner::new(gap);
let trigger = OnWatermarkTrigger::new();
StreamProcessor::new(self.config, assigner, trigger, aggregator_factory)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::aggregation::AverageAggregator;
use tokio_stream::StreamExt;
#[tokio::test]
async fn test_stream_processor_basic() {
let processor = StreamProcessorBuilder::new()
.build_tumbling(
Duration::seconds(1),
|_key| AverageAggregator::new(),
);
let base_time = Utc::now();
processor.process_event("key1".to_string(), base_time, 10.0).await.unwrap();
processor.process_event("key1".to_string(), base_time + Duration::milliseconds(500), 20.0).await.unwrap();
let stats = processor.stats().await;
assert_eq!(stats.events_processed, 2);
}
#[tokio::test]
async fn test_stream_processor_windowing() {
let mut processor = StreamProcessorBuilder::new()
.build_tumbling(
Duration::milliseconds(1000),
|_key| AverageAggregator::new(),
);
let mut results_stream = processor.results().unwrap();
let base_time = Utc::now();
processor.process_event("key1".to_string(), base_time, 10.0).await.unwrap();
processor.process_event("key1".to_string(), base_time + Duration::milliseconds(500), 20.0).await.unwrap();
processor.process_event(
"key1".to_string(),
base_time + Duration::milliseconds(2000),
30.0
).await.unwrap();
tokio::time::timeout(
std::time::Duration::from_secs(1),
results_stream.next()
).await.ok();
}
#[tokio::test]
async fn test_stream_processor_multiple_keys() {
let processor = StreamProcessorBuilder::new()
.build_tumbling(
Duration::seconds(1),
|_key| AverageAggregator::new(),
);
let base_time = Utc::now();
processor.process_event("key1".to_string(), base_time, 10.0).await.unwrap();
processor.process_event("key2".to_string(), base_time, 20.0).await.unwrap();
processor.process_event("key1".to_string(), base_time + Duration::milliseconds(500), 30.0).await.unwrap();
let stats = processor.stats().await;
assert_eq!(stats.events_processed, 3);
}
#[tokio::test]
async fn test_stream_processor_cleanup() {
let processor = StreamProcessorBuilder::new()
.build_tumbling(
Duration::milliseconds(100),
|_key| AverageAggregator::new(),
);
let base_time = Utc::now();
processor.process_event("key1".to_string(), base_time, 10.0).await.unwrap();
processor.process_event(
"key1".to_string(),
base_time + Duration::seconds(1),
20.0
).await.unwrap();
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
let removed = processor.cleanup().await.unwrap();
assert!(removed > 0);
}
}