Turn key websocket broadcasting

fanout.py

"""
Turn key websocket broadcasting: POST a message on a channel it will be
broadcasted to all the websockets listening on that channel.

Channels are identified by the URL path e.g. connecting a websocket to
/listen/to/a/channel registers the websocket on the channel to-a-channel.

Symmetrically a POST on /broadcast/to/a/channel will send a message to all
listeners on to-a-channel.

Third party dependencies: the almighty CherryPy and ws4py
"""


import weakref
from collections import defaultdict
import threading

import cherrypy
from ws4py.server.cherrypyserver import WebSocketPlugin, WebSocketTool
from ws4py.websocket import WebSocket


def main():
    """
    configure, instantiate and start the service
    """
    cherrypy.config.update({'server.socket_port': 9000})
    WebSocketPlugin(cherrypy.engine).subscribe()
    cherrypy.tools.websocket = WebSocketTool()
    cherrypy.quickstart(App(), '/', config={'/listen': {'tools.websocket.on': True,
                                                        'tools.websocket.handler_cls': WebSocket}})


class Channels:
    """
    broadcasting service layer
    """

    channels = defaultdict(weakref.WeakSet)
    lock = threading.RLock()

    @classmethod
    def add(cls, channel, websocket):
        """
        add a websocket listener to a channel, the later os created if it
        does not exist
        """
        with cls.lock:
            cls.channels[channel].add(websocket)

    @classmethod
    def send(cls, channel, message):
        """
        send message to all listeners on a channel, remove disconnected listeners
        and delete empty channels
        """
        with cls.lock:

            sockets = cls.channels.get(channel)
            if sockets is None:
                return

            cleanup = []
            for ws_handler in sockets:
                try:
                    ws_handler.send(message)
                except AttributeError:
                    cleanup.append(ws_handler)

            for ws_handler in cleanup:
                sockets.remove(ws_handler)

            if len(cls.channels[channel]) == 0:
                del cls.channels[channel]


class App(object):
    """
    the service itself, listen on a channel, broadcast to all listeners on
    a channel
    """

    @cherrypy.expose
    def index(self):
        """
        just to demonstrate how it works, open a browser and point to root
        then, in a terminal window:

        $ curl -d "message=$(date)" http://localhost:9000/broadcast/test/channel

        the date should appear in the console of the browser
        """
        return """
            <script type="text/javascript">
                var connection = new WebSocket('ws://localhost:9000/listen/test/channel')
                connection.onmessage = function(e) {console.log(e.data)}
            </script>
               """

    @cherrypy.expose
    def broadcast(self, *args, **kvargs):
        """
        broadcasting end point, POST a message
        """
        if cherrypy.request.method == 'POST':
            channel = '-'.join(args)
            Channels.send(channel, kvargs['message'])

    @cherrypy.expose
    def listen(self, *args):
        """
        register a listener on a channel by connecting a websocket
        """
        handler = cherrypy.request.ws_handler
        Channels.add('-'.join(args), handler)


if __name__ == '__main__':
    main()

Looks like it will do what it says on the label. 😸

You only lock when mutating the channels dict, but what about reading?

The Go version https://gist.github.com/nf/7c03729770315c05570f (untested)