zmq_socket_pool.py

from eventlet.pools import Pool

class SocketPool(Pool):
    """A pool of sockets connected to a component

    If a socket times out in use, simply close if before handing it back to the
    pool and it will be discarded and a replacement inserted into the pool.
    """

    def __init__(self, address, **kwargs):
        self.address = address
        super(SocketPool, self).__init__(**kwargs)

    def __str__(self):
        return self.address

    def __repr__(self):
        return "<SocketPool: %s>" % self

    def create(self):
        socket = zmq.Context().socket(zmq.REQ)
        socket.connect(self.address)
        socket.setsockopt(zmq.LINGER, 0)
        return socket

    def put(self, socket):
        """Wrapper around superclass put, replacing socket if closed"""
        if socket.closed:
            socket = self.create()
        super(SocketPool, self).put(socket)


class SocketPoolManager(object):
    """A container for socket pools"""

    def __init__(self, components):
        """Creates a :class:`SocketPool` for each :class:`Component` in components"""
        self._pools = dict((component.name, SocketPool(component.address, max_size=20))
                           for component in components)

    def __repr__(self):
        return "\n".join("%s -> %r" % p for p in self._pools.items())

    def __getitem__(self, item):
        """Sockets are keyed off and fetched by the component's ``name``"""
        socket = self._pools[item].item()
        return socket

# Component could be anything you want here, the idea is that when I want to communicate with one I can say
# with pools['component_name'] as socket:
#     try:
#         do_some_stuff(socket)
#     except (SomeExceptionIndicatingIWantToThrowAwayTheSocket, Timeout, Whatever):
#         socket.close()
# zmq.LINGER is in there as I don't want the message to hang around after I throw away a socket