Fix DB scheduling on mariaDB

Utils.py

@@ -6,12 +6,15 @@ def tuplize_version(version: str) -> typing.Tuple[int, ...]:
     return tuple(int(piece, 10) for piece in version.split("."))
 
 
-__version__ = "2.5.3"
+__version__ = "2.5.4"
 _version_tuple = tuplize_version(__version__)
 
 import os
 import subprocess
 import sys
+import pickle
+import io
+import builtins
 
 import functools
 
@@ -294,3 +297,23 @@ def get_unique_identifier():
     uuid = uuid.getnode()
     persistent_store("client", "uuid", uuid)
     return uuid
+
+
+safe_builtins = {
+    'set',
+    'frozenset',
+}
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == "builtins" and name in safe_builtins:
+            return getattr(builtins, name)
+        # Forbid everything else.
+        raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
+                                     (module, name))
+
+
+def restricted_loads(s):
+    """Helper function analogous to pickle.loads()."""
+    return RestrictedUnpickler(io.BytesIO(s)).load()

WebHostLib/__init__.py

@@ -32,12 +32,14 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 app.config['MAX_CONTENT_LENGTH'] = 4 * 1024 * 1024  # 4 megabyte limit
 # if you want persistent sessions on your server, make sure you make this a constant in your config.yaml
 app.config["SECRET_KEY"] = os.urandom(32)
+# at what amount of worlds should scheduling be used, instead of rolling in the webthread
+app.config["JOB_THRESHOLD"] = 2
 app.config['SESSION_PERMANENT'] = True
 
 # waitress uses one thread for I/O, these are for processing of views that then get sent
 # berserkermulti.world uses gunicorn + nginx; ignoring this option
 app.config["WAITRESS_THREADS"] = 10
-#a default that just works. berserkermulti.world runs on mariadb
+# a default that just works. berserkermulti.world runs on mariadb
 app.config["PONY"] = {
     'provider': 'sqlite',
     'filename': os.path.abspath('db.db3'),

WebHostLib/autolauncher.py

@@ -8,6 +8,7 @@ import time
 
 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"""
@@ -77,10 +78,12 @@ def handle_generation_failure(result: BaseException):
 
 
 def launch_generator(pool: multiprocessing.pool.Pool, generation: Generation):
-    logging.info(f"Generating {generation.id} for {len(generation.options)} players")
+    options = restricted_loads(generation.options)
+    logging.info(f"Generating {generation.id} for {len(options)} players")
 
-    pool.apply_async(gen_game, (generation.options,),
-                     {"race": generation.meta["race"], "sid": generation.id, "owner": generation.owner},
+    meta = restricted_loads(generation.meta)
+    pool.apply_async(gen_game, (options,),
+                     {"race": meta["race"], "sid": generation.id, "owner": generation.owner},
                      handle_generation_success, handle_generation_failure)
     generation.state = STATE_STARTED
 

WebHostLib/generate.py

@@ -3,13 +3,13 @@ import tempfile
 import random
 import zlib
 import json
-import multiprocessing
 
 from flask import request, flash, redirect, url_for, session, render_template
 
 from EntranceRandomizer import parse_arguments
 from Main import main as ERmain
 from Main import get_seed, seeddigits
+import pickle
 
 from .models import *
 from WebHostLib import app
@@ -35,15 +35,20 @@ def generate(race=False):
                 elif len(gen_options) > app.config["MAX_ROLL"]:
                     flash(f"Sorry, generating of multiworlds is limited to {app.config['MAX_ROLL']} players for now. "
                           f"If you have a larger group, please generate it yourself and upload it.")
-                else:
-
-                    gen = Generation(options={name: vars(options) for name, options in gen_options.items()},
-                                     # convert to json compatible
-                                     meta={"race": race, "owner": session["_id"].int}, state=STATE_QUEUED,
-                                     owner=session["_id"])
+                elif len(gen_options) >= app.config["JOB_THRESHOLD"]:
+                    gen = Generation(
+                        options=pickle.dumps({name: vars(options) for name, options in gen_options.items()}),
+                        # convert to json compatible
+                        meta=pickle.dumps({"race": race}), state=STATE_QUEUED,
+                        owner=session["_id"])
                     commit()
 
                     return redirect(url_for("wait_seed", seed=gen.id))
+                else:
+                    seed_id = gen_game({name: vars(options) for name, options in gen_options.items()},
+                                       race=race, owner=session["_id"].int)
+                    return redirect(url_for("view_seed", seed=seed_id))
+
     return render_template("generate.html", race=race)
 
 
@@ -89,8 +94,11 @@ def gen_game(gen_options, race=False, owner=None, sid=None):
 
         return upload_to_db(target.name, owner, sid)
     except BaseException:
-        with db_session:
-            Generation.get(id=sid).state = STATE_ERROR
+        if sid:
+            with db_session:
+                gen = Generation.get(id=sid)
+                if gen is not None:
+                    gen.state = STATE_ERROR
         raise
 
 
@@ -127,8 +135,14 @@ def upload_to_db(folder, owner, sid):
                 flash(e)
     if multidata:
         with db_session:
-            seed = Seed(multidata=multidata, spoiler=spoiler, patches=patches, owner=owner, id=sid)
+            if sid:
+                seed = Seed(multidata=multidata, spoiler=spoiler, patches=patches, owner=owner, id=sid)
+            else:
+                seed = Seed(multidata=multidata, spoiler=spoiler, patches=patches, owner=owner)
             for patch in patches:
                 patch.seed = seed
-            Generation.get(id=sid).delete()
+            if sid:
+                gen = Generation.get(id=sid)
+                if gen is not None:
+                    gen.delete()
         return seed.id

WebHostLib/models.py

@@ -49,6 +49,6 @@ class Command(db.Entity):
 class Generation(db.Entity):
     id = PrimaryKey(UUID, default=uuid4)
     owner = Required(UUID)
-    options = Required(Json, lazy=True)
-    meta = Required(Json)
+    options = Required(bytes, lazy=True)  # these didn't work as JSON on mariaDB, so they're getting pickled now
+    meta = Required(bytes, lazy=True)
     state = Required(int, default=0, index=True)

WebHostLib/requirements.txt

@@ -4,4 +4,4 @@ waitress>=1.4.4
 flask-caching>=1.9.0
 Flask-Autoversion>=0.2.0
 Flask-Compress>=1.5.0
-Flask-Limiter>=1.3.1
+Flask-Limiter>=1.4