use std::collections::HashMap;
use std::time::Duration;
use crate::automation::event_api;
use crate::envelope::EventEnvelope;
use crate::function::AppError;
use crate::platform::Platform;
use crate::trace;
pub const BUSINESS_CID_TAG: &str = "my_cid";
pub(crate) const RPC_TAG: &str = "rpc";
pub(crate) fn apply_current_trace(mut event: EventEnvelope) -> EventEnvelope {
let snapshot = trace::with_current(|state| {
(
state.route.clone(),
state.trace_id.clone(),
state.trace_path.clone(),
state.span_id.clone(),
state.cid.clone(),
state.zero_traced,
)
});
if let Some((route, trace_id, trace_path, span_id, cid, zero_traced)) = snapshot {
let effective_id = event.trace_id().unwrap_or(&trace_id).to_string();
let effective_path = event.trace_path().unwrap_or(&trace_path).to_string();
event = event.set_trace(&effective_id, &effective_path);
if !zero_traced {
event = event.set_span_id(&span_id);
}
if event.from().is_none() {
event = event.set_from(&route);
}
if event.tag(BUSINESS_CID_TAG).is_none() {
if let Some(cid) = cid {
event = event.add_tag(BUSINESS_CID_TAG, &cid);
}
}
}
event
}
fn scheduled_events(
) -> &'static std::sync::Mutex<std::collections::HashMap<String, tokio::task::AbortHandle>> {
static TIMERS: std::sync::OnceLock<
std::sync::Mutex<std::collections::HashMap<String, tokio::task::AbortHandle>>,
> = std::sync::OnceLock::new();
TIMERS.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()))
}
#[derive(Clone)]
pub struct PostOffice {
platform: Platform,
}
impl PostOffice {
pub fn new(platform: &Platform) -> Self {
PostOffice {
platform: platform.clone(),
}
}
pub async fn send(&self, event: EventEnvelope) -> Result<(), AppError> {
let event = apply_current_trace(event);
let Some(route) = event.to().map(str::to_string) else {
return Err(AppError::new(400, "Missing routing path ('to')"));
};
if event.header(event_api::X_EVENT_API).is_none() {
if let Some(entry) = event_api::get_event_http_target(&route) {
return event_api::send_with_event_http(&self.platform, event, &route, entry);
}
}
self.platform.deliver(&route, event).await
}
pub fn send_later(&self, event: EventEnvelope, delay: std::time::Duration) -> String {
let event = apply_current_trace(event);
let timer_id = uuid::Uuid::new_v4().simple().to_string();
let platform = self.platform.clone();
let id_for_task = timer_id.clone();
let handle = tokio::spawn(async move {
tokio::time::sleep(delay).await;
scheduled_events()
.lock()
.expect("timer registry")
.remove(&id_for_task);
if let Some(route) = event.to().map(str::to_string) {
if let Err(e) = PostOffice::new(&platform).send(event).await {
log::warn!(
"Unable to deliver scheduled event to {route} - {}",
e.message()
);
}
}
});
scheduled_events()
.lock()
.expect("timer registry")
.insert(timer_id.clone(), handle.abort_handle());
timer_id
}
pub fn cancel_future_event(&self, timer_id: &str) -> bool {
match scheduled_events()
.lock()
.expect("timer registry")
.remove(timer_id)
{
Some(handle) => {
handle.abort();
true
}
None => false,
}
}
pub fn my_correlation_id(&self) -> Option<String> {
trace::with_current(|state| state.cid.clone()).flatten()
}
pub fn my_trace_id(&self) -> Option<String> {
trace::with_current(|state| state.trace_id.clone())
}
pub fn my_trace_path(&self) -> Option<String> {
trace::with_current(|state| state.trace_path.clone())
}
pub fn annotate_trace(&self, key: &str, value: impl serde::Serialize) -> &Self {
if let Ok(value) = serde_json::to_value(value) {
trace::with_current_mut(|state| {
state.annotations.insert(key.to_string(), value);
});
}
self
}
pub fn update_context(&self, key: &str, value: impl serde::Serialize) -> Result<(), AppError> {
if crate::trace::RESERVED_KEYS.contains(&key) {
return Err(AppError::new(
400,
format!("'{key}' is a reserved log context key"),
));
}
let value = serde_json::to_value(value)
.map_err(|e| AppError::new(400, format!("unable to serialize context value: {e}")))?;
trace::with_current_mut(|state| {
if value.is_null() {
state.custom_log_keys.remove(key);
} else {
state.custom_log_keys.insert(key.to_string(), value);
}
});
Ok(())
}
pub async fn request(
&self,
event: EventEnvelope,
timeout: Duration,
) -> Result<EventEnvelope, AppError> {
let event = apply_current_trace(event);
if event.header(event_api::X_EVENT_API).is_none() {
if let Some(entry) = event.to().and_then(event_api::get_event_http_target) {
let forward = event.set_header(event_api::X_EVENT_API, "request");
return event_api::event_over_http_with_headers(
self,
&entry.target,
forward,
timeout,
true,
&entry.headers,
)
.await;
}
}
self.request_direct(event, timeout).await
}
pub(crate) async fn request_direct(
&self,
event: EventEnvelope,
timeout: Duration,
) -> Result<EventEnvelope, AppError> {
let (inbox_cid, rx) = crate::inbox::open();
let original_cid = event.correlation_id().map(str::to_string);
let mut event = event;
if event.tag(BUSINESS_CID_TAG).is_none() {
if let Some(cid) = &original_cid {
event = event.add_tag(BUSINESS_CID_TAG, cid);
}
}
let rpc_trace = RpcTraceCapture::of(&event);
let begin = std::time::Instant::now();
let event = event
.set_reply_to(crate::inbox::TEMPORARY_INBOX)
.set_correlation_id(&inbox_cid)
.add_tag(RPC_TAG, &timeout.as_millis().to_string());
if let Err(e) = self.send(event).await {
crate::inbox::close(&inbox_cid);
return Err(e);
}
let outcome = tokio::time::timeout(timeout, rx).await;
match outcome {
Ok(Ok(response)) => {
let diff = begin.elapsed().as_secs_f32() * 1000.0;
let diff = (diff.max(0.0) * 1000.0).round() / 1000.0;
let mut response = response.set_round_trip(diff);
if original_cid.is_some() {
response.set_cid_internal(original_cid);
}
let annotations = response.annotations().clone();
let response = response.clear_annotations();
self.record_rpc_trace(&rpc_trace, &response, annotations);
Ok(response)
}
Ok(Err(_)) => Err(AppError::new(500, "Reply channel closed unexpectedly")),
Err(_) => {
crate::inbox::close(&inbox_cid);
Err(AppError::new(
408,
format!("Request timeout for {} ms", timeout.as_millis()),
))
}
}
}
fn record_rpc_trace(
&self,
rpc: &RpcTraceCapture,
reply: &EventEnvelope,
annotations: HashMap<String, rmpv::Value>,
) {
let (Some(to), Some(trace_id), Some(trace_path)) =
(&rpc.to, &rpc.trace_id, &rpc.trace_path)
else {
return; };
let service = trim_origin(to).to_string();
if crate::platform::in_skip_rpc_tracing_list(&service) {
return;
}
if !self
.platform
.has_route(crate::telemetry::DISTRIBUTED_TRACING)
{
return; }
let mut metrics = serde_json::Map::new();
let mut put = |k: &str, v: serde_json::Value| {
metrics.insert(k.to_string(), v);
};
put(
"origin",
serde_json::Value::String(Platform::origin().to_string()),
);
put("id", serde_json::Value::String(trace_id.clone()));
put("service", serde_json::Value::String(service));
if let Some(from) = &rpc.from {
put(
"from",
serde_json::Value::String(trim_origin(from).to_string()),
);
}
if let Some(span_id) = span_id_from_responder(to, reply) {
put("span_id", serde_json::Value::String(span_id.to_string()));
}
if let Some(parent) = &rpc.parent_span {
put("parent_span_id", serde_json::Value::String(parent.clone()));
}
if let Some(exec_time) = reply.exec_time() {
put(
"exec_time",
serde_json::Value::from(((exec_time as f64) * 1000.0).round() / 1000.0),
);
}
if let Some(round_trip) = reply.round_trip() {
put(
"round_trip",
serde_json::Value::from(((round_trip as f64) * 1000.0).round() / 1000.0),
);
}
put("start", serde_json::Value::String(rpc.start.clone()));
put("path", serde_json::Value::String(trace_path.clone()));
let status = reply.status();
put("status", serde_json::Value::from(status));
if status >= 400 {
put("success", serde_json::Value::Bool(false));
let message = match reply.body() {
rmpv::Value::String(s) => s.as_str().unwrap_or("***").to_string(),
_ => "***".to_string(),
};
put("exception", serde_json::Value::String(message));
} else {
put("success", serde_json::Value::Bool(true));
}
let mut dataset = serde_json::Map::new();
dataset.insert("trace".to_string(), serde_json::Value::Object(metrics));
if !annotations.is_empty() {
let folded: serde_json::Map<String, serde_json::Value> = annotations
.into_iter()
.filter_map(|(k, v)| serde_json::to_value(&v).ok().map(|value| (k, value)))
.collect();
if !folded.is_empty() {
dataset.insert("annotations".to_string(), serde_json::Value::Object(folded));
}
}
let platform = self.platform.clone();
tokio::spawn(async move {
match EventEnvelope::new()
.set_to(crate::telemetry::DISTRIBUTED_TRACING)
.set_body(serde_json::Value::Object(dataset))
{
Ok(event) => {
if let Err(e) = platform
.deliver(crate::telemetry::DISTRIBUTED_TRACING, event)
.await
{
log::error!("Unable to send to distributed.tracing - {}", e.message());
}
}
Err(e) => log::error!("Unable to send to distributed.tracing - {}", e.message()),
}
});
}
}
struct RpcTraceCapture {
to: Option<String>,
from: Option<String>,
trace_id: Option<String>,
trace_path: Option<String>,
parent_span: Option<String>,
start: String,
}
impl RpcTraceCapture {
fn of(event: &EventEnvelope) -> Self {
RpcTraceCapture {
to: event.to().map(str::to_string),
from: event.from().map(str::to_string),
trace_id: event.trace_id().map(str::to_string),
trace_path: event.trace_path().map(str::to_string),
parent_span: event.span_id().map(str::to_string),
start: trace::iso8601_utc_now(),
}
}
}
fn trim_origin(route: &str) -> &str {
match route.find('@') {
Some(at) => &route[..at],
None => route,
}
}
fn span_id_from_responder<'a>(to: &str, reply: &'a EventEnvelope) -> Option<&'a str> {
match reply.from() {
Some(from) if trim_origin(to) == from => reply.span_id(),
_ => None,
}
}