Improve local cache and prefer local reads (#1545)

avinassh · web-flow · commit f563a741ca4e · 2024-07-12T06:39:37.000Z
* update WIP list

* prefer local max_frame_no

this assumes that cache is always up to date and uses that instead
of calling the storage server. later we will add replication and make
sure that cache is upto date.

* Prefer local reads and also save page versions locally

Let's store `page_no` along with the `frame_no`. And then while reading,
we will do a range read, read frame which is less than txn's
`max_frame_no`

* minor improvements: index creation, fix max_frame_num query

* update local cache on inserts

* avoid unnecessary call to ss in `find_frame`

* remove unused `get_frame`

* eliminiate all calls to storage server during read path
diff --git a/libsql-storage/src/lib.rs b/libsql-storage/src/lib.rs
@@ -13,7 +13,7 @@ use libsql_sys::rusqlite;
 use libsql_sys::wal::{Result, Vfs, Wal, WalManager};
 use rpc::storage_client::StorageClient;
 use tonic::transport::Channel;
-use tracing::{error, trace};
+use tracing::{error, trace, warn};
 
 pub mod rpc {
     #![allow(clippy::all)]
@@ -23,12 +23,7 @@ pub mod rpc {
 // What does (not) work:
 // - there are no read txn locks nor upgrades
 // - no lock stealing
-// - write set is kept in mem
-// - txns don't use max_frame_no yet
 // - no savepoints, yet
-// - no multi tenancy, uses `default` namespace
-// - txn can read new frames after it started (since there are no read locks)
-// - requires huge memory as it assumes all the txn data fits in the memory
 
 #[derive(Clone, Default)]
 pub struct DurableWalConfig {
@@ -178,17 +173,6 @@ impl DurableWal {
             .flatten();
         Ok(frame_no)
     }
-
-    async fn frames_count(&self) -> u64 {
-        let req = rpc::FramesInWalRequest {
-            namespace: self.namespace.to_string(),
-        };
-        let mut binding = self.client.clone();
-        let resp = binding.frames_in_wal(req).await.unwrap();
-        let count = resp.into_inner().count;
-        trace!("DurableWal::frames_in_wal() = {}", count);
-        count
-    }
 }
 
 impl Wal for DurableWal {
@@ -200,9 +184,7 @@ impl Wal for DurableWal {
         // - create a read lock
         // - save the current max_frame_no for this txn
         trace!("DurableWal::begin_read_txn()");
-        let rt = tokio::runtime::Handle::current();
-        let frame_no = tokio::task::block_in_place(|| rt.block_on(self.frames_count()));
-        self.max_frame_no = frame_no;
+        self.max_frame_no = self.local_cache.get_max_frame_num().unwrap();
         Ok(true)
     }
 
@@ -232,21 +214,14 @@ impl Wal for DurableWal {
         if self.max_frame_no == 0 {
             return Ok(None);
         }
-        let rt = tokio::runtime::Handle::current();
-        // TODO: find_frame should account for `max_frame_no` of this txn
-        let frame_no =
-            tokio::task::block_in_place(|| rt.block_on(self.find_frame_by_page_no(page_no)))
-                .unwrap();
-        if frame_no.is_none() {
-            return Ok(None);
-        }
         return Ok(Some(page_no));
     }
 
     #[tracing::instrument(skip_all, fields(page_no))]
     fn read_frame(&mut self, page_no: std::num::NonZeroU32, buffer: &mut [u8]) -> Result<()> {
         trace!("DurableWal::read_frame()");
-        let rt = tokio::runtime::Handle::current();
+        // to read a frame, first we check in transaction cache, then frames cache and lastly
+        // storage server
         if let Ok(Some(frame)) = self
             .local_cache
             .get_page(self.conn_id.as_str(), u32::from(page_no))
@@ -258,20 +233,23 @@ impl Wal for DurableWal {
             buffer.copy_from_slice(&frame);
             return Ok(());
         }
-        // TODO: this call is unnecessary since `read_frame` is always called after `find_frame`
-        let frame_no =
-            tokio::task::block_in_place(|| rt.block_on(self.find_frame_by_page_no(page_no)))
-                .unwrap()
-                .unwrap();
         // check if the frame exists in the local cache
-        if let Ok(Some(frame)) = self.local_cache.get_frame(frame_no.into()) {
+        if let Ok(Some(frame)) = self
+            .local_cache
+            .get_frame_by_page(u32::from(page_no), self.max_frame_no)
+        {
             trace!(
                 "DurableWal::read_frame(page_no: {:?}) -- read cache hit",
                 page_no
             );
             buffer.copy_from_slice(&frame);
             return Ok(());
         }
+        let rt = tokio::runtime::Handle::current();
+        let frame_no =
+            tokio::task::block_in_place(|| rt.block_on(self.find_frame_by_page_no(page_no)))
+                .unwrap()
+                .unwrap();
         let req = rpc::ReadFrameRequest {
             namespace: self.namespace.to_string(),
             frame_no: frame_no.get(),
@@ -281,17 +259,17 @@ impl Wal for DurableWal {
         let resp = tokio::task::block_in_place(|| rt.block_on(resp)).unwrap();
         let frame = resp.into_inner().frame.unwrap();
         buffer.copy_from_slice(&frame);
-        let _ = self.local_cache.insert_frame(frame_no.into(), &frame);
+        let _ = self
+            .local_cache
+            .insert_frame(frame_no.into(), u32::from(page_no), &frame);
         Ok(())
     }
 
     fn db_size(&self) -> u32 {
-        let rt = tokio::runtime::Handle::current();
-        let size = tokio::task::block_in_place(|| rt.block_on(self.frames_count()))
-            .try_into()
-            .unwrap();
+        let size = self.local_cache.get_max_frame_num().unwrap();
         trace!("DurableWal::db_size() => {}", size);
-        size
+        // TODO: serve the db size from the meta table
+        size as u32
     }
 
     fn begin_write_txn(&mut self) -> Result<()> {
@@ -379,7 +357,7 @@ impl Wal for DurableWal {
 
         let req = rpc::InsertFramesRequest {
             namespace: self.namespace.to_string(),
-            frames,
+            frames: frames.clone(),
             max_frame_no: self.max_frame_no,
         };
         let mut binding = self.client.clone();
@@ -392,6 +370,10 @@ impl Wal for DurableWal {
             }
             return Err(rusqlite::ffi::Error::new(SQLITE_ABORT));
         }
+        // TODO: fix parity with storage server frame num with local cache
+        self.local_cache
+            .insert_frames(self.max_frame_no, frames)
+            .unwrap();
         Ok(resp.unwrap().into_inner().num_frames as usize)
     }
 
diff --git a/libsql-storage/src/local_cache.rs b/libsql-storage/src/local_cache.rs
@@ -1,6 +1,7 @@
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
+use crate::rpc::Frame;
 use libsql_sys::rusqlite::{ffi, params, Connection, Error, Result};
 
 /// We use LocalCache to cache frames and transaction state. Each namespace gets its own cache
@@ -36,10 +37,15 @@ impl LocalCache {
         self.conn.execute(
             "CREATE TABLE IF NOT EXISTS frames (
                 frame_no INTEGER PRIMARY KEY NOT NULL,
+                page_no INTEGER NOT NULL,
                 data BLOB NOT NULL
             )",
             [],
         )?;
+        self.conn.execute(
+            "CREATE INDEX IF NOT EXISTS idx_page_no_frame_no ON frames (page_no, frame_no)",
+            [],
+        )?;
 
         self.conn.execute(
             "CREATE TABLE IF NOT EXISTS transactions (
@@ -53,10 +59,10 @@ impl LocalCache {
         Ok(())
     }
 
-    pub fn insert_frame(&self, frame_no: u64, frame_data: &[u8]) -> Result<()> {
+    pub fn insert_frame(&self, frame_no: u64, page_no: u32, frame_data: &[u8]) -> Result<()> {
         match self.conn.execute(
-            "INSERT INTO frames (frame_no, data) VALUES (?1, ?2)",
-            params![frame_no, frame_data],
+            "INSERT INTO frames (frame_no, page_no, data) VALUES (?1, ?2, ?3)",
+            params![frame_no, page_no, frame_data],
         ) {
             Ok(_) => Ok(()),
             Err(Error::SqliteFailure(e, _)) if e.code == ffi::ErrorCode::ConstraintViolation => {
@@ -66,17 +72,45 @@ impl LocalCache {
         }
     }
 
-    pub fn get_frame(&self, frame_no: u64) -> Result<Option<Vec<u8>>> {
-        let mut stmt = self
-            .conn
-            .prepare("SELECT data FROM frames WHERE frame_no = ?1")?;
-        match stmt.query_row(params![frame_no], |row| row.get(0)) {
+    pub fn insert_frames(&mut self, frame_no: u64, frames: Vec<Frame>) -> Result<()> {
+        let tx = self.conn.transaction().unwrap();
+        {
+            let mut stmt =
+                tx.prepare("INSERT INTO frames (frame_no, page_no, data) VALUES (?1, ?2, ?3)")?;
+            let mut frame_no = frame_no;
+            for f in frames {
+                frame_no += 1;
+                stmt.execute(params![frame_no, f.page_no, f.data]).unwrap();
+            }
+        }
+        tx.commit().unwrap();
+        Ok(())
+    }
+
+    pub fn get_frame_by_page(&self, page_no: u32, max_frame_no: u64) -> Result<Option<Vec<u8>>> {
+        let mut stmt = self.conn.prepare(
+            "SELECT data FROM frames WHERE page_no=?1 AND frame_no <= ?2
+            ORDER BY frame_no DESC LIMIT 1",
+        )?;
+        match stmt.query_row(params![page_no, max_frame_no], |row| row.get(0)) {
             Ok(frame_data) => Ok(Some(frame_data)),
             Err(Error::QueryReturnedNoRows) => Ok(None),
             Err(e) => Err(e),
         }
     }
 
+    pub fn get_max_frame_num(&self) -> Result<u64> {
+        match self
+            .conn
+            .query_row("SELECT MAX(frame_no) from frames", (), |row| {
+                row.get::<_, Option<u64>>(0)
+            }) {
+            Ok(Some(frame_no)) => Ok(frame_no),
+            Ok(None) | Err(Error::QueryReturnedNoRows) => Ok(0),
+            Err(e) => Err(e),
+        }
+    }
+
     pub fn insert_page(&self, txn_id: &str, page_no: u32, frame_data: &[u8]) -> Result<()> {
         self.conn.execute(
             "INSERT INTO transactions (txn_id, page_no, data) VALUES (?1, ?2, ?3)
