Cache encoded compilation config (#7236)

edb/server/compiler/compiler.py

@@ -411,7 +411,9 @@ def state_serializer_factory(self) -> sertypes.StateSerializerFactory:
         )
 
     @functools.cached_property
-    def compilation_config_serializer(self) -> sertypes.InputShapeSerializer:
+    def compilation_config_serializer(
+        self
+    ) -> sertypes.CompilationConfigSerializer:
         return (
             self.state_serializer_factory.make_compilation_config_serializer()
         )
@@ -1132,7 +1134,7 @@ def make_state_serializer(
 
     def make_compilation_config_serializer(
         self,
-    ) -> sertypes.InputShapeSerializer:
+    ) -> sertypes.CompilationConfigSerializer:
         return self.state.compilation_config_serializer
 
     def describe_database_dump(

edb/server/compiler/rpc.pyi

@@ -40,7 +40,8 @@ class CompilationRequest:
     session_config: immutables.Map[str, config.SettingValue] | None
 
     def __init__(
-        self, compilation_config_serializer: sertypes.InputShapeSerializer
+        self,
+        compilation_config_serializer: sertypes.CompilationConfigSerializer,
     ):
         ...
 

edb/server/compiler/rpc.pyx

@@ -76,7 +76,8 @@ cdef deserialize_output_format(char mode):
 @cython.final
 cdef class CompilationRequest:
     def __cinit__(
-        self, compilation_config_serializer: sertypes.InputShapeSerializer
+        self,
+        compilation_config_serializer: sertypes.CompilationConfigSerializer,
     ):
         self._serializer = compilation_config_serializer
 
@@ -221,25 +222,19 @@ cdef class CompilationRequest:
         hash_obj = hashlib.blake2b(memoryview(out), digest_size=16)
         hash_obj.update(self.source.cache_key())
 
-        # Build config that affects compilation: session -> database -> system.
-        # comp_config is only used for calculating cache_key, while session
-        # config itself is separately stored in the serialized format.
-        comp_config = {}
-        if self.system_config is not None:
-            comp_config.update(self.system_config)
-        if self.database_config is not None:
-            comp_config.update(self.database_config)
         if self.session_config is None:
             session_config = b""
         else:
-            comp_config.update(self.session_config)
-            session_config = self._serializer.encode(
-                {k: v.value for k, v in self.session_config.items()}
+            session_config = self._serializer.encode_configs(
+                self.session_config
             )
         out.write_len_prefixed_bytes(session_config)
 
-        serialized_comp_config = self._serializer.encode(
-            {k: v.value for k, v in comp_config.items()}
+        # Build config that affects compilation: session -> database -> system.
+        # This is only used for calculating cache_key, while session
+        # config itself is separately stored above in the serialized format.
+        serialized_comp_config = self._serializer.encode_configs(
+            self.system_config, self.database_config, self.session_config
         )
         hash_obj.update(serialized_comp_config)
 
@@ -324,9 +319,10 @@ cdef class CompilationRequest:
             serializer = self._serializer
             buf.read_len_prefixed_bytes()
         else:
-            serializer = sertypes.InputShapeSerializer(
+            serializer = sertypes.CompilationConfigSerializer(
                 type_id, buf.read_len_prefixed_bytes(), defines.CURRENT_PROTOCOL
             )
+            self._serializer = serializer
 
         data = buf.read_len_prefixed_bytes()
         if data:

edb/server/compiler/sertypes.py

@@ -1771,7 +1771,7 @@ def make(
             type_id, type_data, global_reps, protocol_version
         )
 
-    def make_compilation_config_serializer(self) -> InputShapeSerializer:
+    def make_compilation_config_serializer(self) -> CompilationConfigSerializer:
         ctx = Context(
             schema=self._schema,
             protocol_version=edbdef.CURRENT_PROTOCOL,
@@ -1782,7 +1782,7 @@ def make_compilation_config_serializer(self) -> InputShapeSerializer:
             ctx=ctx
         )
         type_data = b''.join(ctx.buffer)
-        return InputShapeSerializer(
+        return CompilationConfigSerializer(
             type_id,
             type_data,
             edbdef.CURRENT_PROTOCOL
@@ -1851,6 +1851,18 @@ def get_global_type_rep(
         return self._global_reps.get(global_name)
 
 
+class CompilationConfigSerializer(InputShapeSerializer):
+    @functools.lru_cache(64)
+    def encode_configs(
+        self, *configs: immutables.Map[str, config.SettingValue] | None
+    ) -> bytes:
+        state: dict[str, Any] = {}
+        for conf in configs:
+            if conf is not None:
+                state.update((k, v.value) for k, v in conf.items())
+        return self.encode(state)
+
+
 def derive_alias(
     schema: s_schema.Schema,
     parent: s_objtypes.ObjectType,

edb/server/server.py

@@ -128,7 +128,7 @@ class BaseServer:
     _stmt_cache_size: int | None = None
 
     _compiler_pool: compiler_pool.AbstractPool | None
-    compilation_config_serializer: sertypes.InputShapeSerializer
+    compilation_config_serializer: sertypes.CompilationConfigSerializer
     _http_request_logger: asyncio.Task | None
     _auth_gc: asyncio.Task | None