Core, WebHost: lazy-load worlds in unpickler, WebHost and WebHostLib (A…

…rchipelagoMW#2156)

* Core: lazy-load worlds in unpickler

this should hopefully fix customserver's memory consumption

* WebHost: move imports around to save memory in MP

* MultiServer: prefer loading _speedups without pyximport

This saves ~15MB per MP and speeds up module import if it was built in-place.

* Tests: fix tests for changed WebHost imports

* CustomServer: run GC after setup

* CustomServer: cleanup exception handling

NetUtils.py

@@ -407,14 +407,22 @@ def get_remaining(self, state: typing.Dict[typing.Tuple[int, int], typing.Set[in
 if typing.TYPE_CHECKING:  # type-check with pure python implementation until we have a typing stub
     LocationStore = _LocationStore
 else:
-    try:
-        import pyximport
-        pyximport.install()
-    except ImportError:
-        pyximport = None
     try:
         from _speedups import LocationStore
+        import _speedups
+        import os.path
+        if os.path.getctime(_speedups.__file__) < os.path.getctime("_speedups.pyx"):
+            warnings.warn(f"{_speedups.__file__} outdated! "
+                          f"Please rebuild with `cythonize -b -i _speedups.pyx` or delete it!")
     except ImportError:
-        warnings.warn("_speedups not available. Falling back to pure python LocationStore. "
-                      "Install a matching C++ compiler for your platform to compile _speedups.")
-        LocationStore = _LocationStore
+        try:
+            import pyximport
+            pyximport.install()
+        except ImportError:
+            pyximport = None
+        try:
+            from _speedups import LocationStore
+        except ImportError:
+            warnings.warn("_speedups not available. Falling back to pure python LocationStore. "
+                          "Install a matching C++ compiler for your platform to compile _speedups.")
+            LocationStore = _LocationStore

Utils.py

@@ -359,11 +359,13 @@ def get_unique_identifier():
 
 
 class RestrictedUnpickler(pickle.Unpickler):
+    generic_properties_module: Optional[object]
+
     def __init__(self, *args, **kwargs):
         super(RestrictedUnpickler, self).__init__(*args, **kwargs)
         self.options_module = importlib.import_module("Options")
         self.net_utils_module = importlib.import_module("NetUtils")
-        self.generic_properties_module = importlib.import_module("worlds.generic")
+        self.generic_properties_module = None
 
     def find_class(self, module, name):
         if module == "builtins" and name in safe_builtins:
@@ -373,6 +375,8 @@ def find_class(self, module, name):
             return getattr(self.net_utils_module, name)
         # Options and Plando are unpickled by WebHost -> Generate
         if module == "worlds.generic" and name in {"PlandoItem", "PlandoConnection"}:
+            if not self.generic_properties_module:
+                self.generic_properties_module = importlib.import_module("worlds.generic")
             return getattr(self.generic_properties_module, name)
         # pep 8 specifies that modules should have "all-lowercase names" (options, not Options)
         if module.lower().endswith("options"):

WebHost.py

@@ -13,22 +13,16 @@
 import settings
 
 Utils.local_path.cached_path = os.path.dirname(__file__) or "."  # py3.8 is not abs. remove "." when dropping 3.8
-
-from WebHostLib import register, cache, app as raw_app
-from waitress import serve
-
-from WebHostLib.models import db
-from WebHostLib.autolauncher import autohost, autogen
-from WebHostLib.lttpsprites import update_sprites_lttp
-from WebHostLib.options import create as create_options_files
-
 settings.no_gui = True
 configpath = os.path.abspath("config.yaml")
 if not os.path.exists(configpath):  # fall back to config.yaml in home
     configpath = os.path.abspath(Utils.user_path('config.yaml'))
 
 
 def get_app():
+    from WebHostLib import register, cache, app as raw_app
+    from WebHostLib.models import db
+
     register()
     app = raw_app
     if os.path.exists(configpath) and not app.config["TESTING"]:
@@ -121,6 +115,11 @@ def create_ordered_tutorials_file() -> typing.List[typing.Dict[str, typing.Any]]
     multiprocessing.freeze_support()
     multiprocessing.set_start_method('spawn')
     logging.basicConfig(format='[%(asctime)s] %(message)s', level=logging.INFO)
+
+    from WebHostLib.lttpsprites import update_sprites_lttp
+    from WebHostLib.autolauncher import autohost, autogen
+    from WebHostLib.options import create as create_options_files
+
     try:
         update_sprites_lttp()
     except Exception as e:
@@ -137,4 +136,5 @@ def create_ordered_tutorials_file() -> typing.List[typing.Dict[str, typing.Any]]
         if app.config["DEBUG"]:
             app.run(debug=True, port=app.config["PORT"])
         else:
+            from waitress import serve
             serve(app, port=app.config["PORT"], threads=app.config["WAITRESS_THREADS"])

WebHostLib/autolauncher.py

@@ -3,8 +3,6 @@
 import json
 import logging
 import multiprocessing
-import os
-import sys
 import threading
 import time
 import typing
@@ -13,55 +11,7 @@
 from pony.orm import db_session, select, commit
 
 from Utils import restricted_loads
-
-
-class CommonLocker():
-    """Uses a file lock to signal that something is already running"""
-    lock_folder = "file_locks"
-
-    def __init__(self, lockname: str, folder=None):
-        if folder:
-            self.lock_folder = folder
-        os.makedirs(self.lock_folder, exist_ok=True)
-        self.lockname = lockname
-        self.lockfile = os.path.join(self.lock_folder, f"{self.lockname}.lck")
-
-
-class AlreadyRunningException(Exception):
-    pass
-
-
-if sys.platform == 'win32':
-    class Locker(CommonLocker):
-        def __enter__(self):
-            try:
-                if os.path.exists(self.lockfile):
-                    os.unlink(self.lockfile)
-                self.fp = os.open(
-                    self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
-            except OSError as e:
-                raise AlreadyRunningException() from e
-
-        def __exit__(self, _type, value, tb):
-            fp = getattr(self, "fp", None)
-            if fp:
-                os.close(self.fp)
-                os.unlink(self.lockfile)
-else:  # unix
-    import fcntl
-
-
-    class Locker(CommonLocker):
-        def __enter__(self):
-            try:
-                self.fp = open(self.lockfile, "wb")
-                fcntl.flock(self.fp.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
-            except OSError as e:
-                raise AlreadyRunningException() from e
-
-        def __exit__(self, _type, value, tb):
-            fcntl.flock(self.fp.fileno(), fcntl.LOCK_UN)
-            self.fp.close()
+from .locker import Locker, AlreadyRunningException
 
 
 def launch_room(room: Room, config: dict):

WebHostLib/customserver.py

@@ -19,6 +19,7 @@
 
 from MultiServer import Context, server, auto_shutdown, ServerCommandProcessor, ClientMessageProcessor, load_server_cert
 from Utils import restricted_loads, cache_argsless
+from .locker import Locker
 from .models import Command, GameDataPackage, Room, db
 
 
@@ -163,16 +164,19 @@ def run_server_process(room_id, ponyconfig: dict, static_server_data: dict,
     db.generate_mapping(check_tables=False)
 
     async def main():
+        import gc
+
         Utils.init_logging(str(room_id), write_mode="a")
         ctx = WebHostContext(static_server_data)
         ctx.load(room_id)
         ctx.init_save()
         ssl_context = load_server_cert(cert_file, cert_key_file) if cert_file else None
+        gc.collect()  # free intermediate objects used during setup
         try:
             ctx.server = websockets.serve(functools.partial(server, ctx=ctx), ctx.host, ctx.port, ssl=ssl_context)
 
             await ctx.server
-        except Exception:  # likely port in use - in windows this is OSError, but I didn't check the others
+        except OSError:  # likely port in use
             ctx.server = websockets.serve(functools.partial(server, ctx=ctx), ctx.host, 0, ssl=ssl_context)
 
             await ctx.server
@@ -198,16 +202,15 @@ async def main():
         await ctx.shutdown_task
         logging.info("Shutting down")
 
-    from .autolauncher import Locker
     with Locker(room_id):
         try:
             asyncio.run(main())
-        except KeyboardInterrupt:
+        except (KeyboardInterrupt, SystemExit):
             with db_session:
                 room = Room.get(id=room_id)
                 # ensure the Room does not spin up again on its own, minute of safety buffer
                 room.last_activity = datetime.datetime.utcnow() - datetime.timedelta(minutes=1, seconds=room.timeout)
-        except:
+        except Exception:
             with db_session:
                 room = Room.get(id=room_id)
                 room.last_port = -1

WebHostLib/locker.py

@@ -0,0 +1,51 @@
+import os
+import sys
+
+
+class CommonLocker:
+    """Uses a file lock to signal that something is already running"""
+    lock_folder = "file_locks"
+
+    def __init__(self, lockname: str, folder=None):
+        if folder:
+            self.lock_folder = folder
+        os.makedirs(self.lock_folder, exist_ok=True)
+        self.lockname = lockname
+        self.lockfile = os.path.join(self.lock_folder, f"{self.lockname}.lck")
+
+
+class AlreadyRunningException(Exception):
+    pass
+
+
+if sys.platform == 'win32':
+    class Locker(CommonLocker):
+        def __enter__(self):
+            try:
+                if os.path.exists(self.lockfile):
+                    os.unlink(self.lockfile)
+                self.fp = os.open(
+                    self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
+            except OSError as e:
+                raise AlreadyRunningException() from e
+
+        def __exit__(self, _type, value, tb):
+            fp = getattr(self, "fp", None)
+            if fp:
+                os.close(self.fp)
+                os.unlink(self.lockfile)
+else:  # unix
+    import fcntl
+
+
+    class Locker(CommonLocker):
+        def __enter__(self):
+            try:
+                self.fp = open(self.lockfile, "wb")
+                fcntl.flock(self.fp.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except OSError as e:
+                raise AlreadyRunningException() from e
+
+        def __exit__(self, _type, value, tb):
+            fcntl.flock(self.fp.fileno(), fcntl.LOCK_UN)
+            self.fp.close()

test/webhost/TestAPIGenerate.py

@@ -5,7 +5,8 @@
 class TestDocs(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None:
-        from WebHost import get_app, raw_app
+        from WebHostLib import app as raw_app
+        from WebHost import get_app
         raw_app.config["PONY"] = {
             "provider": "sqlite",
             "filename": ":memory:",

test/webhost/TestFileGeneration.py

@@ -14,7 +14,8 @@ def setUpClass(cls) -> None:
         cls.incorrect_path = os.path.join(os.path.split(os.path.dirname(__file__))[0], "WebHostLib")
 
     def testOptions(self):
-        WebHost.create_options_files()
+        from WebHostLib.options import create as create_options_files
+        create_options_files()
         target = os.path.join(self.correct_path, "static", "generated", "configs")
         self.assertTrue(os.path.exists(target))
         self.assertFalse(os.path.exists(os.path.join(self.incorrect_path, "static", "generated", "configs")))