feat: schema cache to prevent thundering herd

integration/go/go_pgx/pg_tests_test.go

@@ -15,8 +15,16 @@ import (
 )
 
 func assertShowField(t *testing.T, query string, field string, value int64, user string, database string, shard int64, role string) {
-	expected_value := value
+	actualValue := showField(t, query, field, user, database, shard, role)
+	assert.Equal(t, value, actualValue, fmt.Sprintf("\"%s\" in %s is not %d", field, query, value))
+}
+
+func assertShowFieldOneOf(t *testing.T, query string, field string, values []int64, user string, database string, shard int64, role string) {
+	actualValue := showField(t, query, field, user, database, shard, role)
+	assert.Contains(t, values, actualValue, fmt.Sprintf("\"%s\" in %s is not one of %v", field, query, values))
+}
 
+func showField(t *testing.T, query string, field string, user string, database string, shard int64, role string) int64 {
 	conn, err := pgx.Connect(context.Background(), "postgres://admin:pgdog@127.0.0.1:6432/admin")
 	if err != nil {
 		panic(err)
@@ -41,10 +49,9 @@ func assertShowField(t *testing.T, query string, field string, value int64, user
 		if row_db == database && row_user == user && row_shard.Int.Int64() == shard && row_role == role {
 			for i, description := range rows.FieldDescriptions() {
 				if description.Name == field {
-					actual_value, err := values[i].(pgtype.Numeric).Int64Value()
+					actualValue, err := values[i].(pgtype.Numeric).Int64Value()
 					assert.NoError(t, err)
-					assert.Equal(t, expected_value, actual_value.Int64, fmt.Sprintf("\"%s\" in %s is not %d", field, query, value))
-					return
+					return actualValue.Int64
 				}
 			}
 		}

integration/go/go_pgx/sharded_test.go

@@ -119,7 +119,9 @@ func rangeCases() []shardCase {
 }
 
 // TestShardedListDeprecated tests list mapping via the legacy [[sharded_mappings]] config format.
-func TestShardedListDeprecated(t *testing.T) { runShardingTest(t, "sharded_list_deprecated", listCases()) }
+func TestShardedListDeprecated(t *testing.T) {
+	runShardingTest(t, "sharded_list_deprecated", listCases())
+}
 
 // TestShardedList tests list mapping via the inline mapping = [...] config format.
 func TestShardedList(t *testing.T) { runShardingTest(t, "sharded_list", listCases()) }
@@ -228,11 +230,13 @@ func TestShardedTwoPc(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
-	// +5 is for schema sync
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 0, "primary")
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 1, "primary")
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+5, "pgdog_2pc", "pgdog_sharded", 0, "primary") // PREPARE, COMMIT for each transaction + TRUNCATE
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+5, "pgdog_2pc", "pgdog_sharded", 1, "primary")
+	// The schema cache is shared across users. If pgdog_2pc wins the cache-fill
+	// race after RELOAD, the five schema queries are counted in its pool stats.
+	expectedXactCounts := []int64{401, 406}
+	assertShowFieldOneOf(t, "SHOW STATS", "total_xact_count", expectedXactCounts, "pgdog_2pc", "pgdog_sharded", 0, "primary")
+	assertShowFieldOneOf(t, "SHOW STATS", "total_xact_count", expectedXactCounts, "pgdog_2pc", "pgdog_sharded", 1, "primary")
 
 	for i := range 200 {
 		rows, err := conn.Query(

pgdog/src/backend/databases.rs

@@ -19,6 +19,7 @@ use tracing::{debug, error, info, warn};
 
 use crate::auth::AuthResult;
 use crate::backend::replication::ShardedSchemas;
+use crate::backend::schema::SchemaCache;
 use crate::config::PoolerMode;
 use crate::frontend::PreparedStatements;
 use crate::frontend::client::query_engine::two_pc::Manager;
@@ -57,6 +58,9 @@ pub fn databases() -> Arc<Databases> {
 
 /// Replace databases pooler-wide.
 pub fn replace_databases(new_databases: Databases, reload: bool) -> Result<(), Error> {
+    // Bust the schema cache before launching pools.
+    SchemaCache::global().clear();
+
     // Order of operations is important
     // to ensure zero downtime for clients.
     //

pgdog/src/backend/pool/shard/mod.rs

@@ -14,6 +14,7 @@ use crate::backend::Schema;
 use crate::backend::databases::User;
 use crate::backend::pool::lb::ban::Ban;
 use crate::backend::pub_sub::listener::Listener;
+use crate::backend::schema::SchemaCache;
 use crate::config::{LoadBalancingStrategy, ReadWriteSplit, Role};
 use crate::net::Parameters;
 use crate::net::messages::FrontendPid;
@@ -37,7 +38,7 @@ pub(super) struct ShardConfig<'a> {
     pub(super) lb_strategy: LoadBalancingStrategy,
     /// Primary/replica read/write split strategy.
     pub(super) rw_split: ReadWriteSplit,
-    /// Cluster identifier (user/password).
+    /// Cluster identifier (user/database).
     pub(super) identifier: Arc<User>,
     /// LSN check interval
     pub(super) lsn_check_interval: Duration,
@@ -120,23 +121,41 @@ impl Shard {
         }
     }
 
-    /// Load schema from the shard's primary.
-    pub async fn load_schema(&self) -> Result<bool, crate::backend::Error> {
+    /// Load schema from the shard's primary database.
+    ///
+    /// Uses the global schema cache, so most requests will not actually touch
+    /// the database.
+    ///
+    pub(crate) async fn load_schema(&self) -> Result<bool, crate::backend::Error> {
         if self.schema.initialized() {
             return Ok(false);
         }
 
+        // This is syncrhonized by database/shard number, so this prevents
+        // a thundering herd with 100s of users, for example, all fetching
+        // the same schema.
+        let schema = SchemaCache::global().get(self).await?;
+        let _ = self.schema.set(schema);
+        self.schema_waiter.notify_one();
+
+        Ok(true)
+    }
+
+    /// Fetch schema from the shard. This does not use the
+    /// cache and returns the freshed schema available.
+    ///
+    pub(crate) async fn fetch_schema(&self) -> Result<Schema, crate::backend::Error> {
         let mut server = self.primary_or_replica(&Request::default()).await?;
         let schema = Schema::load(&mut server).await?;
+
         info!(
             "loaded schema for {} tables on shard {} [{}]",
             schema.tables().len(),
             self.number(),
             server.addr()
         );
-        let _ = self.schema.set(schema);
-        self.schema_waiter.notify_one();
-        Ok(true)
+
+        Ok(schema)
     }
 
     /// Set the schema to its default value.

pgdog/src/backend/schema/cache/mod.rs

@@ -0,0 +1,63 @@
+//! Cache the schema per database, so we don't have to fetch
+//! it for each [`crate::backend::pool::Cluster`].
+
+use dashmap::DashMap;
+use once_cell::sync::Lazy;
+use std::sync::Arc;
+use tokio::sync::Mutex;
+
+use crate::backend::{Schema, Shard};
+
+type Entry = Arc<Mutex<Schema>>;
+
+static CACHE: Lazy<SchemaCache> = Lazy::new(SchemaCache::default);
+
+/// Schema cache.
+#[derive(Debug, Default)]
+pub(crate) struct SchemaCache {
+    // Database => shard => Schema
+    cache: DashMap<String, DashMap<usize, Entry>>,
+}
+
+impl SchemaCache {
+    /// Get a schema entry from the cache or load it from
+    /// the server and store it in the cache.
+    ///
+    /// The loading is synchronized with a mutex, so only one user
+    /// can load a schema at a time, preventing a thundering herd situtation.
+    ///
+    pub(crate) async fn get(&self, shard: &Shard) -> Result<Schema, super::Error> {
+        // This is synchronized.
+        let entry = self
+            .cache
+            .entry(shard.identifier().database.clone())
+            .or_default()
+            .entry(shard.number())
+            .or_default()
+            .clone();
+
+        // This is syncrhonized too,
+        // so only one shard/user can fetch the schema at a time.
+        let mut guard = entry.lock().await;
+
+        if guard.is_loaded() {
+            return Ok(guard.clone());
+        }
+
+        let schema = shard.fetch_schema().await?;
+
+        *guard = schema.clone();
+
+        Ok(schema)
+    }
+
+    /// Remove all entries from the schema cache.
+    pub(crate) fn clear(&self) {
+        self.cache.clear();
+    }
+
+    /// Get a reference the the global schema cache.
+    pub(crate) fn global() -> &'static SchemaCache {
+        &CACHE
+    }
+}

pgdog/src/backend/schema/mod.rs

@@ -1,4 +1,5 @@
 //! Schema operations.
+pub mod cache;
 pub mod columns;
 pub mod relation;
 pub mod sync;
@@ -11,6 +12,7 @@ use std::ops::DerefMut;
 use std::{collections::HashMap, ops::Deref};
 use tracing::info;
 
+pub(crate) use cache::SchemaCache;
 pub use relation::Relation;
 
 use super::{Cluster, Error, Server, pool::Request};