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;
@@ -553,7 +554,20 @@ fn resolve_table_mapping_deprecated(
     )
 }
 
-fn new_pool(user: &crate::config::User, config: &crate::config::Config) -> Option<(User, Cluster)> {
+// Create new Cluster from user and databases in `pgdog.toml`.
+//
+// # Arguments
+//
+// - `user`: `[[users]]` entry in `users.toml`
+// - `config`: all of `pgdog.toml`
+// - `schema_cache`: A cache of database tables, shared between all clusters. This is passed here
+//                   to ensure all clusters share the same schema cache, and to make sure a new one
+//                   is created on each config reload.
+fn new_pool(
+    user: &crate::config::User,
+    config: &crate::config::Config,
+    schema_cache: SchemaCache,
+) -> Option<(User, Cluster)> {
     let omnisharded_tables = config.omnisharded_tables();
     let sharded_mappings = config.sharded_mappings();
     let sharded_schemas = config.sharded_schemas();
@@ -624,6 +638,7 @@ fn new_pool(user: &crate::config::User, config: &crate::config::Config) -> Optio
         sharded_tables,
         sharded_schemas,
         query_parser,
+        schema_cache,
     );
 
     Some((
@@ -638,6 +653,8 @@ fn new_pool(user: &crate::config::User, config: &crate::config::Config) -> Optio
 /// Load databases from config.
 pub fn from_config(config: &ConfigAndUsers) -> Databases {
     let mut databases = HashMap::new();
+    // The schema cache is shared between all databases.
+    let schema_cache = SchemaCache::default();
 
     for user in &config.users.users {
         let users = if user.databases.is_empty() && !user.all_databases {
@@ -673,7 +690,7 @@ pub fn from_config(config: &ConfigAndUsers) -> Databases {
         };
 
         for user in users {
-            if let Some((user, cluster)) = new_pool(&user, &config.config) {
+            if let Some((user, cluster)) = new_pool(&user, &config.config, schema_cache.clone()) {
                 databases.insert(user, cluster);
             }
         }
@@ -1696,7 +1713,7 @@ mod tests {
             ..Default::default()
         };
 
-        let result = new_pool(&user, &config);
+        let result = new_pool(&user, &config, SchemaCache::default());
         assert!(
             result.is_none(),
             "new_pool should return None when database doesn't exist"

pgdog/src/backend/pool/cluster.rs

@@ -8,6 +8,7 @@ use pgdog_config::{
 };
 use std::{sync::Arc, time::Duration};
 
+use crate::backend::schema::SchemaCache;
 use crate::frontend::router::sharding::ShardedTable;
 use crate::{
     backend::{
@@ -126,54 +127,60 @@ impl ClusterShardConfig {
 /// Cluster creation config.
 #[derive(Debug)]
 pub struct ClusterConfig<'a> {
-    pub name: &'a str,
-    pub shards: &'a [ClusterShardConfig],
-    pub lb_strategy: LoadBalancingStrategy,
-    pub user: &'a str,
-    pub passwords: Vec<PasswordKind>,
-    pub pooler_mode: PoolerMode,
-    pub sharded_tables: ShardedTables,
-    pub replication_sharding: Option<String>,
-    pub multi_tenant: &'a Option<MultiTenant>,
-    pub rw_strategy: ReadWriteStrategy,
-    pub rw_split: ReadWriteSplit,
-    pub schema_admin: bool,
-    pub cross_shard_disabled: bool,
-    pub two_pc: bool,
-    pub two_pc_auto: bool,
-    pub sharded_schemas: ShardedSchemas,
-    pub rewrite: &'a Rewrite,
-    pub prepared_statements: &'a PreparedStatements,
-    pub dry_run: bool,
-    pub expanded_explain: bool,
-    pub pub_sub_channel_size: usize,
-    pub query_parser: QueryParserLevel,
-    pub query_parser_engine: QueryParserEngine,
-    pub log_min_duration_parse: Option<Duration>,
-    pub log_query_sample_length: usize,
-    pub connection_recovery: ConnectionRecovery,
-    pub client_connection_recovery: ConnectionRecovery,
-    pub lsn_check_interval: Duration,
-    pub reload_schema_on_ddl: bool,
-    pub load_schema: LoadSchema,
-    pub resharding_parallel_copies: usize,
-    pub resharding_copy_retry_max_attempts: usize,
-    pub resharding_copy_retry_min_delay: u64,
-    pub resharding_replication_retry_max_attempts: usize,
-    pub resharding_replication_retry_min_delay: u64,
-    pub regex_parser_limit: usize,
-    pub pub_sub_enabled: bool,
-    pub identity: &'a Option<String>,
+    name: &'a str,
+    shards: &'a [ClusterShardConfig],
+    lb_strategy: LoadBalancingStrategy,
+    user: &'a str,
+    passwords: Vec<PasswordKind>,
+    pooler_mode: PoolerMode,
+    sharded_tables: ShardedTables,
+    replication_sharding: Option<String>,
+    multi_tenant: &'a Option<MultiTenant>,
+    rw_strategy: ReadWriteStrategy,
+    rw_split: ReadWriteSplit,
+    schema_admin: bool,
+    cross_shard_disabled: bool,
+    two_pc: bool,
+    two_pc_auto: bool,
+    sharded_schemas: ShardedSchemas,
+    rewrite: &'a Rewrite,
+    prepared_statements: &'a PreparedStatements,
+    dry_run: bool,
+    expanded_explain: bool,
+    pub_sub_channel_size: usize,
+    query_parser: QueryParserLevel,
+    query_parser_engine: QueryParserEngine,
+    log_min_duration_parse: Option<Duration>,
+    log_query_sample_length: usize,
+    connection_recovery: ConnectionRecovery,
+    client_connection_recovery: ConnectionRecovery,
+    lsn_check_interval: Duration,
+    reload_schema_on_ddl: bool,
+    load_schema: LoadSchema,
+    resharding_parallel_copies: usize,
+    resharding_copy_retry_max_attempts: usize,
+    resharding_copy_retry_min_delay: u64,
+    resharding_replication_retry_max_attempts: usize,
+    resharding_replication_retry_min_delay: u64,
+    regex_parser_limit: usize,
+    pub_sub_enabled: bool,
+    identity: &'a Option<String>,
+    schema_cache: SchemaCache,
 }
 
 impl<'a> ClusterConfig<'a> {
+    /// Dependencies for creating a Cluster.
+    ///
+    /// TODO(lev): This is getting unruly. We may need a struct for a struct :)
+    ///
     pub(crate) fn new(
         config: &'a crate::config::Config,
         user: &'a User,
         shards: &'a [ClusterShardConfig],
         sharded_tables: ShardedTables,
         sharded_schemas: ShardedSchemas,
         query_parser: QueryParser,
+        schema_cache: SchemaCache,
     ) -> Self {
         let general = &config.general;
         let multi_tenant = config.multi_tenant();
@@ -228,6 +235,7 @@ impl<'a> ClusterConfig<'a> {
             regex_parser_limit: general.regex_parser_limit,
             pub_sub_enabled: general.pub_sub_enabled(),
             identity: &user.identity,
+            schema_cache,
         }
     }
 }
@@ -274,6 +282,7 @@ impl Cluster {
             regex_parser_limit,
             pub_sub_enabled,
             identity,
+            schema_cache,
         } = config;
 
         let identifier = Arc::new(DatabaseUser {
@@ -296,6 +305,7 @@ impl Cluster {
                         identifier: identifier.clone(),
                         lsn_check_interval,
                         pub_sub_enabled,
+                        schema_cache: schema_cache.clone(),
                     })
                 })
                 .collect(),
@@ -661,6 +671,7 @@ mod test {
         ConfigAndUsers, OmnishardedTable, PoolerMode, QueryParserLevel, ShardedSchema,
     };
 
+    use crate::backend::schema::SchemaCache;
     use crate::frontend::router::sharding::ShardedTable;
     use crate::{
         backend::{
@@ -707,6 +718,7 @@ mod test {
                         identifier: identifier.clone(),
                         lsn_check_interval: Duration::MAX,
                         pub_sub_enabled: false,
+                        schema_cache: SchemaCache::default(),
                     })
                 })
                 .collect::<Vec<_>>();
@@ -834,6 +846,7 @@ mod test {
                 identifier: cluster.identifier.clone(),
                 lsn_check_interval: Duration::MAX,
                 pub_sub_enabled: false,
+                schema_cache: SchemaCache::default(),
             });
             cluster
         }
@@ -857,6 +870,7 @@ mod test {
                     identifier: identifier.clone(),
                     lsn_check_interval: Duration::default(),
                     pub_sub_enabled: false,
+                    schema_cache: SchemaCache::default(),
                 })],
                 prepared_statements: config.config.general.prepared_statements,
                 dry_run: config.config.general.dry_run,
@@ -891,6 +905,7 @@ mod test {
                 identifier,
                 lsn_check_interval: Duration::default(),
                 pub_sub_enabled: false,
+                schema_cache: SchemaCache::default(),
             });
 
             cluster