Omni routing and change to convergence for schema sharding (#620)

* Omni routing and change to convergence for schema sharding

* save

* Fix test

* Add config test

* Rename

Author: levkk

integration/d_plus/pgdog.toml

@@ -23,3 +23,4 @@ tables = [
     "pg_namespace",
     "pg_am"
 ]
+sticky = true

pgdog-config/src/core.rs

@@ -5,7 +5,9 @@ use std::path::PathBuf;
 use tracing::{info, warn};
 
 use crate::sharding::ShardedSchema;
-use crate::{EnumeratedDatabase, Memory, PassthoughAuth, PreparedStatements, RewriteMode};
+use crate::{
+    EnumeratedDatabase, Memory, OmnishardedTable, PassthoughAuth, PreparedStatements, RewriteMode,
+};
 
 use super::database::Database;
 use super::error::Error;
@@ -227,15 +229,18 @@ impl Config {
         tables
     }
 
-    pub fn omnisharded_tables(&self) -> HashMap<String, Vec<String>> {
+    pub fn omnisharded_tables(&self) -> HashMap<String, Vec<OmnishardedTable>> {
         let mut tables = HashMap::new();
 
         for table in &self.omnisharded_tables {
             let entry = tables
                 .entry(table.database.clone())
                 .or_insert_with(Vec::new);
             for t in &table.tables {
-                entry.push(t.clone());
+                entry.push(OmnishardedTable {
+                    name: t.clone(),
+                    sticky_routing: table.sticky,
+                });
             }
         }
 
@@ -647,4 +652,58 @@ password = "users_admin_password"
         );
         assert!(config_and_users.users.admin.is_none());
     }
+
+    #[test]
+    fn test_omnisharded_tables() {
+        let source = r#"
+[general]
+host = "0.0.0.0"
+port = 6432
+
+[[databases]]
+name = "db1"
+host = "127.0.0.1"
+port = 5432
+
+[[databases]]
+name = "db2"
+host = "127.0.0.1"
+port = 5433
+
+[[omnisharded_tables]]
+database = "db1"
+tables = ["table_a", "table_b"]
+
+[[omnisharded_tables]]
+database = "db1"
+tables = ["table_c"]
+sticky = true
+
+[[omnisharded_tables]]
+database = "db2"
+tables = ["table_x"]
+"#;
+
+        let config: Config = toml::from_str(source).unwrap();
+
+        assert_eq!(config.omnisharded_tables.len(), 3);
+
+        let tables = config.omnisharded_tables();
+
+        assert_eq!(tables.len(), 2);
+
+        let db1_tables = tables.get("db1").unwrap();
+        assert_eq!(db1_tables.len(), 3);
+        assert_eq!(db1_tables[0].name, "table_a");
+        assert!(!db1_tables[0].sticky_routing);
+        assert_eq!(db1_tables[1].name, "table_b");
+        assert!(!db1_tables[1].sticky_routing);
+        assert_eq!(db1_tables[2].name, "table_c");
+        assert!(db1_tables[2].sticky_routing);
+
+        let db2_tables = tables.get("db2").unwrap();
+        assert_eq!(db2_tables.len(), 1);
+        assert_eq!(db2_tables[0].name, "table_x");
+        assert!(!db2_tables[0].sticky_routing);
+    }
 }

pgdog-config/src/sharding.rs

@@ -177,6 +177,14 @@ impl From<String> for FlexibleType {
 pub struct OmnishardedTables {
     pub database: String,
     pub tables: Vec<String>,
+    #[serde(default)]
+    pub sticky: bool,
+}
+
+#[derive(PartialEq, Debug, Clone, Default)]
+pub struct OmnishardedTable {
+    pub name: String,
+    pub sticky_routing: bool,
 }
 
 /// Queries with manual routing rules.

pgdog/src/backend/pool/cluster.rs

@@ -538,6 +538,8 @@ impl Cluster {
 mod test {
     use std::{sync::Arc, time::Duration};
 
+    use pgdog_config::OmnishardedTable;
+
     use crate::{
         backend::{
             pool::{Address, Config, PoolConfig, ShardConfig},
@@ -597,7 +599,16 @@ mod test {
                         hasher: Hasher::Postgres,
                         ..Default::default()
                     }],
-                    vec!["sharded_omni".into()],
+                    vec![
+                        OmnishardedTable {
+                            name: "sharded_omni".into(),
+                            sticky_routing: false,
+                        },
+                        OmnishardedTable {
+                            name: "sharded_omni_sticky".into(),
+                            sticky_routing: true,
+                        },
+                    ],
                 ),
                 shards,
                 identifier,

pgdog/src/backend/replication/logical/subscriber/context.rs

@@ -56,6 +56,7 @@ impl<'a> StreamContext<'a> {
             &mut self.prepared_statements,
             &self.params,
             None,
+            1,
         )?)
     }
 }

pgdog/src/backend/replication/sharded_tables.rs

@@ -1,15 +1,20 @@
 //! Tables sharded in the database.
+use pgdog_config::OmnishardedTable;
+
 use crate::{
     config::{DataType, ShardedTable},
     frontend::router::sharding::Mapping,
     net::messages::Vector,
 };
-use std::{collections::HashSet, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
 
 #[derive(Default, Debug)]
 struct Inner {
     tables: Vec<ShardedTable>,
-    omnisharded: HashSet<String>,
+    omnisharded: HashMap<String, bool>, // Name <-> sticky routing
     /// This is set only if we have the same sharding scheme
     /// across all tables, i.e., 3 tables with the same data type
     /// and list/range/hash function.
@@ -46,7 +51,7 @@ impl From<&[ShardedTable]> for ShardedTables {
 }
 
 impl ShardedTables {
-    pub fn new(tables: Vec<ShardedTable>, omnisharded_tables: Vec<String>) -> Self {
+    pub fn new(tables: Vec<ShardedTable>, omnisharded_tables: Vec<OmnishardedTable>) -> Self {
         let mut common_mapping = HashSet::new();
         for table in &tables {
             common_mapping.insert((
@@ -69,7 +74,10 @@ impl ShardedTables {
         Self {
             inner: Arc::new(Inner {
                 tables,
-                omnisharded: omnisharded_tables.into_iter().collect(),
+                omnisharded: omnisharded_tables
+                    .into_iter()
+                    .map(|table| (table.name, table.sticky_routing))
+                    .collect(),
                 common_mapping,
             }),
         }
@@ -79,7 +87,7 @@ impl ShardedTables {
         &self.inner.tables
     }
 
-    pub fn omnishards(&self) -> &HashSet<String> {
+    pub fn omnishards(&self) -> &HashMap<String, bool> {
         &self.inner.omnisharded
     }
 

pgdog/src/frontend/client/mod.rs

@@ -4,6 +4,7 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
+use rand::{thread_rng, Rng};
 use timeouts::Timeouts;
 use tokio::{select, spawn, time::timeout};
 use tracing::{debug, enabled, error, info, trace, Level as LogLevel};
@@ -28,7 +29,6 @@ use crate::state::State;
 use crate::stats::memory::MemoryUsage;
 use crate::util::user_database_from_params;
 
-// pub mod counter;
 pub mod query_engine;
 pub mod timeouts;
 
@@ -50,6 +50,7 @@ pub struct Client {
     client_request: ClientRequest,
     stream_buffer: MessageBuffer,
     passthrough_password: Option<String>,
+    omni_sticky_index: usize,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
@@ -315,6 +316,7 @@ impl Client {
             stream_buffer: MessageBuffer::new(config.config.memory.message_buffer),
             shutdown: false,
             passthrough_password,
+            omni_sticky_index: thread_rng().gen_range(1..usize::MAX),
         }))
     }
 
@@ -342,6 +344,7 @@ impl Client {
             stream_buffer: MessageBuffer::new(4096),
             shutdown: false,
             passthrough_password: None,
+            omni_sticky_index: 1,
         }
     }
 

pgdog/src/frontend/client/query_engine/context.rs

@@ -1,3 +1,5 @@
+use rand::{thread_rng, Rng};
+
 use crate::{
     backend::pool::{connection::mirror::Mirror, stats::MemoryStats},
     frontend::{
@@ -34,6 +36,8 @@ pub struct QueryEngineContext<'a> {
     pub(super) admin: bool,
     /// Executing rollback statement.
     pub(super) rollback: bool,
+    /// Omnisharded modulo.
+    pub(super) omni_sticky_index: usize,
 }
 
 impl<'a> QueryEngineContext<'a> {
@@ -53,6 +57,7 @@ impl<'a> QueryEngineContext<'a> {
             admin: client.admin,
             requests_left: 0,
             rollback: false,
+            omni_sticky_index: client.omni_sticky_index,
         }
     }
 
@@ -77,6 +82,7 @@ impl<'a> QueryEngineContext<'a> {
             admin: false,
             requests_left: 0,
             rollback: false,
+            omni_sticky_index: thread_rng().gen_range(1..usize::MAX),
         }
     }
 

pgdog/src/frontend/client/query_engine/route_query.rs

@@ -50,6 +50,7 @@ impl QueryEngine {
             context.prepared_statements,
             context.params,
             context.transaction,
+            context.omni_sticky_index,
         )?;
         match self.router.query(router_context) {
             Ok(cmd) => {

pgdog/src/frontend/router/cli.rs

@@ -55,6 +55,7 @@ impl RouterCli {
                 &mut stmt,
                 &mut params,
                 None,
+                1,
             )?)?;
             result.push(cmd);
         }