otel_bootstrap/
profiling.rs1#![cfg(feature = "profiling")]
2
3use std::error::Error;
4use std::sync::OnceLock;
5
6#[cfg(feature = "profiling-bridge-pyroscope-rs")]
7use opentelemetry::trace::TraceContextExt;
8
9fn validate_pyroscope_endpoint(endpoint: &str) -> Result<(), Box<dyn Error>> {
13 use url::Url;
14
15 if endpoint.starts_with("unix://") {
17 return Ok(());
18 }
19
20 if endpoint.starts_with("http://") || endpoint.starts_with("https://") {
22 let url = Url::parse(endpoint)?;
23
24 if !url.username().is_empty() || url.password().is_some() {
26 return Err(format!(
27 "pyroscope endpoint must not contain userinfo; got: {endpoint} (ADR platform/0203 AC1)"
28 ).into());
29 }
30
31 let host = url.host_str().unwrap_or("");
32
33 match host {
34 "127.0.0.1" | "::1" | "[::1]" | "localhost" => Ok(()),
35 _ => Err(format!(
36 "pyroscope endpoint must target loopback (127.0.0.1, ::1, localhost, or unix socket); \
37 got: {endpoint} (ADR platform/0203 AC1)"
38 ).into()),
39 }
40 } else {
41 Err(
42 format!("pyroscope endpoint must be http://, https://, or unix://; got: {endpoint}")
43 .into(),
44 )
45 }
46}
47
48pub struct ProfilingHandle {
51 #[cfg(feature = "profiling-bridge-pyroscope-rs")]
52 agent: Option<pyroscope::PyroscopeAgent<pyroscope::pyroscope::PyroscopeAgentRunning>>,
53}
54
55#[cfg(feature = "profiling-bridge-pyroscope-rs")]
56impl Drop for ProfilingHandle {
57 fn drop(&mut self) {
58 if let Some(agent) = self.agent.take() {
59 let _ = agent.stop();
60 }
61 }
62}
63
64#[cfg(feature = "profiling-bridge-pyroscope-rs")]
65type BoxedTagFn = Box<dyn Fn(String, String) -> pyroscope::Result<()> + Send + Sync>;
66
67#[cfg(feature = "profiling-bridge-pyroscope-rs")]
70static PROFILING_TAG_FNS: OnceLock<(BoxedTagFn, BoxedTagFn)> = OnceLock::new();
71
72#[cfg(feature = "profiling-bridge-pyroscope-rs")]
77static PROFILING_STARTED: OnceLock<()> = OnceLock::new();
78
79#[cfg(feature = "profiling-bridge-pyroscope-rs")]
90pub(crate) fn start_pyroscope_bridge(
91 service_name: &str,
92 pyroscope_endpoint: &str,
93) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
94 use pyroscope_pprofrs::{PprofConfig, pprof_backend};
95
96 validate_pyroscope_endpoint(pyroscope_endpoint)?;
98
99 if PROFILING_STARTED.set(()).is_err() {
102 return Ok(None);
103 }
104
105 let agent = pyroscope::PyroscopeAgent::builder(pyroscope_endpoint, service_name)
106 .backend(pprof_backend(PprofConfig::new().sample_rate(100)))
107 .build()?
108 .start()?;
109
110 let (add_tag, remove_tag) = agent.tag_wrapper();
111 PROFILING_TAG_FNS
112 .set((Box::new(add_tag), Box::new(remove_tag)))
113 .ok();
114
115 Ok(Some(ProfilingHandle { agent: Some(agent) }))
116}
117
118#[cfg(all(feature = "profiling", not(feature = "profiling-bridge-pyroscope-rs")))]
120pub(crate) fn start_pyroscope_bridge(
121 _service_name: &str,
122 _pyroscope_endpoint: &str,
123) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
124 Ok(None)
125}
126
127#[cfg(feature = "profiling-bridge-pyroscope-rs")]
130pub struct ProfilingTagLayer;
131
132#[cfg(feature = "profiling-bridge-pyroscope-rs")]
133impl<S> tracing_subscriber::Layer<S> for ProfilingTagLayer
134where
135 S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
136{
137 fn on_enter(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
138 if let Some((add_tag, _)) = PROFILING_TAG_FNS.get() {
139 let cx = opentelemetry::Context::current();
140 let span_ref = cx.span();
141 let span_context = span_ref.span_context();
142 if span_context.is_valid() {
143 let trace_id = span_context.trace_id();
144 let span_id = span_context.span_id();
145 let _ = add_tag("trace_id".to_string(), format!("{trace_id:x}"));
146 let _ = add_tag("span_id".to_string(), format!("{span_id:x}"));
147 }
148 }
149 }
150
151 fn on_exit(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
152 if let Some((_, remove_tag)) = PROFILING_TAG_FNS.get() {
153 let cx = opentelemetry::Context::current();
154 let span_ref = cx.span();
155 let span_context = span_ref.span_context();
156 if span_context.is_valid() {
157 let trace_id = span_context.trace_id();
158 let span_id = span_context.span_id();
159 let _ = remove_tag("trace_id".to_string(), format!("{trace_id:x}"));
160 let _ = remove_tag("span_id".to_string(), format!("{span_id:x}"));
161 }
162 }
163 }
164}
165
166#[cfg(all(test, feature = "profiling-bridge-pyroscope-rs"))]
167mod tests {
168 use super::*;
169
170 #[test]
171 fn start_bridge_with_nonexistent_server() {
172 let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
173 assert!(
174 result.is_ok(),
175 "pyroscope agent start() is lazy and does not eagerly connect"
176 );
177 if let Ok(Some(_handle)) = result {
178 }
180 }
181
182 #[test]
183 fn start_bridge_multiple_times_ignores_second() {
184 let result1 = start_pyroscope_bridge("test-svc-1", "http://localhost:4040");
185 assert!(result1.is_ok());
186 let result2 = start_pyroscope_bridge("test-svc-2", "http://localhost:4041");
187 assert!(result2.is_ok());
188 assert!(result2.unwrap().is_none());
191 }
192
193 #[test]
194 fn validate_endpoint_accepts_loopback_ipv4() {
195 assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040").is_ok());
196 }
197
198 #[test]
199 fn validate_endpoint_accepts_loopback_ipv6() {
200 assert!(validate_pyroscope_endpoint("http://[::1]:4040").is_ok());
202 }
203
204 #[test]
205 fn validate_endpoint_accepts_localhost() {
206 assert!(validate_pyroscope_endpoint("http://localhost:4040").is_ok());
207 }
208
209 #[test]
210 fn validate_endpoint_accepts_https_loopback() {
211 assert!(validate_pyroscope_endpoint("https://127.0.0.1:4040").is_ok());
212 }
213
214 #[test]
215 fn validate_endpoint_rejects_routable_ipv4() {
216 assert!(validate_pyroscope_endpoint("http://10.0.0.1:4040").is_err());
217 }
218
219 #[test]
220 fn validate_endpoint_rejects_userinfo_bypass() {
221 assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040@evil.com/").is_err());
223 }
224
225 #[test]
226 fn validate_endpoint_rejects_userinfo_with_password() {
227 assert!(validate_pyroscope_endpoint("http://user:pass@localhost:4040").is_err());
228 }
229
230 #[test]
231 fn validate_endpoint_rejects_unix_socket_check() {
232 assert!(validate_pyroscope_endpoint("unix:///var/run/profiling.sock").is_ok());
233 }
234}
235
236#[cfg(all(
237 test,
238 feature = "profiling",
239 not(feature = "profiling-bridge-pyroscope-rs")
240))]
241mod tests_no_bridge {
242 use super::*;
243
244 #[test]
245 fn start_bridge_returns_none() {
246 let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
247 assert!(result.is_ok());
248 if let Ok(handle) = result {
249 assert!(handle.is_none());
250 }
251 }
252}