1use 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
27pub const QUEUE: usize = 20_000;
29pub const BATCH: usize = 500;
31pub const FLUSH: Duration = Duration::from_secs(10);
33
34#[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#[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 pub kind: &'static str,
62 pub name: Option<String>,
63 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
92pub fn enabled() -> bool {
94 ON.load(Ordering::Relaxed)
95}
96
97pub fn stats() -> (u64, u64, u64) {
99 (WRITTEN.load(Ordering::Relaxed), DROPPED.load(Ordering::Relaxed), FAILED.load(Ordering::Relaxed))
100}
101
102#[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#[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]), ¶ms).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]), ¶ms).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
257pub 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
276pub fn stop() {
278 ON.store(false, Ordering::Relaxed);
279 sender().lock().unwrap_or_else(|e| e.into_inner()).take();
281}
282
283#[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#[derive(Clone, Debug, Default, PartialEq)]
302pub struct Query {
303 pub from: u64,
305 pub to: u64,
307 pub host: Option<String>,
308 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#[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 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 #[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
386pub 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, ¶ms).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}