# Source code for ws4py.manager

    # -*- coding: utf-8 -*-
    __doc__ = """
    The manager module provides a selected classes to
    handle websocket's execution.

    Initially the rationale was to:

    - Externalize the way the CherryPy server had been setup
      as its websocket management was too tightly coupled with
      the plugin implementation.
    - Offer a management that could be used by other
      server or client implementations.
    - Move away from the threaded model to the event-based
      model by relying on `select` or `epoll` (when available).

    A simple usage for handling websocket clients:

    .. code-block:: python

        from ws4py.client import WebSocketBaseClient
        from ws4py.manager import WebSocketManager

        m = WebSocketManager()

        class EchoClient(WebSocketBaseClient):
            def handshake_ok(self):
                m.add(self)  # register the client once the handshake is done

            def received_message(self, msg):
                print str(msg)

        m.start()

        client = EchoClient('ws://localhost:9000/ws')
        client.connect()

        m.join()  # blocks forever

    Managers are not compulsory but hopefully will help your
    workflow. For clients, you can still rely on threaded, gevent or
    tornado based implementations of course.
    """
    import logging
    import select
    import threading
    import time

    from ws4py import format_addresses
    from ws4py.compat import py3k

    logger = logging.getLogger('ws4py')

    class SelectPoller(object):
        def __init__(self, timeout=0.1):
            """
            A socket poller that uses the `select`
            implementation to determines which
            file descriptors have data available to read.

            It is available on all platforms.
            """
            self._fds = []
            self.timeout = timeout

        def release(self):
            """
            Cleanup resources.
            """
            self._fds = []

        def register(self, fd):
            """
            Register a new file descriptor to be
            part of the select polling next time around.
            """
            if fd not in self._fds:
                self._fds.append(fd)

        def unregister(self, fd):
            """
            Unregister the given file descriptor.
            """
            if fd in self._fds:
                self._fds.remove(fd)

        def poll(self):
            """
            Polls once and returns a list of
            ready-to-be-read file descriptors.
            """
            if not self._fds:
                time.sleep(self.timeout)
                return []
            try:
                r, w, x = select.select(self._fds, [], [], self.timeout)
            except IOError as e:
                return []
            return r

    class EPollPoller(object):
        def __init__(self, timeout=0.1):
            """
            An epoll poller that uses the ``epoll``
            implementation to determines which
            file descriptors have data available to read.

            Available on Unix flavors mostly.
            """
            self.poller = select.epoll()
            self.timeout = timeout

        def release(self):
            """
            Cleanup resources.
            """
            self.poller.close()

        def register(self, fd):
            """
            Register a new file descriptor to be
            part of the select polling next time around.
            """
            try:
                self.poller.register(fd, select.EPOLLIN | select.EPOLLPRI)
            except IOError:
                pass

        def unregister(self, fd):
            """
            Unregister the given file descriptor.
            """
            self.poller.unregister(fd)

        def poll(self):
            """
            Polls once and yields each ready-to-be-read
            file-descriptor
            """
            try:
                events = self.poller.poll(timeout=self.timeout)
            except IOError:
                events = []

            for fd, event in events:
                if event | select.EPOLLIN | select.EPOLLPRI:
                    yield fd

    class KQueuePoller(object):
        def __init__(self, timeout=0.1):
            """
            An epoll poller that uses the ``epoll``
            implementation to determines which
            file descriptors have data available to read.

            Available on Unix flavors mostly.
            """
            self.poller = select.epoll()
            self.timeout = timeout

        def release(self):
            """
            Cleanup resources.
            """
            self.poller.close()

        def register(self, fd):
            """
            Register a new file descriptor to be
            part of the select polling next time around.
            """
            try:
                self.poller.register(fd, select.EPOLLIN | select.EPOLLPRI)
            except IOError:
                pass

        def unregister(self, fd):
            """
            Unregister the given file descriptor.
            """
            self.poller.unregister(fd)

        def poll(self):
            """
            Polls once and yields each ready-to-be-read
            file-descriptor
            """
            try:
                events = self.poller.poll(timeout=self.timeout)
            except IOError:
                events = []
            for fd, event in events:
                if event | select.EPOLLIN | select.EPOLLPRI:
                    yield fd

    class WebSocketManager(threading.Thread):
        def __init__(self, poller=None):
            """
            An event-based websocket manager. By event-based, we mean
            that the websockets will be called when their
            sockets have data to be read from.

            The manager itself runs in its own thread as not to
            be the blocking mainloop of your application.

            The poller's implementation is automatically chosen
            with ``epoll`` if available else ``select`` unless you
            provide your own ``poller``.
            """
            threading.Thread.__init__(self)
            self.name = "WebSocketManager"
            self.lock = threading.Lock()
            self.websockets = {}
            self.running = False

            if poller:
                self.poller = poller
            else:
                if hasattr(select, "epoll"):
                    self.poller = EPollPoller()
                    logger.info("Using epoll")
                else:
                    self.poller = SelectPoller()
                    logger.info("Using select as epoll is not available")

        def __len__(self):
            return len(self.websockets)

        def __iter__(self):
            if py3k:
                return iter(self.websockets.values())
            else:
                return self.websockets.itervalues()

        def __contains__(self, ws):
            fd = ws.sock.fileno()
            # just in case the file descriptor was reused
            # we actually check the instance (well, this might
            # also have been reused...)
            return self.websockets.get(fd) is ws

        def add(self, websocket):
            """
            Manage a new websocket.

            First calls its :meth:`opened() <ws4py.websocket.WebSocket.opened>`
            method and register its socket against the poller
            for reading events.
            """
            if websocket in self:
                return

            logger.info("Managing websocket %s" % format_addresses(websocket))
            websocket.opened()
            with self.lock:
                fd = websocket.sock.fileno()
                self.websockets[fd] = websocket
                self.poller.register(fd)

        def remove(self, websocket):
            """
            Remove the given ``websocket`` from the manager.

            This does not call its :meth:`closed() <ws4py.websocket.WebSocket.closed>`
            method as it's out-of-band by your application
            or from within the manager's run loop.
            """
            if websocket not in self:
                return

            logger.info("Removing websocket %s" % format_addresses(websocket))
            with self.lock:
                fd = websocket.sock.fileno()
                self.websockets.pop(fd, None)
                self.poller.unregister(fd)

        def stop(self):
            """
            Mark the manager as terminated and
            releases its resources.
            """
            self.running = False
            with self.lock:
                self.websockets.clear()
                self.poller.release()

        def run(self):
            """
            Manager's mainloop executed from within a thread.

            Constantly poll for read events and, when available,
            call related websockets' `once` method to
            read and process the incoming data.

            If the :meth:`once() <ws4py.websocket.WebSocket.once>`
            method returns a `False` value, its :meth:`terminate() <ws4py.websocket.WebSocket.terminate>`
            method is also applied to properly close
            the websocket and its socket is unregistered from the poller.

            Note that websocket shouldn't take long to process
            their data or they will block the remaining
            websockets with data to be handled. As for what long means,
            it's up to your requirements.
            """
            self.running = True
            while self.running:
                with self.lock:
                    polled = self.poller.poll()
                if not self.running:
                    break

                for fd in polled:
                    if not self.running:
                        break

                    ws = self.websockets.get(fd)
                    if ws and not ws.terminated:
                        # I don't know what kind of errors might spew out of here
                        # but they probably shouldn't crash the entire server.
                        try:
                            x = ws.once()
                        # Treat the error as if once() had returned None
                        except Exception as e:
                            x = None
                            logger.error("Terminating websocket %s due to exception: %s in once method" % (format_addresses(ws), repr(e)) )
                        if not x:
                            with self.lock:
                                self.websockets.pop(fd, None)
                                self.poller.unregister(fd)

                            if not ws.terminated:
                                logger.info("Terminating websocket %s" % format_addresses(ws))
                                ws.terminate()

        def close_all(self, code=1001, message='Server is shutting down'):
            """
            Execute the :meth:`close() <ws4py.websocket.WebSocket.close>`
            method of each registered websockets to initiate the closing handshake.
            It doesn't wait for the handshake to complete properly.
            """
            with self.lock:
                logger.info("Closing all websockets with [%d] '%s'" % (code, message))
                for ws in iter(self):
                    ws.close(code=code, reason=message)

        def broadcast(self, message, binary=False):
            """
            Broadcasts the given message to all registered
            websockets, at the time of the call.

            Broadcast may fail on a given registered peer
            but this is silent as it's not the method's
            purpose to handle websocket's failures.
            """
            with self.lock:
                websockets = self.websockets.copy()
                if py3k:
                    ws_iter = iter(websockets.values())
                else:
                    ws_iter = websockets.itervalues()

            for ws in ws_iter:
                if not ws.terminated:
                    try:
                        ws.send(message, binary)
                    except:
                        pass