lightswitch_metadata/
metadata_provider.rs

1use crate::system_metadata::SystemMetadata;
2use crate::task_metadata::TaskMetadata;
3use crate::types::{MetadataLabel, SystemMetadataProvider, TaskKey, TaskMetadataProvider};
4
5use lru::LruCache;
6use std::num::NonZeroUsize;
7use std::sync::{Arc, Mutex};
8use tracing::warn;
9
10pub struct GlobalMetadataProvider {
11    process_label_cache: LruCache<TaskKey, Vec<MetadataLabel>>,
12    default_task_metadata: TaskMetadata,
13    default_system_metadata: SystemMetadata,
14    custom_system_metadata_providers: Vec<Box<dyn SystemMetadataProvider + Send>>,
15    custom_task_metadata_providers: Vec<Box<dyn TaskMetadataProvider + Send>>,
16}
17
18pub type ThreadSafeGlobalMetadataProvider = Arc<Mutex<GlobalMetadataProvider>>;
19
20impl Default for GlobalMetadataProvider {
21    fn default() -> Self {
22        Self::new(NonZeroUsize::new(5000).unwrap(), Vec::new(), Vec::new())
23    }
24}
25
26impl GlobalMetadataProvider {
27    pub fn new(
28        metadata_cache_size: NonZeroUsize,
29        system_metadata_providers: Vec<Box<dyn SystemMetadataProvider + Send>>,
30        task_metadata_providers: Vec<Box<dyn TaskMetadataProvider + Send>>,
31    ) -> Self {
32        Self {
33            process_label_cache: LruCache::new(metadata_cache_size),
34            default_task_metadata: TaskMetadata {},
35            default_system_metadata: SystemMetadata {},
36            custom_system_metadata_providers: system_metadata_providers,
37            custom_task_metadata_providers: task_metadata_providers,
38        }
39    }
40    pub fn register_task_metadata_providers(
41        &mut self,
42        providers: Vec<Box<dyn TaskMetadataProvider + Send>>,
43    ) {
44        self.custom_task_metadata_providers.extend(providers);
45    }
46
47    pub fn register_system_metadata_providers(
48        &mut self,
49        providers: Vec<Box<dyn SystemMetadataProvider + Send>>,
50    ) {
51        self.custom_system_metadata_providers.extend(providers);
52    }
53
54    fn get_labels(&mut self, task_key: TaskKey) -> Vec<MetadataLabel> {
55        let mut labels = self
56            .default_task_metadata
57            .get_metadata(task_key)
58            .map_err(|err| warn!("{}", err))
59            .unwrap_or_default();
60
61        labels.extend(
62            self.default_system_metadata
63                .get_metadata()
64                .map_err(|err| {
65                    warn!(
66                        "Failed to retrieve default system metadata, error = {}",
67                        err
68                    )
69                })
70                .unwrap_or_default(),
71        );
72
73        for provider in &self.custom_system_metadata_providers {
74            match provider.get_metadata() {
75                Ok(custom_system_labels) => {
76                    labels.extend(custom_system_labels.into_iter());
77                }
78                Err(err) => {
79                    warn!("Failed to retrieve custom system metadata, error = {}", err);
80                }
81            }
82        }
83
84        for provider in &self.custom_task_metadata_providers {
85            match provider.get_metadata(task_key) {
86                Ok(custom_task_labels) => {
87                    labels.extend(custom_task_labels.into_iter());
88                }
89                Err(err) => {
90                    warn!("Failed to retrieve custom task metadata, error = {}", err);
91                }
92            }
93        }
94        labels
95    }
96
97    pub fn get_metadata(&mut self, task_key: TaskKey) -> Vec<MetadataLabel> {
98        if let Some(cached_labels) = self.process_label_cache.get(&task_key) {
99            cached_labels.to_vec()
100        } else {
101            let labels = self.get_labels(task_key);
102            self.process_label_cache.push(task_key, labels.clone());
103            labels
104        }
105    }
106
107    pub fn register_task(&mut self, task_key: TaskKey) {
108        if !self.process_label_cache.contains(&task_key) {
109            let labels = self.get_labels(task_key);
110            self.process_label_cache.push(task_key, labels);
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::taskname::TaskName;
119    use crate::types::MetadataLabelValue;
120    use nix::unistd;
121
122    #[test]
123    fn test_get_metadata_returns_minimal_labels() {
124        // Given
125        let tid = unistd::gettid().as_raw();
126        let pid = unistd::getpgrp().as_raw();
127        let mut metadata_provider = GlobalMetadataProvider::default();
128        let expected = TaskName::for_task(tid).unwrap();
129
130        // When
131        let labels = metadata_provider.get_metadata(TaskKey { tid, pid });
132
133        // Then
134        assert_eq!(labels[0].key, "pid");
135        assert_eq!(
136            labels[0].value,
137            MetadataLabelValue::Number(tid.into(), "task-id".into())
138        );
139        assert_eq!(labels[1].key, "thread.name");
140        assert_eq!(
141            labels[1].value,
142            MetadataLabelValue::String(expected.current_thread)
143        );
144        assert_eq!(labels[2].key, "process.name");
145        assert_eq!(
146            labels[2].value,
147            MetadataLabelValue::String(expected.main_thread)
148        );
149        assert_eq!(labels[3].key, "pid");
150        assert_eq!(
151            labels[3].value,
152            MetadataLabelValue::Number(pid.into(), "task-tgid".into())
153        );
154    }
155}
lightswitch_metadata/metadata_provider.rs

lightswitch_metadata/
metadata_provider.rs
