use std::collections::HashMap;
use std::sync::Arc;
use serde_json::json;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use crate::infra::io::state::DynStateBus;
use crate::pipeline::engines::BasicPipelineEngine;
use crate::pipeline::interface::PipelineEngine;
use crate::pipeline::model::PipelineContext;
use crate::platform::adapters::data::DataAdapter;
use crate::platform::model::PipelineInvocationEntry;
use crate::platform::services::pipeline_hits::PipelineHitsService;
use crate::platform::services::pipeline_runtime::PipelineRuntimeService;
use crate::platform::services::project_config::ProjectConfigurationService;
pub struct KvSubscriber {
tasks: Arc<Mutex<HashMap<String, JoinHandle<()>>>>,
state_bus: DynStateBus,
runtime: Arc<PipelineRuntimeService>,
engine: Arc<BasicPipelineEngine>,
hits: Arc<PipelineHitsService>,
data: Arc<dyn DataAdapter>,
zebflow_cfg: Arc<ProjectConfigurationService>,
}
impl KvSubscriber {
pub fn new(
state_bus: DynStateBus,
runtime: Arc<PipelineRuntimeService>,
engine: Arc<BasicPipelineEngine>,
hits: Arc<PipelineHitsService>,
data: Arc<dyn DataAdapter>,
zebflow_cfg: Arc<ProjectConfigurationService>,
) -> Self {
Self {
tasks: Arc::new(Mutex::new(HashMap::new())),
state_bus,
runtime,
engine,
hits,
data,
zebflow_cfg,
}
}
pub async fn register_all(&self) {
for pipeline in self.runtime.list_all() {
for trigger in &pipeline.kv_subscribe_triggers {
self.register_task(
&pipeline.owner,
&pipeline.project,
&pipeline.file_rel_path,
&pipeline.graph.id,
&trigger.node_id,
&trigger.channel,
)
.await;
}
}
}
pub async fn sync_pipeline(&self, owner: &str, project: &str, file_rel_path: &str) {
let key_prefix = format!("{}/{}/{}", owner, project, file_rel_path);
let stale: Vec<String> = {
let tasks = self.tasks.lock().await;
tasks
.keys()
.filter(|k| k.starts_with(&key_prefix))
.cloned()
.collect()
};
for key in stale {
if let Some(handle) = self.tasks.lock().await.remove(&key) {
handle.abort();
}
}
if let Some(pipeline) = self.runtime.get(owner, project, file_rel_path) {
for trigger in &pipeline.kv_subscribe_triggers {
self.register_task(
owner,
project,
file_rel_path,
&pipeline.graph.id,
&trigger.node_id,
&trigger.channel,
)
.await;
}
}
}
async fn register_task(
&self,
owner: &str,
project: &str,
file_rel_path: &str,
graph_id: &str,
node_id: &str,
channel: &str,
) {
let task_key = format!("{}/{}/{}:{}", owner, project, file_rel_path, node_id);
let Ok(mut rx) = self.state_bus.subscribe(owner, project, channel) else {
eprintln!(
"KvSubscriber: failed subscribing to channel '{}' for {}/{}/{}",
channel, owner, project, file_rel_path
);
return;
};
let runtime = self.runtime.clone();
let engine = self.engine.clone();
let hits = self.hits.clone();
let data = self.data.clone();
let zebflow_cfg = self.zebflow_cfg.clone();
let owner_s = owner.to_string();
let project_s = project.to_string();
let file_rel_path_s = file_rel_path.to_string();
let graph_id_s = graph_id.to_string();
let node_id_s = node_id.to_string();
let channel_s = channel.to_string();
let handle = tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(message) => {
let Some(compiled) = runtime.get(&owner_s, &project_s, &file_rel_path_s)
else {
eprintln!(
"KvSubscriber: pipeline no longer active — stopping listener {}",
file_rel_path_s
);
break;
};
let project_cfg = match zebflow_cfg.read_or_default(&owner_s, &project_s) {
Ok(config) => config,
Err(err) => {
eprintln!(
"KvSubscriber: invalid project configuration for {}/{}: {}",
owner_s, project_s, err
);
break;
}
};
let pipeline_retention =
compiled.graph.metadata.as_ref().and_then(|metadata| {
metadata.settings.invocation_retention.as_ref()
});
let log_max_n = pipeline_retention
.and_then(|retention| retention.max_invocations)
.map(|value| value.max(1) as usize)
.unwrap_or_else(|| {
project_cfg
.configs
.pipelines
.logging
.effective_max_invocations()
});
let max_age_secs = pipeline_retention
.and_then(|retention| retention.max_age_secs)
.map(|value| value.max(1) as i64);
let fired_at = chrono::Utc::now();
let ctx = PipelineContext {
owner: owner_s.clone(),
project: project_s.clone(),
pipeline: graph_id_s.clone(),
request_id: format!(
"kvsub-{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
),
route: String::new(),
input: json!({
"trigger": "kv.subscribe",
"channel": channel_s,
"node_id": node_id_s,
"message": message,
}),
trigger: None,
placeholder: None,
};
let exec_start = std::time::Instant::now();
match engine.execute_async(&compiled.graph, &ctx).await {
Ok(output) => {
let duration_ms = exec_start.elapsed().as_millis() as u64;
hits.record_success(&owner_s, &project_s, &file_rel_path_s);
let _ = data.log_pipeline_invocation(
&owner_s,
&project_s,
&file_rel_path_s,
&PipelineInvocationEntry {
run_id: ctx.request_id.clone(),
at: fired_at.timestamp(),
duration_ms,
status: "ok".to_string(),
trigger: "kv.subscribe".to_string(),
error: None,
trace: output.node_trace,
},
log_max_n,
max_age_secs,
);
}
Err(e) => {
let duration_ms = exec_start.elapsed().as_millis() as u64;
hits.record_failure(
&owner_s,
&project_s,
&file_rel_path_s,
"kv.subscribe",
e.code,
&e.message,
);
let _ = data.log_pipeline_invocation(
&owner_s,
&project_s,
&file_rel_path_s,
&PipelineInvocationEntry {
run_id: ctx.request_id.clone(),
at: fired_at.timestamp(),
duration_ms,
status: "error".to_string(),
trigger: "kv.subscribe".to_string(),
error: Some(e.message.clone()),
trace: e.node_trace.clone(),
},
log_max_n,
max_age_secs,
);
eprintln!(
"❌ KvSubscriber failed {}/{}/{}: {}",
owner_s, project_s, file_rel_path_s, e
);
}
}
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => {
eprintln!(
"KvSubscriber: channel '{}' lagged by {} messages — some events dropped",
channel_s, n
);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
break;
}
}
}
});
self.tasks.lock().await.insert(task_key, handle);
}
}