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::backend::{BackendConfig, 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::pyroscope::PyroscopeAgentBuilder::new(
106 pyroscope_endpoint,
107 service_name,
108 100,
109 "pyroscope-rs",
110 env!("CARGO_PKG_VERSION"),
111 pprof_backend(PprofConfig { sample_rate: 100 }, BackendConfig::default()),
112 )
113 .build()?
114 .start()?;
115
116 let (add_tag, remove_tag) = agent.tag_wrapper();
117 PROFILING_TAG_FNS
118 .set((Box::new(add_tag), Box::new(remove_tag)))
119 .ok();
120
121 Ok(Some(ProfilingHandle { agent: Some(agent) }))
122}
123
124#[cfg(all(feature = "profiling", not(feature = "profiling-bridge-pyroscope-rs")))]
126pub(crate) fn start_pyroscope_bridge(
127 _service_name: &str,
128 _pyroscope_endpoint: &str,
129) -> Result<Option<ProfilingHandle>, Box<dyn Error>> {
130 Ok(None)
131}
132
133#[cfg(feature = "profiling-bridge-pyroscope-rs")]
136pub struct ProfilingTagLayer;
137
138#[cfg(feature = "profiling-bridge-pyroscope-rs")]
139impl<S> tracing_subscriber::Layer<S> for ProfilingTagLayer
140where
141 S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
142{
143 fn on_enter(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
144 if let Some((add_tag, _)) = PROFILING_TAG_FNS.get() {
145 let cx = opentelemetry::Context::current();
146 let span_ref = cx.span();
147 let span_context = span_ref.span_context();
148 if span_context.is_valid() {
149 let trace_id = span_context.trace_id();
150 let span_id = span_context.span_id();
151 let _ = add_tag("trace_id".to_string(), format!("{trace_id:x}"));
152 let _ = add_tag("span_id".to_string(), format!("{span_id:x}"));
153 }
154 }
155 }
156
157 fn on_exit(&self, _id: &tracing::span::Id, _ctx: tracing_subscriber::layer::Context<'_, S>) {
158 if let Some((_, remove_tag)) = PROFILING_TAG_FNS.get() {
159 let cx = opentelemetry::Context::current();
160 let span_ref = cx.span();
161 let span_context = span_ref.span_context();
162 if span_context.is_valid() {
163 let trace_id = span_context.trace_id();
164 let span_id = span_context.span_id();
165 let _ = remove_tag("trace_id".to_string(), format!("{trace_id:x}"));
166 let _ = remove_tag("span_id".to_string(), format!("{span_id:x}"));
167 }
168 }
169 }
170}
171
172#[cfg(all(test, feature = "profiling-bridge-pyroscope-rs"))]
173mod tests {
174 use super::*;
175
176 #[test]
177 fn start_bridge_with_nonexistent_server() {
178 let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
179 assert!(
180 result.is_ok(),
181 "pyroscope agent start() is lazy and does not eagerly connect"
182 );
183 if let Ok(Some(_handle)) = result {
184 }
186 }
187
188 #[test]
189 fn start_bridge_multiple_times_ignores_second() {
190 let result1 = start_pyroscope_bridge("test-svc-1", "http://localhost:4040");
191 assert!(result1.is_ok());
192 let result2 = start_pyroscope_bridge("test-svc-2", "http://localhost:4041");
193 assert!(result2.is_ok());
194 assert!(result2.unwrap().is_none());
197 }
198
199 #[test]
200 fn validate_endpoint_accepts_loopback_ipv4() {
201 assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040").is_ok());
202 }
203
204 #[test]
205 fn validate_endpoint_accepts_loopback_ipv6() {
206 assert!(validate_pyroscope_endpoint("http://[::1]:4040").is_ok());
208 }
209
210 #[test]
211 fn validate_endpoint_accepts_localhost() {
212 assert!(validate_pyroscope_endpoint("http://localhost:4040").is_ok());
213 }
214
215 #[test]
216 fn validate_endpoint_accepts_https_loopback() {
217 assert!(validate_pyroscope_endpoint("https://127.0.0.1:4040").is_ok());
218 }
219
220 #[test]
221 fn validate_endpoint_rejects_routable_ipv4() {
222 assert!(validate_pyroscope_endpoint("http://10.0.0.1:4040").is_err());
223 }
224
225 #[test]
226 fn validate_endpoint_rejects_userinfo_bypass() {
227 assert!(validate_pyroscope_endpoint("http://127.0.0.1:4040@evil.com/").is_err());
229 }
230
231 #[test]
232 fn validate_endpoint_rejects_userinfo_with_password() {
233 assert!(validate_pyroscope_endpoint("http://user:pass@localhost:4040").is_err());
234 }
235
236 #[test]
237 fn validate_endpoint_rejects_unix_socket_check() {
238 assert!(validate_pyroscope_endpoint("unix:///var/run/profiling.sock").is_ok());
239 }
240}
241
242#[cfg(all(
243 test,
244 feature = "profiling",
245 not(feature = "profiling-bridge-pyroscope-rs")
246))]
247mod tests_no_bridge {
248 use super::*;
249
250 #[test]
251 fn start_bridge_returns_none() {
252 let result = start_pyroscope_bridge("test-svc", "http://localhost:4040");
253 assert!(result.is_ok());
254 if let Ok(handle) = result {
255 assert!(handle.is_none());
256 }
257 }
258}