From c3723aaf272c8b699b522c60a07b5b994a0906c9 Mon Sep 17 00:00:00 2001
From: Ellie Huxtable <ellie@elliehuxtable.com>
Date: Tue, 2 Jul 2024 14:47:41 +0100
Subject: [PATCH] feat: monitor idx cache consistency before switching (#2229)

---
 Cargo.lock                              |  1 +
 crates/atuin-server-postgres/Cargo.toml |  1 +
 crates/atuin-server-postgres/src/lib.rs | 32 ++++++++++++++++++++++---
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6956f54d..d5a911c6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -447,6 +447,7 @@ dependencies = [
  "atuin-server-database",
  "eyre",
  "futures-util",
+ "metrics",
  "serde",
  "sqlx",
  "time",
diff --git a/crates/atuin-server-postgres/Cargo.toml b/crates/atuin-server-postgres/Cargo.toml
index 10753272..0d31ef6f 100644
--- a/crates/atuin-server-postgres/Cargo.toml
+++ b/crates/atuin-server-postgres/Cargo.toml
@@ -20,5 +20,6 @@ serde = { workspace = true }
 sqlx = { workspace = true }
 async-trait = { workspace = true }
 uuid = { workspace = true }
+metrics = "0.21.1"
 futures-util = "0.3"
 url = "2.5.2"
diff --git a/crates/atuin-server-postgres/src/lib.rs b/crates/atuin-server-postgres/src/lib.rs
index f12e4801..6c3572bc 100644
--- a/crates/atuin-server-postgres/src/lib.rs
+++ b/crates/atuin-server-postgres/src/lib.rs
@@ -8,6 +8,7 @@ use atuin_common::utils::crypto_random_string;
 use atuin_server_database::models::{History, NewHistory, NewSession, NewUser, Session, User};
 use atuin_server_database::{Database, DbError, DbResult};
 use futures_util::TryStreamExt;
+use metrics::counter;
 use serde::{Deserialize, Serialize};
 use sqlx::postgres::PgPoolOptions;
 use sqlx::Row;
@@ -643,16 +644,41 @@ impl Database for Postgres {
         const STATUS_SQL: &str =
             "select host, tag, max(idx) from store where user_id = $1 group by host, tag";
 
-        let res: Vec<(Uuid, String, i64)> = sqlx::query_as(STATUS_SQL)
+        let mut res: Vec<(Uuid, String, i64)> = sqlx::query_as(STATUS_SQL)
             .bind(user.id)
             .fetch_all(&self.pool)
             .await
             .map_err(fix_error)?;
+        res.sort();
+
+        // We're temporarily increasing latency in order to improve confidence in the cache
+        // If it runs for a few days, and we confirm that cached values are equal to realtime, we
+        // can replace realtime with cached.
+        //
+        // But let's check so sync doesn't do Weird Things.
+
+        let mut cached_res: Vec<(Uuid, String, i64)> =
+            sqlx::query_as("select host, tag, idx from store_idx_cache where user_id = $1")
+                .bind(user.id)
+                .fetch_all(&self.pool)
+                .await
+                .map_err(fix_error)?;
+        cached_res.sort();
 
         let mut status = RecordStatus::new();
 
-        for i in res {
-            status.set_raw(HostId(i.0), i.1, i.2 as u64);
+        let equal = res == cached_res;
+
+        if equal {
+            counter!("atuin_store_idx_cache_consistent", 1);
+        } else {
+            // log the values if we have an inconsistent cache
+            tracing::debug!(user = user.username, cache_match = equal, res = ?res, cached = ?cached_res, "record store index request");
+            counter!("atuin_store_idx_cache_inconsistent", 1);
+        };
+
+        for i in res.iter() {
+            status.set_raw(HostId(i.0), i.1.clone(), i.2 as u64);
         }
 
         Ok(status)