diff options
Diffstat (limited to 'src/leap/bitmask/backend/backend_proxy.py')
| -rw-r--r-- | src/leap/bitmask/backend/backend_proxy.py | 239 |
1 files changed, 135 insertions, 104 deletions
diff --git a/src/leap/bitmask/backend/backend_proxy.py b/src/leap/bitmask/backend/backend_proxy.py index 04046d3d..8d6930d6 100644 --- a/src/leap/bitmask/backend/backend_proxy.py +++ b/src/leap/bitmask/backend/backend_proxy.py @@ -21,11 +21,13 @@ to the backend. # XXX should document the relationship to the API here. import functools -import Queue import threading -import time import zmq +from zmq.eventloop import ioloop +from zmq.eventloop import zmqstream + +from taskthread import TimerTask from leap.bitmask.backend.api import API, STOP_REQUEST, PING_REQUEST from leap.bitmask.backend.settings import Settings @@ -37,35 +39,39 @@ import logging logger = logging.getLogger(__name__) -class BackendProxy(object): +class ZmqREQConnection(threading.Thread): """ - The BackendProxy handles calls from the GUI and forwards (through ZMQ) - to the backend. + A threaded zmq req connection. """ - if flags.ZMQ_HAS_CURVE: - PORT = '5556' - SERVER = "tcp://localhost:%s" % PORT - else: - SERVER = "ipc:///tmp/bitmask.socket.0" - - POLL_TIMEOUT = 4000 # ms - POLL_TRIES = 3 - - PING_INTERVAL = 2 # secs - - def __init__(self): - generate_zmq_certificates_if_needed() - - self._socket = None + def __init__(self, server_address, on_recv): + """ + Initialize the connection. - self.settings = Settings() + :param server_address: The address of the backend zmq server. + :type server: str + :param on_recv: The callback to be executed when a message is + received. + :type on_recv: callable(msg) + """ + threading.Thread.__init__(self) + self._server_address = server_address + self._on_recv = on_recv + self._stream = None + self._init_zmq() - # initialize ZMQ stuff: + def _init_zmq(self): + """ + Configure the zmq components and connection. + """ + logger.debug("Setting up ZMQ connection to server...") context = zmq.Context() - logger.debug("Connecting to server...") socket = context.socket(zmq.REQ) + # we use zmq's eventloop in order to asynchronously send requests + loop = ioloop.ZMQIOLoop.current() + self._stream = zmqstream.ZMQStream(socket, loop) + if flags.ZMQ_HAS_CURVE: # public, secret = zmq.curve_keypair() client_keys = zmq.curve_keypair() @@ -79,66 +85,128 @@ class BackendProxy(object): socket.setsockopt(zmq.RCVTIMEO, 1000) socket.setsockopt(zmq.LINGER, 0) # Terminate early - socket.connect(self.SERVER) - self._socket = socket - self._ping_at = 0 + self._stream.on_recv(self._on_recv) + + def run(self): + """ + Run the threaded stream connection loop. + """ + self._stream.socket.connect(self._server_address) + logger.debug("Starting ZMQ loop.") + self._stream.io_loop.start() + logger.debug("Finished ZMQ loop.") + + def stop(self): + """ + Stop the threaded connection loop. + """ + self._stream.io_loop.stop() + + def send(self, *args, **kwargs): + """ + Send a message through this connection. + """ + # Important note: calling send on the zmqstream from another + # thread doesn’t properly tell the IOLoop thread that there’s an + # event to process. This could cuase small delays if the IOLoop is + # already processing lots of events, but it can cause the message + # to never send if the zmq socket is the only one it’s handling. + # + # Because of that, we want ZmqREQConnection.send to hand off the + # stream.send to the IOLoop’s thread via IOLoop.add_callback: + self._stream.io_loop.add_callback( + lambda: self._stream.send(*args, **kwargs)) + + +class BackendProxy(object): + """ + The BackendProxy handles calls from the GUI and forwards (through ZMQ) + to the backend. + """ + + if flags.ZMQ_HAS_CURVE: + PORT = '5556' + SERVER = "tcp://localhost:%s" % PORT + else: + SERVER = "ipc:///tmp/bitmask.socket.0" + + PING_INTERVAL = 2 # secs + + def __init__(self): + """ + Initialize the backend proxy. + """ + generate_zmq_certificates_if_needed() + self._do_work = threading.Event() + self._work_lock = threading.Lock() + self._connection = ZmqREQConnection(self.SERVER, self._set_online) + self._heartbeat = TimerTask(self._ping, delay=self.PING_INTERVAL) + self._ping_event = threading.Event() self.online = False + self.settings = Settings() - self._call_queue = Queue.Queue() - self._worker_caller = threading.Thread(target=self._worker) + def _set_online(self, _): + """ + Mark the backend as being online. - def start(self): - self._worker_caller.start() + This is used as the zmq connection's on_recv callback, and so it is + passed the received message as a parameter. Because we currently don't + use that message, we just ignore it for now. + """ + self.online = True + # the following event is used when checking whether the backend is + # online + self._ping_event.set() + + def _set_offline(self): + """ + Mark the backend as being offline. + """ + self.online = False def check_online(self): """ Return whether the backend is accessible or not. You don't need to do `run` in order to use this. - :rtype: bool """ - # we use a small timeout in order to response quickly if the backend is - # offline - self._send_request(PING_REQUEST, retry=False, timeout=500) - self._socket.close() + logger.debug("Checking whether backend is online...") + self._send_request(PING_REQUEST) + # self._ping_event will eventually be set by the zmq connection's + # on_recv callback, so we use a small timeout in order to response + # quickly if the backend is offline + if not self._ping_event.wait(0.5): + logger.warning("Backend is offline!") + self._set_offline() return self.online - def _worker(self): + def start(self): """ - Worker loop that processes the Queue of pending requests to do. + Start the backend proxy. """ - while True: - try: - request = self._call_queue.get(block=False) - # break the loop after sending the 'stop' action to the - # backend. - if request == STOP_REQUEST: - break - - self._send_request(request) - except Queue.Empty: - pass - time.sleep(0.01) - self._ping() + logger.debug("Starting backend proxy...") + self._do_work.set() + self._connection.start() + self.check_online() + self._heartbeat.start() - logger.debug("BackendProxy worker stopped.") - - def _reset_ping(self): + def _stop(self): """ - Reset the ping timeout counter. - This is called for every ping and request. + Stop the backend proxy. """ - self._ping_at = time.time() + self.PING_INTERVAL + with self._work_lock: # avoid sending after connection was closed + self._do_work.clear() + self._heartbeat.stop() + self._connection.stop() + logger.debug("BackendProxy worker stopped.") def _ping(self): """ Heartbeat helper. Sends a PING request just to know that the server is alive. """ - if time.time() > self._ping_at: - self._send_request(PING_REQUEST) - self._reset_ping() + self._send_request(PING_REQUEST) def _api_call(self, *args, **kwargs): """ @@ -162,6 +230,8 @@ class BackendProxy(object): 'arguments': kwargs, } + request_json = None + try: request_json = zmq.utils.jsonapi.dumps(request) except Exception as e: @@ -172,12 +242,12 @@ class BackendProxy(object): raise # queue the call in order to handle the request in a thread safe way. - self._call_queue.put(request_json) + self._send_request(request_json) if api_method == STOP_REQUEST: - self._call_queue.put(STOP_REQUEST) + self._stop() - def _send_request(self, request, retry=True, timeout=None): + def _send_request(self, request): """ Send the given request to the server. This is used from a thread safe loop in order to avoid sending a @@ -185,49 +255,10 @@ class BackendProxy(object): :param request: the request to send. :type request: str - :param retry: whether we should retry or not in case of timeout. - :type retry: bool - :param timeout: a custom timeout (milliseconds) to wait for a response. - :type timeout: int """ - # logger.debug("Sending request to backend: {0}".format(request)) - self._socket.send(request) - - poll = zmq.Poller() - poll.register(self._socket, zmq.POLLIN) - - reply = None - - tries = 0 - if not retry: - tries = self.POLL_TRIES + 1 # this means: no retries left - - if timeout is None: - timeout = self.POLL_TIMEOUT - - while True: - socks = dict(poll.poll(timeout)) - if socks.get(self._socket) == zmq.POLLIN: - reply = self._socket.recv() - break - - tries += 1 - if tries < self.POLL_TRIES: - logger.warning('Retrying receive... {0}/{1}'.format( - tries, self.POLL_TRIES)) - else: - break - - if reply is None: - msg = "Timeout error contacting backend." - logger.critical(msg) - self.online = False - else: - # msg = "Received reply for '{0}' -> '{1}'".format(request, reply) - # logger.debug(msg) - self.online = True - # request received, no ping needed for other interval. - self._reset_ping() + with self._work_lock: # avoid sending after connection was closed + if self._do_work.is_set(): + self._connection.send(request) def __getattribute__(self, name): """ |