diff --git a/wptserve/request.py b/wptserve/request.py index d81119616d6c0d..0d4b4196cba36b 100644 --- a/wptserve/request.py +++ b/wptserve/request.py @@ -28,9 +28,15 @@ class Server(object): config = None def __init__(self, request): - host, port, authkey = stash.load_env_config() - address = (host, port) - self.stash = stash.Stash(request.url_parts.path, address, authkey) + self._stash = None + self._request = request + + @property + def stash(self): + if self._stash is None: + address, authkey = stash.load_env_config() + self._stash = stash.Stash(self._request.url_parts.path, address, authkey) + return self._stash class InputFile(object): diff --git a/wptserve/stash.py b/wptserve/stash.py index 7d2d6dd55b52d0..b5461919322400 100644 --- a/wptserve/stash.py +++ b/wptserve/stash.py @@ -1,32 +1,57 @@ +import base64 +import json +import os import uuid from multiprocessing import Process -from multiprocessing.managers import SyncManager, DictProxy -import os -import json +from multiprocessing.managers import BaseManager, DictProxy +class ServerDictManager(BaseManager): + shared_data = {} -WPT_STASH_CONFIG = "WPT_STASH_CONFIG" +def _get_shared(): + return ServerDictManager.shared_data -def load_env_config(): - return json.loads(os.environ[WPT_STASH_CONFIG]) +ServerDictManager.register("get_dict", + callable=_get_shared, + proxytype=DictProxy) -def store_env_config(config): - os.environ[WPT_STASH_CONFIG] = json.dumps(config) +class ClientDictManager(BaseManager): + pass -def start_server(address=None, authkey=None): - shared_data = {} - class DictManager(SyncManager): - pass +ClientDictManager.register("get_dict") - DictManager.register("get_dict", - callable=lambda:shared_data, - proxytype=DictProxy) - manager = DictManager(address, authkey) - server = manager.get_server() - server_process = Process(target=server.serve_forever) - server_process.start() +class StashServer(object): + def __init__(self, address=None, authkey=None): + self.address = address + self.authkey = authkey + self.manager = None - return (server_process, manager._address, manager._authkey) + def __enter__(self): + self.manager, self.address, self.authkey = start_server(self.address, self.authkey) + store_env_config(self.address, self.authkey) + + def __exit__(self, *args, **kwargs): + if self.manager is not None: + self.manager.shutdown() + +def load_env_config(): + address, authkey = json.loads(os.environ["WPT_STASH_CONFIG"]) + if isinstance(address, list): + address = tuple(address) + else: + address = str(address) + authkey = base64.decodestring(authkey) + return address, authkey + +def store_env_config(address, authkey): + authkey = base64.encodestring(authkey) + os.environ["WPT_STASH_CONFIG"] = json.dumps((address, authkey)) + +def start_server(address=None, authkey=None): + manager = ServerDictManager(address, authkey) + manager.start() + + return (manager, manager._address, manager._authkey) #TODO: Consider expiring values after some fixed time for long-running @@ -66,11 +91,7 @@ def _get_proxy(self, address=None, authkey=None): Stash._proxy = {} if Stash._proxy is None: - class DictManager(SyncManager): - pass - - DictManager.register("get_dict") - manager = DictManager(address, authkey) + manager = ClientDictManager(address, authkey) manager.connect() Stash._proxy = manager.get_dict()