Skip to content

Commit 647d55d

Browse files
authored
Track unique shortids seen (#1459)
2 parents c528143 + 03f4dba commit 647d55d

6 files changed

Lines changed: 273 additions & 2 deletions

File tree

Cargo-minimal.lock

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2004,6 +2004,15 @@ dependencies = [
20042004
"tracing",
20052005
]
20062006

2007+
[[package]]
2008+
name = "hyperloglogplus"
2009+
version = "0.4.1"
2010+
source = "registry+https://github.com/rust-lang/crates.io-index"
2011+
checksum = "621debdf94dcac33e50475fdd76d34d5ea9c0362a834b9db08c3024696c1fbe3"
2012+
dependencies = [
2013+
"serde",
2014+
]
2015+
20072016
[[package]]
20082017
name = "iana-time-zone"
20092018
version = "0.1.65"
@@ -2863,6 +2872,7 @@ dependencies = [
28632872
"hyper",
28642873
"hyper-rustls",
28652874
"hyper-util",
2875+
"hyperloglogplus",
28662876
"ipnet",
28672877
"maxminddb",
28682878
"mockito",

Cargo-recent.lock

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1976,6 +1976,15 @@ dependencies = [
19761976
"tracing",
19771977
]
19781978

1979+
[[package]]
1980+
name = "hyperloglogplus"
1981+
version = "0.4.1"
1982+
source = "registry+https://github.com/rust-lang/crates.io-index"
1983+
checksum = "621debdf94dcac33e50475fdd76d34d5ea9c0362a834b9db08c3024696c1fbe3"
1984+
dependencies = [
1985+
"serde",
1986+
]
1987+
19791988
[[package]]
19801989
name = "iana-time-zone"
19811990
version = "0.1.64"
@@ -2831,6 +2840,7 @@ dependencies = [
28312840
"hyper",
28322841
"hyper-rustls",
28332842
"hyper-util",
2843+
"hyperloglogplus",
28342844
"ipnet",
28352845
"maxminddb",
28362846
"mockito",

payjoin-mailroom/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ hyper-rustls = { version = "0.27.7", default-features = false, features = [
4747
"ring",
4848
] }
4949
hyper-util = { version = "0.1.16", features = ["client-legacy", "service"] }
50+
hyperloglogplus = "0.4.1"
5051
ipnet = { version = "2", optional = true }
5152
maxminddb = { version = "0.27", optional = true }
5253
ohttp = { package = "bitcoin-ohttp", version = "0.6.0" }

payjoin-mailroom/src/db/mod.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,7 @@ impl<D: Db> Db for MetricsDb<D> {
256256
mailbox_id: &ShortId,
257257
data: Vec<u8>,
258258
) -> Result<Option<()>, Error<Self::OperationalError>> {
259+
self.metrics.record_short_id(mailbox_id);
259260
let result = self.inner.post_v2_payload(mailbox_id, data).await?;
260261
if result.is_some() {
261262
self.metrics.record_db_entry(PayjoinVersion::Two);
@@ -267,6 +268,7 @@ impl<D: Db> Db for MetricsDb<D> {
267268
&self,
268269
mailbox_id: &ShortId,
269270
) -> Result<Arc<Vec<u8>>, Error<Self::OperationalError>> {
271+
self.metrics.record_short_id(mailbox_id);
270272
self.inner.wait_for_v2_payload(mailbox_id).await
271273
}
272274

@@ -275,6 +277,7 @@ impl<D: Db> Db for MetricsDb<D> {
275277
mailbox_id: &ShortId,
276278
data: Vec<u8>,
277279
) -> Result<(), Error<Self::OperationalError>> {
280+
self.metrics.record_short_id(mailbox_id);
278281
self.inner.post_v1_response(mailbox_id, data).await
279282
}
280283

@@ -283,6 +286,7 @@ impl<D: Db> Db for MetricsDb<D> {
283286
mailbox_id: &ShortId,
284287
data: Vec<u8>,
285288
) -> Result<Arc<Vec<u8>>, Error<Self::OperationalError>> {
289+
self.metrics.record_short_id(mailbox_id);
286290
let result = self.inner.post_v1_request_and_wait_for_response(mailbox_id, data).await?;
287291
self.metrics.record_db_entry(PayjoinVersion::One);
288292
Ok(result)

payjoin-mailroom/src/directory.rs

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -877,4 +877,64 @@ mod tests {
877877
other => panic!("expected U64 Sum, got {other:?}"),
878878
}
879879
}
880+
881+
#[tokio::test]
882+
async fn post_mailbox_records_short_id_cardinality() {
883+
use opentelemetry_sdk::metrics::{
884+
InMemoryMetricExporter, PeriodicReader, SdkMeterProvider,
885+
};
886+
887+
use crate::db::MetricsDb;
888+
use crate::metrics::{MetricsService, UNIQUE_SHORT_IDS};
889+
890+
let exporter = InMemoryMetricExporter::default();
891+
let reader = PeriodicReader::builder(exporter.clone()).build();
892+
let provider = SdkMeterProvider::builder().with_reader(reader).build();
893+
let metrics = MetricsService::new(Some(provider.clone()));
894+
895+
let dir = tempfile::tempdir().expect("tempdir");
896+
let db = FilesDb::init(
897+
Duration::from_millis(100),
898+
dir.keep(),
899+
Duration::from_secs(60 * 60 * 24 * 7),
900+
)
901+
.await
902+
.expect("db init");
903+
let db = MetricsDb::new(db, metrics);
904+
let ohttp: ohttp::Server =
905+
crate::key_config::gen_ohttp_server_config().expect("ohttp config").into();
906+
let svc = Service::new(db, ohttp, SentinelTag::new([0u8; 32]), None);
907+
908+
let id = valid_short_id_path();
909+
let res = svc
910+
.post_mailbox(&id, Body::from(b"payload".to_vec()))
911+
.await
912+
.expect("post_mailbox should succeed");
913+
assert_eq!(res.status(), StatusCode::OK);
914+
915+
provider.force_flush().expect("flush failed");
916+
let finished = exporter.get_finished_metrics().expect("metrics");
917+
let gauge = finished
918+
.iter()
919+
.flat_map(|rm| rm.scope_metrics())
920+
.flat_map(|sm| sm.metrics())
921+
.find(|m| m.name() == UNIQUE_SHORT_IDS)
922+
.expect("missing unique_short_ids metric");
923+
924+
use opentelemetry::KeyValue;
925+
use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData};
926+
927+
match gauge.data() {
928+
AggregatedMetrics::U64(MetricData::Gauge(g)) => {
929+
let points: Vec<_> = g.data_points().collect();
930+
assert!(!points.is_empty(), "expected at least one data point");
931+
let hourly = points.iter().find(|dp| {
932+
dp.attributes().any(|kv| kv == &KeyValue::new("interval", "hourly"))
933+
});
934+
assert!(hourly.is_some(), "expected hourly data point");
935+
assert!(hourly.unwrap().value() >= 1, "expected at least 1 unique short ID");
936+
}
937+
other => panic!("expected U64 Gauge, got {other:?}"),
938+
}
939+
}
880940
}

payjoin-mailroom/src/metrics.rs

Lines changed: 188 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,155 @@
1+
use std::collections::hash_map::RandomState;
2+
use std::collections::BTreeMap;
13
use std::fmt;
4+
use std::sync::{Arc, Mutex};
5+
use std::time::{Duration, SystemTime, UNIX_EPOCH};
26

3-
use opentelemetry::metrics::{Counter, MeterProvider, UpDownCounter};
7+
use hyperloglogplus::{HyperLogLog, HyperLogLogPlus};
8+
use opentelemetry::metrics::{Counter, MeterProvider, ObservableGauge, UpDownCounter};
49
use opentelemetry::KeyValue;
510
use opentelemetry_sdk::metrics::SdkMeterProvider;
11+
use payjoin::directory::ShortId;
612

713
pub(crate) const TOTAL_CONNECTIONS: &str = "total_connections";
814
pub(crate) const ACTIVE_CONNECTIONS: &str = "active_connections";
915
pub(crate) const HTTP_REQUESTS: &str = "http_request_total";
1016
pub(crate) const DB_ENTRIES: &str = "db_entries_total";
17+
pub(crate) const UNIQUE_SHORT_IDS: &str = "unique_short_ids";
18+
19+
const HLL_PRECISION: u8 = 14;
20+
const HOURLY_RETENTION_HOURS: u64 = 168; // 7 days
21+
const DAILY_RETENTION_DAYS: u64 = 90;
22+
23+
type HllSketch = HyperLogLogPlus<[u8; 8], RandomState>;
24+
25+
const HOUR: Duration = Duration::from_secs(3600);
26+
const DAY: Duration = Duration::from_secs(86400);
27+
28+
/// Convenience helpers for SystemTime to get integer intervals
29+
/// since the UNIX epoch.
30+
trait SystemTimeExt {
31+
fn intervals_since_epoch(&self, interval: Duration) -> u64;
32+
33+
fn hours_since_epoch(&self) -> u64 { self.intervals_since_epoch(HOUR) }
34+
35+
fn days_since_epoch(&self) -> u64 { self.intervals_since_epoch(DAY) }
36+
}
37+
38+
impl SystemTimeExt for SystemTime {
39+
fn intervals_since_epoch(&self, interval: Duration) -> u64 {
40+
self.duration_since(UNIX_EPOCH).expect("system clock before UNIX epoch").as_secs()
41+
/ interval.as_secs()
42+
}
43+
}
44+
45+
fn new_sketch() -> HllSketch {
46+
HyperLogLogPlus::new(HLL_PRECISION, RandomState::new()).expect("precision 14 is always valid")
47+
}
48+
49+
/// Estimates the number of unique ShortIds seen per time window.
50+
/// Two tiers of HLL sketches:
51+
/// - **Hourly** -- one sketch per hour, pruned after 7 days.
52+
/// - **Daily** -- one sketch per day, pruned after 90 days.
53+
struct HllSketches {
54+
hourly: BTreeMap<u64, HllSketch>,
55+
daily: BTreeMap<u64, HllSketch>,
56+
}
57+
58+
impl HllSketches {
59+
fn new() -> Self { Self { hourly: BTreeMap::new(), daily: BTreeMap::new() } }
60+
61+
// NOTE: ShortIds are passed directly to the HyperLogLog sketches held in mailroom's
62+
// RAM; only the resulting cardinality estimates are published to Prometheus. If
63+
// Prometheus ever gains native cardinality-estimation support and raw IDs are
64+
// exposed to it, they should be hashed (e.g. with a keyed PRF) before being
65+
// forwarded to avoid leaking session identifiers through the metrics pipeline.
66+
fn add_id(&mut self, id: &ShortId) {
67+
let now = SystemTime::now();
68+
let hour = now.hours_since_epoch();
69+
let day = now.days_since_epoch();
70+
71+
self.hourly.entry(hour).or_insert_with(new_sketch).insert(&id.0);
72+
self.daily.entry(day).or_insert_with(new_sketch).insert(&id.0);
73+
74+
let hourly_cutoff = hour.saturating_sub(HOURLY_RETENTION_HOURS);
75+
while let Some((&k, _)) = self.hourly.first_key_value() {
76+
if k < hourly_cutoff {
77+
self.hourly.pop_first();
78+
} else {
79+
break;
80+
}
81+
}
82+
let daily_cutoff = day.saturating_sub(DAILY_RETENTION_DAYS);
83+
while let Some((&k, _)) = self.daily.first_key_value() {
84+
if k < daily_cutoff {
85+
self.daily.pop_first();
86+
} else {
87+
break;
88+
}
89+
}
90+
}
91+
92+
fn hourly_count(&mut self) -> u64 {
93+
let hour = SystemTime::now().hours_since_epoch();
94+
self.hourly.get_mut(&hour).map(|hll| hll.count().trunc() as u64).unwrap_or(0)
95+
}
96+
97+
fn daily_count(&mut self) -> u64 {
98+
let day = SystemTime::now().days_since_epoch();
99+
self.daily.get_mut(&day).map(|hll| hll.count().trunc() as u64).unwrap_or(0)
100+
}
101+
102+
fn weekly_count(&mut self) -> u64 {
103+
let today = SystemTime::now().days_since_epoch();
104+
self.daily_union_count(today.saturating_sub(6)..=today)
105+
}
106+
107+
fn monthly_count(&mut self) -> u64 {
108+
let today = SystemTime::now().days_since_epoch();
109+
self.daily_union_count(today.saturating_sub(30)..=today)
110+
}
111+
112+
fn daily_union_count<R: std::ops::RangeBounds<u64>>(&self, range: R) -> u64 {
113+
let mut union = new_sketch();
114+
for sketch in self.daily.range(range).map(|(_, v)| v) {
115+
union.merge(sketch).expect("same precision");
116+
}
117+
union.count().trunc() as u64
118+
}
119+
}
120+
121+
#[derive(Clone)]
122+
pub struct UniqueShortIdTracker {
123+
inner: Arc<Mutex<HllSketches>>,
124+
}
125+
126+
impl UniqueShortIdTracker {
127+
pub fn new() -> Self { Self { inner: Arc::new(Mutex::new(HllSketches::new())) } }
128+
129+
pub fn add_id(&self, id: &ShortId) {
130+
self.inner.lock().expect("tracker lock poisoned").add_id(id);
131+
}
132+
133+
pub fn hourly_count(&self) -> u64 {
134+
self.inner.lock().expect("tracker lock poisoned").hourly_count()
135+
}
136+
137+
pub fn daily_count(&self) -> u64 {
138+
self.inner.lock().expect("tracker lock poisoned").daily_count()
139+
}
140+
141+
pub fn weekly_count(&self) -> u64 {
142+
self.inner.lock().expect("tracker lock poisoned").weekly_count()
143+
}
144+
145+
pub fn monthly_count(&self) -> u64 {
146+
self.inner.lock().expect("tracker lock poisoned").monthly_count()
147+
}
148+
}
149+
150+
impl Default for UniqueShortIdTracker {
151+
fn default() -> Self { Self::new() }
152+
}
11153

12154
#[derive(Clone)]
13155
pub struct MetricsService {
@@ -19,6 +161,8 @@ pub struct MetricsService {
19161
active_connections: UpDownCounter<i64>,
20162
/// Total v1/v2 mailbox entries written, labelled by `version`
21163
db_entries_total: Counter<u64>,
164+
tracker: UniqueShortIdTracker,
165+
_unique_ids_gauge: Option<Arc<ObservableGauge<u64>>>,
22166
}
23167

24168
#[repr(u8)]
@@ -36,6 +180,7 @@ impl fmt::Display for PayjoinVersion {
36180

37181
impl MetricsService {
38182
pub fn new(provider: Option<SdkMeterProvider>) -> Self {
183+
let has_reader = provider.is_some();
39184
let provider = provider.unwrap_or_default();
40185
let meter = provider.meter("payjoin-mailroom");
41186

@@ -59,7 +204,46 @@ impl MetricsService {
59204
.with_description("Total mailbox entries stored by protocol version")
60205
.build();
61206

62-
Self { http_requests_total, total_connections, active_connections, db_entries_total }
207+
let tracker = UniqueShortIdTracker::new();
208+
209+
let unique_ids_gauge = if has_reader {
210+
let gauge_tracker = tracker.clone();
211+
Some(Arc::new(
212+
meter
213+
.u64_observable_gauge(UNIQUE_SHORT_IDS)
214+
.with_description("Estimated unique short IDs")
215+
.with_callback(move |observer| {
216+
observer.observe(
217+
gauge_tracker.hourly_count(),
218+
&[KeyValue::new("interval", "hourly")],
219+
);
220+
observer.observe(
221+
gauge_tracker.daily_count(),
222+
&[KeyValue::new("interval", "daily")],
223+
);
224+
observer.observe(
225+
gauge_tracker.weekly_count(),
226+
&[KeyValue::new("interval", "weekly")],
227+
);
228+
observer.observe(
229+
gauge_tracker.monthly_count(),
230+
&[KeyValue::new("interval", "monthly")],
231+
);
232+
})
233+
.build(),
234+
))
235+
} else {
236+
None
237+
};
238+
239+
Self {
240+
http_requests_total,
241+
total_connections,
242+
active_connections,
243+
db_entries_total,
244+
tracker,
245+
_unique_ids_gauge: unique_ids_gauge,
246+
}
63247
}
64248

65249
pub fn record_http_request(&self, endpoint: &str, method: &str, status_code: u16) {
@@ -83,4 +267,6 @@ impl MetricsService {
83267
pub fn record_db_entry(&self, version: PayjoinVersion) {
84268
self.db_entries_total.add(1, &[KeyValue::new("version", version.to_string())]);
85269
}
270+
271+
pub fn record_short_id(&self, id: &ShortId) { self.tracker.add_id(id); }
86272
}

0 commit comments

Comments
 (0)