1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
//! Tenant identity for log lines: the field that says whose request
//! this was.
//!
//! `Tenant` is an axum extractor, so the resolved
//! [`Org`](crate::tenancy::Org) is moved into the handler and the
//! logging middleware outside it cannot see the tenant. This module is
//! the missing slot. [`scope`] opens one per request, the resolver
//! calls [`record`], and outer middleware reads it with [`current`].
//!
//! ```ignore
//! // Middleware, outside the handler. Read back INSIDE the scope:
//! // once the `scope` future resolves the slot is gone, and a read
//! // after it returns `None` for every request.
//! let (response, who) = tenant_log::scope(async {
//! let response = next.run(req).await;
//! (response, tenant_log::current()) // Some(acme)
//! })
//! .await;
//! ```
//!
//! [`record`] also puts `tenant` and `org_id` on the current `tracing`
//! span, so with [`crate::tracing_layer::TracingLayer`] installed every
//! event in the request carries the tenant, the ORM's events included.
//!
//! ## Scope and safety
//!
//! The slot is per-task and set once. `tokio::spawn` does not inherit
//! it, so work that leaves the request leaves the scope. Keep it that
//! way: a slot shared or reused between requests would label one
//! tenant's log lines with another tenant's slug, which leaks customer
//! names across tenants and corrupts any per-tenant audit trail. Open a
//! fresh [`scope`] per request and never hold a [`TenantLabel`] past
//! it. Background jobs need their own mechanism.
//!
//! [`scope`]: crate::tenant_log::scope
//! [`record`]: crate::tenant_log::record
//! [`current`]: crate::tenant_log::current
//! [`TenantLabel`]: crate::tenant_log::TenantLabel
use std::future::Future;
use std::sync::{Arc, OnceLock};
/// The resolved tenant, as logs should name it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TenantLabel {
/// `Org.slug`. Often a customer name, so treat logs holding it as
/// customer data. See [`crate::access_log::TenantField`] to log the
/// id instead.
pub slug: String,
/// `Org.id`. `None` only for an unsaved org, because `Auto<i64>` is
/// optional; no request path produces that.
pub id: Option<i64>,
}
tokio::task_local! {
/// Per-request slot for the resolved tenant.
///
/// It holds `Arc<OnceLock<_>>`, not a plain value, because the write
/// happens below the read: the resolver runs inside the handler and
/// the logging middleware sits outside it. Set-once matches the
/// rule of one request, one tenant.
static TENANT: Arc<OnceLock<TenantLabel>>;
}
/// Run `fut` with a fresh tenant slot in scope.
///
/// Nesting is safe: an inner call reuses the outer slot rather than
/// shadowing it, so two layers that both open a scope still agree on
/// one tenant per request.
pub async fn scope<F, T>(fut: F) -> T
where
F: Future<Output = T>,
{
if TENANT.try_with(|_| ()).is_ok() {
return fut.await;
}
TENANT.scope(Arc::new(OnceLock::new()), fut).await
}
/// Publish the request's tenant. The resolver calls this as soon as it
/// identifies an `Org`.
///
/// Fills the [`scope`] slot if one is open, and puts `tenant` and
/// `org_id` on the current span. Both halves do nothing outside their
/// context, so a CLI verb or test can call it safely. First call wins,
/// so a later resolver cannot relabel the request.
pub fn record(slug: &str, id: Option<i64>) {
// Check the slot first: it is what stops a chain of resolvers from
// stamping the same field once each, which is a duplicate key in
// JSON logs.
//
// `Ok(false)` means a resolver already filled it, so stop. `Err`
// means no scope is open, so there is nothing to dedupe against and
// we record anyway.
let first = TENANT
.try_with(|slot| {
slot.set(TenantLabel {
slug: slug.to_owned(),
id,
})
.is_ok()
})
.unwrap_or(true);
if !first {
return;
}
let span = tracing::Span::current();
span.record("tenant", slug);
if let Some(id) = id {
span.record("org_id", id);
}
}
/// The tenant resolved for the current request, if any.
///
/// `None` means no tenant was identified: an apex-domain or operator
/// console request, a single-tenant deployment, or code outside a
/// [`scope`]. It never means one was found and then lost.
#[must_use]
pub fn current() -> Option<TenantLabel> {
TENANT.try_with(|slot| slot.get().cloned()).ok().flatten()
}
#[cfg(test)]
mod tests {
use super::*;
/// One `tenant=` per line, however many resolvers run. More than
/// one is a duplicate key in JSON logs.
#[cfg(feature = "runtime")]
#[tokio::test]
async fn a_field_is_stamped_once_however_many_resolvers_run() {
let buf = crate::testkit::CaptureWriter::default();
let subscriber = tracing_subscriber::fmt()
.with_writer(buf.clone())
.with_ansi(false)
.finish();
let _g = tracing::subscriber::set_default(subscriber);
scope(async {
let span = tracing::info_span!(
"http.request",
tenant = tracing::field::Empty,
org_id = tracing::field::Empty,
);
let _e = span.enter();
// Three resolvers publish the same tenant.
record("acme", Some(7));
record("acme", Some(7));
record("acme", Some(7));
tracing::info!("handled");
})
.await;
let out = buf.contents();
assert!(out.contains("handled"), "nothing rendered:\n{out}");
assert_eq!(
out.matches("tenant=").count(),
1,
"tenant was stamped onto the span once per resolver instead of once \
per request:\n{out}"
);
assert_eq!(
out.matches("org_id=").count(),
1,
"org_id was stamped more than once:\n{out}"
);
}
#[tokio::test]
async fn records_and_reads_back_inside_a_scope() {
let seen = scope(async {
record("acme", Some(7));
current()
})
.await;
assert_eq!(
seen,
Some(TenantLabel {
slug: "acme".to_owned(),
id: Some(7),
})
);
}
#[tokio::test]
async fn current_is_none_without_a_resolution() {
let seen = scope(async { current() }).await;
assert_eq!(seen, None, "an apex request resolves no tenant");
}
#[tokio::test]
async fn record_outside_a_scope_is_a_noop() {
record("acme", Some(7));
assert_eq!(current(), None);
}
#[tokio::test]
async fn first_resolution_wins() {
let seen = scope(async {
record("acme", Some(1));
record("globex", Some(2));
current()
})
.await;
assert_eq!(seen.unwrap().slug, "acme");
}
/// The slot does not outlive its scope. A read placed after it
/// fails silently, so this needs its own test.
#[tokio::test]
async fn a_read_after_the_scope_closes_sees_nothing() {
scope(async { record("acme", Some(7)) }).await;
assert_eq!(current(), None);
}
/// Two layers opening a scope must not shadow each other, or the
/// resolver fills the inner slot and the outer read comes back
/// empty.
#[tokio::test]
async fn a_nested_scope_reuses_the_outer_slot() {
let seen = scope(async {
scope(async { record("acme", Some(7)) }).await;
current()
})
.await;
assert_eq!(seen.map(|t| t.slug), Some("acme".to_owned()));
}
/// A spawned task is a different task: it neither sees the parent's
/// tenant nor leaks one back.
#[tokio::test]
async fn a_spawned_task_does_not_inherit_the_scope() {
let seen = scope(async {
record("acme", Some(7));
tokio::spawn(async { current() }).await.unwrap()
})
.await;
assert_eq!(seen, None);
}
}