1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::path::{Path, PathBuf};
4use std::time::Duration;
5use tokio::sync::mpsc;
6use tokio::time::{interval, Instant, MissedTickBehavior};
7use tracing::{debug, error, warn};
8use uuid::Uuid;
9
10const MAX_IN_MEMORY_RETRIES: u32 = 3;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13#[serde(tag = "type", rename_all = "snake_case")]
14pub enum WebSocketUsageEvent {
15 ConnectionEstablished {
16 client_id: String,
17 remote_addr: String,
18 deployment_id: Option<String>,
19 metering_key: Option<String>,
20 subject: Option<String>,
21 key_class: Option<String>,
22 },
23 ConnectionClosed {
24 client_id: String,
25 deployment_id: Option<String>,
26 metering_key: Option<String>,
27 subject: Option<String>,
28 duration_secs: Option<f64>,
29 subscription_count: u32,
30 },
31 SubscriptionCreated {
32 client_id: String,
33 deployment_id: Option<String>,
34 metering_key: Option<String>,
35 subject: Option<String>,
36 view_id: String,
37 },
38 SubscriptionRemoved {
39 client_id: String,
40 deployment_id: Option<String>,
41 metering_key: Option<String>,
42 subject: Option<String>,
43 view_id: String,
44 },
45 SnapshotSent {
46 client_id: String,
47 deployment_id: Option<String>,
48 metering_key: Option<String>,
49 subject: Option<String>,
50 view_id: String,
51 rows: u32,
52 messages: u32,
53 bytes: u64,
54 },
55 UpdateSent {
56 client_id: String,
57 deployment_id: Option<String>,
58 metering_key: Option<String>,
59 subject: Option<String>,
60 view_id: String,
61 messages: u32,
62 bytes: u64,
63 },
64}
65
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct WebSocketUsageEnvelope {
68 pub event_id: String,
69 pub occurred_at_ms: u64,
70 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub build_id: Option<String>,
72 pub event: WebSocketUsageEvent,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct WebSocketUsageBatch {
77 pub events: Vec<WebSocketUsageEnvelope>,
78}
79
80#[async_trait]
81pub trait WebSocketUsageEmitter: Send + Sync {
82 async fn emit(&self, event: WebSocketUsageEvent);
83}
84
85#[derive(Clone)]
86pub struct ChannelUsageEmitter {
87 sender: mpsc::UnboundedSender<WebSocketUsageEvent>,
88}
89
90impl ChannelUsageEmitter {
91 pub fn new(sender: mpsc::UnboundedSender<WebSocketUsageEvent>) -> Self {
92 Self { sender }
93 }
94}
95
96#[async_trait]
97impl WebSocketUsageEmitter for ChannelUsageEmitter {
98 async fn emit(&self, event: WebSocketUsageEvent) {
99 let _ = self.sender.send(event);
100 }
101}
102
103pub struct HttpUsageEmitter {
104 sender: mpsc::UnboundedSender<WebSocketUsageEvent>,
105}
106
107#[derive(Debug, thiserror::Error)]
108#[error("WebSocket usage build ID must be a positive integer")]
109pub struct InvalidUsageBuildId;
110
111#[derive(Debug, Clone)]
112struct RetryState {
113 batch: WebSocketUsageBatch,
114 attempts: u32,
115 next_retry_at: Instant,
116}
117
118impl HttpUsageEmitter {
119 pub fn new(endpoint: String, auth_token: Option<String>) -> Self {
120 Self::with_config(endpoint, auth_token, 50, Duration::from_secs(2))
121 }
122
123 pub fn new_attributed(
124 endpoint: String,
125 auth_token: Option<String>,
126 build_id: impl Into<String>,
127 ) -> Result<Self, InvalidUsageBuildId> {
128 let build_id = validate_build_id(build_id.into())?;
129 Ok(Self::with_full_config(
130 endpoint,
131 auth_token,
132 50,
133 Duration::from_secs(2),
134 None,
135 Some(build_id),
136 ))
137 }
138
139 pub fn with_spool_dir(
140 endpoint: String,
141 auth_token: Option<String>,
142 spool_dir: impl Into<PathBuf>,
143 ) -> Self {
144 Self::with_full_config(
145 endpoint,
146 auth_token,
147 50,
148 Duration::from_secs(2),
149 Some(spool_dir.into()),
150 None,
151 )
152 }
153
154 pub fn with_attributed_spool_dir(
155 endpoint: String,
156 auth_token: Option<String>,
157 spool_dir: impl Into<PathBuf>,
158 build_id: impl Into<String>,
159 ) -> Result<Self, InvalidUsageBuildId> {
160 let build_id = validate_build_id(build_id.into())?;
161 Ok(Self::with_full_config(
162 endpoint,
163 auth_token,
164 50,
165 Duration::from_secs(2),
166 Some(spool_dir.into()),
167 Some(build_id),
168 ))
169 }
170
171 pub fn with_config(
172 endpoint: String,
173 auth_token: Option<String>,
174 batch_size: usize,
175 flush_interval: Duration,
176 ) -> Self {
177 Self::with_full_config(endpoint, auth_token, batch_size, flush_interval, None, None)
178 }
179
180 fn with_full_config(
181 endpoint: String,
182 auth_token: Option<String>,
183 batch_size: usize,
184 flush_interval: Duration,
185 spool_dir: Option<PathBuf>,
186 build_id: Option<String>,
187 ) -> Self {
188 let (sender, mut receiver) = mpsc::unbounded_channel::<WebSocketUsageEvent>();
189 let client = reqwest::Client::new();
190
191 tokio::spawn(async move {
192 let mut ticker = interval(flush_interval);
193 ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
194 let mut pending: Vec<WebSocketUsageEnvelope> = Vec::new();
195 let mut retry_state: Option<RetryState> = None;
196
197 if let Some(dir) = spool_dir.as_ref() {
198 if let Err(error) = ensure_spool_dir(dir) {
199 warn!(error = %error, path = %dir.display(), "failed to initialize websocket usage spool directory");
200 }
201 }
202
203 loop {
204 tokio::select! {
205 maybe_event = receiver.recv() => {
206 match maybe_event {
207 Some(event) => {
208 pending.push(WebSocketUsageEnvelope {
209 event_id: Uuid::new_v4().to_string(),
210 occurred_at_ms: current_time_ms(),
211 build_id: build_id.clone(),
212 event,
213 });
214
215 if retry_state.is_none() && pending.len() >= batch_size {
216 flush_pending_batch(
217 &client,
218 &endpoint,
219 auth_token.as_deref(),
220 &mut pending,
221 &mut retry_state,
222 spool_dir.as_deref(),
223 ).await;
224 }
225 }
226 None => {
227 if retry_state.is_none() && !pending.is_empty() {
228 flush_pending_batch(
229 &client,
230 &endpoint,
231 auth_token.as_deref(),
232 &mut pending,
233 &mut retry_state,
234 spool_dir.as_deref(),
235 ).await;
236 }
237
238 if let Some(state) = retry_state.take() {
239 if let Err(retry_state_failed) = flush_existing_batch(
240 &client,
241 &endpoint,
242 auth_token.as_deref(),
243 state,
244 ).await {
245 if let Some(dir) = spool_dir.as_deref() {
246 if let Err(error) = spool_retry_state(dir, &retry_state_failed) {
247 warn!(error = %error, count = retry_state_failed.batch.events.len(), "failed to spool websocket usage batch during shutdown");
248 }
249 } else {
250 warn!(
251 count = retry_state_failed.batch.events.len(),
252 attempts = retry_state_failed.attempts,
253 "dropping websocket usage batch during shutdown after failed retry"
254 );
255 }
256 }
257 }
258
259 if !pending.is_empty() {
260 if let Some(dir) = spool_dir.as_deref() {
261 let batch = WebSocketUsageBatch { events: std::mem::take(&mut pending) };
262 if let Err(error) = spool_batch(dir, &batch) {
263 warn!(error = %error, count = batch.events.len(), "failed to spool pending websocket usage batch during shutdown");
264 }
265 } else {
266 warn!(count = pending.len(), "dropping pending websocket usage events during shutdown without spool directory");
267 }
268 }
269 break;
270 }
271 }
272 }
273 _ = ticker.tick() => {
274 if let Some(dir) = spool_dir.as_deref() {
275 if retry_state.is_none() {
276 if let Err(error) = flush_one_spooled_batch(&client, &endpoint, auth_token.as_deref(), dir).await {
277 warn!(error = %error, path = %dir.display(), "failed to process spooled websocket usage batch");
278 }
279 }
280 }
281
282 if let Some(state) = retry_state.take() {
283 if Instant::now() >= state.next_retry_at {
284 match flush_existing_batch(
285 &client,
286 &endpoint,
287 auth_token.as_deref(),
288 state,
289 ).await {
290 Ok(()) => {
291 if !pending.is_empty() {
292 flush_pending_batch(
293 &client,
294 &endpoint,
295 auth_token.as_deref(),
296 &mut pending,
297 &mut retry_state,
298 spool_dir.as_deref(),
299 ).await;
300 }
301 }
302 Err(state) => {
303 if state.attempts >= MAX_IN_MEMORY_RETRIES {
304 if let Some(dir) = spool_dir.as_deref() {
305 if let Err(error) = spool_retry_state(dir, &state) {
306 warn!(error = %error, count = state.batch.events.len(), "failed to spool websocket usage batch after retries");
307 retry_state = Some(state);
308 }
309 } else {
310 retry_state = Some(state);
311 }
312 } else {
313 retry_state = Some(state)
314 }
315 }
316 }
317 } else {
318 retry_state = Some(state);
319 }
320 } else if !pending.is_empty() {
321 flush_pending_batch(
322 &client,
323 &endpoint,
324 auth_token.as_deref(),
325 &mut pending,
326 &mut retry_state,
327 spool_dir.as_deref(),
328 ).await;
329 }
330 }
331 }
332 }
333 });
334
335 Self { sender }
336 }
337}
338
339fn validate_build_id(value: String) -> Result<String, InvalidUsageBuildId> {
340 if value.trim() != value || value.parse::<i32>().ok().is_none_or(|value| value <= 0) {
341 return Err(InvalidUsageBuildId);
342 }
343 Ok(value)
344}
345
346#[async_trait]
347impl WebSocketUsageEmitter for HttpUsageEmitter {
348 async fn emit(&self, event: WebSocketUsageEvent) {
349 if let Err(error) = self.sender.send(event) {
350 warn!(error = %error, "failed to queue websocket usage event");
351 }
352 }
353}
354
355fn current_time_ms() -> u64 {
356 std::time::SystemTime::now()
357 .duration_since(std::time::UNIX_EPOCH)
358 .unwrap_or_default()
359 .as_millis() as u64
360}
361
362async fn flush_batch(
363 client: &reqwest::Client,
364 endpoint: &str,
365 auth_token: Option<&str>,
366 batch: &WebSocketUsageBatch,
367) -> bool {
368 if batch.events.is_empty() {
369 return true;
370 }
371
372 let mut request = client.post(endpoint).json(batch);
373 if let Some(token) = auth_token {
374 request = request.header("Authorization", format!("Bearer {}", token));
375 }
376
377 match request.send().await {
378 Ok(response) if response.status().is_success() => {
379 debug!(count = batch.events.len(), "flushed websocket usage batch");
380 true
381 }
382 Ok(response) => {
383 error!(status = %response.status(), count = batch.events.len(), "failed to ingest websocket usage batch");
384 false
385 }
386 Err(error) => {
387 error!(error = %error, count = batch.events.len(), "failed to post websocket usage batch");
388 false
389 }
390 }
391}
392
393async fn flush_pending_batch(
394 client: &reqwest::Client,
395 endpoint: &str,
396 auth_token: Option<&str>,
397 pending: &mut Vec<WebSocketUsageEnvelope>,
398 retry_state: &mut Option<RetryState>,
399 spool_dir: Option<&Path>,
400) {
401 let batch = WebSocketUsageBatch {
402 events: std::mem::take(pending),
403 };
404
405 if !flush_batch(client, endpoint, auth_token, &batch).await {
406 let state = RetryState {
407 batch,
408 attempts: 1,
409 next_retry_at: Instant::now() + retry_delay(1),
410 };
411
412 if let Some(dir) = spool_dir.filter(|_| MAX_IN_MEMORY_RETRIES <= 1) {
413 if let Err(error) = spool_retry_state(dir, &state) {
414 warn!(error = %error, count = state.batch.events.len(), "failed to spool websocket usage batch after first failure");
415 *retry_state = Some(state);
416 }
417 } else {
418 *retry_state = Some(state);
419 }
420 }
421}
422
423async fn flush_existing_batch(
424 client: &reqwest::Client,
425 endpoint: &str,
426 auth_token: Option<&str>,
427 mut state: RetryState,
428) -> Result<(), RetryState> {
429 if flush_batch(client, endpoint, auth_token, &state.batch).await {
430 Ok(())
431 } else {
432 state.attempts += 1;
433 state.next_retry_at = Instant::now() + retry_delay(state.attempts);
434 Err(state)
435 }
436}
437
438fn retry_delay(attempt: u32) -> Duration {
439 let capped_attempt = attempt.min(6);
440 Duration::from_secs(1_u64 << capped_attempt)
441}
442
443fn ensure_spool_dir(path: &Path) -> std::io::Result<()> {
444 std::fs::create_dir_all(path)
445}
446
447fn spool_retry_state(path: &Path, state: &RetryState) -> std::io::Result<PathBuf> {
448 spool_batch(path, &state.batch)
449}
450
451fn spool_batch(path: &Path, batch: &WebSocketUsageBatch) -> std::io::Result<PathBuf> {
452 ensure_spool_dir(path)?;
453
454 let file_name = format!(
455 "ws-usage-{}-{}.json",
456 current_time_ms(),
457 Uuid::new_v4().simple()
458 );
459 let final_path = path.join(file_name);
460 let temp_path = final_path.with_extension("tmp");
461 let data = serde_json::to_vec(batch).map_err(std::io::Error::other)?;
462 std::fs::write(&temp_path, data)?;
463 std::fs::rename(&temp_path, &final_path)?;
464 Ok(final_path)
465}
466
467fn load_batch_from_file(path: &Path) -> std::io::Result<WebSocketUsageBatch> {
468 let data = std::fs::read(path)?;
469 serde_json::from_slice(&data).map_err(std::io::Error::other)
470}
471
472fn oldest_spooled_batch(path: &Path) -> std::io::Result<Option<PathBuf>> {
473 if !path.exists() {
474 return Ok(None);
475 }
476
477 let mut entries: Vec<PathBuf> = std::fs::read_dir(path)?
478 .filter_map(|entry| entry.ok().map(|entry| entry.path()))
479 .filter(|entry| entry.extension().and_then(|ext| ext.to_str()) == Some("json"))
480 .collect();
481 entries.sort();
482 Ok(entries.into_iter().next())
483}
484
485async fn flush_one_spooled_batch(
486 client: &reqwest::Client,
487 endpoint: &str,
488 auth_token: Option<&str>,
489 spool_dir: &Path,
490) -> std::io::Result<()> {
491 let Some(path) = oldest_spooled_batch(spool_dir)? else {
492 return Ok(());
493 };
494
495 let batch = load_batch_from_file(&path)?;
496 if flush_batch(client, endpoint, auth_token, &batch).await {
497 std::fs::remove_file(path)?;
498 }
499
500 Ok(())
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506 use std::fs;
507
508 fn temp_spool_dir() -> PathBuf {
509 let dir = std::env::temp_dir().join(format!("arete-usage-test-{}", Uuid::new_v4()));
510 fs::create_dir_all(&dir).expect("temp dir should be created");
511 dir
512 }
513
514 #[tokio::test]
515 async fn channel_usage_emitter_forwards_events() {
516 let (tx, mut rx) = mpsc::unbounded_channel();
517 let emitter = ChannelUsageEmitter::new(tx);
518
519 emitter
520 .emit(WebSocketUsageEvent::SubscriptionCreated {
521 client_id: "client-1".to_string(),
522 deployment_id: Some("deployment-1".to_string()),
523 metering_key: Some("meter-1".to_string()),
524 subject: Some("subject-1".to_string()),
525 view_id: "OreRound/latest".to_string(),
526 })
527 .await;
528
529 let event = rx.recv().await.expect("event should be forwarded");
530 match event {
531 WebSocketUsageEvent::SubscriptionCreated { view_id, .. } => {
532 assert_eq!(view_id, "OreRound/latest");
533 }
534 other => panic!("unexpected event: {other:?}"),
535 }
536 }
537
538 #[test]
539 fn retry_delay_grows_and_caps() {
540 assert_eq!(retry_delay(1), Duration::from_secs(2));
541 assert_eq!(retry_delay(2), Duration::from_secs(4));
542 assert_eq!(retry_delay(6), Duration::from_secs(64));
543 assert_eq!(retry_delay(9), Duration::from_secs(64));
544 }
545
546 #[test]
547 fn spooled_batches_round_trip() {
548 let dir = temp_spool_dir();
549 let batch = WebSocketUsageBatch {
550 events: vec![WebSocketUsageEnvelope {
551 event_id: "evt_1".to_string(),
552 occurred_at_ms: 123,
553 build_id: Some("7".to_string()),
554 event: WebSocketUsageEvent::UpdateSent {
555 client_id: "client-1".to_string(),
556 deployment_id: Some("1".to_string()),
557 metering_key: Some("api_key:1".to_string()),
558 subject: Some("user:1".to_string()),
559 view_id: "OreRound/latest".to_string(),
560 messages: 1,
561 bytes: 42,
562 },
563 }],
564 };
565
566 let path = spool_batch(&dir, &batch).expect("batch should spool");
567 let loaded = load_batch_from_file(&path).expect("batch should load");
568 assert_eq!(loaded.events.len(), 1);
569 assert_eq!(loaded.events[0].build_id.as_deref(), Some("7"));
570
571 fs::remove_dir_all(dir).expect("temp dir should be removed");
572 }
573
574 #[test]
575 fn legacy_and_attributed_envelopes_are_compatible() {
576 let legacy = serde_json::json!({
577 "event_id": "legacy",
578 "occurred_at_ms": 123,
579 "event": {
580 "type": "update_sent",
581 "client_id": "client",
582 "deployment_id": "1",
583 "metering_key": null,
584 "subject": null,
585 "view_id": "view",
586 "messages": 1,
587 "bytes": 2
588 }
589 });
590 let decoded: WebSocketUsageEnvelope = serde_json::from_value(legacy).unwrap();
591 assert_eq!(decoded.build_id, None);
592
593 let mut attributed = serde_json::to_value(decoded).unwrap();
594 assert!(attributed.get("build_id").is_none());
595 attributed["build_id"] = serde_json::json!("42");
596 let decoded: WebSocketUsageEnvelope = serde_json::from_value(attributed).unwrap();
597 assert_eq!(decoded.build_id.as_deref(), Some("42"));
598 }
599
600 #[test]
601 fn attributed_constructor_validates_build_identity() {
602 assert!(
603 HttpUsageEmitter::new_attributed("http://localhost".to_string(), None, "0").is_err()
604 );
605 assert!(
606 HttpUsageEmitter::new_attributed("http://localhost".to_string(), None, " 7").is_err()
607 );
608 }
609
610 #[test]
611 fn oldest_spooled_batch_prefers_lexicographically_oldest_file() {
612 let dir = temp_spool_dir();
613 fs::write(dir.join("ws-usage-100-a.json"), b"{\"events\":[]}").expect("first batch");
614 fs::write(dir.join("ws-usage-200-b.json"), b"{\"events\":[]}").expect("second batch");
615
616 let oldest = oldest_spooled_batch(&dir)
617 .expect("listing should succeed")
618 .expect("batch should exist");
619 assert!(oldest.ends_with("ws-usage-100-a.json"));
620
621 fs::remove_dir_all(dir).expect("temp dir should be removed");
622 }
623}