use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::Instant,
};
use chrono::Utc;
use serde_json::Value;
use traverse_contracts::{CapabilityContract, ServiceType, ViolationRecord};
use crate::{
events::types::{EventBroker, TraverseEvent},
executor::{ArtifactType, CapabilityExecutor, ExecutorCapability},
placement::{
PlacementConstraintEvaluator, PlacementDecision, PlacementError, PlacementRequest,
RuntimeSnapshot,
},
trace::{PrivateTraceEntry, PublicTraceEntry, TraceOutcome, TraceStore, new_trace_id_and_time},
};
use traverse_contracts::ExecutionTarget;
pub type CapabilityExecutorRegistry = HashMap<ArtifactType, Box<dyn CapabilityExecutor>>;
pub struct RouterRequest {
pub capability_id: String,
pub artifact_type: ArtifactType,
pub contract: CapabilityContract,
pub target_hint: Option<ExecutionTarget>,
pub runtime_snapshot: RuntimeSnapshot,
pub input: Value,
pub executor_capability: ExecutorCapability,
pub trace_id_override: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RouterError {
PlacementFailed(PlacementError),
ExecutorNotFound(String),
ExecutionFailed(String),
ContractViolation(Vec<ViolationRecord>),
TraceLockPoisoned,
}
impl std::fmt::Display for RouterError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::PlacementFailed(e) => write!(f, "placement failed: {e:?}"),
Self::ExecutorNotFound(t) => write!(f, "no executor registered for artifact type: {t}"),
Self::ExecutionFailed(msg) => write!(f, "execution failed: {msg}"),
Self::ContractViolation(violations) => {
write!(f, "contract violation: {} violation(s)", violations.len())
}
Self::TraceLockPoisoned => write!(f, "trace store lock is poisoned"),
}
}
}
impl std::error::Error for RouterError {}
#[derive(Debug)]
pub struct RouterResponse {
pub output: Value,
pub emitted_events: Vec<TraverseEvent>,
pub trace_id: String,
pub placement_decision: PlacementDecision,
}
pub struct PlacementRouter {
evaluator: PlacementConstraintEvaluator,
executor_registry: CapabilityExecutorRegistry,
trace_store: Arc<Mutex<TraceStore>>,
event_broker: Arc<dyn EventBroker>,
}
impl PlacementRouter {
#[must_use]
pub fn new(
evaluator: PlacementConstraintEvaluator,
executor_registry: CapabilityExecutorRegistry,
trace_store: Arc<Mutex<TraceStore>>,
event_broker: Arc<dyn EventBroker>,
) -> Self {
Self {
evaluator,
executor_registry,
trace_store,
event_broker,
}
}
pub fn execute(&self, request: RouterRequest) -> Result<RouterResponse, RouterError> {
let executor = self
.executor_registry
.get(&request.artifact_type)
.ok_or_else(|| RouterError::ExecutorNotFound(format!("{:?}", request.artifact_type)))?;
self.execute_with_executor(request, executor.as_ref())
}
pub fn execute_with_executor(
&self,
request: RouterRequest,
executor: &dyn CapabilityExecutor,
) -> Result<RouterResponse, RouterError> {
let placement_req = PlacementRequest {
capability_id: request.capability_id.clone(),
target_hint: request.target_hint,
runtime_snapshot: request.runtime_snapshot,
};
let decision = self
.evaluator
.evaluate(&placement_req, &request.contract)
.map_err(RouterError::PlacementFailed)?;
let placement_target_str = format!("{:?}", decision.target);
let start = Instant::now();
let exec_result = executor.execute(&request.executor_capability, &request.input);
let duration_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
let (output, emitted_events, outcome) = match exec_result {
Ok(exec_output) => (
exec_output.value,
exec_output.emitted_events,
TraceOutcome::Success,
),
Err(e) => return Err(RouterError::ExecutionFailed(format!("{e}"))),
};
let (trace_id, time) = match request.trace_id_override {
Some(override_id) => (override_id, Utc::now().to_rfc3339()),
None => new_trace_id_and_time(),
};
let public_entry = PublicTraceEntry::new(
trace_id.clone(),
request.capability_id.clone(),
placement_target_str,
outcome,
duration_ms,
time,
);
let input_str = serde_json::to_string(&request.input).unwrap_or_default();
let output_str = serde_json::to_string(&output).unwrap_or_default();
let private_entry =
PrivateTraceEntry::new(trace_id.clone(), &input_str, &output_str, duration_ms);
{
let mut store = self
.trace_store
.lock()
.map_err(|_| RouterError::TraceLockPoisoned)?;
store.insert(public_entry, Some(private_entry));
}
let published_events = if request.contract.service_type == ServiceType::Subscribable {
for event in &emitted_events {
let _ = self.event_broker.publish(event.clone());
}
emitted_events
} else {
Vec::new()
};
Ok(RouterResponse {
output,
emitted_events: published_events,
trace_id,
placement_decision: decision,
})
}
}