1use std::{
4 future::Future,
5 pin::Pin,
6 sync::Arc,
7 task::{Context, Poll},
8 time::Instant,
9};
10
11use bytes::BytesMut;
12use http::Request;
13use obs_proto::obs::v1::{ObsEnvelope, ObsHttpClientCompleted};
14use pin_project_lite::pin_project;
15use tower::Service;
16
17use crate::propagator::{TraceContext, W3cPropagator, fresh_span_id, fresh_trace_id, status_class};
18
19type StatusFn = Arc<dyn Fn(u16) -> &'static str + Send + Sync>;
20type RouteFn<B> = Arc<dyn Fn(&Request<B>) -> String + Send + Sync>;
21
22pub struct ObsHttpClientLayer<B = ()> {
24 propagator: Arc<W3cPropagator>,
25 target_extractor: RouteFn<B>,
26 status_classifier: StatusFn,
27}
28
29impl<B> std::fmt::Debug for ObsHttpClientLayer<B> {
30 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31 f.debug_struct("ObsHttpClientLayer").finish_non_exhaustive()
32 }
33}
34
35impl<B> Clone for ObsHttpClientLayer<B> {
36 fn clone(&self) -> Self {
37 Self {
38 propagator: Arc::clone(&self.propagator),
39 target_extractor: Arc::clone(&self.target_extractor),
40 status_classifier: Arc::clone(&self.status_classifier),
41 }
42 }
43}
44
45impl<B> ObsHttpClientLayer<B> {
46 #[must_use]
48 pub fn new() -> Self {
49 Self {
50 propagator: Arc::new(W3cPropagator::new()),
51 target_extractor: Arc::new(|req: &Request<B>| {
52 req.uri()
53 .host()
54 .map(ToString::to_string)
55 .unwrap_or_else(|| req.uri().to_string())
56 }),
57 status_classifier: Arc::new(|s| status_class(s)),
58 }
59 }
60
61 #[must_use]
63 pub fn with_target_extractor<F>(mut self, f: F) -> Self
64 where
65 F: Fn(&Request<B>) -> String + Send + Sync + 'static,
66 {
67 self.target_extractor = Arc::new(f);
68 self
69 }
70}
71
72impl<B> Default for ObsHttpClientLayer<B> {
73 fn default() -> Self {
74 Self::new()
75 }
76}
77
78impl<S, B> tower::Layer<S> for ObsHttpClientLayer<B>
79where
80 S: Service<Request<B>>,
81{
82 type Service = ObsHttpClientService<S, B>;
83 fn layer(&self, inner: S) -> Self::Service {
84 ObsHttpClientService {
85 inner,
86 layer: self.clone(),
87 }
88 }
89}
90
91pub struct ObsHttpClientService<S, B> {
93 inner: S,
94 layer: ObsHttpClientLayer<B>,
95}
96
97impl<S, B> std::fmt::Debug for ObsHttpClientService<S, B> {
98 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99 f.debug_struct("ObsHttpClientService")
100 .field("layer", &self.layer)
101 .finish_non_exhaustive()
102 }
103}
104
105impl<S, B> Clone for ObsHttpClientService<S, B>
106where
107 S: Clone,
108{
109 fn clone(&self) -> Self {
110 Self {
111 inner: self.inner.clone(),
112 layer: self.layer.clone(),
113 }
114 }
115}
116
117impl<S, B, ResBody> Service<Request<B>> for ObsHttpClientService<S, B>
118where
119 S: Service<Request<B>, Response = http::Response<ResBody>>,
120 S::Future: Send + 'static,
121 B: Send + 'static,
122{
123 type Response = S::Response;
124 type Error = S::Error;
125 type Future = ObsHttpClientFuture<S::Future>;
126
127 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
128 self.inner.poll_ready(cx)
129 }
130
131 fn call(&mut self, mut req: Request<B>) -> Self::Future {
132 let started = Instant::now();
133 let target = (self.layer.target_extractor)(&req);
134 let method = req.method().as_str().to_string();
135 let propagator = Arc::clone(&self.layer.propagator);
136 let status_classifier = Arc::clone(&self.layer.status_classifier);
137
138 let sampled = obs_core::scope::active_sampled().unwrap_or(true);
144 let flags = if sampled { "01" } else { "00" };
145 let ctx = match obs_core::scope::active_correlation() {
146 Some((trace_id, parent_span)) => TraceContext {
147 trace_id,
148 span_id: fresh_span_id(),
149 flags: flags.to_string(),
150 tracestate: format!("parent={parent_span}"),
151 },
152 None => TraceContext {
153 trace_id: fresh_trace_id(),
154 span_id: fresh_span_id(),
155 flags: flags.to_string(),
156 tracestate: String::new(),
157 },
158 };
159 propagator.inject(req.headers_mut(), &ctx);
160 let trace_id = ctx.trace_id.clone();
161 let span_id = ctx.span_id.clone();
162 emit_client_started(&target, &method, &trace_id);
163
164 ObsHttpClientFuture {
165 inner: self.inner.call(req),
166 started: Some(started),
167 target,
168 method,
169 trace_id,
170 span_id,
171 status_classifier,
172 }
173 }
174}
175
176pin_project! {
177 pub struct ObsHttpClientFuture<F> {
179 #[pin]
180 inner: F,
181 started: Option<Instant>,
182 target: String,
183 method: String,
184 trace_id: String,
185 span_id: String,
186 status_classifier: StatusFn,
187 }
188}
189
190impl<F, ResBody, E> Future for ObsHttpClientFuture<F>
191where
192 F: Future<Output = Result<http::Response<ResBody>, E>>,
193{
194 type Output = F::Output;
195 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
196 let this = self.project();
197 match this.inner.poll(cx) {
198 Poll::Pending => Poll::Pending,
199 Poll::Ready(out) => {
200 let started = this.started.take().unwrap_or_else(Instant::now);
201 let elapsed_ms = started.elapsed().as_millis() as u64;
202 let class = match &out {
203 Ok(resp) => (this.status_classifier)(resp.status().as_u16()),
204 Err(_) => "err",
205 };
206 emit_client_completed(
207 this.target,
208 this.method,
209 class,
210 elapsed_ms,
211 this.trace_id,
212 this.span_id,
213 );
214 Poll::Ready(out)
215 }
216 }
217 }
218}
219
220fn encode_into<M: ::buffa::Message>(msg: &M, out: &mut Vec<u8>) {
223 let mut cache = ::buffa::SizeCache::default();
224 let size = msg.compute_size(&mut cache);
225 let mut buf = BytesMut::with_capacity(size as usize);
226 msg.write_to(&mut cache, &mut buf);
227 out.clear();
228 out.extend_from_slice(&buf);
229}
230
231fn emit_client_started(target: &str, method: &str, trace_id: &str) {
232 let typed = obs_proto::obs::v1::ObsHttpClientStarted {
233 method: method.to_string(),
234 host: target.to_string(),
235 __buffa_unknown_fields: Default::default(),
236 };
237 let mut env = ObsEnvelope {
238 full_name: "obs.v1.ObsHttpClientStarted".to_string(),
239 tier: ::buffa::EnumValue::Known(obs_proto::obs::v1::Tier::TIER_LOG),
240 sev: ::buffa::EnumValue::Known(obs_proto::obs::v1::Severity::SEVERITY_INFO),
241 trace_id: trace_id.to_string(),
242 ..Default::default()
243 };
244 encode_into(&typed, &mut env.payload);
245 env.labels.insert("host".to_string(), target.to_string());
246 env.labels.insert("method".to_string(), method.to_string());
247 obs_core::observer().emit_envelope(env);
248}
249
250fn emit_client_completed(
251 target: &str,
252 method: &str,
253 status_class: &str,
254 latency_ms: u64,
255 trace_id: &str,
256 span_id: &str,
257) {
258 let typed = ObsHttpClientCompleted {
264 method: method.to_string(),
265 host: target.to_string(),
266 status_class: status_class.to_string(),
267 latency_ms,
268 __buffa_unknown_fields: Default::default(),
269 };
270 let mut env = ObsEnvelope {
271 full_name: "obs.v1.ObsHttpClientCompleted".to_string(),
272 tier: ::buffa::EnumValue::Known(obs_proto::obs::v1::Tier::TIER_LOG),
273 sev: ::buffa::EnumValue::Known(obs_proto::obs::v1::Severity::SEVERITY_INFO),
274 trace_id: trace_id.to_string(),
275 span_id: span_id.to_string(),
276 ..Default::default()
277 };
278 encode_into(&typed, &mut env.payload);
279 env.labels.insert("host".to_string(), target.to_string());
280 env.labels.insert("method".to_string(), method.to_string());
281 env.labels
282 .insert("status_class".to_string(), status_class.to_string());
283 obs_core::observer().emit_envelope(env);
284}