From cc6b3bed453e77315ab00c63ffc8cbb9c5570d73 Mon Sep 17 00:00:00 2001
From: Piotr Sarna <sarna@turso.tech>
Date: Fri, 20 Oct 2023 12:53:13 +0200
Subject: [PATCH] libsql-server,admin: add top-k queries to stats

Top 10 most expensive queries are now stored and available
in the admin api:

```
$ curl -s http://localhost:8081/v1/namespaces/default/stats | jq '.top_queries'
[
  {
    "weight": 1,
    "rows_written": 0,
    "rows_read": 1,
    "query": "SELECT 1;"
  },
  {
    "weight": 1,
    "rows_written": 0,
    "rows_read": 1,
    "query": "SELECT 42;"
  },
  {
    "weight": 3,
    "rows_written": 0,
    "rows_read": 3,
    "query": "SELECT * FROM sqlite_master;"
  }
]
```
---
 libsql-server/src/connection/libsql.rs |  6 +++
 libsql-server/src/http/admin/stats.rs  | 10 ++++-
 libsql-server/src/stats.rs             | 53 +++++++++++++++++++++++++-
 3 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/libsql-server/src/connection/libsql.rs b/libsql-server/src/connection/libsql.rs
index a4f1ba80b6..a52d89f73a 100644
--- a/libsql-server/src/connection/libsql.rs
+++ b/libsql-server/src/connection/libsql.rs
@@ -650,6 +650,12 @@ impl<W: WalHook> Connection<W> {
         };
         self.stats.inc_rows_read(rows_read as u64);
         self.stats.inc_rows_written(rows_written as u64);
+        let weight = (rows_read + rows_written) as i64;
+        if self.stats.qualifies_as_top_query(weight) {
+            let query = stmt.expanded_sql().unwrap_or("<unknown>".into());
+            self.stats
+                .add_top_query(crate::stats::TopQuery::new(query, rows_read, rows_written));
+        }
     }
 
     fn describe(&self, sql: &str) -> DescribeResult {
diff --git a/libsql-server/src/http/admin/stats.rs b/libsql-server/src/http/admin/stats.rs
index f10c481797..8a93fee9a6 100644
--- a/libsql-server/src/http/admin/stats.rs
+++ b/libsql-server/src/http/admin/stats.rs
@@ -7,7 +7,7 @@ use axum::Json;
 
 use crate::namespace::{MakeNamespace, NamespaceName};
 use crate::replication::FrameNo;
-use crate::stats::Stats;
+use crate::stats::{Stats, TopQuery};
 
 use super::AppState;
 
@@ -18,6 +18,7 @@ pub struct StatsResponse {
     pub storage_bytes_used: u64,
     pub write_requests_delegated: u64,
     pub replication_index: FrameNo,
+    pub top_queries: Vec<TopQuery>,
 }
 
 impl From<&Stats> for StatsResponse {
@@ -28,6 +29,13 @@ impl From<&Stats> for StatsResponse {
             storage_bytes_used: stats.storage_bytes_used(),
             write_requests_delegated: stats.write_requests_delegated(),
             replication_index: stats.get_current_frame_no(),
+            top_queries: stats
+                .top_queries()
+                .read()
+                .unwrap()
+                .iter()
+                .cloned()
+                .collect(),
         }
     }
 }
diff --git a/libsql-server/src/stats.rs b/libsql-server/src/stats.rs
index 500c8fb948..e63970a4e3 100644
--- a/libsql-server/src/stats.rs
+++ b/libsql-server/src/stats.rs
@@ -1,14 +1,34 @@
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::{Arc, Weak};
+use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
+use std::sync::{Arc, RwLock, Weak};
 
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeSet;
 use tokio::io::AsyncWriteExt;
 use tokio::task::JoinSet;
 use tokio::time::Duration;
 
 use crate::replication::FrameNo;
 
+#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
+pub struct TopQuery {
+    pub weight: i64,
+    pub rows_written: i32,
+    pub rows_read: i32,
+    pub query: String,
+}
+
+impl TopQuery {
+    pub fn new(query: String, rows_read: i32, rows_written: i32) -> Self {
+        Self {
+            weight: rows_read as i64 + rows_written as i64,
+            rows_read,
+            rows_written,
+            query,
+        }
+    }
+}
+
 #[derive(Debug, Default, Serialize, Deserialize)]
 pub struct Stats {
     #[serde(default)]
@@ -22,6 +42,11 @@ pub struct Stats {
     write_requests_delegated: AtomicU64,
     #[serde(default)]
     current_frame_no: AtomicU64,
+    // Lowest value in currently stored top queries
+    #[serde(default)]
+    top_query_threshold: AtomicI64,
+    #[serde(default)]
+    top_queries: Arc<RwLock<BTreeSet<TopQuery>>>,
 }
 
 impl Stats {
@@ -91,6 +116,30 @@ impl Stats {
     pub(crate) fn get_current_frame_no(&self) -> FrameNo {
         self.current_frame_no.load(Ordering::Relaxed)
     }
+
+    pub(crate) fn add_top_query(&self, query: TopQuery) {
+        let mut top_queries = self.top_queries.write().unwrap();
+        tracing::debug!(
+            "top query: {},{}:{}",
+            query.rows_read,
+            query.rows_written,
+            query.query
+        );
+        top_queries.insert(query);
+        if top_queries.len() > 10 {
+            top_queries.pop_first();
+        }
+        self.top_query_threshold
+            .store(top_queries.first().unwrap().weight, Ordering::Relaxed);
+    }
+
+    pub(crate) fn qualifies_as_top_query(&self, weight: i64) -> bool {
+        weight >= self.top_query_threshold.load(Ordering::Relaxed)
+    }
+
+    pub(crate) fn top_queries(&self) -> &Arc<RwLock<BTreeSet<TopQuery>>> {
+        &self.top_queries
+    }
 }
 
 async fn spawn_stats_persist_thread(stats: Weak<Stats>, path: PathBuf) -> anyhow::Result<()> {