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 extract_context(
159        &self,
160        _carrier: &BTreeMap<String, String>,
161    ) -> Box<dyn RuntimeTelemetryPropagationContext> {
162        Box::new(NoopRuntimeTelemetryPropagationContext)
163    }
164    fn flush(&self) {}
165    fn shutdown(&self) {}
166}
167
168pub trait RuntimeTelemetryPropagationContext: Send + Sync {
169    fn with_context(&self, callback: &mut dyn FnMut()) {
170        callback();
171    }
172}
173
174struct NoopRuntimeTelemetryPropagationContext;
175impl RuntimeTelemetryPropagationContext for NoopRuntimeTelemetryPropagationContext {}
176
177pub struct FailOpenRuntimeTelemetryPropagationContext {
178    delegate: Option<Box<dyn RuntimeTelemetryPropagationContext>>,
179}
180
181impl FailOpenRuntimeTelemetryPropagationContext {
182    pub async fn run<F: Future>(&self, future: F) -> F::Output {
183        futures_util::pin_mut!(future);
184        futures_util::future::poll_fn(|task_context| {
185            let Some(delegate) = self.delegate.as_ref() else {
186                return future.as_mut().poll(task_context);
187            };
188            let mut result = None;
189            let mut invoked = false;
190            let mut callback = || {
191                invoked = true;
192                result = Some(future.as_mut().poll(task_context));
193            };
194            let context_result = catch_unwind(AssertUnwindSafe(|| {
195                delegate.with_context(&mut callback);
196            }));
197            match (context_result, result) {
198                (_, Some(result)) => result,
199                (Err(payload), None) if invoked => std::panic::resume_unwind(payload),
200                _ => future.as_mut().poll(task_context),
201            }
202        })
203        .await
204    }
205}
206
207pub fn extract_runtime_context(
208    telemetry: &Arc<dyn RuntimeTelemetry>,
209    carrier: &BTreeMap<String, String>,
210) -> FailOpenRuntimeTelemetryPropagationContext {
211    let delegate = catch_unwind(AssertUnwindSafe(|| telemetry.extract_context(carrier))).ok();
212    FailOpenRuntimeTelemetryPropagationContext { delegate }
213}
214
215#[derive(Default)]
216pub struct NoopRuntimeTelemetry;
217
218impl RuntimeTelemetry for NoopRuntimeTelemetry {
219    fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
220        Box::new(NoopRuntimeTelemetryScope)
221    }
222}
223
224struct NoopRuntimeTelemetryScope;
225
226impl RuntimeTelemetryScope for NoopRuntimeTelemetryScope {
227    fn success(&mut self, _attributes: BTreeMap<String, RuntimeAttributeValue>) {}
228    fn failure(&mut self, _error_type: &str) {}
229}
230
231pub struct FailOpenRuntimeTelemetryScope {
232    delegate: Mutex<Option<Box<dyn RuntimeTelemetryScope>>>,
233}
234
235impl FailOpenRuntimeTelemetryScope {
236    pub async fn run<F: Future>(&self, future: F) -> F::Output {
237        futures_util::pin_mut!(future);
238        futures_util::future::poll_fn(|task_context| {
239            let mut result = None;
240            let Ok(delegate) = self.delegate.lock() else {
241                return future.as_mut().poll(task_context);
242            };
243            let Some(scope) = delegate.as_ref() else {
244                return future.as_mut().poll(task_context);
245            };
246            let mut invoked = false;
247            let mut callback = || {
248                invoked = true;
249                result = Some(future.as_mut().poll(task_context));
250            };
251            let context_result = catch_unwind(AssertUnwindSafe(|| {
252                scope.with_context(&mut callback);
253            }));
254            match (context_result, result) {
255                (_, Some(result)) => result,
256                (Err(payload), None) if invoked => std::panic::resume_unwind(payload),
257                _ => future.as_mut().poll(task_context),
258            }
259        })
260        .await
261    }
262
263    pub fn success(&self, attributes: BTreeMap<String, RuntimeAttributeValue>) {
264        self.finish(|scope| scope.success(attributes));
265    }
266
267    pub fn failure(&self, error_type: &str) {
268        self.finish(|scope| scope.failure(error_type));
269    }
270
271    fn finish(&self, action: impl FnOnce(&mut dyn RuntimeTelemetryScope)) {
272        let Ok(mut delegate) = self.delegate.lock() else {
273            return;
274        };
275        let Some(mut scope) = delegate.take() else {
276            return;
277        };
278        let _ = catch_unwind(AssertUnwindSafe(|| action(scope.as_mut())));
279    }
280}
281
282pub fn start_runtime_operation(
283    telemetry: &Arc<dyn RuntimeTelemetry>,
284    operation: RuntimeOperation,
285) -> FailOpenRuntimeTelemetryScope {
286    let delegate = catch_unwind(AssertUnwindSafe(|| telemetry.start(operation))).ok();
287    FailOpenRuntimeTelemetryScope {
288        delegate: Mutex::new(delegate),
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn classifies_native_error_types_without_messages() {
298        assert_eq!(runtime_error_category("DatabaseTimeoutError"), "timeout");
299        assert_eq!(runtime_error_category("PermissionError"), "authorization");
300        assert_eq!(runtime_error_category("UnknownTeaQLError"), "internal");
301    }
302
303    struct BrokenTelemetry;
304    impl RuntimeTelemetry for BrokenTelemetry {
305        fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
306            panic!("adapter failed")
307        }
308    }
309
310    struct BrokenContextTelemetry;
311    impl RuntimeTelemetry for BrokenContextTelemetry {
312        fn start(&self, _operation: RuntimeOperation) -> Box<dyn RuntimeTelemetryScope> {
313            Box::new(BrokenContextScope)
314        }
315    }
316    struct BrokenContextScope;
317    impl RuntimeTelemetryScope for BrokenContextScope {
318        fn with_context(&self, _callback: &mut dyn FnMut()) {
319            panic!("context adapter failed")
320        }
321        fn success(&mut self, _attributes: BTreeMap<String, RuntimeAttributeValue>) {}
322        fn failure(&mut self, _error_type: &str) {}
323    }
324
325    #[test]
326    fn strips_forbidden_attributes_and_is_fail_open() {
327        let operation = RuntimeOperation::new("query", "School.list")
328            .attribute("teaql.entity.type", "School")
329            .attribute("teaql.entity.id", 42_i64);
330        assert_eq!(
331            operation.attributes.get("teaql.entity.type"),
332            Some(&"School".into())
333        );
334        assert!(!operation.attributes.contains_key("teaql.entity.id"));
335
336        let telemetry: Arc<dyn RuntimeTelemetry> = Arc::new(BrokenTelemetry);
337        let scope = start_runtime_operation(&telemetry, operation);
338        scope.success(BTreeMap::new());
339        scope.failure("late");
340    }
341
342    #[tokio::test]
343    async fn context_activation_is_fail_open() {
344        let telemetry: Arc<dyn RuntimeTelemetry> = Arc::new(BrokenContextTelemetry);
345        let scope =
346            start_runtime_operation(&telemetry, RuntimeOperation::new("query", "School.list"));
347        assert_eq!(scope.run(async { 42 }).await, 42);
348    }
349}
teaql_runtime/telemetry.rs

teaql_runtime/
telemetry.rs
