use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use sofie_core::{Date, Error, Result, SqlValue, Timestamp, Uuid};
use sofie_telemetry::{ProductRecord, keep_product_events};
use tokio::sync::{mpsc, oneshot};
use crate::{Db, Dialect, query_values};
const QUEUED: usize = 8192;
const BATCH: usize = 512;
pub struct ProductSql {
pub insert: &'static str,
pub cohorts: &'static str,
pub prune: &'static str,
pub forget: &'static str,
pub sweep: &'static str,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Period {
Day,
Week,
Month,
}
impl Period {
#[must_use]
pub fn parse(text: &str) -> Option<Self> {
match text {
"day" => Some(Self::Day),
"week" => Some(Self::Week),
"month" => Some(Self::Month),
_ => None,
}
}
#[must_use]
pub fn name(self) -> &'static str {
match self {
Self::Day => "day",
Self::Week => "week",
Self::Month => "month",
}
}
}
pub struct Retention<'a> {
pub start: &'a str,
pub returning: &'a str,
pub period: Period,
pub since: Timestamp,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Cohort {
pub start: Date,
pub size: u64,
pub returned: Vec<u64>,
}
pub async fn cohorts(db: &Db, retention: &Retention<'_>) -> Result<Vec<Cohort>> {
let sql = product_sql(db.driver().dialect())?;
let params = [
SqlValue::Text(retention.period.name().to_owned()),
SqlValue::Text(retention.start.to_owned()),
SqlValue::Text(retention.returning.to_owned()),
SqlValue::Timestamp(retention.since),
];
let rows = query_values(db.driver(), sql.cohorts, ¶ms).await?;
let mut cohorts: Vec<Cohort> = Vec::new();
for row in rows {
let [
SqlValue::DateTime(start),
SqlValue::Int(offset),
SqlValue::Int(count),
] = row.as_slice()
else {
return Err(Error::internal_message(format!(
"a cohort row has unexpected columns: {:?}",
row.iter().map(SqlValue::kind).collect::<Vec<_>>()
)));
};
let count = u64::try_from(*count).unwrap_or_default();
let start = start.date();
if cohorts.last().is_none_or(|cohort| cohort.start != start) {
cohorts.push(Cohort {
start,
size: 0,
returned: Vec::new(),
});
}
let cohort = cohorts.last_mut().expect("a cohort was just pushed");
match usize::try_from(*offset) {
Err(_) => cohort.size = count,
Ok(offset) => {
if cohort.returned.len() <= offset {
cohort.returned.resize(offset + 1, 0);
}
cohort.returned[offset] = count;
}
}
}
Ok(cohorts)
}
pub async fn forget_product_events(db: &Db, principal: Uuid) -> Result<u64> {
let sql = product_sql(db.driver().dialect())?.forget;
let rows = query_values(db.driver(), sql, &[SqlValue::Uuid(principal)]).await?;
let Some([SqlValue::Int(forgotten)]) = rows.first().map(Vec::as_slice) else {
return Err(Error::internal_message(
"forgetting product events answered no count",
));
};
u64::try_from(*forgotten).map_err(|error| Error::internal("counting forgotten events", error))
}
impl crate::Prune {
pub async fn product_events(&self, before: sofie_core::Instant) -> Result<u64> {
let sql = product_sql(self.db().driver().dialect())?;
query_values(self.db().driver(), sql.sweep, &[]).await?;
let params = [SqlValue::Timestamp(before.timestamp())];
let mut deleted = 0;
loop {
let batch = query_values(self.db().driver(), sql.prune, ¶ms)
.await?
.len();
if batch == 0 {
return Ok(deleted);
}
deleted += u64::try_from(batch)
.map_err(|error| Error::internal("counting deleted events", error))?;
}
}
}
pub(crate) fn product_sql(dialect: &dyn Dialect) -> Result<&'static ProductSql> {
dialect
.product_events()
.ok_or_else(|| Error::internal_message("this database keeps no product events"))
}
enum Kept {
Event(ProductRecord),
Flush(oneshot::Sender<()>),
}
pub struct ProductEvents {
queue: mpsc::Sender<Kept>,
}
impl ProductEvents {
pub fn keep(db: &Db) -> Result<Self> {
let sql = product_sql(db.driver().dialect())?;
let (queue, mut received) = mpsc::channel(QUEUED);
let raised = queue.clone();
let dropped = Arc::new(AtomicU64::new(0));
let installed = keep_product_events(move |record| {
if raised.try_send(Kept::Event(record)).is_ok() {
return;
}
let lost = dropped.fetch_add(1, Ordering::Relaxed) + 1;
if lost.is_power_of_two() {
tracing::warn!(
product.dropped = lost,
"product events dropped: the database keeps up with fewer"
);
}
});
if !installed {
return Err(Error::internal_message(
"product events are kept already in this process",
));
}
let db = db.clone();
tokio::spawn(async move {
let mut batch = Vec::with_capacity(BATCH);
while let Some(first) = received.recv().await {
let mut flushed = Vec::new();
let mut next = Some(first);
while let Some(kept) = next {
match kept {
Kept::Event(record) => batch.push(record),
Kept::Flush(done) => flushed.push(done),
}
next = (batch.len() < BATCH)
.then(|| received.try_recv().ok())
.flatten();
}
write(&db, sql, &mut batch).await;
for done in flushed {
let _ = done.send(());
}
}
});
Ok(Self { queue })
}
pub async fn flush(&self) {
let (done, written) = oneshot::channel();
if self.queue.send(Kept::Flush(done)).await.is_ok() {
let _ = written.await;
}
}
}
async fn write(db: &Db, sql: &ProductSql, batch: &mut Vec<ProductRecord>) {
if batch.is_empty() {
return;
}
let mut at = Vec::with_capacity(batch.len());
let mut names = Vec::with_capacity(batch.len());
let mut principals = Vec::with_capacity(batch.len());
for record in batch.drain(..) {
at.push(SqlValue::Timestamp(record.at));
names.push(SqlValue::Text(record.name));
principals.push(record.principal.map_or(SqlValue::Null, SqlValue::Uuid));
}
let count = at.len();
let params = [
SqlValue::Array(at),
SqlValue::Array(names),
SqlValue::Array(principals),
];
if let Err(error) = query_values(db.driver(), sql.insert, ¶ms).await {
tracing::warn!(
error = %error,
product.lost = count,
"writing product events failed"
);
}
}