feat(nunciod): per-RPC tracing span with request-id correlation and auth logging #371

Merged
KofTwentyTwo merged 1 commits from story/obs-2-rpc-spans-correlation into dev 2026-07-31 23:39:01 +00:00
8 changed files with 711 additions and 22 deletions

No files matched your search

Generated
+4
View File
@@ -2679,6 +2679,9 @@ dependencies = [
"bytes",
"fs4",
"hex",
"http",
"http-body",
"http-body-util",
"nuncio-cal",
"nuncio-cli",
"nuncio-contacts",
@@ -2696,6 +2699,7 @@ dependencies = [
"tokio",
"tokio-stream",
"tonic",
"tower",
"tracing",
"tracing-appender",
"tracing-subscriber",
+4
View File
@@ -80,6 +80,10 @@ flate2 = "1.0"
tar = "0.4"
zip = { version = "8.6", default-features = false, features = ["deflate"] }
tonic = "0.14"
tower = "0.5"
http = "1"
http-body = "1"
http-body-util = "0.1"
tonic-prost = "0.14"
prost = "0.14"
prost-types = "0.14"
+4
View File
@@ -30,6 +30,9 @@ tracing-subscriber.workspace = true
tracing-appender.workspace = true
tempfile.workspace = true
tonic.workspace = true
tower.workspace = true
http.workspace = true
http-body.workspace = true
async-trait.workspace = true
subtle.workspace = true
hex.workspace = true
@@ -41,4 +44,5 @@ fs4.workspace = true
nuncio-cli = { path = "../nuncio-cli" }
sqlx.workspace = true
wiremock.workspace = true
http-body-util.workspace = true
tokio = { workspace = true, features = ["test-util"] }
+126 -6
View File
@@ -2458,25 +2458,41 @@ impl BearerAuthInterceptor {
impl tonic::service::Interceptor for BearerAuthInterceptor {
fn call(&mut self, request: Request<()>) -> Result<Request<()>, Status> {
// Only ever log the peer and a fixed reason string. The token value,
// the raw `authorization` header, and anything derived from either
// (length, prefix, substring) are NEVER logged -- doing so would leak
// the bearer credential into the logs. This runs inside the per-RPC
// span (see `rpc_trace`), so these events carry the request_id.
let peer = request
.remote_addr()
.map(|addr| addr.to_string())
.unwrap_or_else(|| "unknown".to_string());
let reject = |reason: &'static str| {
tracing::warn!(peer = %peer, reason, "gRPC auth rejected");
Status::unauthenticated(reason)
};
let header = request
.metadata()
.get("authorization")
.ok_or_else(|| Status::unauthenticated("missing authorization metadata"))?;
.ok_or_else(|| reject("missing authorization metadata"))?;
let header_str = header
.to_str()
.map_err(|_| Status::unauthenticated("authorization metadata is not valid ASCII"))?;
let provided = header_str.strip_prefix("Bearer ").ok_or_else(|| {
Status::unauthenticated("expected 'Bearer <token>' authorization scheme")
})?;
.map_err(|_| reject("authorization metadata is not valid ASCII"))?;
let provided = header_str
.strip_prefix("Bearer ")
.ok_or_else(|| reject("expected 'Bearer <token>' authorization scheme"))?;
let matches: bool = provided
.as_bytes()
.ct_eq(self.expected_token.as_bytes())
.into();
if matches {
tracing::debug!(peer = %peer, "gRPC auth accepted");
Ok(request)
} else {
Err(Status::unauthenticated("invalid bearer token"))
Err(reject("invalid bearer token"))
}
}
}
@@ -2861,6 +2877,12 @@ pub async fn serve_on_listener_with_overrides_and_shutdown(
);
Server::builder()
// Wrap every RPC in a correlation span (request_id + method + peer)
// before the resource-limit tuning below, so the span brackets the
// whole request -- routing, the per-service `BearerAuthInterceptor`,
// and the handler -- and its request_id is inherited by every event
// emitted while the request is handled.
.layer(crate::rpc_trace::RpcTraceLayer)
.concurrency_limit_per_connection(GRPC_CONCURRENCY_LIMIT_PER_CONNECTION)
.timeout(GRPC_REQUEST_TIMEOUT)
.max_concurrent_streams(GRPC_MAX_CONCURRENT_STREAMS)
@@ -2888,6 +2910,104 @@ mod tests {
use nuncio_proto::v1::system_client::SystemClient;
use tonic::Code;
use crate::test_tracing::with_recorder;
use tonic::service::Interceptor;
use tracing::Level;
/// The bearer token used by the auth-logging tests. A test asserts this
/// exact value never appears in any captured log field.
const TEST_BEARER_TOKEN: &str = "s3cr3t-bearer-token-value";
fn bearer_request(header: Option<&str>) -> Request<()> {
let mut request = Request::new(());
if let Some(value) = header {
request.metadata_mut().insert(
"authorization",
value.parse().expect("valid ASCII metadata value"),
);
}
request
}
#[test]
fn auth_accept_logs_debug_without_leaking_the_token() {
let (recorder, result) = with_recorder(|| {
let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN);
interceptor.call(bearer_request(Some(&format!("Bearer {TEST_BEARER_TOKEN}"))))
});
assert!(result.is_ok(), "a valid bearer token is accepted");
let events = recorder.events();
let accept = events
.iter()
.find(|e| e.message() == "gRPC auth accepted")
.expect("an accept must be logged");
assert_eq!(accept.level, Level::DEBUG, "accept is logged at DEBUG");
assert!(
recorder
.all_field_values()
.iter()
.all(|v| !v.contains(TEST_BEARER_TOKEN)),
"the bearer token must never appear in any log field"
);
}
#[test]
fn auth_reject_bad_token_logs_warn_without_leaking_the_token() {
let presented = "attacker-guess-token";
let (recorder, result) = with_recorder(|| {
let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN);
interceptor.call(bearer_request(Some(&format!("Bearer {presented}"))))
});
let status = result.expect_err("a mismatched token is rejected");
assert_eq!(status.code(), Code::Unauthenticated);
let events = recorder.events();
let reject = events
.iter()
.find(|e| e.message() == "gRPC auth rejected")
.expect("a reject must be logged");
assert_eq!(reject.level, Level::WARN, "reject is logged at WARN");
assert_eq!(
reject.fields.get("reason").map(String::as_str),
Some("invalid bearer token"),
"the reject must carry a fixed reason string"
);
// Neither the real token nor the presented guess may appear anywhere.
assert!(
recorder
.all_field_values()
.iter()
.all(|v| !v.contains(TEST_BEARER_TOKEN) && !v.contains(presented)),
"no token material may appear in any log field"
);
}
#[test]
fn auth_reject_missing_header_logs_warn_reason() {
let (recorder, result) = with_recorder(|| {
let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN);
interceptor.call(bearer_request(None))
});
assert!(
result.is_err(),
"a missing authorization header is rejected"
);
let events = recorder.events();
let reject = events
.iter()
.find(|e| e.message() == "gRPC auth rejected")
.expect("a reject must be logged");
assert_eq!(reject.level, Level::WARN);
assert_eq!(
reject.fields.get("reason").map(String::as_str),
Some("missing authorization metadata")
);
}
/// Proves a stored second-granularity instant survives the proto
/// boundary exactly: `map_email_to_proto`'s `received_at` decodes back
/// to the identical unix-seconds value via `timestamp_to_unix_secs`.
+3
View File
@@ -11,9 +11,12 @@ pub mod logging;
pub mod orchestrator;
pub mod outbox;
pub mod pagination;
pub mod rpc_trace;
pub mod send;
pub mod sync;
pub mod sync_dispatcher;
#[cfg(test)]
pub(crate) mod test_tracing;
pub use orchestrator::SelfHealingSyncOrchestrator;
/// Default PERSISTENT database path for `nunciod`: `~/.nuncio/nuncio.db`
+401
View File
@@ -0,0 +1,401 @@
//! Per-RPC tracing span and request-id correlation for the gRPC transport.
//!
//! This is a [`tower::Layer`] mounted on the tonic `Server` so that every
//! inbound RPC is wrapped in a single root span carrying:
//!
//! * a generated `request_id` -- a lightweight per-process identifier
//! (process-start tag + monotonic counter) that lets an operator follow one
//! request across every log line it produced, without pulling in a heavier
//! UUID dependency; and
//! * the gRPC method path and the peer socket address.
//!
//! The inner service (routing + the per-service `BearerAuthInterceptor` + the
//! handler future) is entered inside that span, so *every* event emitted while
//! the request is handled -- including the auth accept/reject logs and any
//! event a handler emits -- inherits the `request_id` automatically. The span
//! opens with a DEBUG "received" event and closes, once the response body has
//! been fully drained, with an INFO event carrying the terminal gRPC status
//! code and the elapsed wall-clock time.
//!
//! ## Why the response body is wrapped
//!
//! A gRPC call's final status lives in the HTTP/2 *trailers* (`grpc-status`),
//! which are only produced once the response body finishes streaming (an
//! immediate error is instead a trailers-only response with `grpc-status` in
//! the leading headers). The layer therefore wraps the response body in
//! [`TracedBody`], which observes the terminal status as it flows past and
//! emits the completion event when the stream ends (or when the body is
//! dropped, e.g. on client disconnect), so a long-lived server-streaming RPC
//! is not logged as "complete" until it actually ends.
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;
use std::task::{Context, Poll};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use http::{HeaderMap, Request, Response};
use http_body::{Body, Frame, SizeHint};
use tonic::body::Body as TonicBody;
use tonic::transport::server::TcpConnectInfo;
use tonic::Code;
use tower::{Layer, Service};
use tracing::{Instrument, Span};
/// HTTP/2 trailer (or leading header, for a trailers-only error response)
/// carrying the final gRPC status code.
const GRPC_STATUS_HEADER: &str = "grpc-status";
/// Generates a correlation id unique within this daemon process.
///
/// The id is `<process-tag>-<counter>`: the process tag is the daemon's start
/// time in unix nanoseconds (captured once), and the counter increments per
/// call. This is deliberately not a UUID -- it needs only to be unique within
/// a single running daemon so its logs can be correlated, which a per-process
/// tag plus a monotonic counter satisfies without a new dependency.
fn next_request_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(0);
static PROCESS_TAG: OnceLock<u64> = OnceLock::new();
let tag = *PROCESS_TAG.get_or_init(|| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos() as u64)
.unwrap_or(0)
});
let n = COUNTER.fetch_add(1, Ordering::Relaxed);
format!("{tag:016x}-{n:08x}")
}
/// Reads the `grpc-status` code from a header/trailer map, if present.
fn grpc_status_code(headers: &HeaderMap) -> Option<Code> {
headers
.get(GRPC_STATUS_HEADER)
.map(|value| Code::from_bytes(value.as_bytes()))
}
/// [`tower::Layer`] that wraps each RPC in a correlation span. Mounted on the
/// tonic `Server` via `Server::builder().layer(RpcTraceLayer)`.
#[derive(Clone, Copy, Debug, Default)]
pub struct RpcTraceLayer;
impl<S> Layer<S> for RpcTraceLayer {
type Service = RpcTrace<S>;
fn layer(&self, inner: S) -> Self::Service {
RpcTrace { inner }
}
}
/// The [`RpcTraceLayer`]'s service: opens the span, runs the inner service
/// (routing, auth interceptor, handler) inside it, and wraps the response body
/// so the terminal status can be observed at stream end.
#[derive(Clone, Debug)]
pub struct RpcTrace<S> {
inner: S,
}
impl<S> Service<Request<TonicBody>> for RpcTrace<S>
where
S: Service<Request<TonicBody>, Response = Response<TonicBody>> + Clone + Send + 'static,
S::Future: Send + 'static,
S::Error: Send + 'static,
{
type Response = Response<TracedBody<TonicBody>>;
type Error = S::Error;
type Future =
Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<TonicBody>) -> Self::Future {
// `poll_ready` was called on `self.inner`, so it -- not the fresh
// clone -- is the readied service that must handle this request. Swap
// the readied service out and leave the clone behind for the next
// call (the standard tower clone-and-replace pattern that avoids
// calling an un-readied clone).
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
let method = req.uri().path().to_string();
let peer = req
.extensions()
.get::<TcpConnectInfo>()
.and_then(TcpConnectInfo::remote_addr)
.map(|addr| addr.to_string())
.unwrap_or_else(|| "unknown".to_string());
let request_id = next_request_id();
let span = tracing::info_span!(
"grpc.request",
request_id = %request_id,
rpc.method = %method,
peer = %peer,
);
let start = Instant::now();
Box::pin(async move {
// Enter the span for the synchronous routing + interceptor call so
// the auth accept/reject logs inherit the request_id, and emit the
// open event under it.
let call = span.in_scope(|| {
tracing::debug!("gRPC request received");
inner.call(req)
});
// Instrument the handler future so events emitted across its await
// points also inherit the span.
let response = call.instrument(span.clone()).await?;
// A trailers-only error response carries `grpc-status` in the
// leading headers; a normal response carries it in the trailers,
// observed by `TracedBody` as the body drains.
let initial_status = grpc_status_code(response.headers());
let (parts, body) = response.into_parts();
let traced = TracedBody::new(body, span, start, initial_status);
Ok(Response::from_parts(parts, traced))
})
}
}
/// Response body wrapper that observes the terminal `grpc-status` and logs the
/// request's completion (status + elapsed) exactly once, when the body stream
/// ends or is dropped.
///
/// All fields are `Unpin`, so the pin projection in [`Body::poll_frame`] is a
/// safe `Pin::new` over the inner body rather than an `unsafe` projection.
#[derive(Debug)]
pub struct TracedBody<B> {
inner: B,
span: Span,
start: Instant,
status: Option<Code>,
logged: bool,
}
impl<B> TracedBody<B> {
fn new(inner: B, span: Span, start: Instant, initial_status: Option<Code>) -> Self {
Self {
inner,
span,
start,
status: initial_status,
logged: false,
}
}
/// Emits the terminal completion event once. Idempotent: subsequent calls
/// (e.g. a `Drop` after the stream already ended) are no-ops.
fn log_completion(&mut self) {
if self.logged {
return;
}
self.logged = true;
let elapsed_ms = self.start.elapsed().as_millis() as u64;
let status = self.status.unwrap_or(Code::Unknown);
self.span.in_scope(|| {
tracing::info!(
grpc.status = ?status,
elapsed_ms,
"gRPC request completed"
);
});
}
}
impl<B> Body for TracedBody<B>
where
B: Body + Unpin,
{
type Data = B::Data;
type Error = B::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
// Safe because `Self: Unpin` (all fields are `Unpin`).
let this = self.get_mut();
let polled = Pin::new(&mut this.inner).poll_frame(cx);
match &polled {
Poll::Ready(Some(Ok(frame))) => {
if let Some(trailers) = frame.trailers_ref() {
if let Some(code) = grpc_status_code(trailers) {
this.status = Some(code);
}
}
}
Poll::Ready(Some(Err(_))) | Poll::Ready(None) => this.log_completion(),
Poll::Pending => {}
}
polled
}
fn is_end_stream(&self) -> bool {
self.inner.is_end_stream()
}
fn size_hint(&self) -> SizeHint {
self.inner.size_hint()
}
}
impl<B> Drop for TracedBody<B> {
fn drop(&mut self) {
// Covers stream ends that do not surface a final frame here (client
// disconnect, an outer layer dropping the body): the request is still
// logged as completed exactly once.
self.log_completion();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_tracing::with_recorder;
use http_body_util::BodyExt;
use std::collections::HashSet;
use std::convert::Infallible;
use tower::ServiceExt;
use tracing::Level;
/// Minimal inner service that returns a response with `grpc-status` set in
/// the leading headers (a trailers-only response, tonic's shape for an
/// immediate status), so the layer's terminal-status extraction can be
/// exercised without a live server.
#[derive(Clone)]
struct StatusService {
code: i32,
}
impl Service<Request<TonicBody>> for StatusService {
type Response = Response<TonicBody>;
type Error = Infallible;
type Future =
Pin<Box<dyn Future<Output = Result<Self::Response, Infallible>> + Send + 'static>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Infallible>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, _req: Request<TonicBody>) -> Self::Future {
let code = self.code;
Box::pin(async move {
// A handler-emitted event: it must inherit the span so it
// carries the request_id set by the layer.
tracing::info!("handler ran");
let mut response = Response::new(TonicBody::empty());
response
.headers_mut()
.insert(GRPC_STATUS_HEADER, code.to_string().parse().unwrap());
Ok(response)
})
}
}
fn request() -> Request<TonicBody> {
Request::builder()
.uri("/nuncio.v1.Mail/Sync")
.body(TonicBody::empty())
.unwrap()
}
#[test]
fn rpc_produces_span_with_request_id_and_logs_terminal_status() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let (recorder, ()) = with_recorder(|| {
runtime.block_on(async {
let mut svc = RpcTraceLayer.layer(StatusService { code: 5 });
let response = svc.ready().await.unwrap().call(request()).await.unwrap();
// Drain the body so the completion event fires at stream end.
response.into_body().collect().await.unwrap();
});
});
let spans = recorder.spans();
let rpc_span = spans
.iter()
.find(|s| s.name == "grpc.request")
.expect("a grpc.request span must be created");
let request_id = rpc_span
.fields
.get("request_id")
.expect("the span must carry a request_id field");
assert!(!request_id.is_empty(), "request_id must be populated");
assert_eq!(
rpc_span.fields.get("rpc.method").map(String::as_str),
Some("/nuncio.v1.Mail/Sync"),
"the span must carry the gRPC method path"
);
let events = recorder.events();
let completion = events
.iter()
.find(|e| e.message() == "gRPC request completed")
.expect("the terminal completion event must be logged");
assert_eq!(
completion.level,
Level::INFO,
"the completion event is logged at INFO"
);
assert_eq!(
completion.fields.get("grpc.status").map(String::as_str),
Some("NotFound"),
"the completion event must carry the terminal gRPC status code"
);
assert!(
completion.fields.contains_key("elapsed_ms"),
"the completion event must carry the elapsed time"
);
// The handler ran inside the span, so its own event was captured too
// (proving handler events inherit the correlation context).
assert!(
events.iter().any(|e| e.message() == "handler ran"),
"a handler-emitted event must be observed under the request span"
);
}
#[test]
fn ok_status_is_recorded_at_completion() {
let runtime = tokio::runtime::Builder::new_current_thread()
.build()
.unwrap();
let (recorder, ()) = with_recorder(|| {
runtime.block_on(async {
let mut svc = RpcTraceLayer.layer(StatusService { code: 0 });
let response = svc.ready().await.unwrap().call(request()).await.unwrap();
response.into_body().collect().await.unwrap();
});
});
let events = recorder.events();
let completion = events
.iter()
.find(|e| e.message() == "gRPC request completed")
.expect("completion event");
assert_eq!(
completion.fields.get("grpc.status").map(String::as_str),
Some("Ok")
);
}
#[test]
fn grpc_status_code_parses_known_codes_and_ignores_absent() {
let mut headers = HeaderMap::new();
assert!(grpc_status_code(&headers).is_none());
headers.insert(GRPC_STATUS_HEADER, "3".parse().unwrap());
assert_eq!(grpc_status_code(&headers), Some(Code::InvalidArgument));
}
#[test]
fn request_ids_are_unique_within_the_process() {
let ids: HashSet<String> = (0..1000).map(|_| next_request_id()).collect();
assert_eq!(ids.len(), 1000, "request ids must not collide");
}
}
+36 -16
View File
@@ -48,6 +48,7 @@ use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use thiserror::Error;
use tokio::sync::{watch, Semaphore};
use tracing::Instrument;
/// Environment variable bounding how many account syncs may run at once.
pub const CONCURRENCY_ENV_VAR: &str = "NUNCIO_SYNC_CONCURRENCY";
@@ -285,18 +286,30 @@ impl SyncDispatcher {
let inner = self.inner.clone();
let id = account_id.to_string();
tokio::spawn(async move {
let outcome = inner.run_guarded(&id).await;
// Remove the in-flight entry before publishing so a request
// arriving after completion starts a fresh sync rather than
// joining an already-finished one.
inner
.inflight
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&id);
let _ = tx.send(Some(outcome));
});
// The actual sync runs in a detached task, which does NOT
// inherit the caller's span automatically. Capture the current
// span (the RPC's correlation span when a `Mail/Sync` triggered
// this) and instrument the task with it, so the sync's logs are
// correlatable to the request that kicked it off. A request
// that COALESCES onto an already-running sync joins the owner's
// task and is therefore correlated to the owner's request_id,
// not its own.
let span = tracing::Span::current();
tokio::spawn(
async move {
let outcome = inner.run_guarded(&id).await;
// Remove the in-flight entry before publishing so a
// request arriving after completion starts a fresh sync
// rather than joining an already-finished one.
inner
.inflight
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&id);
let _ = tx.send(Some(outcome));
}
.instrument(span),
);
rx
}
};
@@ -319,12 +332,19 @@ impl SyncDispatcher {
/// production syncer additionally reports them on the event bus).
pub async fn sync_all(&self, account_ids: Vec<String>) -> usize {
let mut set = tokio::task::JoinSet::new();
// Propagate the caller's span (the RPC correlation span for an
// all-accounts `Mail/Sync`) into each per-account task so the fan-out
// stays correlated to the triggering request.
let span = tracing::Span::current();
for account_id in account_ids {
let dispatcher = self.clone();
set.spawn(async move {
let result = dispatcher.sync_account(&account_id).await;
(account_id, result)
});
set.spawn(
async move {
let result = dispatcher.sync_account(&account_id).await;
(account_id, result)
}
.instrument(span.clone()),
);
}
let mut total = 0usize;
+133
View File
@@ -0,0 +1,133 @@
//! Shared in-test `tracing` capture helper.
//!
//! Installs a thread-local subscriber that records every span (with its
//! fields) and every event (with its level, fields, and message) into an
//! in-memory buffer, so tests can assert on what the code logged -- including
//! the negative assertion that a secret never appears in any recorded field.
//! Deterministic and fully offline: no global subscriber, no I/O.
#![cfg(test)]
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tracing::field::{Field, Visit};
use tracing::span::Attributes;
use tracing::{Event, Id, Level, Subscriber};
use tracing_subscriber::layer::Context as LayerContext;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
/// A captured span: its name and the string-rendered fields it opened with.
#[derive(Debug, Clone)]
pub(crate) struct RecordedSpan {
pub name: String,
pub fields: HashMap<String, String>,
}
/// A captured event: its level and its string-rendered fields (including the
/// `message`).
#[derive(Debug, Clone)]
pub(crate) struct RecordedEvent {
pub level: Level,
pub fields: HashMap<String, String>,
}
impl RecordedEvent {
/// The event's `message` field, or the empty string if it had none.
pub fn message(&self) -> &str {
self.fields.get("message").map(String::as_str).unwrap_or("")
}
}
/// In-memory sink of everything the subscriber captured during a test.
#[derive(Debug, Default)]
pub(crate) struct Recorder {
spans: Mutex<Vec<RecordedSpan>>,
events: Mutex<Vec<RecordedEvent>>,
}
impl Recorder {
/// A snapshot of every captured span.
pub fn spans(&self) -> Vec<RecordedSpan> {
self.spans.lock().unwrap_or_else(|e| e.into_inner()).clone()
}
/// A snapshot of every captured event.
pub fn events(&self) -> Vec<RecordedEvent> {
self.events
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
}
/// Every field value recorded across all spans and events, so a test can
/// assert that a secret never appears anywhere in the telemetry.
pub fn all_field_values(&self) -> Vec<String> {
let mut values = Vec::new();
for span in self.spans().iter() {
values.extend(span.fields.values().cloned());
}
for event in self.events().iter() {
values.extend(event.fields.values().cloned());
}
values
}
}
struct FieldVisitor(HashMap<String, String>);
impl Visit for FieldVisitor {
fn record_str(&mut self, field: &Field, value: &str) {
self.0.insert(field.name().to_string(), value.to_string());
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
self.0
.entry(field.name().to_string())
.or_insert_with(|| format!("{value:?}"));
}
}
struct CaptureLayer(Arc<Recorder>);
impl<S> tracing_subscriber::Layer<S> for CaptureLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_new_span(&self, attrs: &Attributes<'_>, _id: &Id, _ctx: LayerContext<'_, S>) {
let mut visitor = FieldVisitor(HashMap::new());
attrs.record(&mut visitor);
self.0
.spans
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(RecordedSpan {
name: attrs.metadata().name().to_string(),
fields: visitor.0,
});
}
fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) {
let mut visitor = FieldVisitor(HashMap::new());
event.record(&mut visitor);
self.0
.events
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(RecordedEvent {
level: *event.metadata().level(),
fields: visitor.0,
});
}
}
/// Runs `f` with a thread-local capturing subscriber installed and returns the
/// recorder alongside `f`'s result. Because the subscriber is thread-local,
/// any async work driven inside `f` must run on the calling thread (e.g. a
/// current-thread runtime's `block_on`) for its telemetry to be captured.
pub(crate) fn with_recorder<T>(f: impl FnOnce() -> T) -> (Arc<Recorder>, T) {
let recorder = Arc::new(Recorder::default());
let subscriber = tracing_subscriber::registry().with(CaptureLayer(recorder.clone()));
let value = tracing::subscriber::with_default(subscriber, f);
(recorder, value)
}