witchcraft_server/debug/
mod.rs

1// Copyright 2022 Palantir Technologies, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//! Debug endpoints.
15use std::collections::hash_map::Entry;
16use std::collections::HashMap;
17use std::sync::Arc;
18
19use bytes::Bytes;
20use conjure_error::Error;
21use http::HeaderValue;
22use once_cell::sync::Lazy;
23use parking_lot::Mutex;
24use regex::Regex;
25
26pub(crate) mod diagnostic_types;
27pub(crate) mod endpoint;
28#[cfg(feature = "jemalloc")]
29pub(crate) mod heap_profile;
30#[cfg(feature = "jemalloc")]
31pub(crate) mod heap_stats;
32pub(crate) mod metric_names;
33#[cfg(target_os = "linux")]
34pub(crate) mod thread_dump;
35
36static TYPE_PATTERN: Lazy<Regex> = Lazy::new(|| Regex::new(r"([a-z0-9]+\.)+v[0-9]+").unwrap());
37
38/// An SLS diagnostic. See the [SLS debug spec](https://github.palantir.build/deployability/sls-spec/blob/develop/docs/debug.md)
39pub trait Diagnostic {
40    /// The type of the diagnostic. Must be lower cased, dot delimited, and end with a version.
41    ///
42    /// Example: "my.diagnostic.v1"
43    fn type_(&self) -> &str;
44
45    /// The value of the "Content-Type" header.
46    ///
47    /// Example: "application/json"
48    fn content_type(&self) -> HeaderValue;
49
50    /// Whether the value is safe to log
51    fn safe_loggable(&self) -> bool;
52
53    /// The bytes of the response to send
54    fn result(&self) -> Result<Bytes, Error>;
55}
56
57/// A registry of diagnostics for the server.
58pub struct DiagnosticRegistry {
59    diagnostics: Mutex<HashMap<String, Arc<dyn Diagnostic + Sync + Send>>>,
60}
61
62impl Default for DiagnosticRegistry {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68impl DiagnosticRegistry {
69    /// Create a new diagnostic registry
70    pub fn new() -> Self {
71        DiagnosticRegistry {
72            diagnostics: Mutex::new(HashMap::new()),
73        }
74    }
75
76    /// Register a new diagnostic.
77    ///
78    /// # Panics
79    ///
80    /// Panics if the diagnostic type is not valid. See [Diagnostic::type_] for valid diagnostic
81    /// types.
82    ///
83    /// Panics if another diagnostic of the same type is already registered.
84    pub fn register<T>(&self, diagnostic: T)
85    where
86        T: Diagnostic + 'static + Sync + Send,
87    {
88        self.register_inner(Arc::new(diagnostic));
89    }
90
91    fn register_inner(&self, diagnostic: Arc<dyn Diagnostic + Sync + Send>) {
92        let type_ = diagnostic.type_();
93
94        assert!(
95            TYPE_PATTERN.is_match(type_),
96            "{type_} must be `lower.case.dot.delimited.v1`",
97        );
98
99        match self.diagnostics.lock().entry(type_.to_string()) {
100            Entry::Occupied(_) => {
101                panic!("a diagnostic has already been registered for type {type_}")
102            }
103            Entry::Vacant(e) => {
104                e.insert(diagnostic);
105            }
106        }
107    }
108
109    /// Get the diagnostic with the specified type.
110    pub fn get(&self, type_: &str) -> Option<Arc<dyn Diagnostic + Sync + Send>> {
111        self.diagnostics.lock().get(type_).cloned()
112    }
113}
witchcraft_server/debug/mod.rs

witchcraft_server/debug/
mod.rs
