Add cfg::Config._query_cache_mode (#7158)

Author: fantix

edb/buildmeta.py

@@ -55,7 +55,7 @@
 
 
 # Increment this whenever the database layout or stdlib changes.
-EDGEDB_CATALOG_VERSION = 2024_02_27_13_20
+EDGEDB_CATALOG_VERSION = 2024_04_04_00_00
 EDGEDB_MAJOR_VERSION = 5
 
 

edb/common/debug.py

@@ -35,7 +35,6 @@
 
 import builtins
 import contextlib
-import platform
 import os
 import sys
 import time
@@ -181,19 +180,6 @@ class flags(metaclass=FlagsMeta):
 
     zombodb = Flag(doc="Enabled zombodb and disables postgres FTS")
 
-    disable_persistent_cache = Flag(
-        doc="Don't use persistent cache",
-        # Persistent cache disabled for now by default on arm64 linux
-        # because of observed problems in CI test runs.
-        default=(
-            platform.system() == 'Linux'
-            and platform.machine() == 'arm64'
-        ),
-    )
-
-    # Function cache is an experimental feature that may not fully work
-    func_cache = Flag(doc="Use stored functions for persistent cache")
-
 
 @contextlib.contextmanager
 def timeit(title='block'):

edb/lib/cfg.edgeql

@@ -40,6 +40,8 @@ CREATE SCALAR TYPE cfg::memory EXTENDING std::anyscalar;
 CREATE SCALAR TYPE cfg::AllowBareDDL EXTENDING enum<AlwaysAllow, NeverAllow>;
 CREATE SCALAR TYPE cfg::ConnectionTransport EXTENDING enum<
     TCP, TCP_PG, HTTP, SIMPLE_HTTP, HTTP_METRICS, HTTP_HEALTH>;
+CREATE SCALAR TYPE cfg::QueryCacheMode EXTENDING enum<
+    InMemory, RegInline, PgFunc, Default>;
 
 CREATE ABSTRACT TYPE cfg::ConfigObject EXTENDING std::BaseObject;
 
@@ -192,6 +194,13 @@ ALTER TYPE cfg::AbstractConfig {
             'Recompile all cached queries on DDL if enabled.';
     };
 
+    CREATE PROPERTY query_cache_mode -> cfg::QueryCacheMode {
+        SET default := cfg::QueryCacheMode.Default;
+        CREATE ANNOTATION cfg::affects_compilation := 'true';
+        CREATE ANNOTATION std::description :=
+            'Where the query cache is finally stored';
+    };
+
     # Exposed backend settings follow.
     # When exposing a new setting, remember to modify
     # the _read_sys_config function to select the value

edb/pgsql/patches.py

@@ -108,5 +108,17 @@ def _setup_patches(patches: list[tuple[str, str]]) -> list[tuple[str, str]]:
 ) {
     set volatility := 'Stable';
 };
+'''),
+    ('edgeql+schema+config', '''
+CREATE SCALAR TYPE cfg::QueryCacheMode EXTENDING enum<
+    InMemory, RegInline, PgFunc, Default>;
+ALTER TYPE cfg::AbstractConfig {
+    CREATE PROPERTY query_cache_mode -> cfg::QueryCacheMode {
+        SET default := cfg::QueryCacheMode.Default;
+        CREATE ANNOTATION cfg::affects_compilation := 'true';
+        CREATE ANNOTATION std::description :=
+            'Where the query cache is finally stored';
+    };
+}
 '''),
 ])

edb/server/compiler/compiler.py

@@ -141,9 +141,13 @@ class CompileContext:
     log_ddl_as_migrations: bool = True
     dump_restore_mode: bool = False
     notebook: bool = False
-    persistent_cache: bool = False
     cache_key: Optional[uuid.UUID] = None
 
+    def get_cache_mode(self) -> config.QueryCacheMode:
+        return config.QueryCacheMode.effective(
+            _get_config_val(self, 'query_cache_mode')
+        )
+
     def _assert_not_in_migration_block(self, ql: qlast.Base) -> None:
         """Check that a START MIGRATION block is *not* active."""
         current_tx = self.state.current_tx()
@@ -861,7 +865,6 @@ def compile_request(
             request.protocol_version,
             request.inline_objectids,
             request.json_parameters,
-            persistent_cache=not debug.flags.disable_persistent_cache,
             cache_key=request.get_cache_key(),
         )
         return units, cstate
@@ -884,7 +887,6 @@ def compile(
         protocol_version: defines.ProtocolVersion,
         inline_objectids: bool = True,
         json_parameters: bool = False,
-        persistent_cache: bool = False,
         cache_key: Optional[uuid.UUID] = None,
     ) -> Tuple[dbstate.QueryUnitGroup,
                Optional[dbstate.CompilerConnectionState]]:
@@ -923,7 +925,6 @@ def compile(
             json_parameters=json_parameters,
             source=source,
             protocol_version=protocol_version,
-            persistent_cache=persistent_cache,
             cache_key=cache_key,
         )
 
@@ -967,7 +968,6 @@ def compile_in_tx_request(
             request.inline_objectids,
             request.json_parameters,
             expect_rollback=expect_rollback,
-            persistent_cache=not debug.flags.disable_persistent_cache,
             cache_key=request.get_cache_key(),
         )
         return units, cstate
@@ -986,7 +986,6 @@ def compile_in_tx(
         inline_objectids: bool = True,
         json_parameters: bool = False,
         expect_rollback: bool = False,
-        persistent_cache: bool = False,
         cache_key: Optional[uuid.UUID] = None,
     ) -> Tuple[dbstate.QueryUnitGroup, dbstate.CompilerConnectionState]:
         if (
@@ -1013,7 +1012,6 @@ def compile_in_tx(
             protocol_version=protocol_version,
             json_parameters=json_parameters,
             expect_rollback=expect_rollback,
-            persistent_cache=persistent_cache,
             cache_key=cache_key,
         )
 
@@ -1834,30 +1832,32 @@ def _compile_ql_query(
     # This low-hanging-fruit is temporary; persistent cache should cover all
     # cacheable cases properly in future changes.
     use_persistent_cache = (
-        ctx.persistent_cache
-        and cacheable
+        cacheable
         and not ctx.bootstrap_mode
         and script_info is None
         and ctx.cache_key is not None
     )
+    cache_mode = ctx.get_cache_mode()
 
     sql_res = pg_compiler.compile_ir_to_sql_tree(
         ir,
         expected_cardinality_one=ctx.expected_cardinality_one,
         output_format=_convert_format(ctx.output_format),
         backend_runtime_params=ctx.backend_runtime_params,
         expand_inhviews=options.expand_inhviews,
-        detach_params=bool(use_persistent_cache and debug.flags.func_cache),
+        detach_params=(use_persistent_cache
+                       and cache_mode is config.QueryCacheMode.PgFunc),
     )
 
     pg_debug.dump_ast_and_query(sql_res.ast, ir)
 
-    if use_persistent_cache:
-        if debug.flags.func_cache:
-            cache_sql, sql_ast = _build_cache_function(ctx, ir, sql_res)
-        else:
-            sql_ast = sql_res.ast
-            cache_sql = (b"", b"")
+    if use_persistent_cache and cache_mode is config.QueryCacheMode.PgFunc:
+        cache_sql, sql_ast = _build_cache_function(ctx, ir, sql_res)
+    elif (
+        use_persistent_cache and cache_mode is config.QueryCacheMode.RegInline
+    ):
+        sql_ast = sql_res.ast
+        cache_sql = (b"", b"")
     else:
         sql_ast = sql_res.ast
         cache_sql = None

edb/server/config/__init__.py

@@ -34,6 +34,7 @@
     load_ext_settings_from_schema,
 )
 from .types import ConfigType, CompositeConfigType
+from .types import QueryCacheMode
 
 
 __all__ = (
@@ -47,6 +48,7 @@
     'load_ext_settings_from_schema',
     'get_compilation_config',
     'coerce_single_value',
+    'QueryCacheMode',
 )
 
 

edb/server/config/types.py

@@ -21,6 +21,9 @@
 
 from typing import Any, TypeVar, TYPE_CHECKING, TypeGuard
 
+import enum
+import platform
+
 from edb import errors
 from edb.common import typeutils
 from edb.common import typing_inspect
@@ -255,3 +258,25 @@ def _err(
     ) -> errors.ConfigurationError:
         return errors.ConfigurationError(
             f'invalid {tspec.name.lower()!r} value: {msg}')
+
+
+class QueryCacheMode(enum.StrEnum):
+    InMemory = "InMemory"
+    RegInline = "RegInline"
+    PgFunc = "PgFunc"
+    Default = "Default"
+
+    @classmethod
+    def effective(cls, value: str | None) -> QueryCacheMode:
+        if value is None:
+            rv = cls.Default
+        else:
+            rv = cls(value)
+        if rv is QueryCacheMode.Default:
+            # Persistent cache disabled for now by default on arm64 linux
+            # because of observed problems in CI test runs.
+            if platform.system() == 'Linux' and platform.machine() == 'arm64':
+                rv = QueryCacheMode.InMemory
+            else:
+                rv = QueryCacheMode.RegInline
+        return rv

edb/server/dbview/dbview.pyx

@@ -149,13 +149,12 @@ cdef class Database:
 
         self._cache_worker_task = self._cache_queue = None
         self._cache_notify_task = self._cache_notify_queue = None
-        if not debug.flags.disable_persistent_cache:
-            self._cache_queue = asyncio.Queue()
-            self._cache_worker_task = asyncio.create_task(
-                self.monitor(self.cache_worker, 'cache_worker'))
-            self._cache_notify_queue = asyncio.Queue()
-            self._cache_notify_task = asyncio.create_task(
-                self.monitor(self.cache_notifier, 'cache_notifier'))
+        self._cache_queue = asyncio.Queue()
+        self._cache_worker_task = asyncio.create_task(
+            self.monitor(self.cache_worker, 'cache_worker'))
+        self._cache_notify_queue = asyncio.Queue()
+        self._cache_notify_task = asyncio.create_task(
+            self.monitor(self.cache_notifier, 'cache_notifier'))
 
     @property
     def server(self):
@@ -365,6 +364,9 @@ cdef class Database:
                 group.cache_state = CacheState.Present
                 self._eql_to_compiled[query_req] = group
 
+    def clear_query_cache(self):
+        self._eql_to_compiled.clear()
+
     def iter_views(self):
         yield from self._views
 
@@ -375,10 +377,18 @@ cdef class Database:
         if self.user_schema_pickle is None:
             async with self._introspection_lock:
                 if self.user_schema_pickle is None:
-                    await self.tenant.introspect_db(
-                        self.name,
-                        hydrate_cache=not debug.flags.disable_persistent_cache,
-                    )
+                    await self.tenant.introspect_db(self.name)
+
+    def lookup_config(self, name: str):
+        spec = self._index._sys_config_spec
+        if self.user_config_spec is not None:
+            spec = config.ChainedSpec(spec, self.user_config_spec)
+        return config.lookup(
+            name,
+            self.db_config or DEFAULT_CONFIG,
+            self._index._sys_config,
+            spec=spec,
+        )
 
 
 cdef class DatabaseConnectionView:
@@ -1579,3 +1589,10 @@ cdef class DatabaseIndex:
                 dbs, self._global_schema_pickle, self._comp_sys_config
             )
         return self._cached_compiler_args
+
+    def lookup_config(self, name: str):
+        return config.lookup(
+            name,
+            self._sys_config,
+            spec=self._sys_config_spec,
+        )

edb/server/tenant.py

@@ -848,9 +848,7 @@ async def _introspect_extensions(
 
         return extensions
 
-    async def introspect_db(
-        self, dbname: str, hydrate_cache: bool = False
-    ) -> None:
+    async def introspect_db(self, dbname: str) -> bool:
         """Use this method to (re-)introspect a DB.
 
         If the DB is already registered in self._dbindex, its
@@ -864,12 +862,21 @@ async def introspect_db(
         dropped and recreated again. It's safer to refresh the entire state
         than refreshing individual components of it. Besides, DDL and
         database-level config modifications are supposed to be rare events.
+
+        Returns True if the query cache mode changed.
         """
         logger.info("introspecting database '%s'", dbname)
 
+        assert self._dbindex is not None
+        if db := self._dbindex.maybe_get_db(dbname):
+            cache_mode_val = db.lookup_config('query_cache_mode')
+        else:
+            cache_mode_val = self._dbindex.lookup_config('query_cache_mode')
+        old_cache_mode = config.QueryCacheMode.effective(cache_mode_val)
+
         conn = await self._acquire_intro_pgcon(dbname)
         if not conn:
-            return
+            return False
 
         try:
             user_schema_json = (
@@ -917,7 +924,7 @@ async def introspect_db(
             extensions = await self._introspect_extensions(conn)
 
             query_cache: list[tuple[bytes, ...]] | None = None
-            if hydrate_cache:
+            if old_cache_mode is not config.QueryCacheMode.InMemory:
                 query_cache = await self._load_query_cache(conn)
         finally:
             self.release_pgcon(dbname, conn)
@@ -926,7 +933,6 @@ async def introspect_db(
         parsed_db = await compiler_pool.parse_user_schema_db_config(
             user_schema_json, db_config_json, self.get_global_schema_pickle()
         )
-        assert self._dbindex is not None
         db = self._dbindex.register_db(
             dbname,
             user_schema_pickle=parsed_db.user_schema_pickle,
@@ -941,8 +947,12 @@ async def introspect_db(
             parsed_db.protocol_version,
             parsed_db.state_serializer,
         )
-        if query_cache:
+        cache_mode = config.QueryCacheMode.effective(
+            db.lookup_config('query_cache_mode')
+        )
+        if query_cache and cache_mode is not config.QueryCacheMode.InMemory:
             db.hydrate_cache(query_cache)
+        return old_cache_mode is not cache_mode
 
     async def _early_introspect_db(self, dbname: str) -> None:
         """We need to always introspect the extensions for each database.
@@ -1436,7 +1446,16 @@ def on_local_database_config_change(self, dbname: str) -> None:
         # of the DB and update all components of it.
         async def task():
             try:
-                await self.introspect_db(dbname)
+                if await self.introspect_db(dbname):
+                    logger.info(
+                        "clearing query cache for database '%s'", dbname)
+                    conn = await self.acquire_pgcon(dbname)
+                    try:
+                        await conn.sql_execute(
+                            b'SELECT edgedb._clear_query_cache()')
+                        self._dbindex.get_db(dbname).clear_query_cache()
+                    finally:
+                        self.release_pgcon(dbname, conn)
             except Exception:
                 metrics.background_errors.inc(
                     1.0, self._instance_name, "on_local_database_config_change"