Example 4-1. Message protocol: read and write

# msgproto.py
import asyncio
from asyncio import StreamReader, StreamWriter

async def read_msg(reader: StreamReader) -> bytes:
    # Raises asyncio.streams.IncompleteReadError
    size_bytes = await reader.readexactly(4)  
    size = int.from_bytes(size_bytes, byteorder='big') 
    data = await reader.readexactly(size)  
    return data

async def send_msg(writer: StreamWriter, data: bytes):
    writer.write(len(data).to_bytes(4, byteorder='big'))  
    writer.write(data)  
    await writer.drain()  

def run_server(client, host='127.0.0.1', port=25000):  
    loop = asyncio.get_event_loop()
    coro = asyncio.start_server(client, '127.0.0.1', 25000)
    server = loop.run_until_complete(coro)
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        print('Bye!')
    server.close()
    loop.run_until_complete(server.wait_closed())
    tasks = asyncio.Task.all_tasks()
    for t in tasks:
        t.cancel()
    group = asyncio.gather(*tasks, return_exceptions=True
    loop.run_until_complete(group)
    loop.close()

Get the first 4 bytes. This is the size prefix.

Those four bytes must be converted into an integer.

Now we know the payload size, read that off the stream.

Write is the inverse of read: first send the length of the data, encoded as 4 bytes.

Then send the data.

drain() ensures that the data is fully sent. Without drain(), the data may still be waiting in the send buffer when this coroutine function exits.

It doesn’t belong here, but I also snuck in a boilerplate function to run a TCP server. The shutdown sequence has been discussed before in a previous section, and I’m including it here only to save space in the code samples that follow. Server shutdown will begin on SIGINT or Ctrl-C.

Example 4-2. A 35-line prototype

# mq_server.py
import asyncio
from asyncio import StreamReader, StreamWriter, gather
from collections import deque, defaultdict
from typing import Deque, DefaultDict
from msgproto import read_msg, send_msg, run_server  

SUBSCRIBERS: DefaultDict[bytes, Deque] = defaultdict(deque) 

async def client(reader: StreamReader, writer: StreamWriter):
  peername = writer.transport.get_extra_info('peername')  
  subscribe_chan = await read_msg(reader)  
  SUBSCRIBERS[subscribe_chan].append(writer)  
  print(f'Remote {peername} subscribed to {subscribe_chan}')
  try:
    while True:
      channel_name = await read_msg(reader)  
      data = await read_msg(reader)  
      print(f'Sending to {channel_name}: {data[:19]}...')
      writers = SUBSCRIBERS[channel_name]  
      if writers and channel_name.startswith(b'/queue'):  
          writers.rotate()  
          writers = [writers[0]]  
      await gather(*[send_msg(w, data) for w in writers]) 
  except asyncio.CancelledError:
    print(f'Remote {peername} closing connection.')
    writer.close()
  except asyncio.streams.IncompleteReadError:
    print(f'Remote {peername} disconnected')
  finally:
    print(f'Remote {peername} closed')
    SUBSCRIBERS[subscribe_chan].remove(writer)  

if __name__ == '__main__':
    run_server(client)

Imports from our msgproto.py module.

A global collection of currently active subscribers. Every time a client connects, they must first send a channel name they’re subscribing to. A deque will hold all the subscribers for a particular channel.

The client() coroutine function will produce a long-lived coroutine for each new connection. Think of it as a callback for the TCP server started in run_server(). On this line, I’ve shown how the host and port of the remote peer can be obtained, e.g., for logging.

Our protocol for clients is the following:

• On first connect, a client must send a message containing the channel to subscribe to (here, subscribe_chan).

• Thereafter, for the life of the connection, a client sends a message to a channel by first sending a message containing the destination channel name, followed by a message containing the data. Our broker will send such data-messages to every client subscribed to that channel name.

Add the StreamWriter instance to the global collection of subscribers.

An infinite loop, waiting for data from this client. The first message from a client must be the destination channel name.

Next comes the actual data to distribute to the channel.

Get the deque of subscribers on the target channel.

Some special handling if the channel name begins with the magic word “/queue”: in this case, we send the data to only one of the subscribers, not all of them. This can be used for sharing work between a bunch of workers, rather than the usual pub-sub notification scheme where all subscribers on a channel get all the messages.

Here is why we use a deque and not a list: rotation of the deque is how we keep track of which client is next in line for “/queue” distribution. This seems expensive until you realize that a single deque rotation is an O(1) operation.

Target only whichever client is first; this changes after every rotation.

Create a list of coroutines for sending the message to each writer, and then unpack these into gather() so we can wait for all of the sending to complete.

Note: This line is a bad flaw in our program, but it may not be obvious why: though it may be true that all of the sending to each subscriber will happen concurrently, what happens if we have one very slow client? In this case the gather() will only finish when the slowest subscriber has received their data. We can’t receive any more data from the sending client until all these send_msg() coroutines finish. This slows down all message distribution to the speed of the slowest subscriber.

When leaving the client() coroutine, make sure to remove ourselves from the global SUBSCRIBERS collection. Unfortunately, this is an O(n) operation which can be a little expensive for very large n. A different data structure would fix this, but for now we console ourselves with the understanding that connections are intended to be long-lived thus few disconnection events; and n is unlikely to be very large (say ~10,000 as a rough order-of-magnitude estimate); and this code is at least very easy to understand!

Example 4-3. Listener: a toolkit for listening for messages on our message broker

# mq_client_listen.py
import asyncio
import argparse, uuid
from msgproto import read_msg, send_msg

async def main(args):
  me = uuid.uuid4().hex[:8] 
  print(f'Starting up {me}')
  reader, writer = await asyncio.open_connection(
    args.host, args.port) 
  print(f'I am {writer.transport.get_extra_info("sockname")}')
  channel = args.listen.encode() 
  await send_msg(writer, channel) 
  try:
    while True:
      data = await read_msg(reader)  
      if not data:
          print('Connection ended.')
          break
      print(f'Received by {me}: {data[:20]}')
  except asyncio.streams.IncompleteReadError:
    print('Server closed.')

if __name__ == '__main__':
  parser = argparse.ArgumentParser() 
  parser.add_argument('--host', default='localhost')
  parser.add_argument('--port', default=25000)
  parser.add_argument('--listen', default='/topic/foo')
  loop = asyncio.get_event_loop()
  try:
      loop.run_until_complete(main(parser.parse_args()))
  except KeyboardInterrupt:
      print('Bye!')
  loop.close()

The uuid standard library module is a convenient way of creating an “identity” for this listener. If you start up multiple instances of these, each will have their own identity, and you’ll be able to track what is happening in the logs.

Open a connection to the server.

The channel to subscribe to is an input parameter, captured in args.listen. Encode it into bytes before sending.

By our protocol rules (as discussed in the broker code analysis previously), the first thing to do after connecting is to send the channel name to subscribe to.

This loop does nothing else but wait for data to appear on the socket.

The command-line arguments for this program make it easy to point to a host, a port, and a channel name to listen to.

Example 4-4. Sender: a toolkit for sending data to our message broker

# mq_client_sender.py
import asyncio
import argparse, uuid
from itertools import count
from msgproto import send_msg

async def main(args):
  me = uuid.uuid4().hex[:8]  
  print(f'Starting up {me}')
  reader, writer = await asyncio.open_connection(
      host=args.host, port=args.port)  
  print(f'I am {writer.transport.get_extra_info("sockname")}')

  channel = b'/null'  
  await send_msg(writer, channel) 

  chan = args.channel.encode()  
  for i in count():  
    await asyncio.sleep(args.interval)  
    data = b'X'*args.size or f'Msg {i} from {me}'.encode()
    try:
        await send_msg(writer, chan)
        await send_msg(writer, data) 
    except ConnectionResetError:
        print('Connection ended.')
        break
  writer.close()

if __name__ == '__main__':
  parser = argparse.ArgumentParser()  
  parser.add_argument('--host', default='localhost')
  parser.add_argument('--port', default=25000, type=int)
  parser.add_argument('--channel', default='/topic/foo')
  parser.add_argument('--interval', default=1, type=float)
  parser.add_argument('--size', default=0, type=int)
  loop = asyncio.get_event_loop()
  try:
      loop.run_until_complete(main(parser.parse_args()))
  except KeyboardInterrupt:
      print('Bye!')
  loop.close()

As with the listener, claim an identity.

As with the listener, reach out and make a connection.

According to our protocol rules, the first thing to do after connecting to the server is to give the name of the channel to subscribe to; however, since we are a sender, we don’t really care about subscribing to any channels; nevertheless, the protocol requires it so just provide a null channel to subscribe to (we won’t actually listen for anything).

Send the channel to subscribe to.

The command-line parameter args.channel provides the channel to which we want to send messages. Note that it must be converted to bytes first before sending.

Using itertools.count() is like a while True loop, except that you get an iteration variable to use. We use this in the debugging messages since it makes it a bit easier to track which message got sent from where.

The delay between sent messages is an input parameter, args.interval. The next line generates the message payload. It’s either a bytestring of specified size (args.size), or it’s a descriptive message. This flexibility is just for testing.

Send! Note that there are two messages here: the first is the destination channel name and the second is the payload.

As with the listener, there are a bunch of command-line options for tweaking the sender: “channel” determines the target channel to send to, while “interval” controls the delay between sends. The “size” parameter controls the size of each message payload.

Example 4-5. Message broker (server) output

$ python mq_server.py
Remote ('127.0.0.1', 55382) subscribed to b'/queue/blah'
Remote ('127.0.0.1', 55386) subscribed to b'/queue/blah'
Remote ('127.0.0.1', 55390) subscribed to b'/null'
Sending to b'/queue/blah': b'Msg 0 from 6b5a8e1d'...
Sending to b'/queue/blah': b'Msg 1 from 6b5a8e1d'...
Sending to b'/queue/blah': b'Msg 2 from 6b5a8e1d'...
Sending to b'/queue/blah': b'Msg 3 from 6b5a8e1d'...
Sending to b'/queue/blah': b'Msg 4 from 6b5a8e1d'...
Sending to b'/queue/blah': b'Msg 5 from 6b5a8e1d'...
^CBye!
Remote ('127.0.0.1', 55382) closing connection.
Remote ('127.0.0.1', 55382) closed
Remote ('127.0.0.1', 55390) closing connection.
Remote ('127.0.0.1', 55390) closed
Remote ('127.0.0.1', 55386) closing connection.
Remote ('127.0.0.1', 55386) closed

Example 4-6. Sender (client) output

$ python mq_client_sender.py --channel /queue/blah
Starting up 6b5a8e1d
I am ('127.0.0.1', 55390)
Connection ended.

Example 4-7. Listener 1 (client) output

$ python mq_client_listen.py --listen /queue/blah
Starting up 9ae04690
I am ('127.0.0.1', 55382)
Received by 9ae04690: b'Msg 1 from 6b5a8e1d'
Received by 9ae04690: b'Msg 3 from 6b5a8e1d'
Received by 9ae04690: b'Msg 5 from 6b5a8e1d'
Server closed.

Example 4-8. Listener 2 (client) output

$ python mq_client_listen.py --listen /queue/blah
Starting up bd4e3baa
I am ('127.0.0.1', 55386)
Received by bd4e3baa: b'Msg 0 from 6b5a8e1d'
Received by bd4e3baa: b'Msg 2 from 6b5a8e1d'
Received by bd4e3baa: b'Msg 4 from 6b5a8e1d'
Server closed.

Example 4-9. Message broker: improved design

# mq_server_plus.py
import asyncio
from asyncio import StreamReader, StreamWriter, Queue
from collections import deque, defaultdict
from contextlib import suppress
from typing import Deque, DefaultDict, Dict
from msgproto import read_msg, send_msg, run_server

SUBSCRIBERS: DefaultDict[bytes, Deque] = defaultdict(deque)
SEND_QUEUES: DefaultDict[StreamWriter, Queue] = defaultdict(Queue)
CHAN_QUEUES: Dict[bytes, Queue] = {}  

async def client(reader: StreamReader, writer: StreamWriter):
  peername = writer.transport.get_extra_info('peername')
  subscribe_chan = await read_msg(reader)
  SUBSCRIBERS[subscribe_chan].append(writer)  
  loop = asyncio.get_event_loop()
  send_task = loop.create_task(
      send_client(writer, SEND_QUEUES[writer]))  
  print(f'Remote {peername} subscribed to {subscribe_chan}')
  try:
    while True:
      channel_name = await read_msg(reader)
      data = await read_msg(reader)
      if channel_name not in CHAN_QUEUES:  
        CHAN_QUEUES[channel_name] = Queue(maxsize=10)  
        loop.create_task(chan_sender(channel_name))  
      await CHAN_QUEUES[channel_name].put(data)  
  except asyncio.CancelledError:
    print(f'Remote {peername} connection cancelled.')
  except asyncio.streams.IncompleteReadError:
    print(f'Remote {peername} disconnected')
  finally:
    print(f'Remote {peername} closed')
    await SEND_QUEUES[writer].put(None)  
    await send_task  
    del SEND_QUEUES[writer]  
    SUBSCRIBERS[subscribe_chan].remove(writer)

async def send_client(writer: StreamWriter, queue: Queue):  
    while True:
        with suppress(asyncio.CancelledError):
            data = await queue.get()
            if not data:
                writer.close()
                break
            await send_msg(writer, data)

async def chan_sender(name: bytes):
    with suppress(asyncio.CancelledError):
        while True:
            writers = SUBSCRIBERS[name]
            if not writers:
                await asyncio.sleep(1)
                continue  
            if name.startswith(b'/queue'):  
                writers.rotate()
                writers = [writers[0]]
            msg = await CHAN_QUEUES[name].get()  
            if not msg:
                break
            for writer in writers:
                if not SEND_QUEUES[writer].full():
                    print(f'Sending to {name}: {msg[:19]}...')
                    await SEND_QUEUES[writer].put(msg)  

if __name__ == '__main__':
    run_server(client)

In the previous implementation, there were only SUBSCRIBERS; now there are SEND_QUEUES and CHAN_QUEUES as global collections. This is a consequence of completely decoupling the receiving and sending of data. SEND_QUEUES has one queue entry for each client connection: all data that must be sent to that client must be placed onto that queue. (If you peek ahead, the send_client() coroutine will pull data off SEND_QUEUES and send it.)

Up till this point in the client() coroutine function, the code is the same as the simple server: the subscribed channel name is received and we add the StreamWriter instance for the new client to the global SUBSCRIBERS collection.

This is new: we create a long-lived task that will do all the sending of data to this client. The task will run independently as a separate coroutine, and will pull messages off the supplied queue, SEND_QUEUES[writer], for sending.

Now we’re inside the loop where we receive data. Remember that we always receive two messages: one for the destination channel name, and one for the data. We’re going to create a new, dedicated Queue for every destination channel, and that’s what CHAN_QUEUES is for: when any client wants to push data to a channel, we’re going to put that data onto the appropriate queue and then go immediately back to listening for more data. This approach decouples the distribution of messages from the receiving of messages from this client.

If there isn’t already a queue for the target channel, make one.

Create a dedicated, long-lived task for that channel. The coroutine, chan_sender(), will be responsible for taking data off the channel queue and distributing that data to subscribers.

Place the newly received data onto the specific channel’s queue. Note that if the queue fills up, we’ll wait here until there’s space for the new data. By waiting here, we won’t be reading any new data off the socket, which means that the client will have to wait on sending new data into the socket on their side. This isn’t necessarily a bad thing, since it communicates so-called back-pressure to this client. (Alternatively, you could choose to drop messages here if the use-case is OK with that.)

When the connection is closed, it’s time to clean up! The long-lived task we created for sending data to this client, send_task, can be shut down by placing None onto its queue, SEND_QUEUES[writer] (check the code for send_client()). It’s important to use a value on the queue, rather than outright cancellation, because there may already be data on that queue and we want that data to be sent out before send_client() is ended.

Wait for that sender task to finish.

Remove the entry in the SEND_QUEUES collection (and in the next line we also remove the sock from the SUBSCRIBERS collection as before).

The send_client() coroutine function is very nearly a textbook example of pulling work off a queue. Note how the coroutine will exit if None is placed onto the queue. Note also how we suppress CancelledError inside the loop: this is because we want this task to only be closed by receiving a None on the queue. This way, all pending data on the queue can be sent out before shutdown.

chan_sender() is the distribution logic for a channel: it sends data from a dedicated channel Queue instance to all the subscribers on that channel. But what happens if there are no subscribers for this channel yet? We’ll just wait a bit and try again. (Note that the queue for this channel, i.e., CHAN_QUEUES[name] will keep filling up though.)

As before in our previous broker implementation, we do something special for channels whose name begins with “/queue”: we rotate the deque and send only to the first entry. This acts like a crude load-balancing system because each subscriber gets different messages off the same queue. For all other channels, all subscribers get all the messages.

We’ll wait here for data on the queue. On the next line, exit if None is received. Currently this isn’t triggered anywhere (so these chan_sender() coroutines live forever); but if logic were added to clean up these channel tasks after, say, some period of inactivity, that’s how it would be done.

Data has been received, so it’s time to send to subscribers. Note that we do not do the sending here: instead, we place the data onto each subscriber’s own send queue. This decoupling is necessary to make sure that a slow subscriber doesn’t slow down anyone else receiving data. And furthermore, if the subscriber is so slow that their send queue fills up, we don’t put that data on their queue, i.e., it is lost.

Tip

To obtain and run the Splash container, run these commands in your shell:

$ docker pull scrapinghub/splash
$ docker run --rm -p 8050:8050 scrapinghub/splash

Our server backend will call the Splash API at http://localhost:8050.

Example 4-12. The traditional approach

# poller.py
import zmq

context = zmq.Context()
receiver = context.socket(zmq.PULL)  
receiver.connect("tcp://localhost:5557")

subscriber = context.socket(zmq.SUB)  
subscriber.connect("tcp://localhost:5556")
subscriber.setsockopt_string(zmq.SUBSCRIBE, '')

poller = zmq.Poller()  
poller.register(receiver, zmq.POLLIN)
poller.register(subscriber, zmq.POLLIN)

while True:
    try:
        socks = dict(poller.poll())  
    except KeyboardInterrupt:
        break

    if receiver in socks:
        message = receiver.recv_json()
        print(f'Via PULL: {message}')

    if subscriber in socks:
        message = subscriber.recv_json()
        print(f'Via SUB: {message}')

ØMQ sockets have types! This is a PULL socket. You can think of it as a “receive-only” kind of socket, that will be fed by some other “send-only” socket which will be a PUSH type.

The SUB socket type is another kind of “receive-only” socket, and will be fed a PUB type socket which is send-only.

If you need to move data between multiple sockets in a threaded ØMQ application, you’re going to need a poller. This is because these sockets are not thread-safe, so you cannot recv() on different sockets in different threads.⁴

It works similar to the select() system call. The poller will unblock when there is data ready to be received on one of the registered sockets, and then it’s up to you to pull the data off and do something with it. The big if block is how you have to detect the correct socket.

Example 4-13. Server code

# poller_srv.py
import zmq, itertools, time

context = zmq.Context()
pusher = context.socket(zmq.PUSH)
pusher.bind("tcp://*:5557")

publisher = context.socket(zmq.PUB)
publisher.bind("tcp://*:5556")

for i in itertools.count():
    time.sleep(1)
    pusher.send_json(i)
    publisher.send_json(i)

Note

For the code examples that follow, it is necessary to use pyzmq >= 17.0.0. At the time of writing, version 17 wasn’t released yet, so if necessary you will have to install the latest beta of pyzmq with a major version of 17.

Example 4-14. Clean separation with asyncio

# poller_aio.py
import asyncio
import zmq
from zmq.asyncio import Context

context = Context()

async def do_receiver():
    receiver = context.socket(zmq.PULL)  
    receiver.connect("tcp://localhost:5557")
    while True:
        message = await receiver.recv_json()  
        print(f'Via PULL: {message}')

async def do_subscriber():
    subscriber = context.socket(zmq.SUB)  
    subscriber.connect("tcp://localhost:5556")
    subscriber.setsockopt_string(zmq.SUBSCRIBE, '')
    while True:
        message = await subscriber.recv_json()  
        print(f'Via SUB: {message}')

loop = asyncio.get_event_loop()
loop.create_task(do_receiver())  
loop.create_task(do_subscriber())
loop.run_forever()

This code sample does the same as before, except that now we’re taking advantage of coroutines to restructure everything. Now we can deal with each socket in isolation. We’ve created two coroutine functions, one for each socket, and this one is for the PULL socket.

We’re using the asyncio support in pyzmq, which means that all send() and recv() calls must use the await keyword. The Poller no longer appears anywhere, because it’s been integrated into the asyncio event loop itself.

This is the handler for the SUB socket. The structure is very similar to the PULL socket’s handler, but that need not have been the case. If more complex logic had been required, we’d have been able to easily add it here, fully encapsulated within the SUB-handler code only.

Again: the asyncio-compatible sockets require the await keyword to send and receive.

The extra lines required to start the asyncio event loop and create the tasks for each socket. I’ve cut a few corners here, and omitted all error-handling and cleanup, because I want to emphasize the impact on code layout further up.

Example 4-15. The application layer: producing metrics

import argparse
from asyncio import get_event_loop, gather, sleep, CancelledError
from random import randint, uniform
from datetime import datetime as dt
from datetime import timezone as tz
from contextlib import suppress
import zmq, zmq.asyncio, psutil
from signal import SIGINT

# zmq.asyncio.install()  
ctx = zmq.asyncio.Context()

async def stats_reporter(color: str):  
    p = psutil.Process()
    sock = ctx.socket(zmq.PUB)  
    sock.setsockopt(zmq.LINGER, 1)
    sock.connect('tcp://localhost:5555')  
    with suppress(CancelledError):  
        while True:  
            await sock.send_json(dict(  
                color=color,
                timestamp=dt.now(tz=tz.utc).isoformat(),  
                cpu=p.cpu_percent(),
                mem=p.memory_full_info().rss / 1024 / 1024
            ))
            await sleep(1)
    sock.close()  

async def main(args):
    leak = []
    with suppress(CancelledError):
        while True:
            sum(range(randint(1_000, 10_000_000)))  
            await sleep(uniform(0, 1))
            leak += [0] * args.leak

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--color', type=str)  
    parser.add_argument('--leak', type=int, default=0)
    args = parser.parse_args()
    loop = get_event_loop()
    loop.add_signal_handler(SIGINT, loop.call_soon, loop.stop)  
    tasks = gather(main(args), stats_reporter(args.color))  
    loop.run_forever()
    print('Leaving...')
    for t in asyncio.Task.all_tasks():
        t.cancel()  
    loop.run_until_complete(tasks)
    ctx.term()  

In versions of pyzmq below 17.0.0, it was necessary to use this explicit zmq.asyncio.install() command to enable Asyncio support. At the time of writing, version 17 is currently in beta but hopefully it will have a stable release by the time you read this.

This coroutine function will run as a long-lived coroutine, continually sending out data to the server process.

Create a ØMQ socket! There are different flavors of socket. This one is a PUB socket type, which allows one-way messages to be sent to another ØMQ socket. This socket has—as the ØMQ guide says—superpowers. It will automatically handle all reconnect and buffering logic for us.

Connect to the server.

Our shutdown sequence is driven by KeyboardInterrupt, further down. When the signal is received, we’ll cancel all the tasks. Here we handle the raised CancelledError with the handy suppress() context manager from the contextlib standard library module.

Iterate forever, sending out data to the server.

Since ØMQ knows how to work with complete messages, and not just chunks off a bytestream, it opens the door to a bunch of useful wrappers around the usual sock.send() idiom: here, we use one of those helper methods, send_json(), which will automatically serialize the argument into JSON. This allows us to use a dict() directly.

A reliable way to transmit datetime information is via the ISO 8601 format. This is especially true if you have to pass datetime data between software written in different languages, since the vast majority of language implementations will be able to work with this standard.

To end up here, we must have received the CancelledError exception resulting from task cancellation. The ØMQ socket must be closed to allow program shutdown.

The main() function symbolizes the actual microservice application. Fake work is produced with this sum over random numbers, just to give us some non-zero data to view in the visualization layer a bit later.

We’re going to create multiple instances of this application, so it would be convenient to be able to distinguish between them (later, in the graphs) with a --color parameter.

When a SIGINT signal is received (e.g., pressing Ctrl-C), schedule a call to stop the loop.

Create and gather tasks for each of the coroutine functions.

Having received the shutdown signal, cancel the tasks. This will raise a CancelledError inside all of the coroutines represented in the tasks group. After cancellation, it is still necessary to run the tasks to completion, by allowing them the chance to handle the cancellation appropriately. For example, we must close the ØMQ socket in order to shut down at all.

Finally, the ØMQ context can be terminated.

Example 4-16. The collection layer: this server collects process stats

# metric-server.py
import asyncio
from contextlib import suppress
import zmq
import zmq.asyncio
import aiohttp
from aiohttp import web
from aiohttp_sse import sse_response
from weakref import WeakSet
import json

# zmq.asyncio.install()
ctx = zmq.asyncio.Context()
connections = WeakSet()  

async def collector():
    sock = ctx.socket(zmq.SUB)  
    sock.setsockopt_string(zmq.SUBSCRIBE, '')  
    sock.bind('tcp://*:5555')  
    with suppress(asyncio.CancelledError):
        while True:
            data = await sock.recv_json()  
            print(data)
            for q in connections:
                await q.put(data)  
    sock.close()

async def feed(request):  
    queue = asyncio.Queue()
    connections.add(queue)  
    with suppress(asyncio.CancelledError):
        async with sse_response(request) as resp:  
            while True:
                data = await queue.get()  
                print('sending data:', data)
                resp.send(json.dumps(data))  
    return resp

async def index(request):  
    return aiohttp.web.FileResponse('./charts.html')

async def start_collector(app):  
    app['collector'] = app.loop.create_task(collector())

async def stop_collector(app):
    print('Stopping collector...')
    app['collector'].cancel()  
    await app['collector']
    ctx.term()

if __name__ == '__main__':
    app = web.Application()
    app.router.add_route('GET', '/', index)
    app.router.add_route('GET', '/feed', feed)
    app.on_startup.append(start_collector)  
    app.on_cleanup.append(stop_collector)
    web.run_app(app, host='127.0.0.1', port=8088)

One half of this program will receive data from other applications, and the other half will provide data to browser clients via server-sent events (SSE). We use a WeakSet() to keep track of all the currently connected web clients. Each connected client will have an associated Queue() instance, so this connections identifier is really a set of queues.

Recall that in the application layer, we used a zmq.PUB socket; here in the collection layer we use its partner, the zmq.SUB socket type. This ØMQ socket can only receive, not send.

For the zmq.SUB socket type, it is required to provide a subscription name, but for our goals we’ll just take everything that comes in, hence the empty topic name.

Here we bind the zmq.SUB socket. Think about that for second! In “pubsub” configurations you usually have to make the pub end the server (bind()) and the sub end the client (connect()). ØMQ is different: either end can be the server. For our use-case this is important, because each of our application-layer instances will be connecting to the same collection server domain name and not the other way round.

The support for asyncio in pyzmq allows us to await on data from our connected apps. And not only that, but the incoming data will be automatically deserialized from JSON (yes, this means data is a dict()).

Recall that our connections set holds a queue for every connected web client? Now that data has been received, it’s time to send it to all the clients: the data is placed onto each queue.

The feed() coroutine function will create coroutines for each connected web client. Internally, server-sent events are used to push data to the web clients.

As described earlier, each web client will have its own queue instance, in order to receive data from the collector() coroutine. The queue instance is added to the connections set, but because connections is a weak set, the entry will automatically be removed from connections when the queue goes out of scope, i.e., when a web client disconnects. Weakrefs are really great for simplifying these kinds of bookkeeping tasks.

The aiohttp_sse package provides the sse_response() context manager. This gives us a scope inside which to push data to the web client.

We remain connected to the web client, and wait for data on this specific client’s queue.

As soon as the data comes in (inside collector()) it will be sent to the connected web client. Note that we reserialize the data dict here. An optimization to the code shown here would be to avoid deserializing JSON in collector(), and instead use sock.recv_string to avoid the serialization round-trip. Of course, in a real scenario you might want to deserialize in the collector anyway, and perform some validation on the data before sending to the browser client. So many choices!

The index() endpoint is the primary page-load, and here we serve a static file called charts.html.

The aiohttp library provides facilities for you to hook in additional long-lived coroutines you might need. With the collector() coroutine, we have exactly that situation, so we create a startup coroutine start_collector(), and a shutdown coroutine. These will be called during specific phases of aiohttp’s startup and shutdown sequence. Note that we add the collector task to the app itself, which implements a mapping protocol so that you can use it like a dict.

Here you can see that we obtain our collector() coroutine off the app identifier and call cancel() on that.

Finally, you can see where the custom startup and shutdown coroutines are hooked in: the app instance provides hooks to which your custom coroutines may be appended.

Example 4-17. The visualization layer, which is a fancy way of saying “the browser”

<snip>
var evtSource = new EventSource("/feed");  
evtSource.onmessage = function(e) {
    var obj = JSON.parse(e.data);  
    if (!(obj.color in cpu)) {
        add_timeseries(cpu, cpu_chart, obj.color);
    }
    if (!(obj.color in mem)) {
        add_timeseries(mem, mem_chart, obj.color);
    }
    cpu[obj.color].append(
        Date.parse(obj.timestamp), obj.cpu);  
    mem[obj.color].append(
        Date.parse(obj.timestamp), obj.mem);
};
<snip>

Create a new EventSource instance on the /feed URL. The browser will connect to /feed on our server, metric-server.py. Note that the browser will automatically try to reconnect if the connection is lost. Server-sent events are often overlooked, but there are many situations where the simplicity of SSE might be preferred over websockets.

The onmessage() event will fire every time the server sends data. Here the data is parsed as JSON.

Recall that the cpu identifier is a mapping of color to a TimeSeries() instance. Here, we obtain that time series and append data to it. We also obtain the timestamp and parse it to get the correct format required by the chart.

Example 4-18. Starting the collector

$ python metric-server.py
======== Running on http://127.0.0.1:8088 ========
(Press CTRL+C to quit)

Example 4-19. Starting the monitored applications

$ python backend-app.py --color red &
$ python backend-app.py --color blue --leak 10000 &
$ python backend-app.py --color green --leak 100000 &