test: add pooler SELECT 1 microbench; mimics existing e2e pooler

pgdog/src/backend/fake_server.rs

@@ -181,7 +181,7 @@ impl Reply {
                 let columns = select_columns(select);
                 let row_count = if reads_catalog(select) {
                     0
-                } else if returns_one_row(select) {
+                } else if returns_one_row(select) || select.from_clause().is_empty() {
                     1
                 } else {
                     limit_count(select)

pgdog/src/frontend/client/query_engine/test/bench.rs

@@ -33,7 +33,7 @@ use pgdog_config::{DataType, ShardedTableConfig};
 use super::prelude::*;
 use crate::backend::databases::{databases, reload_from_existing};
 use crate::backend::fake_server;
-use crate::config::{config, load_test_sharded, set};
+use crate::config::{config, load_test, load_test_sharded, set};
 use crate::frontend::router::parser::Cache;
 
 struct Case(&'static str, &'static str);
@@ -50,6 +50,8 @@ impl Case {
     }
 }
 
+const POOLER_CASES: &[Case] = &[Case("select_1", "SELECT 1")];
+
 const CASES: &[Case] = &[
     Case("select_by_key", "SELECT * FROM sharded WHERE id = 1"),
     Case(
@@ -136,54 +138,91 @@ enum Scope {
     Full,
 }
 
+#[derive(Clone, Copy)]
+enum BenchConfig {
+    Sharded,
+    Pooler,
+}
+
+impl BenchConfig {
+    fn load(self) {
+        match self {
+            Self::Sharded => bench_config(),
+            Self::Pooler => pooler_config(),
+        }
+    }
+
+    fn cases(self) -> &'static [Case] {
+        match self {
+            Self::Sharded => CASES,
+            Self::Pooler => POOLER_CASES,
+        }
+    }
+}
+
 struct Group {
     name: &'static str,
     protocol: ClientProtocol,
     caches: Caches,
     scope: Scope,
+    config: BenchConfig,
 }
 
 const SIMPLE_COLD: Group = Group {
     name: "simple_cold",
     protocol: ClientProtocol::Simple,
     caches: Caches::Cold,
     scope: Scope::Route,
+    config: BenchConfig::Sharded,
 };
 const EXTENDED_COLD: Group = Group {
     name: "extended_cold",
     protocol: ClientProtocol::Unnamed,
     caches: Caches::Cold,
     scope: Scope::Route,
+    config: BenchConfig::Sharded,
 };
 const EXTENDED_CACHED: Group = Group {
     name: "extended_cached",
     protocol: ClientProtocol::Unnamed,
     caches: Caches::Warm,
     scope: Scope::Route,
+    config: BenchConfig::Sharded,
 };
 const EXTENDED_EVICTING: Group = Group {
     name: "extended_evicting",
     protocol: ClientProtocol::Unnamed,
     caches: Caches::Evicting,
     scope: Scope::Route,
+    config: BenchConfig::Sharded,
 };
 const FULL_SIMPLE: Group = Group {
     name: "full_simple",
     protocol: ClientProtocol::Simple,
     caches: Caches::Cold,
     scope: Scope::Full,
+    config: BenchConfig::Sharded,
 };
 const FULL_EXTENDED_CACHED: Group = Group {
     name: "full_extended_cached",
     protocol: ClientProtocol::Unnamed,
     caches: Caches::Warm,
     scope: Scope::Full,
+    config: BenchConfig::Sharded,
 };
 const FULL_NAMED_CACHED: Group = Group {
     name: "full_named_cached",
     protocol: ClientProtocol::Named,
     caches: Caches::Warm,
     scope: Scope::Full,
+    config: BenchConfig::Sharded,
+};
+const POOLER_SIMPLE: Group = Group {
+    name: "pooler_simple",
+    protocol: ClientProtocol::Simple,
+    caches: Caches::Cold,
+    scope: Scope::Full,
+    config: BenchConfig::Pooler,
 };
 
 const GROUPS: &[Group] = &[
@@ -194,6 +233,7 @@ const GROUPS: &[Group] = &[
     FULL_SIMPLE,
     FULL_EXTENDED_CACHED,
     FULL_NAMED_CACHED,
+    POOLER_SIMPLE,
 ];
 
 const NAMED_STATEMENT: &str = "bench";
@@ -374,14 +414,26 @@ fn bench_config() {
     reload_from_existing().unwrap();
 }
 
+fn pooler_config() {
+    load_test();
+
+    let mut config = config().deref().clone();
+    for database in &mut config.config.databases {
+        database.host = fake_server::HOST.into();
+    }
+
+    set(config).unwrap();
+    reload_from_existing().unwrap();
+}
+
 fn bench(bencher: divan::Bencher, case: &Case, group: &Group) {
     let rt = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
         .unwrap();
 
     let _guard = rt.enter();
-    bench_config();
+    group.config.load();
 
     let runner = RefCell::new(rt.block_on(Runner::new(case, group)));
 
@@ -439,6 +491,14 @@ fn full_named_cached(bencher: divan::Bencher, case: &Case) {
     std::hint::black_box(());
 }
 
+// Mimics our `pooler` end-to-end benchmark.
+#[divan::bench(args = POOLER_CASES, sample_count = 1000)]
+#[inline(never)]
+fn pooler_simple(bencher: divan::Bencher, case: &Case) {
+    bench(bencher, case, &POOLER_SIMPLE);
+    std::hint::black_box(());
+}
+
 #[test]
 #[cfg(all(feature = "bench", target_os = "linux"))]
 fn count_instructions() {
@@ -456,8 +516,8 @@ fn count_instructions() {
     let module = module_path!().trim_start_matches("pgdog::");
 
     for group in GROUPS {
-        for case in CASES {
-            bench_config();
+        for case in group.config.cases() {
+            group.config.load();
             let mut runner = rt.block_on(Runner::new(case, group));
 
             for _ in 0..WARMUP {
@@ -495,8 +555,8 @@ fn count_instructions() {
 #[tokio::test]
 async fn test_bench_groups() {
     for group in GROUPS {
-        for case in CASES {
-            bench_config();
+        for case in group.config.cases() {
+            group.config.load();
             let mut runner = Runner::new(case, group).await;
             let mut input = runner.input();
             let ok = Box::pin(runner.request(&mut input)).await;