feat: Remove query comments to improve cache hit rate

pgdog/src/frontend/router/parser/cache/ast.rs

@@ -7,9 +7,7 @@ use std::{collections::HashSet, ops::Deref};
 use parking_lot::Mutex;
 use std::sync::Arc;
 
-use super::super::{
-    comment::comment, Error, Route, Shard, StatementRewrite, StatementRewriteContext, Table,
-};
+use super::super::{Error, Route, Shard, StatementRewrite, StatementRewriteContext, Table};
 use super::{Fingerprint, Stats};
 use crate::backend::schema::Schema;
 use crate::frontend::router::parser::rewrite::statement::RewritePlan;
@@ -72,6 +70,8 @@ impl Ast {
         schema: &ShardingSchema,
         db_schema: &Schema,
         prepared_statements: &mut PreparedStatements,
+        comment_shard: Option<Shard>,
+        comment_role: Option<Role>,
         user: &str,
         search_path: Option<&ParameterValue>,
     ) -> Result<Self, Error> {
@@ -81,7 +81,6 @@ impl Ast {
             QueryParserEngine::PgQueryRaw => parse_raw(query),
         }
         .map_err(Error::PgQuery)?;
-        let (comment_shard, comment_role) = comment(query, schema)?;
         let fingerprint =
             Fingerprint::new(query, schema.query_parser_engine).map_err(Error::PgQuery)?;
 
@@ -125,12 +124,16 @@ impl Ast {
         query: &BufferedQuery,
         ctx: &super::AstContext<'_>,
         prepared_statements: &mut PreparedStatements,
+        shard: Option<Shard>,
+        role: Option<Role>,
     ) -> Result<Self, Error> {
         Self::new(
             query,
             &ctx.sharding_schema,
             &ctx.db_schema,
             prepared_statements,
+            shard,
+            role,
             ctx.user,
             ctx.search_path,
         )

pgdog/src/frontend/router/parser/cache/cache_impl.rs

@@ -2,6 +2,7 @@ use lru::LruCache;
 use once_cell::sync::Lazy;
 use pg_query::normalize;
 use pgdog_config::QueryParserEngine;
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::time::Duration;
 
@@ -11,6 +12,7 @@ use tracing::debug;
 
 use super::super::{Error, Route};
 use super::{Ast, AstContext};
+use crate::frontend::router::parser::comment;
 use crate::frontend::{BufferedQuery, PreparedStatements};
 
 static CACHE: Lazy<Cache> = Lazy::new(Cache::new);
@@ -97,6 +99,10 @@ impl Cache {
     /// Parse a statement by either getting it from cache
     /// or using pg_query parser.
     ///
+    /// In the event of cache miss, we retry after removing all comments except
+    /// for pgdog metadata. We retain it for correctness, since a query with
+    /// that metadata must not map to an identical query without it.
+    ///
     /// N.B. There is a race here that allows multiple threads to
     /// parse the same query. That's better imo than locking the data structure
     /// while we parse the query.
@@ -118,12 +124,38 @@ impl Cache {
             }
         }
 
+        let (maybe_shard, maybe_role, maybe_filtered_query) =
+            comment::parse_comment(&query, &ctx.sharding_schema)?;
+
+        let query_to_cache: Cow<'_, str>;
+
+        if let Some(filtered_query) = maybe_filtered_query {
+            query_to_cache = Cow::Owned(filtered_query);
+
+            // Check cache again after removing comments from query
+            let mut guard = self.inner.lock();
+
+            let ast = guard.queries.get_mut(&*query_to_cache).map(|entry| {
+                entry.stats.lock().hits += 1;
+                entry.clone()
+            });
+
+            if let Some(ast) = ast {
+                guard.stats.hits += 1;
+                return Ok(ast);
+            }
+        } else {
+            query_to_cache = Cow::Borrowed(query.query());
+        }
+
         // Parse query without holding lock.
-        let entry = Ast::with_context(query, ctx, prepared_statements)?;
+        let entry = Ast::with_context(query, ctx, prepared_statements, maybe_shard, maybe_role)?;
         let parse_time = entry.stats.lock().parse_time;
 
         let mut guard = self.inner.lock();
-        guard.queries.put(query.query().to_string(), entry.clone());
+        guard
+            .queries
+            .put(query_to_cache.into_owned(), entry.clone());
         guard.stats.misses += 1;
         guard.stats.parse_time += parse_time;
 
@@ -138,7 +170,10 @@ impl Cache {
         ctx: &AstContext<'_>,
         prepared_statements: &mut PreparedStatements,
     ) -> Result<Ast, Error> {
-        let mut entry = Ast::with_context(query, ctx, prepared_statements)?;
+        let (maybe_shard, maybe_role, _) = comment::parse_comment(&query, &ctx.sharding_schema)?;
+
+        let mut entry =
+            Ast::with_context(query, ctx, prepared_statements, maybe_shard, maybe_role)?;
         entry.cached = false;
 
         let parse_time = entry.stats.lock().parse_time;

pgdog/src/frontend/router/parser/cache/test.rs

@@ -121,7 +121,7 @@ fn test_tables_list() {
         "DELETE FROM private_schema.test",
         "DROP TABLE private_schema.test",
     ] {
-        let ast = Ast::new(&BufferedQuery::Query(Query::new(q)), &ShardingSchema::default(), &db_schema, &mut prepared_statements, "", None).unwrap();
+        let ast = Ast::new(&BufferedQuery::Query(Query::new(q)), &ShardingSchema::default(), &db_schema, &mut prepared_statements, None, None, "", None).unwrap();
         let tables = ast.tables();
         println!("{:?}", tables);
     }