From eba82e3fc12cf503c25aa2bbe0fb3ae51dd4cdc2 Mon Sep 17 00:00:00 2001 From: Julius de Bruijn Date: Wed, 24 Jun 2020 19:43:51 +0200 Subject: [PATCH] LRU caching for SQLite --- sqlx-core/src/common/statement_cache.rs | 19 ++++++++----- sqlx-core/src/mysql/connection/mod.rs | 2 +- sqlx-core/src/postgres/connection/mod.rs | 2 +- sqlx-core/src/sqlite/connection/establish.rs | 8 ++++-- sqlx-core/src/sqlite/connection/executor.rs | 8 ++++-- sqlx-core/src/sqlite/connection/mod.rs | 17 +++++++++++- sqlx-core/src/sqlite/options.rs | 29 ++++++++++++++++---- tests/sqlite/sqlite.rs | 25 ++++++++++++++++- 8 files changed, 89 insertions(+), 21 deletions(-) diff --git a/sqlx-core/src/common/statement_cache.rs b/sqlx-core/src/common/statement_cache.rs index 7ba19a61..2740bde7 100644 --- a/sqlx-core/src/common/statement_cache.rs +++ b/sqlx-core/src/common/statement_cache.rs @@ -3,11 +3,11 @@ use lru_cache::LruCache; /// A cache for prepared statements. When full, the least recently used /// statement gets removed. #[derive(Debug)] -pub struct StatementCache { - inner: LruCache, +pub struct StatementCache { + inner: LruCache, } -impl StatementCache { +impl StatementCache { /// Create a new cache with the given capacity. pub fn new(capacity: usize) -> Self { Self { @@ -17,19 +17,19 @@ impl StatementCache { /// Returns a mutable reference to the value corresponding to the given key /// in the cache, if any. - pub fn get_mut(&mut self, k: &str) -> Option<&mut u32> { + pub fn get_mut(&mut self, k: &str) -> Option<&mut T> { self.inner.get_mut(k) } /// Inserts a new statement to the cache, returning the least recently used /// statement id if the cache is full, or if inserting with an existing key, /// the replaced existing statement. - pub fn insert(&mut self, k: &str, v: u32) -> Option { + pub fn insert(&mut self, k: &str, v: T) -> Option { let mut lru_item = None; if self.inner.capacity() == self.len() && !self.inner.contains_key(k) { lru_item = self.remove_lru(); - } else if self.inner.contains_key(k) { + } else if self.contains_key(k) { lru_item = self.inner.remove(k); } @@ -44,7 +44,7 @@ impl StatementCache { } /// Removes the least recently used item from the cache. - pub fn remove_lru(&mut self) -> Option { + pub fn remove_lru(&mut self) -> Option { self.inner.remove_lru().map(|(_, v)| v) } @@ -52,4 +52,9 @@ impl StatementCache { pub fn clear(&mut self) { self.inner.clear(); } + + /// True if cache has a value for the given key. + pub fn contains_key(&mut self, k: &str) -> bool { + self.inner.contains_key(k) + } } diff --git a/sqlx-core/src/mysql/connection/mod.rs b/sqlx-core/src/mysql/connection/mod.rs index 0711eae4..028c1a90 100644 --- a/sqlx-core/src/mysql/connection/mod.rs +++ b/sqlx-core/src/mysql/connection/mod.rs @@ -37,7 +37,7 @@ pub struct MySqlConnection { pub(crate) stream: MySqlStream, // cache by query string to the statement id - cache_statement: StatementCache, + cache_statement: StatementCache, // working memory for the active row's column information // this allows us to re-use these allocations unless the user is persisting the diff --git a/sqlx-core/src/postgres/connection/mod.rs b/sqlx-core/src/postgres/connection/mod.rs index 987769df..ef1e437b 100644 --- a/sqlx-core/src/postgres/connection/mod.rs +++ b/sqlx-core/src/postgres/connection/mod.rs @@ -48,7 +48,7 @@ pub struct PgConnection { next_statement_id: u32, // cache statement by query string to the id and columns - cache_statement: StatementCache, + cache_statement: StatementCache, // cache user-defined types by id <-> info cache_type_info: HashMap, diff --git a/sqlx-core/src/sqlite/connection/establish.rs b/sqlx-core/src/sqlite/connection/establish.rs index dbebd165..fe2b9d8a 100644 --- a/sqlx-core/src/sqlite/connection/establish.rs +++ b/sqlx-core/src/sqlite/connection/establish.rs @@ -1,7 +1,6 @@ use std::io; use std::ptr::{null, null_mut}; -use hashbrown::HashMap; use libsqlite3_sys::{ sqlite3_busy_timeout, sqlite3_extended_result_codes, sqlite3_open_v2, SQLITE_OK, SQLITE_OPEN_CREATE, SQLITE_OPEN_MEMORY, SQLITE_OPEN_NOMUTEX, SQLITE_OPEN_PRIVATECACHE, @@ -12,7 +11,10 @@ use sqlx_rt::blocking; use crate::error::Error; use crate::sqlite::connection::handle::ConnectionHandle; use crate::sqlite::statement::StatementWorker; -use crate::sqlite::{SqliteConnectOptions, SqliteConnection, SqliteError}; +use crate::{ + common::StatementCache, + sqlite::{SqliteConnectOptions, SqliteConnection, SqliteError}, +}; pub(super) async fn establish(options: &SqliteConnectOptions) -> Result { let mut filename = options @@ -87,7 +89,7 @@ pub(super) async fn establish(options: &SqliteConnectOptions) -> Result( conn: &mut ConnectionHandle, - statements: &'a mut HashMap, + statements: &'a mut StatementCache, statement: &'a mut Option, query: &str, persistent: bool, @@ -28,7 +29,10 @@ fn prepare<'a>( if !statements.contains_key(query) { let statement = SqliteStatement::prepare(conn, query, false)?; - statements.insert(query.to_owned(), statement); + + if let Some(mut statement) = statements.insert(query, statement) { + statement.reset(); + } } let statement = statements.get_mut(query).unwrap(); diff --git a/sqlx-core/src/sqlite/connection/mod.rs b/sqlx-core/src/sqlite/connection/mod.rs index d70c1c89..69042d83 100644 --- a/sqlx-core/src/sqlite/connection/mod.rs +++ b/sqlx-core/src/sqlite/connection/mod.rs @@ -6,6 +6,8 @@ use futures_util::future; use hashbrown::HashMap; use libsqlite3_sys::sqlite3; +use crate::caching_connection::CachingConnection; +use crate::common::StatementCache; use crate::connection::{Connect, Connection}; use crate::error::Error; use crate::ext::ustr::UStr; @@ -25,7 +27,7 @@ pub struct SqliteConnection { pub(crate) worker: StatementWorker, // cache of semi-persistent statements - pub(crate) statements: HashMap, + pub(crate) statements: StatementCache, // most recent non-persistent statement pub(crate) statement: Option, @@ -47,6 +49,19 @@ impl Debug for SqliteConnection { } } +impl CachingConnection for SqliteConnection { + fn cached_statements_count(&self) -> usize { + self.statements.len() + } + + fn clear_cached_statements(&mut self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async move { + self.statements.clear(); + Ok(()) + }) + } +} + impl Connection for SqliteConnection { type Database = Sqlite; diff --git a/sqlx-core/src/sqlite/options.rs b/sqlx-core/src/sqlite/options.rs index 5f925b6b..6408dd68 100644 --- a/sqlx-core/src/sqlite/options.rs +++ b/sqlx-core/src/sqlite/options.rs @@ -1,5 +1,5 @@ use std::path::PathBuf; -use std::str::FromStr; +use std::{io, str::FromStr}; use crate::error::BoxDynError; @@ -10,6 +10,7 @@ use crate::error::BoxDynError; pub struct SqliteConnectOptions { pub(crate) filename: PathBuf, pub(crate) in_memory: bool, + pub(crate) statement_cache_size: usize, } impl Default for SqliteConnectOptions { @@ -23,6 +24,7 @@ impl SqliteConnectOptions { Self { filename: PathBuf::from(":memory:"), in_memory: false, + statement_cache_size: 100, } } } @@ -34,6 +36,7 @@ impl FromStr for SqliteConnectOptions { let mut options = Self { filename: PathBuf::new(), in_memory: false, + statement_cache_size: 100, }; // remove scheme @@ -41,10 +44,26 @@ impl FromStr for SqliteConnectOptions { .trim_start_matches("sqlite://") .trim_start_matches("sqlite:"); - if s == ":memory:" { - options.in_memory = true; - } else { - options.filename = s.parse()?; + let mut splitted = s.split("?"); + + match splitted.next() { + Some(":memory:") => options.in_memory = true, + Some(s) => options.filename = s.parse()?, + None => unreachable!(), + } + + match splitted.next().map(|s| s.split("=")) { + Some(mut splitted) => { + if splitted.next() == Some("statement-cache-size") { + options.statement_cache_size = splitted + .next() + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "Invalid connection string") + })? + .parse()? + } + } + _ => (), } Ok(options) diff --git a/tests/sqlite/sqlite.rs b/tests/sqlite/sqlite.rs index 1630773b..6935f04f 100644 --- a/tests/sqlite/sqlite.rs +++ b/tests/sqlite/sqlite.rs @@ -1,6 +1,7 @@ use futures::TryStreamExt; use sqlx::{ - query, sqlite::Sqlite, Connect, Connection, Executor, Row, SqliteConnection, SqlitePool, + query, sqlite::Sqlite, CachingConnection, Connect, Connection, Executor, Row, SqliteConnection, + SqlitePool, }; use sqlx_test::new; @@ -269,3 +270,25 @@ SELECT id, text FROM _sqlx_test; Ok(()) } + +#[sqlx_macros::test] +async fn it_caches_statements() -> anyhow::Result<()> { + let mut conn = new::().await?; + + for i in 0..2 { + let row = sqlx::query("SELECT ? AS val") + .bind(i) + .fetch_one(&mut conn) + .await?; + + let val: i32 = row.get("val"); + + assert_eq!(i, val); + } + + assert_eq!(1, conn.cached_statements_count()); + conn.clear_cached_statements().await?; + assert_eq!(0, conn.cached_statements_count()); + + Ok(()) +}