Add --compiler-worker-max-rss

Author: fantix

edb/server/args.py

@@ -257,6 +257,8 @@ class ServerConfig(NamedTuple):
     compiler_pool_mode: CompilerPoolMode
     compiler_pool_addr: tuple[str, int]
     compiler_pool_tenant_cache_size: int
+    compiler_worker_max_rss: Optional[int]
+
     echo_runtime_info: bool
     emit_server_status: str
     temp_dir: bool
@@ -1135,6 +1137,12 @@ def resolve_envvar_value(self, ctx: click.Context):
         help='Path to a TOML file to configure the server.',
         hidden=True,
     ),
+    click.option(
+        '--compiler-worker-max-rss',
+        type=int,
+        help='Maximum allowed RSS (in KiB) per compiler worker process. Any '
+             'worker exceeding this limit will be terminated and recreated.',
+    ),
 ])
 
 
@@ -1175,6 +1183,12 @@ def resolve_envvar_value(self, ctx: click.Context):
         '--metrics-port', type=PortType(),
         help=f'Port to listen on for metrics HTTP API.',
     ),
+    click.option(
+        '--worker-max-rss',
+        type=int,
+        help='Maximum allowed RSS (in KiB) per worker process. Any worker '
+             'exceeding this limit will be terminated and recreated.',
+    ),
 ])
 
 
@@ -1409,6 +1423,10 @@ def parse_args(**kwargs: Any):
             kwargs['compiler_pool_addr'] = (
                 "localhost", defines.EDGEDB_REMOTE_COMPILER_PORT
             )
+        if kwargs['compiler_worker_max_rss'] is not None:
+            abort('cannot set --compiler-worker-max-rss when using '
+                  '--compiler-pool-mode=remote')
+
     elif kwargs['compiler_pool_addr'] is not None:
         abort('--compiler-pool-addr is only meaningful '
               'under --compiler-pool-mode=remote')

edb/server/compiler_pool/pool.py

@@ -46,6 +46,7 @@
 import time
 
 import immutables
+import psutil
 
 from edb.common import debug
 from edb.common import lru
@@ -208,6 +209,7 @@ async def _request(
 class Worker(BaseWorker):
 
     _pid: int
+    _proc: psutil.Process
     _manager: BaseLocalPool
     _server: amsg.Server
 
@@ -221,6 +223,7 @@ def __init__(
         super().__init__(*args)
 
         self._pid = pid
+        self._proc = psutil.Process(pid)
         self._manager = manager
         self._server = server
 
@@ -237,6 +240,9 @@ async def _attach(self, init_args_pickled: bytes) -> None:
     def get_pid(self) -> int:
         return self._pid
 
+    def get_rss(self) -> int:
+        return self._proc.memory_info().rss // 1024
+
     def close(self) -> None:
         if self._closed:
             return
@@ -835,6 +841,7 @@ class BaseLocalPool(
 
     _poolsock_name: str
     _pool_size: int
+    _worker_max_rss: Optional[int]
     _server: Optional[amsg.Server]
     _ready_evt: asyncio.Event
     _running: Optional[bool]
@@ -846,6 +853,7 @@ def __init__(
         *,
         runstate_dir: str,
         pool_size: int,
+        worker_max_rss: Optional[int] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -859,6 +867,7 @@ def __init__(
 
         assert pool_size >= 1
         self._pool_size = pool_size
+        self._worker_max_rss = worker_max_rss
         self._workers = {}
 
         self._server = amsg.Server(self._poolsock_name, self._loop, self)
@@ -1036,6 +1045,12 @@ def _release_worker(
     ) -> None:
         # Skip disconnected workers
         if worker.get_pid() in self._workers:
+            if self._worker_max_rss is not None:
+                if worker.get_rss() > self._worker_max_rss:
+                    if debug.flags.server:
+                        print(f"HIT MEMORY LIMIT, KILLING {worker.get_pid()}")
+                    worker.close()
+                    return
             self._workers_queue.release(worker, put_in_front=put_in_front)
 
     def get_debug_info(self) -> dict[str, Any]:

edb/server/compiler_pool/server.py

@@ -648,6 +648,7 @@ async def server_main(
     client_schema_cache_size: int,
     runstate_dir: Optional[str | pathlib.Path],
     metrics_port: Optional[int],
+    worker_max_rss: Optional[int],
 ):
     if listen_port is None:
         listen_port = defines.EDGEDB_REMOTE_COMPILER_PORT
@@ -674,6 +675,7 @@ async def server_main(
             pool_size=pool_size,
             cache_size=client_schema_cache_size,
             secret=secret.encode(),
+            worker_max_rss=worker_max_rss,
         )
         await pool.start()
         try:

edb/server/main.py

@@ -230,6 +230,7 @@ async def _run_server(
             compiler_pool_size=args.compiler_pool_size,
             compiler_pool_mode=args.compiler_pool_mode,
             compiler_pool_addr=args.compiler_pool_addr,
+            compiler_worker_max_rss=args.compiler_worker_max_rss,
             nethosts=args.bind_addresses,
             netport=args.port,
             listen_sockets=tuple(s for ss in sockets.values() for s in ss),

edb/server/multitenant.py

@@ -480,6 +480,7 @@ async def run_server(
             compiler_pool_tenant_cache_size=(
                 args.compiler_pool_tenant_cache_size
             ),
+            compiler_worker_max_rss=args.compiler_worker_max_rss,
             compiler_state=compiler_state,
             use_monitor_fs=args.reload_config_files in [
                 srvargs.ReloadTrigger.Default,

edb/server/server.py

@@ -147,6 +147,7 @@ def __init__(
         compiler_pool_addr,
         nethosts,
         netport,
+        compiler_worker_max_rss: Optional[int] = None,
         listen_sockets: tuple[socket.socket, ...] = (),
         testmode: bool = False,
         daemonized: bool = False,
@@ -192,6 +193,7 @@ def __init__(
         self._compiler_pool_size = compiler_pool_size
         self._compiler_pool_mode = compiler_pool_mode
         self._compiler_pool_addr = compiler_pool_addr
+        self._compiler_worker_max_rss = compiler_worker_max_rss
         self._system_compile_cache = lru.LRUMapping(
             maxsize=defines._MAX_QUERIES_CACHE
         )
@@ -609,6 +611,9 @@ def _get_compiler_args(self) -> dict[str, Any]:
         )
         if self._compiler_pool_mode == srvargs.CompilerPoolMode.Remote:
             args['address'] = self._compiler_pool_addr
+        else:
+            if self._compiler_worker_max_rss is not None:
+                args['worker_max_rss'] = self._compiler_worker_max_rss
         return args
 
     async def _destroy_compiler_pool(self):