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)]
320pub use super::summary::parse_day;
321
322impl Query {
323 #[doc(hidden)]
325 pub fn filter(&self) -> (String, Vec<Value>) {
326 let mut clauses = vec!["at >= to_timestamp($1)".to_string(), "at < to_timestamp($2)".to_string()];
327 let mut params = vec![Value::Int(self.from as i64), Value::Int(self.to as i64)];
328 let mut add = |column: &str, value: &Option<String>, prefix: bool| {
329 if let Some(v) = value.as_ref().filter(|v| !v.is_empty()) {
330 params.push(Value::Text(match (prefix, v.strip_suffix('*')) {
331 (true, Some(p)) => format!("{}%", p.replace('\\', "\\\\").replace('%', "\\%").replace('_', "\\_")),
332 _ => v.clone(),
333 }));
334 let n = params.len();
335 if prefix && v.ends_with('*') {
336 clauses.push(format!("{column} LIKE ${n}"));
337 } else {
338 clauses.push(format!("{column} = ${n}"));
339 }
340 }
341 };
342 add("host", &self.host.as_ref().map(|h| h.to_ascii_lowercase()), false);
343 add("path", &self.path, true);
344 add("country", &self.country.as_ref().map(|c| c.to_ascii_uppercase()), false);
345 add("channel", &self.channel, false);
346 add("utm_source", &self.utm_source, false);
347 add("utm_campaign", &self.utm_campaign, false);
348 add("browser", &self.browser, false);
349 add("os", &self.os, false);
350 add("device", &self.device, false);
351 (clauses.join(" AND "), params)
352 }
353}
354
355fn quote(s: &str) -> String {
356 let mut out = String::with_capacity(s.len() + 2);
357 out.push('"');
358 for c in s.chars() {
359 match c {
360 '"' => out.push_str("\\\""),
361 '\\' => out.push_str("\\\\"),
362 c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)),
363 c => out.push(c),
364 }
365 }
366 out.push('"');
367 out
368}
369
370pub fn query_json(q: &Query) -> Result<String, String> {
373 let pool = pool_slot().lock().unwrap_or_else(|e| e.into_inner()).clone().ok_or("no traffic database (set [traffic] database_url_env)")?;
374 let (filter, params) = q.filter();
375 let (from, to) = (q.from, q.to);
376 rws_server::async_bridge::block_on_shared(async move {
377 let one = |sql: String| {
378 let pool = pool.clone();
379 let params = params.clone();
380 async move { pool.query_rows(&sql, ¶ms).await.map_err(|e| format!("database: {e}")) }
381 };
382 let totals = one(format!(
383 "SELECT count(*)::bigint AS views, count(DISTINCT (visitor, date_trunc('day', at)))::bigint AS visitors, \
384 count(DISTINCT (visitor, session_start))::bigint AS sessions FROM rws_traffic_pageviews WHERE {filter}"
385 ))
386 .await?;
387 let bounces = one(format!(
388 "SELECT count(*)::bigint AS n FROM (SELECT 1 FROM rws_traffic_pageviews WHERE {filter} GROUP BY visitor, session_start HAVING count(*) = 1) s"
389 ))
390 .await?;
391 let series = one(format!(
392 "SELECT (extract(epoch FROM date_trunc('day', at))::bigint) AS day, count(*)::bigint AS views, count(DISTINCT visitor)::bigint AS visitors \
393 FROM rws_traffic_pageviews WHERE {filter} GROUP BY 1 ORDER BY 1"
394 ))
395 .await?;
396 let get = |rows: &[crate::model::ModelRow], col: &str| rows.first().and_then(|r| r.get::<i64>(col).ok()).unwrap_or(0);
397 let mut lists = Vec::new();
398 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")] {
399 let rows = one(format!(
400 "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"
401 ))
402 .await?;
403 let pairs: Vec<String> = rows.iter().filter_map(|r| Some(format!("[{},{}]", quote(&r.get::<String>("k").ok()?), r.get::<i64>("n").ok()?))).collect();
404 lists.push(format!("{}:[{}]", quote(name), pairs.join(",")));
405 }
406 let points: Vec<String> = series
407 .iter()
408 .filter_map(|r| Some(format!(r#"{{"t":{},"views":{},"visitors":{}}}"#, r.get::<i64>("day").ok()?, r.get::<i64>("views").ok()?, r.get::<i64>("visitors").ok()?)))
409 .collect();
410 Ok(format!(
411 r#"{{"from":{},"to":{},"views":{},"visitors":{},"sessions":{},"bounces":{},"series":[{}],{}}}"#,
412 from,
413 to,
414 get(&totals, "views"),
415 get(&totals, "visitors"),
416 get(&totals, "sessions"),
417 get(&bounces, "n"),
418 points.join(","),
419 lists.join(",")
420 ))
421 })
422}