catalyst/zipline/core/monitor.py

import zmq
import time
import gevent
import itertools
import logbook
import gevent_zeromq

from collections import OrderedDict

from zipline.utils.gpoll import _Poller as GeventPoller
from zipline.protocol import CONTROL_PROTOCOL, CONTROL_FRAME, \
    CONTROL_UNFRAME, CONTROL_STATES, INVALID_CONTROL_FRAME \

INIT, SOURCES_READY, RUNNING, TERMINATE = CONTROL_STATES

CONTROLLER_TRANSITIONS = frozenset([
    (-1            , INIT),
    (INIT          , SOURCES_READY),
    (SOURCES_READY , RUNNING),
    (INIT          , TERMINATE),
    (SOURCES_READY , TERMINATE),
    (RUNNING       , TERMINATE),
])

class UnknownChatter(Exception):
    def __init__(self, name):
        self.named = name
    def __str__(self):
        return """Component calling itself "%s" talking on unexpected channel"""\
            % self.named


log = logbook.Logger('Controller')

class Controller(object):
    """
    A N to M messaging system for inter component communication.

    :param pub_socket: Socket to publish messages, the starting
                       point of :func message_listener: .

    :param route_socket: Socket to listen for status updates for
                         the individual components.
                         :func message_sender: .

    Topology is the set of components we expect to show up.
    States are the transitions the sytems go through. The
    simplest is from RUNNING -> NOT RUNNING .

    Usage::

        controller = Controller(
            'tcp://127.0.0.1:5000',
            'tcp://127.0.0.1:5001',
        )

        # typically you'd want to run this async to your main
        # program since it blocks indefinetely.
        controller.manage(
            [ TOPOLOGY ]
            [ STATES ]
        )

    """

    debug = False
    period = 1

    def __init__(self, pub_socket, route_socket):

        self.context    = None
        self.zmq        = None
        self.zmq_poller = None

        self.running = False
        self.polling = False
        self.tracked = set()
        self.responses = set()

        self.ctime    = 0
        self.tic      = time.time()
        self.freeform = False
        self._state   = -1

        self.associated = []

        self.pub_socket   = pub_socket
        self.route_socket = route_socket

        self.error_replay = OrderedDict()

    def init_zmq(self, flavor):

        assert self.zmq_flavor in ['thread', 'mp', 'green']

        if flavor == 'mp':
            self.zmq        = zmq
            self.context    = self.zmq.Context()
            self.zmq_poller = self.zmq.Poller
            return
        if flavor == 'thread':
            self.zmq        = zmq
            self.context    = self.zmq.Context.instance()
            self.zmq_poller = self.zmq.Poller
            return
        if flavor == 'green':
            self.zmq        = gevent_zeromq.zmq
            self.context    = self.zmq.Context.instance()
            self.zmq_poller = GeventPoller
            return
        if flavor == 'pypy':
            self.zmq        = zmq
            self.context    = self.zmq.Context.instance()
            self.zmq_poller = self.zmq.Poller
            return

    def manage(self, topology, states=None, context=None):
        """
        Give the controller a set set of components to manage and
        a set of state transitions for the entire system.
        """

        # A freeform topology is where we heartbeat with anything
        # that shows up.
        if topology == 'freeform':
            self.freeform = True
            self.topology = frozenset([])
        else:
            self.freeform = False
            self.topology = frozenset(topology)

        self.polling = True
        self.state = CONTROL_STATES.INIT

    @property
    def state(self):
        return self._state

    @state.setter
    def state(self, new):
        old, self._state = self._state, new

        if (old, new) not in CONTROLLER_TRANSITIONS:
            raise RuntimeError("Invalid State Transition : %s -> %s" %(old, new))
        else:
            log.error("State Transition : %s -> %s" %(old, new))

    def run(self):
        self.running = True
        self.init_zmq(self.zmq_flavor)

        try:
            return self._poll() # use a python loop
        except KeyboardInterrupt:
            log.debug('Shutdown event loop')

    def log_status(self):
        """
        Snapshot of the tracked components at every period.
        """
        #log.info("Tracking component : %s" % ([c for c in self.tracked],))
        pass

    def replay_errors(self):
        """
        Replay the errors in the order they were reported to the
        controller.
        """
        return [ a for a in sorted(self.replay_errors.keys())]

    # -------------
    # Publications
    # -------------

    def send_heart(self):
        if not self.running:
            return

        heartbeat_frame = CONTROL_FRAME(
            CONTROL_PROTOCOL.HEARTBEAT,
            str(self.ctime)
        )
        self.pub.send(heartbeat_frame)

    def send_hardkill(self):
        if not self.running:
            return

        kill_frame = CONTROL_FRAME(
            CONTROL_PROTOCOL.KILL,
            ''
        )
        self.pub.send(kill_frame)

    def send_softkill(self):
        if not self.running:
            return

        soft_frame = CONTROL_FRAME(
            CONTROL_PROTOCOL.SHUTDOWN,
            ''
        )
        self.pub.send(soft_frame)

    # -----------
    # Event Loops
    # -----------

    def _poll(self):

        assert self.route_socket
        assert self.pub_socket
        assert self.cancel_socket

        # -- Publish --
        # =============
        self.pub = self.context.socket(self.zmq.PUB)
        self.pub.bind(self.pub_socket)

        # -- Cancel --
        # =============
        assert isinstance(self.cancel_socket,basestring), self.cancel_socket
        self.cancel = self.context.socket(self.zmq.REP)
        self.cancel.connect(self.cancel_socket)

        # -- Router --
        # =============
        self.router = self.context.socket(self.zmq.ROUTER)
        self.router.bind(self.route_socket)


        poller = self.zmq.Poller()
        poller.register(self.router, self.zmq.POLLIN)
        poller.register(self.cancel, self.zmq.POLLIN)

        self.associated += [self.pub, self.router, self.cancel]

        # TODO: actually do this
        self.state = CONTROL_STATES.SOURCES_READY

        buffer = []

        for i in itertools.count(0):
            self.log_status()
            self.responses = set()

            self.ctime = time.time()
            self.send_heart()

            while self.polling:
                # Reset the responses for this cycle

                socks = dict(poller.poll(self.period))
                tic = time.time()

                if tic - self.ctime > self.period:
                    break

                if socks.get(self.router) == self.zmq.POLLIN:
                    rawmessage = self.router.recv()

                    if rawmessage:
                        buffer.append(rawmessage)

                    try:
                        if not self.router.getsockopt(self.zmq.RCVMORE):
                            self.handle_recv(buffer[:])
                            buffer = []
                    except INVALID_CONTROL_FRAME:
                        log.error('Invalid frame', rawmessage)
                        pass

                if socks.get(self.cancel) == self.zmq.POLLIN:
                    log.info('Received Cancellation')
                    rawmessage = self.cancel.recv()
                    self.cancel.send('')
                    self.shutdown(soft=True)
                    break

            self.beat()

            if self.zmq_flavor == 'green':
                gevent.sleep(0)

            if self.state is CONTROL_STATES.TERMINATE:
                break

            if not self.polling:
                break

        # After loop exits
        self.terminated = True

    def beat(self):

        # These the set overloaded operations
        # A & B ~ set.intersection
        # A - B ~ set.difference

        # * good - Components we are currently tracking and who just sent
        #          us back the right response.
        # * bad - Components we are currently tracking but who did not
        #         send us back a response.
        # * new - Components we haven't heard from yet, but sent back the
        #         right response.

        good = self.tracked   & self.responses
        bad  = self.tracked   - good
        new  = self.responses - good

        for component in new:
            self.new(component)

        for component in bad:
            self.fail(component)

    # --------------
    # Init Handlers
    # --------------

    def new_source(self):
        if self.state is CONTROL_STATES.RUNNING:
            self.state = SOURCES_READY

    def new_universal(self):
        pass

    # The various "states of being that a component can inform us
    # of
    def new(self, component):
        if self.state is CONTROL_STATES.TERMINATE:
            return

        log.info(' Now Tracking "%s" ' % component)

        universal = self.new_universal
        init_handlers = {
            'FEED' : self.new_source,
        }

        if component in self.topology or self.freeform:
            init_handlers.get(component, universal)()
            self.tracked.add(component)
        else:
            # Some sort of socket collision has occured, this is
            # a very bad failure mode.
            raise UnknownChatter(component)

    # ------------------
    # Epic Fail Handling
    # ------------------

    def fail_universal(self):
        pass
        # TODO: this requires higher order functionality
        #log.error('System in exception state, shutting down')
        #self.shutdown(soft=True)

    def fail(self, component):
        if self.state is CONTROL_STATES.TERMINATE:
            return

        universal = self.fail_universal
        fail_handlers = { }

        if component in self.topology or self.freeform:
            log.error('Component "%s" timed out' % component)
            self.tracked.remove(component)
            fail_handlers.get(component, universal)()

    # -------------------
    # Completion Handling
    # -------------------

    def done(self, component):
        log.info('Component "%s" signaled done.' % component)

    # --------------
    # Error Handling
    # --------------

    def exception_universal(self):
        """
        Shutdown the system on failure.
        """
        log.error('System in exception state, shutting down')
        self.shutdown(soft=True)

    def exception(self, component, failure):
        universal = self.exception_universal
        exception_handlers = { }

        if component in self.topology or self.freeform:
            self.error_replay[(component, time.time())] = failure
            log.error('Component in exception state: %s' % component)

            exception_handlers.get(component, universal)()
        else:
            raise UnknownChatter(component)

    # -----------------
    # Protocol Handling
    # -----------------

    def handle_recv(self, msg):
        """
        Check for proper framing at the transport layer.
        Seperates the proper frames from anything else that might
        be coming over the wire. Which shouldn't happen ...  right?
        """
        identity = msg[0]
        id, status = CONTROL_UNFRAME(msg[1])

        # A component is telling us its alive:
        if id is CONTROL_PROTOCOL.OK:

            if status == str(self.ctime):
                self.responses.add(identity)
            else:
                # Otherwise its something weird and we don't know
                # what to do so just say so, probably line noise
                # from ZeroMQ
                log.error("Weird stuff happened: %s" % msg)

        # A component is telling us it failed, and how
        if id is CONTROL_PROTOCOL.EXCEPTION:
            self.exception(identity, status)

        # A component is telling us its done with work and won't
        # be talking to us anymore
        if id is CONTROL_PROTOCOL.DONE:
            self.done(identity)

    # -------------------
    # Hooks for Endpoints
    # -------------------

    # These are all connects so no complex allocation logic is
    # needed. Dealers and Subscribers can all come and go as a
    # function of time without impacting flow of the whole
    # system.

    def message_sender(self, identity, context = None):
        """
        Spin off a socket used for sending messages to this
        controller.
        """

        if not context:
            context = self.zmq.Context.instance()

        s = context.socket(zmq.DEALER)
        s.setsockopt(zmq.IDENTITY, identity)
        s.connect(self.route_socket)

        self.associated.append(s)
        return s

    def message_listener(self, context = None):
        """
        Spin off a socket used for receiving messages from this
        controller.
        """

        if not context:
            context = self.zmq.Context.instance()

        s = context.socket(zmq.SUB)
        s.connect(self.pub_socket)
        s.setsockopt(zmq.SUBSCRIBE, '')

        self.associated.append(s)
        return s

    def do_error_replay(self):
        for (component, time), error in self.error_replay.iteritems():
            log.info('Error Log for -- %s --:\n%s' %
                (component, error))

    def shutdown(self, hard=False, soft=True, context=None):

        if self.state is CONTROL_STATES.TERMINATE:
            return

        if not self.polling:
            return

        self.polling = False

        assert hard or soft, """ Must specify kill hard or soft """

        if hard:
            self.state = CONTROL_STATES.TERMINATE

            log.info('Hard Shutdown')

            #for asoc in self.associated:
                #asoc.close()

        if soft:
            self.state = CONTROL_STATES.TERMINATE

            log.info('Soft Shutdown')
            self.send_softkill()

            #for asoc in self.associated:
                #asoc.close()

        self.do_error_replay()

if __name__ == '__main__':

    print 'Running on '\
        'tcp://127.0.0.1:5000 '\
        'tcp://127.0.0.1:5001 '

    controller = Controller(
        'tcp://127.0.0.1:5000',
        'tcp://127.0.0.1:5001',
    )
    controller.zmq_flavor = 'green'

    controller.manage(
        'freeform',
        []
    )
    controller.run()