Repository: libsql/libsql
File: libsql-server/src/stats.rs (Rust)

Code before:
            count: self.count + another.count,
            rows_read: self.rows_read + another.rows_read,
            rows_written: self.rows_written + another.rows_written,
        }
    }
}

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct QueriesStats {
    #[serde(default)]
    id: Option<Uuid>,

    #[serde(default)]
    stats_threshold: AtomicU64,
    #[serde(default)]
    stats: HashMap<String, QueryStats>,
}

impl QueriesStats {
    pub fn new() -> Arc<RwLock<Self>> {
        let mut this = QueriesStats::default();
        this.id = Some(Uuid::new_v4());
        Arc::new(RwLock::new(this))
    }

    pub fn register_query(&mut self, sql: &String, stat: QueryStats) {
        let (aggregated, new) = match self.stats.get(sql) {
            Some(stat) => (stat.clone(), false),
            None => (QueryStats::empty(), true),
        };

        let aggregated = aggregated.merge(&stat);
        if aggregated.elapsed_ms < self.stats_threshold.load(Ordering::Relaxed) {
            return;
        }

        self.stats.insert(sql.clone(), aggregated);

        if !new || self.stats.len() <= 30 {
            return;
        }

        let mut vec = self.stats.clone().into_iter().collect_vec();
        vec.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
        let len = vec.len();
        if len <= 30 {
            return;
        }

        for i in 0..len - 30 {
            self.stats.remove(&vec[i].0);

Review comment (refactor):
Instead of using a lock, what about using a background task and a channel. That way the cost of doing this work is offloaded from the request path and we reduce contention to just sending the query over the channel? We need to clone anyways so there is no issue there but we avoid the lock contention. It should also simplify the logic since you don't need to thread this lock everywhere?

Code after:
    }
}

#[derive(Debug)]
pub struct QueriesStats {
    id: Uuid,
    stats: HashMap<String, QueryStats>,
    elapsed: Duration,
    stats_threshold: u64,
    hist: Histogram<u32>,
    expires_at: Option<Instant>,
    created_at: DateTime<Utc>,
}

impl QueriesStats {
    fn new() -> Self {
        Self {
            id: Uuid::new_v4(),
            stats: HashMap::new(),
            elapsed: Duration::default(),
            stats_threshold: 0,
            hist: Histogram::<u32>::new(3).unwrap(),
            expires_at: None,
            created_at: Utc::now(),
        }
    }

    fn with_expiration(expires_at: Instant) -> Self {
        let mut this = Self::new();
        this.expires_at = Some(expires_at);
        this
    }

    fn register_query(&mut self, sql: &String, stat: QueryStats) {
        self.elapsed += stat.elapsed;
        let _ = self.hist.record(stat.elapsed.as_millis() as u64);

        let (aggregated, new) = match self.stats.get(sql) {
            Some(aggregated) => (aggregated.merge(&stat), false),
            None => (stat, true),
        };

        if (aggregated.elapsed.as_micros() as u64) < self.stats_threshold {
            return;
        }

        self.stats.insert(sql.clone(), aggregated);

        if !new || self.stats.len() <= 30 {
            return;
        }