Merge pull request #4161 from fraillt/choose-caching-strategy

Selecting custom caching strategy

CONTRIBUTING.md

@@ -133,7 +133,7 @@ To run rustfmt tests locally:
    rustup component add clippy
    ```
 
-3. Install [typos](https://github.com/crate-ci/typos) via `cargo install typos`
+3. Install [typos](https://github.com/crate-ci/typos) via `cargo install typos-cli`
 
 4. Use `cargo xtask tidy` to check if your changes follow the expected code style.
    This will run `cargo fmt --check`, `typos` and `cargo clippy` internally. See `cargo xtask tidy --help`

diesel/src/connection/mod.rs

@@ -45,6 +45,16 @@ pub use self::instrumentation::{DynInstrumentation, StrQueryHelper};
 ))]
 pub(crate) use self::private::MultiConnectionHelper;
 
+/// Set cache size for a connection
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[non_exhaustive]
+pub enum CacheSize {
+    /// Caches all queries if possible
+    Unbounded,
+    /// Disable statement cache
+    Disabled,
+}
+
 /// Perform simple operations on a backend.
 ///
 /// You should likely use [`Connection`] instead.
@@ -401,6 +411,9 @@ where
 
     /// Set a specific [`Instrumentation`] implementation for this connection
     fn set_instrumentation(&mut self, instrumentation: impl Instrumentation);
+
+    /// Set the prepared statement cache size to [`CacheSize`] for this connection
+    fn set_prepared_statement_cache_size(&mut self, size: CacheSize);
 }
 
 /// The specific part of a [`Connection`] which actually loads data from the database

diesel/src/connection/statement_cache.rs → diesel/src/connection/statement_cache/mod.rs

@@ -10,8 +10,9 @@
 //! statements is [`SimpleConnection::batch_execute`](super::SimpleConnection::batch_execute).
 //!
 //! In order to avoid the cost of re-parsing and planning subsequent queries,
-//! Diesel caches the prepared statement whenever possible. Queries will fall
-//! into one of three buckets:
+//! by default Diesel caches the prepared statement whenever possible, but
+//! this an be customized by calling [`Connection::set_cache_size`](super::Connection::set_cache_size).
+//! Queries will fall into one of three buckets:
 //!
 //! - Unsafe to cache
 //! - Cached by SQL
@@ -94,16 +95,21 @@
 
 use std::any::TypeId;
 use std::borrow::Cow;
-use std::collections::HashMap;
 use std::hash::Hash;
 use std::ops::{Deref, DerefMut};
 
+use strategy::{StatementCacheStrategy, WithCacheStrategy, WithoutCacheStrategy};
+
 use crate::backend::Backend;
 use crate::connection::InstrumentationEvent;
 use crate::query_builder::*;
 use crate::result::QueryResult;
 
-use super::Instrumentation;
+use super::{CacheSize, Instrumentation};
+
+/// Various interfaces and implementations to control connection statement caching.
+#[allow(unreachable_pub)]
+pub mod strategy;
 
 /// A prepared statement cache
 #[allow(missing_debug_implementations, unreachable_pub)]
@@ -112,7 +118,10 @@ use super::Instrumentation;
     doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
 )]
 pub struct StatementCache<DB: Backend, Statement> {
-    pub(crate) cache: HashMap<StatementCacheKey<DB>, Statement>,
+    cache: Box<dyn StatementCacheStrategy<DB, Statement>>,
+    // increment every time a query is cached
+    // some backends might use it to create unique prepared statement names
+    cache_counter: u64,
 }
 
 /// A helper type that indicates if a certain query
@@ -128,45 +137,51 @@ pub struct StatementCache<DB: Backend, Statement> {
 )]
 #[allow(unreachable_pub)]
 pub enum PrepareForCache {
-    /// The statement will be cached
-    Yes,
+    /// The statement will be cached    
+    Yes {
+        /// Counter might be used as unique identifier for prepared statement.
+        #[allow(dead_code)]
+        counter: u64,
+    },
     /// The statement won't be cached
     No,
 }
 
-#[allow(
-    clippy::len_without_is_empty,
-    clippy::new_without_default,
-    unreachable_pub
-)]
+#[allow(clippy::new_without_default, unreachable_pub)]
 impl<DB, Statement> StatementCache<DB, Statement>
 where
-    DB: Backend,
+    DB: Backend + 'static,
+    Statement: 'static,
     DB::TypeMetadata: Clone,
     DB::QueryBuilder: Default,
     StatementCacheKey<DB>: Hash + Eq,
 {
-    /// Create a new prepared statement cache
+    /// Create a new prepared statement cache using [`CacheSize::Unbounded`] as caching strategy.
     #[allow(unreachable_pub)]
     pub fn new() -> Self {
         StatementCache {
-            cache: HashMap::new(),
+            cache: Box::new(WithCacheStrategy::default()),
+            cache_counter: 0,
         }
     }
 
-    /// Get the current length of the statement cache
-    #[allow(unreachable_pub)]
-    #[cfg(any(
-        feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes",
-        feature = "postgres",
-        all(feature = "sqlite", test)
-    ))]
-    #[cfg_attr(
-        docsrs,
-        doc(cfg(feature = "i-implement-a-third-party-backend-and-opt-into-breaking-changes"))
-    )]
-    pub fn len(&self) -> usize {
-        self.cache.len()
+    /// Set caching strategy from predefined implementations
+    pub fn set_cache_size(&mut self, size: CacheSize) {
+        if self.cache.cache_size() != size {
+            self.cache = match size {
+                CacheSize::Unbounded => Box::new(WithCacheStrategy::default()),
+                CacheSize::Disabled => Box::new(WithoutCacheStrategy::default()),
+            }
+        }
+    }
+
+    /// Setting custom caching strategy. It is used in tests, to verify caching logic
+    #[allow(dead_code)]
+    pub(crate) fn set_strategy<Strategy>(&mut self, s: Strategy)
+    where
+        Strategy: StatementCacheStrategy<DB, Statement> + 'static,
+    {
+        self.cache = Box::new(s);
     }
 
     /// Prepare a query as prepared statement
@@ -193,50 +208,44 @@ where
         T: QueryFragment<DB> + QueryId,
         F: FnMut(&str, PrepareForCache) -> QueryResult<Statement>,
     {
-        self.cached_statement_non_generic(
+        Self::cached_statement_non_generic(
+            self.cache.as_mut(),
             T::query_id(),
             source,
             backend,
             bind_types,
-            &mut prepare_fn,
-            instrumentation,
+            &mut |sql, is_cached| {
+                if is_cached {
+                    instrumentation.on_connection_event(InstrumentationEvent::CacheQuery { sql });
+                    self.cache_counter += 1;
+                    prepare_fn(
+                        sql,
+                        PrepareForCache::Yes {
+                            counter: self.cache_counter,
+                        },
+                    )
+                } else {
+                    prepare_fn(sql, PrepareForCache::No)
+                }
+            },
         )
     }
 
     /// Reduce the amount of monomorphized code by factoring this via dynamic dispatch
-    fn cached_statement_non_generic(
-        &mut self,
+    fn cached_statement_non_generic<'a>(
+        cache: &'a mut dyn StatementCacheStrategy<DB, Statement>,
         maybe_type_id: Option<TypeId>,
         source: &dyn QueryFragmentForCachedStatement<DB>,
         backend: &DB,
         bind_types: &[DB::TypeMetadata],
-        prepare_fn: &mut dyn FnMut(&str, PrepareForCache) -> QueryResult<Statement>,
-        instrumentation: &mut dyn Instrumentation,
-    ) -> QueryResult<MaybeCached<'_, Statement>> {
-        use std::collections::hash_map::Entry::{Occupied, Vacant};
-
+        prepare_fn: &mut dyn FnMut(&str, bool) -> QueryResult<Statement>,
+    ) -> QueryResult<MaybeCached<'a, Statement>> {
         let cache_key = StatementCacheKey::for_source(maybe_type_id, source, bind_types, backend)?;
-
         if !source.is_safe_to_cache_prepared(backend)? {
             let sql = cache_key.sql(source, backend)?;
-            return prepare_fn(&sql, PrepareForCache::No).map(MaybeCached::CannotCache);
+            return prepare_fn(&sql, false).map(MaybeCached::CannotCache);
         }
-
-        let cached_result = match self.cache.entry(cache_key) {
-            Occupied(entry) => entry.into_mut(),
-            Vacant(entry) => {
-                let statement = {
-                    let sql = entry.key().sql(source, backend)?;
-                    instrumentation
-                        .on_connection_event(InstrumentationEvent::CacheQuery { sql: &sql });
-                    prepare_fn(&sql, PrepareForCache::Yes)
-                };
-
-                entry.insert(statement?)
-            }
-        };
-
-        Ok(MaybeCached::Cached(cached_result))
+        cache.get(cache_key, backend, source, prepare_fn)
     }
 }