Skip to main content

rws_proxy/traffic/
store.rs

1//! Raw page views and script reports in PostgreSQL (opt-in, 18.31.0):
2//! `[traffic] database_url_env` names the environment variable holding a
3//! `postgres://` URL. A background thread writes in batches; a full queue
4//! drops rows (counted) rather than slowing requests. With it, the admin API
5//! answers any date range and filter combination from SQL
6//! (`GET /traffic/query`), past `retention_days`.
7//!
8//! Rows hold no IP address and no User-Agent: the visitor is the day's
9//! salted hash (it changes every day and at every restart), the browser
10//! only its family.
11//!
12//! ```sql
13//! rws_traffic_pageviews(at timestamptz, host, path, visitor bigint,
14//!   session_start timestamptz, referrer_host, channel, utm_source,
15//!   utm_medium, utm_campaign, country, browser, os, device)
16//! rws_traffic_reports(at timestamptz, host, path, visitor bigint, kind,
17//!   name, props, secs, scroll, width)
18//! ```
19
20use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
21use std::sync::mpsc::{sync_channel, Receiver, RecvTimeoutError, SyncSender, TrySendError};
22use std::sync::{Mutex, OnceLock};
23use std::time::{Duration, Instant};
24
25use crate::model::{DbPool, Value};
26
27/// Rows waiting to be written; past it new rows are dropped.
28pub const QUEUE: usize = 20_000;
29/// Rows written in one statement.
30pub const BATCH: usize = 500;
31/// How long a partial batch waits.
32pub const FLUSH: Duration = Duration::from_secs(10);
33
34/// A person's page view.
35#[derive(Clone, Debug, Default, PartialEq)]
36#[doc(hidden)]
37pub struct PageView {
38    pub at: u64,
39    pub host: String,
40    pub path: String,
41    pub visitor: u64,
42    pub session_start: u64,
43    pub referrer_host: Option<String>,
44    pub channel: Option<&'static str>,
45    pub utm: [Option<String>; 3],
46    pub country: Option<String>,
47    pub browser: Option<&'static str>,
48    pub os: Option<&'static str>,
49    pub device: Option<&'static str>,
50}
51
52/// A report from the first-party script.
53#[derive(Clone, Debug, Default, PartialEq)]
54#[doc(hidden)]
55pub struct Report {
56    pub at: u64,
57    pub host: String,
58    pub path: String,
59    pub visitor: u64,
60    /// `leave` or `event`.
61    pub kind: &'static str,
62    pub name: Option<String>,
63    /// `key=value` pairs, `&`-separated.
64    pub props: Option<String>,
65    pub secs: Option<u64>,
66    pub scroll: Option<u64>,
67    pub width: Option<u64>,
68}
69
70#[derive(Debug)]
71#[doc(hidden)]
72pub enum Row {
73    PageView(PageView),
74    Report(Report),
75}
76
77static ON: AtomicBool = AtomicBool::new(false);
78static DROPPED: AtomicU64 = AtomicU64::new(0);
79static WRITTEN: AtomicU64 = AtomicU64::new(0);
80static FAILED: AtomicU64 = AtomicU64::new(0);
81
82fn sender() -> &'static Mutex<Option<SyncSender<Row>>> {
83    static SENDER: OnceLock<Mutex<Option<SyncSender<Row>>>> = OnceLock::new();
84    SENDER.get_or_init(|| Mutex::new(None))
85}
86
87fn pool_slot() -> &'static Mutex<Option<DbPool>> {
88    static POOL: OnceLock<Mutex<Option<DbPool>>> = OnceLock::new();
89    POOL.get_or_init(|| Mutex::new(None))
90}
91
92/// Whether rows are being written.
93pub fn enabled() -> bool {
94    ON.load(Ordering::Relaxed)
95}
96
97/// Rows written, dropped (queue full) and lost to write errors.
98pub fn stats() -> (u64, u64, u64) {
99    (WRITTEN.load(Ordering::Relaxed), DROPPED.load(Ordering::Relaxed), FAILED.load(Ordering::Relaxed))
100}
101
102/// Queue a row; never blocks.
103#[doc(hidden)]
104pub fn enqueue(row: Row) {
105    if !enabled() {
106        return;
107    }
108    let guard = sender().lock().unwrap_or_else(|e| e.into_inner());
109    if let Some(tx) = guard.as_ref() {
110        if let Err(TrySendError::Full(_) | TrySendError::Disconnected(_)) = tx.try_send(row) {
111            DROPPED.fetch_add(1, Ordering::Relaxed);
112        }
113    }
114}
115
116const SCHEMA: &[&str] = &[
117    "CREATE TABLE IF NOT EXISTS rws_traffic_pageviews (at timestamptz NOT NULL, host text NOT NULL, path text NOT NULL, visitor bigint NOT NULL, \
118     session_start timestamptz NOT NULL, referrer_host text, channel text, utm_source text, utm_medium text, utm_campaign text, country text, \
119     browser text, os text, device text)",
120    "CREATE INDEX IF NOT EXISTS rws_traffic_pageviews_at ON rws_traffic_pageviews (at)",
121    "CREATE TABLE IF NOT EXISTS rws_traffic_reports (at timestamptz NOT NULL, host text NOT NULL, path text NOT NULL, visitor bigint NOT NULL, \
122     kind text NOT NULL, name text, props text, secs bigint, scroll bigint, width bigint)",
123    "CREATE INDEX IF NOT EXISTS rws_traffic_reports_at ON rws_traffic_reports (at)",
124];
125
126fn text(v: &Option<String>) -> Value {
127    v.clone().map_or(Value::Null, Value::Text)
128}
129
130fn word(v: Option<&'static str>) -> Value {
131    v.map_or(Value::Null, |s| Value::Text(s.to_string()))
132}
133
134fn number(v: Option<u64>) -> Value {
135    v.map_or(Value::Null, |n| Value::Int(n as i64))
136}
137
138/// A multi-row INSERT: `VALUES (to_timestamp($1), $2, …), (…)`; columns in
139/// `bigints` get `::bigint` (a NULL parameter is sent as text).
140#[doc(hidden)]
141pub fn insert_sql(table: &str, columns: &[&str], rows: usize, timestamps: &[usize], bigints: &[usize]) -> String {
142    let mut sql = format!("INSERT INTO {table} ({}) VALUES ", columns.join(", "));
143    let mut n = 0;
144    for r in 0..rows {
145        if r > 0 {
146            sql.push_str(", ");
147        }
148        sql.push('(');
149        for c in 0..columns.len() {
150            n += 1;
151            if c > 0 {
152                sql.push_str(", ");
153            }
154            if timestamps.contains(&c) {
155                sql.push_str(&format!("to_timestamp(${n})"));
156            } else if bigints.contains(&c) {
157                sql.push_str(&format!("${n}::bigint"));
158            } else {
159                sql.push_str(&format!("${n}"));
160            }
161        }
162        sql.push(')');
163    }
164    sql
165}
166
167async fn write(pool: &DbPool, rows: Vec<Row>) -> Result<u64, String> {
168    let (mut views, mut reports) = (Vec::new(), Vec::new());
169    for row in rows {
170        match row {
171            Row::PageView(v) => views.push(v),
172            Row::Report(r) => reports.push(r),
173        }
174    }
175    let mut written = 0;
176    if !views.is_empty() {
177        let columns = ["at", "host", "path", "visitor", "session_start", "referrer_host", "channel", "utm_source", "utm_medium", "utm_campaign", "country", "browser", "os", "device"];
178        let mut params = Vec::with_capacity(views.len() * columns.len());
179        for v in &views {
180            params.extend([
181                Value::Int(v.at as i64),
182                Value::Text(v.host.clone()),
183                Value::Text(v.path.clone()),
184                Value::Int(v.visitor as i64),
185                Value::Int(v.session_start as i64),
186                text(&v.referrer_host),
187                word(v.channel),
188                text(&v.utm[0]),
189                text(&v.utm[1]),
190                text(&v.utm[2]),
191                text(&v.country),
192                word(v.browser),
193                word(v.os),
194                word(v.device),
195            ]);
196        }
197        written += pool.execute(&insert_sql("rws_traffic_pageviews", &columns, views.len(), &[0, 4], &[3]), &params).await.map_err(|e| e.to_string())?;
198    }
199    if !reports.is_empty() {
200        let columns = ["at", "host", "path", "visitor", "kind", "name", "props", "secs", "scroll", "width"];
201        let mut params = Vec::with_capacity(reports.len() * columns.len());
202        for r in &reports {
203            params.extend([
204                Value::Int(r.at as i64),
205                Value::Text(r.host.clone()),
206                Value::Text(r.path.clone()),
207                Value::Int(r.visitor as i64),
208                Value::Text(r.kind.to_string()),
209                text(&r.name),
210                text(&r.props),
211                number(r.secs),
212                number(r.scroll),
213                number(r.width),
214            ]);
215        }
216        written += pool.execute(&insert_sql("rws_traffic_reports", &columns, reports.len(), &[0], &[3, 7, 8, 9]), &params).await.map_err(|e| e.to_string())?;
217    }
218    Ok(written)
219}
220
221fn writer(pool: DbPool, rx: Receiver<Row>) {
222    let mut batch = Vec::with_capacity(BATCH);
223    let mut since = Instant::now();
224    loop {
225        let wait = FLUSH.saturating_sub(since.elapsed());
226        let closed = match rx.recv_timeout(wait) {
227            Ok(row) => {
228                batch.push(row);
229                false
230            }
231            Err(RecvTimeoutError::Timeout) => false,
232            Err(RecvTimeoutError::Disconnected) => true,
233        };
234        if batch.len() >= BATCH || (!batch.is_empty() && (since.elapsed() >= FLUSH || closed)) {
235            let rows = std::mem::take(&mut batch);
236            let n = rows.len() as u64;
237            let pool = pool.clone();
238            match rws_server::async_bridge::block_on_shared(async move { write(&pool, rows).await }) {
239                Ok(_) => {
240                    WRITTEN.fetch_add(n, Ordering::Relaxed);
241                }
242                Err(e) => {
243                    FAILED.fetch_add(n, Ordering::Relaxed);
244                    eprintln!("traffic: can't write {n} rows to PostgreSQL: {e}");
245                }
246            }
247        }
248        if batch.is_empty() {
249            since = Instant::now();
250        }
251        if closed {
252            return;
253        }
254    }
255}
256
257/// Connect, create the tables and start the writer. Errors are reported
258/// and leave the store off; counting goes on without it.
259pub fn start(url: &str) -> Result<(), String> {
260    let url = url.to_string();
261    let pool = rws_server::async_bridge::block_on_shared(async move {
262        let pool = DbPool::connect(&url).await.map_err(|e| e.to_string())?;
263        for statement in SCHEMA {
264            pool.execute(statement, &[]).await.map_err(|e| e.to_string())?;
265        }
266        Ok::<_, String>(pool)
267    })?;
268    let (tx, rx) = sync_channel(QUEUE);
269    *sender().lock().unwrap_or_else(|e| e.into_inner()) = Some(tx);
270    *pool_slot().lock().unwrap_or_else(|e| e.into_inner()) = Some(pool.clone());
271    std::thread::Builder::new().name("rws-traffic-db".to_string()).spawn(move || writer(pool, rx)).map_err(|e| e.to_string())?;
272    ON.store(true, Ordering::Relaxed);
273    Ok(())
274}
275
276/// Stop taking rows and write what's queued (at shutdown, in tests).
277pub fn stop() {
278    ON.store(false, Ordering::Relaxed);
279    // Dropping the sender ends the writer after its last batch.
280    sender().lock().unwrap_or_else(|e| e.into_inner()).take();
281}
282
283/// Wait until everything queued so far has been written (tests): polls the
284/// written counter.
285#[doc(hidden)]
286pub fn wait_written(n: u64, timeout: Duration) -> bool {
287    let start = Instant::now();
288    while start.elapsed() < timeout {
289        if WRITTEN.load(Ordering::Relaxed) >= n {
290            return true;
291        }
292        std::thread::sleep(Duration::from_millis(50));
293    }
294    false
295}
296
297// ── querying ──────────────────────────────────────────────────────────────────
298
299/// A query from `GET /traffic/query`: a date range (inclusive, UTC days)
300/// and any of the filters.
301#[derive(Clone, Debug, Default, PartialEq)]
302pub struct Query {
303    /// Seconds since the epoch, start of the first day.
304    pub from: u64,
305    /// Seconds since the epoch, end of the last day (exclusive).
306    pub to: u64,
307    pub host: Option<String>,
308    /// A path, or a prefix ending in `*`.
309    pub path: Option<String>,
310    pub country: Option<String>,
311    pub channel: Option<String>,
312    pub utm_source: Option<String>,
313    pub utm_campaign: Option<String>,
314    pub browser: Option<String>,
315    pub os: Option<String>,
316    pub device: Option<String>,
317}
318
319/// Days since the epoch of `YYYY-MM-DD`.
320#[doc(hidden)]
321pub fn parse_day(text: &str) -> Option<u64> {
322    let mut parts = text.trim().splitn(3, '-');
323    let (y, m, d): (i64, i64, i64) = (parts.next()?.parse().ok()?, parts.next()?.parse().ok()?, parts.next()?.parse().ok()?);
324    if !(1970..=9999).contains(&y) || !(1..=12).contains(&m) || !(1..=31).contains(&d) {
325        return None;
326    }
327    // Howard Hinnant's days_from_civil.
328    let y = if m <= 2 { y - 1 } else { y };
329    let era = y.div_euclid(400);
330    let yoe = y - era * 400;
331    let mp = (m + 9) % 12;
332    let doy = (153 * mp + 2) / 5 + d - 1;
333    let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
334    let days = era * 146_097 + doe - 719_468;
335    (days >= 0).then_some(days as u64)
336}
337
338impl Query {
339    /// The WHERE clause and its parameters (starting at `$1`).
340    #[doc(hidden)]
341    pub fn filter(&self) -> (String, Vec<Value>) {
342        let mut clauses = vec!["at >= to_timestamp($1)".to_string(), "at < to_timestamp($2)".to_string()];
343        let mut params = vec![Value::Int(self.from as i64), Value::Int(self.to as i64)];
344        let mut add = |column: &str, value: &Option<String>, prefix: bool| {
345            if let Some(v) = value.as_ref().filter(|v| !v.is_empty()) {
346                params.push(Value::Text(match (prefix, v.strip_suffix('*')) {
347                    (true, Some(p)) => format!("{}%", p.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_")),
348                    _ => v.clone(),
349                }));
350                let n = params.len();
351                if prefix && v.ends_with('*') {
352                    clauses.push(format!("{column} LIKE ${n}"));
353                } else {
354                    clauses.push(format!("{column} = ${n}"));
355                }
356            }
357        };
358        add("host", &self.host.as_ref().map(|h| h.to_ascii_lowercase()), false);
359        add("path", &self.path, true);
360        add("country", &self.country.as_ref().map(|c| c.to_ascii_uppercase()), false);
361        add("channel", &self.channel, false);
362        add("utm_source", &self.utm_source, false);
363        add("utm_campaign", &self.utm_campaign, false);
364        add("browser", &self.browser, false);
365        add("os", &self.os, false);
366        add("device", &self.device, false);
367        (clauses.join(" AND "), params)
368    }
369}
370
371fn quote(s: &str) -> String {
372    let mut out = String::with_capacity(s.len() + 2);
373    out.push('"');
374    for c in s.chars() {
375        match c {
376            '"' => out.push_str("\\\""),
377            '\\' => out.push_str("\\\\"),
378            c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
379            c => out.push(c),
380        }
381    }
382    out.push('"');
383    out
384}
385
386/// The query's answer: totals (page views, visitors, sessions, bounces),
387/// a daily series, and top lists — as JSON.
388pub fn query_json(q: &Query) -> Result<String, String> {
389    let pool = pool_slot().lock().unwrap_or_else(|e| e.into_inner()).clone().ok_or("no traffic database (set [traffic] database_url_env)")?;
390    let (filter, params) = q.filter();
391    let (from, to) = (q.from, q.to);
392    rws_server::async_bridge::block_on_shared(async move {
393        let one = |sql: String| {
394            let pool = pool.clone();
395            let params = params.clone();
396            async move { pool.query_rows(&sql, &params).await.map_err(|e| format!("database: {e}")) }
397        };
398        let totals = one(format!(
399            "SELECT count(*)::bigint AS views, count(DISTINCT (visitor, date_trunc('day', at)))::bigint AS visitors, \
400             count(DISTINCT (visitor, session_start))::bigint AS sessions FROM rws_traffic_pageviews WHERE {filter}"
401        ))
402        .await?;
403        let bounces = one(format!(
404            "SELECT count(*)::bigint AS n FROM (SELECT 1 FROM rws_traffic_pageviews WHERE {filter} GROUP BY visitor, session_start HAVING count(*) = 1) s"
405        ))
406        .await?;
407        let series = one(format!(
408            "SELECT (extract(epoch FROM date_trunc('day', at))::bigint) AS day, count(*)::bigint AS views, count(DISTINCT visitor)::bigint AS visitors \
409             FROM rws_traffic_pageviews WHERE {filter} GROUP BY 1 ORDER BY 1"
410        ))
411        .await?;
412        let get = |rows: &[crate::model::ModelRow], col: &str| rows.first().and_then(|r| r.get::<i64>(col).ok()).unwrap_or(0);
413        let mut lists = Vec::new();
414        for (name, column) in [("pages", "host || path"), ("countries", "coalesce(country, '(unknown)')"), ("channels", "coalesce(channel, '(none)')"), ("referrers", "referrer_host"), ("utm_source", "utm_source"), ("utm_campaign", "utm_campaign"), ("browsers", "browser"), ("os", "os"), ("devices", "device")] {
415            let rows = one(format!(
416                "SELECT {column} AS k, count(*)::bigint AS n FROM rws_traffic_pageviews WHERE {filter} AND {column} IS NOT NULL GROUP BY 1 ORDER BY 2 DESC, 1 LIMIT 25"
417            ))
418            .await?;
419            let pairs: Vec<String> = rows.iter().filter_map(|r| Some(format!("[{},{}]", quote(&r.get::<String>("k").ok()?), r.get::<i64>("n").ok()?))).collect();
420            lists.push(format!("{}:[{}]", quote(name), pairs.join(",")));
421        }
422        let points: Vec<String> = series
423            .iter()
424            .filter_map(|r| Some(format!(r#"{{"t":{},"views":{},"visitors":{}}}"#, r.get::<i64>("day").ok()?, r.get::<i64>("views").ok()?, r.get::<i64>("visitors").ok()?)))
425            .collect();
426        Ok(format!(
427            r#"{{"from":{},"to":{},"views":{},"visitors":{},"sessions":{},"bounces":{},"series":[{}],{}}}"#,
428            from,
429            to,
430            get(&totals, "views"),
431            get(&totals, "visitors"),
432            get(&totals, "sessions"),
433            get(&bounces, "n"),
434            points.join(","),
435            lists.join(",")
436        ))
437    })
438}