eventlet/examples/distributed_websocket_chat.py

   1 """This is a websocket chat example with many servers. A client can connect to
   2 any of the servers and their messages will be received by all clients connected
   3 to any of the servers.
   4
   5 Run the examples like this:
   6
   7 $ python examples/chat_bridge.py tcp://127.0.0.1:12345 tcp://127.0.0.1:12346
   8
   9 and the servers like this (changing the port for each one obviously):
  10
  11 $ python examples/distributed_websocket_chat.py -p tcp://127.0.0.1:12345 -s tcp://127.0.0.1:12346 7000
  12
  13 So all messages are published to port 12345 and the device forwards all the
  14 messages to 12346 where they are subscribed to
  15 """
  16 import os
  17 import sys
  18 import eventlet
  19 from collections import defaultdict
  20 from eventlet import spawn_n, sleep
  21 from eventlet import wsgi
  22 from eventlet import websocket
  23 from eventlet.green import zmq
  24 from eventlet.hubs import get_hub, use_hub
  25 from uuid import uuid1
  26
  27 use_hub('zeromq')
  28 ctx = zmq.Context()
  29
  30
  31 class IDName(object):
  32
  33     def __init__(self):
  34         self.id = uuid1()
  35         self.name = None
  36
  37     def __str__(self):
  38         if self.name:
  39             return self.name
  40         else:
  41             return str(self.id)
  42
  43     def pack_message(self, msg):
  44         return self, msg
  45
  46     def unpack_message(self, msg):
  47         sender, message = msg
  48         sender_name = 'you said' if sender.id == self.id \
  49             else '%s says' % sender
  50         return "%s: %s" % (sender_name, message)
  51
  52
  53 participants = defaultdict(IDName)
  54
  55
  56 def subscribe_and_distribute(sub_socket):
  57     global participants
  58     while True:
  59         msg = sub_socket.recv_pyobj()
  60         for ws, name_id in participants.items():
  61             to_send = name_id.unpack_message(msg)
  62             if to_send:
  63                 try:
  64                     ws.send(to_send)
  65                 except:
  66                     del participants[ws]
  67
  68
  69 @websocket.WebSocketWSGI
  70 def handle(ws):
  71     global pub_socket
  72     name_id = participants[ws]
  73     ws.send("Connected as %s, change name with 'name: new_name'" % name_id)
  74     try:
  75         while True:
  76             m = ws.wait()
  77             if m is None:
  78                 break
  79             if m.startswith('name:'):
  80                 old_name = str(name_id)
  81                 new_name = m.split(':', 1)[1].strip()
  82                 name_id.name = new_name
  83                 m = 'Changed name from %s' % old_name
  84             pub_socket.send_pyobj(name_id.pack_message(m))
  85             sleep()
  86     finally:
  87         del participants[ws]
  88
  89
  90 def dispatch(environ, start_response):
  91     """Resolves to the web page or the websocket depending on the path."""
  92     global port
  93     if environ['PATH_INFO'] == '/chat':
  94         return handle(environ, start_response)
  95     else:
  96         start_response('200 OK', [('content-type', 'text/html')])
  97         return [open(os.path.join(
  98                      os.path.dirname(__file__),
  99                      'websocket_chat.html')).read() % dict(port=port)]
 100
 101 port = None
 102
 103 if __name__ == "__main__":
 104     usage = 'usage: websocket_chat -p pub address -s sub address port number'
 105     if len(sys.argv) != 6:
 106         print(usage)
 107         sys.exit(1)
 108
 109     pub_addr = sys.argv[2]
 110     sub_addr = sys.argv[4]
 111     try:
 112         port = int(sys.argv[5])
 113     except ValueError:
 114         print("Error port supplied couldn't be converted to int\n", usage)
 115         sys.exit(1)
 116
 117     try:
 118         pub_socket = ctx.socket(zmq.PUB)
 119         pub_socket.connect(pub_addr)
 120         print("Publishing to %s" % pub_addr)
 121         sub_socket = ctx.socket(zmq.SUB)
 122         sub_socket.connect(sub_addr)
 123         sub_socket.setsockopt(zmq.SUBSCRIBE, "")
 124         print("Subscribing to %s" % sub_addr)
 125     except:
 126         print("Couldn't create sockets\n", usage)
 127         sys.exit(1)
 128
 129     spawn_n(subscribe_and_distribute, sub_socket)
 130     listener = eventlet.listen(('127.0.0.1', port))
 131     print("\nVisit http://localhost:%s/ in your websocket-capable browser.\n" % port)
 132     wsgi.server(listener, dispatch)