teaql_runtime/
telemetry.rs

1use std::collections::BTreeMap;
2use std::future::Future;
3use std::panic::{AssertUnwindSafe, catch_unwind};
4use std::sync::{Arc, Mutex};
5
6#[derive(Debug, Clone, PartialEq)]
7pub enum RuntimeAttributeValue {
8    String(String),
9    Integer(i64),
10    Float(f64),
11    Boolean(bool),
12}
13
14impl From<&str> for RuntimeAttributeValue {
15    fn from(value: &str) -> Self {
16        Self::String(value.to_owned())
17    }
18}
19
20impl From<String> for RuntimeAttributeValue {
21    fn from(value: String) -> Self {
22        Self::String(value)
23    }
24}
25
26impl From<i64> for RuntimeAttributeValue {
27    fn from(value: i64) -> Self {
28        Self::Integer(value)
29    }
30}
31
32impl From<usize> for RuntimeAttributeValue {
33    fn from(value: usize) -> Self {
34        Self::Integer(value as i64)
35    }
36}
37
38impl From<bool> for RuntimeAttributeValue {
39    fn from(value: bool) -> Self {
40        Self::Boolean(value)
41    }
42}
43
44#[derive(Debug, Clone, PartialEq)]
45pub struct RuntimeOperation {
46    pub family: String,
47    pub name: String,
48    pub attributes: BTreeMap<String, RuntimeAttributeValue>,
49}
50
51impl RuntimeOperation {
52    pub fn new(family: impl Into<String>, name: impl Into<String>) -> Self {
53        let family = family.into();
54        let name = name.into();
55        let mut attributes = BTreeMap::new();
56        attributes.insert("teaql.operation.family".into(), family.clone().into());
57        attributes.insert("teaql.operation.name".into(), name.clone().into());
58        Self {
59            family,
60            name,
61            attributes,
62        }
63    }
64
65    pub fn attribute(
66        mut self,
67        key: impl Into<String>,
68        value: impl Into<RuntimeAttributeValue>,
69    ) -> Self {
70        let key = key.into();
71        if !is_forbidden_attribute(&key) {
72            self.attributes.insert(key, value.into());
73        }
74        self
75    }
76}
77
78fn is_forbidden_attribute(key: &str) -> bool {
79    matches!(
80        key,
81        "teaql.entity.id"
82            | "teaql.user.id"
83            | "teaql.tenant.id"
84            | "teaql.query.parameters"
85            | "teaql.field.values"
86            | "teaql.audit.reason"
87            | "db.query.parameter_values"
88            | "http.request.body"
89            | "url.full"
90    )
91}
92
93/// Stable category derived from a native error type, never its message.
94pub fn runtime_error_category(error_type: &str) -> &'static str {
95    let error_type = error_type.to_ascii_lowercase();
96    for (category, terms) in [
97        ("timeout", &["timeout", "deadline"][..]),
98        (
99            "authorization",
100            &[
101                "authentication",
102                "authorization",
103                "unauthorized",
104                "forbidden",
105                "permission",
106            ],
107        ),
108        (
109            "validation",
110            &[
111                "validation",
112                "invalidargument",
113                "valueerror",
114                "parse",
115                "format",
116            ],
117        ),
118        (
119            "conflict",
120            &[
121                "conflict",
122                "optimistic",
123                "version",
124                "duplicate",
125                "alreadyexists",
126            ],
127        ),
128        (
129            "transport",
130            &[
131                "transport",
132                "network",
133                "connection",
134                "socket",
135                "http",
136                "ioerror",
137            ],
138        ),
139        ("provider", &["provider", "sql", "database", "jdbc"]),
140    ] {
141        if terms.iter().any(|term| error_type.contains(term)) {
142            return category;
143        }
144    }
145    "internal"
146}
147
148pub trait RuntimeTelemetryScope: Send {
149    fn with_context(&self, callback: &mut dyn FnMut()) {
150        callback();
151    }
152    fn success(&mut self, attributes: BTreeMap<String, RuntimeAttributeValue>);
153    fn failure(&mut self, error_type: &str);
154}
155
156pub trait RuntimeTelemetry: Send + Sync {
157    fn start(&self, operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope>;
158    fn is_noop(&self) -> bool {
159        false
160    }
161    fn extract_context(
162        &self,
163        _carrier: &BTreeMap<String, String>,
164    ) -> Box<dyn RuntimeTelemetryPropagationContext> {
165        Box::new(NoopRuntimeTelemetryPropagationContext)
166    }
167    fn flush(&self) {}
168    fn shutdown(&self) {}
169}
170
171pub trait RuntimeTelemetryPropagationContext: Send + Sync {
172    fn with_context(&self, callback: &mut dyn FnMut()) {
173        callback();
174    }
175}
176
177struct NoopRuntimeTelemetryPropagationContext;
178impl RuntimeTelemetryPropagationContext for NoopRuntimeTelemetryPropagationContext {}
179
180pub struct FailOpenRuntimeTelemetryPropagationContext {
181    delegate: Option<Box<dyn RuntimeTelemetryPropagationContext>>,
182}
183
184impl FailOpenRuntimeTelemetryPropagationContext {
185    pub async fn run<F: Future>(&self, future: F) -> F::Output {
186        futures_util::pin_mut!(future);
187        futures_util::future::poll_fn(|task_context| {
188            let Some(delegate) = self.delegate.as_ref() else {
189                return future.as_mut().poll(task_context);
190            };
191            let mut result = None;
192            let mut invoked = false;
193            let mut callback = || {
194                invoked = true;
195                result = Some(future.as_mut().poll(task_context));
196            };
197            let context_result = catch_unwind(AssertUnwindSafe(|| {
198                delegate.with_context(&mut callback);
199            }));
200            match (context_result, result) {
201                (_, Some(result)) => result,
202                (Err(payload), None) if invoked => std::panic::resume_unwind(payload),
203                _ => future.as_mut().poll(task_context),
204            }
205        })
206        .await
207    }
208}
209
210pub fn extract_runtime_context(
211    telemetry: &Arc<dyn RuntimeTelemetry>,
212    carrier: &BTreeMap<String, String>,
213) -> FailOpenRuntimeTelemetryPropagationContext {
214    let delegate = catch_unwind(AssertUnwindSafe(|| telemetry.extract_context(carrier))).ok();
215    FailOpenRuntimeTelemetryPropagationContext { delegate }
216}
217
218#[derive(Default)]
219pub struct NoopRuntimeTelemetry;
220
221impl RuntimeTelemetry for NoopRuntimeTelemetry {
222    fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
223        Box::new(NoopRuntimeTelemetryScope)
224    }
225
226    fn is_noop(&self) -> bool {
227        true
228    }
229}
230
231struct NoopRuntimeTelemetryScope;
232
233impl RuntimeTelemetryScope for NoopRuntimeTelemetryScope {
234    fn success(&mut self, _attributes: BTreeMap<String, RuntimeAttributeValue>) {}
235    fn failure(&mut self, _error_type: &str) {}
236}
237
238pub struct FailOpenRuntimeTelemetryScope {
239    delegate: Mutex<Option<Box<dyn RuntimeTelemetryScope>>>,
240}
241
242impl FailOpenRuntimeTelemetryScope {
243    pub async fn run<F: Future>(&self, future: F) -> F::Output {
244        futures_util::pin_mut!(future);
245        futures_util::future::poll_fn(|task_context| {
246            let mut result = None;
247            let Ok(delegate) = self.delegate.lock() else {
248                return future.as_mut().poll(task_context);
249            };
250            let Some(scope) = delegate.as_ref() else {
251                return future.as_mut().poll(task_context);
252            };
253            let mut invoked = false;
254            let mut callback = || {
255                invoked = true;
256                result = Some(future.as_mut().poll(task_context));
257            };
258            let context_result = catch_unwind(AssertUnwindSafe(|| {
259                scope.with_context(&mut callback);
260            }));
261            match (context_result, result) {
262                (_, Some(result)) => result,
263                (Err(payload), None) if invoked => std::panic::resume_unwind(payload),
264                _ => future.as_mut().poll(task_context),
265            }
266        })
267        .await
268    }
269
270    pub fn success(&self, attributes: BTreeMap<String, RuntimeAttributeValue>) {
271        self.finish(|scope| scope.success(attributes));
272    }
273
274    pub fn failure(&self, error_type: &str) {
275        self.finish(|scope| scope.failure(error_type));
276    }
277
278    fn finish(&self, action: impl FnOnce(&mut dyn RuntimeTelemetryScope)) {
279        let Ok(mut delegate) = self.delegate.lock() else {
280            return;
281        };
282        let Some(mut scope) = delegate.take() else {
283            return;
284        };
285        let _ = catch_unwind(AssertUnwindSafe(|| action(scope.as_mut())));
286    }
287}
288
289pub fn start_runtime_operation(
290    telemetry: &Arc<dyn RuntimeTelemetry>,
291    operation: RuntimeOperation,
292) -> FailOpenRuntimeTelemetryScope {
293    let delegate = catch_unwind(AssertUnwindSafe(|| telemetry.start(operation))).ok();
294    FailOpenRuntimeTelemetryScope {
295        delegate: Mutex::new(delegate),
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[test]
304    fn classifies_native_error_types_without_messages() {
305        assert_eq!(runtime_error_category("DatabaseTimeoutError"), "timeout");
306        assert_eq!(runtime_error_category("PermissionError"), "authorization");
307        assert_eq!(runtime_error_category("UnknownTeaQLError"), "internal");
308    }
309
310    struct BrokenTelemetry;
311    impl RuntimeTelemetry for BrokenTelemetry {
312        fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
313            panic!("adapter failed")
314        }
315    }
316
317    struct BrokenContextTelemetry;
318    impl RuntimeTelemetry for BrokenContextTelemetry {
319        fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
320            Box::new(BrokenContextScope)
321        }
322    }
323    struct BrokenContextScope;
324    impl RuntimeTelemetryScope for BrokenContextScope {
325        fn with_context(&self, _callback: &mut dyn FnMut()) {
326            panic!("context adapter failed")
327        }
328        fn success(&mut self, _attributes: BTreeMap<String, RuntimeAttributeValue>) {}
329        fn failure(&mut self, _error_type: &str) {}
330    }
331
332    #[test]
333    fn strips_forbidden_attributes_and_is_fail_open() {
334        let operation = RuntimeOperation::new("query", "School.list")
335            .attribute("teaql.entity.type", "School")
336            .attribute("teaql.entity.id", 42_i64);
337        assert_eq!(
338            operation.attributes.get("teaql.entity.type"),
339            Some(&"School".into())
340        );
341        assert!(!operation.attributes.contains_key("teaql.entity.id"));
342
343        let telemetry: Arc<dyn RuntimeTelemetry> = Arc::new(BrokenTelemetry);
344        let scope = start_runtime_operation(&telemetry, operation);
345        scope.success(BTreeMap::new());
346        scope.failure("late");
347    }
348
349    #[tokio::test]
350    async fn context_activation_is_fail_open() {
351        let telemetry: Arc<dyn RuntimeTelemetry> = Arc::new(BrokenContextTelemetry);
352        let scope =
353            start_runtime_operation(&telemetry, RuntimeOperation::new("query", "School.list"));
354        assert_eq!(scope.run(async { 42 }).await, 42);
355    }
356}
teaql_runtime/telemetry.rs

teaql_runtime/
telemetry.rs
