mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-06 13:00:45 +08:00
434 lines
13 KiB
Python
434 lines
13 KiB
Python
"""
|
|
Commonly used messaging components.
|
|
"""
|
|
import datetime
|
|
import ujson as json
|
|
|
|
import zipline.util as qutil
|
|
from zipline.component import Component
|
|
from zipline.protocol import CONTROL_PROTOCOL
|
|
|
|
class ComponentHost(Component):
|
|
"""
|
|
Components that can launch multiple sub-components, synchronize their
|
|
start, and then wait for all components to be finished.
|
|
"""
|
|
|
|
def __init__(self, addresses, gevent_needed=False):
|
|
Component.__init__(self)
|
|
self.addresses = addresses
|
|
self.gevent_needed = gevent_needed
|
|
|
|
self.init()
|
|
|
|
def init(self):
|
|
|
|
# workaround for defect in threaded use of strptime:
|
|
# http://bugs.python.org/issue11108
|
|
qutil.parse_date("2012/02/13-10:04:28.114")
|
|
|
|
self.components = {}
|
|
self.sync_register = {}
|
|
self.timeout = datetime.timedelta(seconds=5)
|
|
|
|
self.feed = ParallelBuffer()
|
|
self.merge = MergedParallelBuffer()
|
|
self.passthrough = PassthroughTransform()
|
|
|
|
#register the feed and the merge
|
|
self.register_components([self.feed, self.merge, self.passthrough])
|
|
|
|
def register_controller(self, controller):
|
|
self.controller = controller
|
|
|
|
for component in self.components.itervalues():
|
|
component.controller = controller
|
|
|
|
def register_components(self, components):
|
|
for component in components:
|
|
component.gevent_needed = self.gevent_needed
|
|
component.addresses = self.addresses
|
|
|
|
if self.controller:
|
|
component.controller = self.controller
|
|
|
|
self.components[component.get_id] = component
|
|
self.sync_register[component.get_id] = datetime.datetime.utcnow()
|
|
|
|
if(isinstance(component, DataSource)):
|
|
self.feed.add_source(component.get_id)
|
|
if(isinstance(component, BaseTransform)):
|
|
self.merge.add_source(component.get_id)
|
|
|
|
def unregister_component(self, component_id):
|
|
del self.components[component_id]
|
|
del self.sync_register[component_id]
|
|
|
|
def setup_sync(self):
|
|
"""
|
|
"""
|
|
qutil.LOGGER.debug("Connecting sync server.")
|
|
|
|
self.sync_socket = self.context.socket(self.zmq.REP)
|
|
self.sync_socket.bind(self.addresses['sync_address'])
|
|
|
|
# There is a namespace collision between three classes
|
|
# which use the self.poller property to mean different
|
|
# things.
|
|
# =====================================================
|
|
self.sync_poller = self.zmq.Poller()
|
|
self.sync_poller.register(self.sync_socket, self.zmq.POLLIN)
|
|
# =====================================================
|
|
|
|
self.sockets.append(self.sync_socket)
|
|
|
|
def open(self):
|
|
for component in self.components.values():
|
|
self.launch_component(component)
|
|
self.launch_controller()
|
|
|
|
def is_timed_out(self):
|
|
cur_time = datetime.datetime.utcnow()
|
|
|
|
if len(self.components) == 0:
|
|
qutil.LOGGER.info("Component register is empty.")
|
|
return True
|
|
|
|
for source, last_dt in self.sync_register.iteritems():
|
|
if (cur_time - last_dt) > self.timeout:
|
|
qutil.LOGGER.info(
|
|
"Time out for {source}. Current component registery: {reg}".
|
|
format(source=source, reg=self.components)
|
|
)
|
|
return True
|
|
|
|
return False
|
|
|
|
def loop(self):
|
|
|
|
while not self.is_timed_out():
|
|
# wait for synchronization request
|
|
socks = dict(self.sync_poller.poll(self.heartbeat_timeout)) #timeout after 2 seconds.
|
|
|
|
if self.sync_socket in socks and socks[self.sync_socket] == self.zmq.POLLIN:
|
|
msg = self.sync_socket.recv()
|
|
parts = msg.split(':')
|
|
sync_id, status = parts
|
|
|
|
# TODO: move into frame protocol
|
|
if len(parts) != 2:
|
|
qutil.LOGGER.info("got bad confirm: {msg}".format(msg=msg))
|
|
continue
|
|
|
|
if status == str(CONTROL_PROTOCOL.DONE): # TODO: other way around
|
|
qutil.LOGGER.info("{id} is DONE".format(id=sync_id))
|
|
self.unregister_component(sync_id)
|
|
else:
|
|
self.sync_register[sync_id] = datetime.datetime.utcnow()
|
|
|
|
#qutil.LOGGER.info("confirmed {id}".format(id=msg))
|
|
# send synchronization reply
|
|
self.sync_socket.send('ack', self.zmq.NOBLOCK)
|
|
|
|
# ------------------
|
|
# Simulation Control
|
|
# ------------------
|
|
|
|
def launch_controller(self, controller):
|
|
raise NotImplementedError
|
|
|
|
def launch_component(self, component):
|
|
raise NotImplementedError
|
|
|
|
def teardown_component(self, component):
|
|
raise NotImplementedError
|
|
|
|
class ParallelBuffer(Component):
|
|
"""
|
|
Connects to N PULL sockets, publishing all messages received to a PUB
|
|
socket. Published messages are guaranteed to be in chronological order
|
|
based on message property dt. Expects to be instantiated in one execution
|
|
context (thread, process, etc) and run in another.
|
|
"""
|
|
|
|
def __init__(self):
|
|
Component.__init__(self)
|
|
self.sent_count = 0
|
|
self.received_count = 0
|
|
self.draining = False
|
|
self.data_buffer = {}
|
|
self.ds_finished_counter = 0
|
|
|
|
def init(self):
|
|
pass
|
|
|
|
@property
|
|
def get_id(self):
|
|
return "FEED"
|
|
|
|
def add_source(self, source_id):
|
|
self.data_buffer[source_id] = []
|
|
|
|
def open(self):
|
|
self.pull_socket = self.bind_data()
|
|
self.feed_socket = self.bind_feed()
|
|
|
|
def do_work(self):
|
|
# wait for synchronization reply from the host
|
|
socks = dict(self.poll.poll(self.heartbeat_timeout)) #timeout after 2 seconds.
|
|
|
|
if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
|
|
msg = self.control_in.recv()
|
|
|
|
if self.pull_socket in socks and socks[self.pull_socket] == self.zmq.POLLIN:
|
|
message = self.pull_socket.recv()
|
|
|
|
if message == str(CONTROL_PROTOCOL.DONE):
|
|
self.ds_finished_counter += 1
|
|
|
|
if len(self.data_buffer) == self.ds_finished_counter:
|
|
#drain any remaining messages in the buffer
|
|
self.drain()
|
|
self.signal_done()
|
|
else:
|
|
event = json.loads(message)
|
|
self.append(event[u'id'], event)
|
|
self.send_next()
|
|
|
|
def append(self, source_id, value):
|
|
"""
|
|
Add an event to the buffer for the source specified by
|
|
source_id.
|
|
"""
|
|
self.data_buffer[source_id].append(value)
|
|
self.received_count += 1
|
|
|
|
def next(self):
|
|
"""
|
|
Get the next message in chronological order.
|
|
"""
|
|
if not(self.is_full() or self.draining):
|
|
return
|
|
|
|
cur = None
|
|
earliest = None
|
|
for events in self.data_buffer.values():
|
|
if len(events) == 0:
|
|
continue
|
|
cur = events
|
|
if (earliest == None) or (cur[0]['dt'] <= earliest[0]['dt']):
|
|
earliest = cur
|
|
|
|
if earliest != None:
|
|
return earliest.pop(0)
|
|
|
|
def is_full(self):
|
|
"""
|
|
Indicates whether the buffer has messages in buffer for
|
|
all un-DONE sources.
|
|
"""
|
|
for events in self.data_buffer.values():
|
|
if len(events) == 0:
|
|
return False
|
|
return True
|
|
|
|
def pending_messages(self):
|
|
"""
|
|
Returns the count of all events from all sources in the
|
|
buffer.
|
|
"""
|
|
total = 0
|
|
for events in self.data_buffer.values():
|
|
total += len(events)
|
|
return total
|
|
|
|
def drain(self):
|
|
"""
|
|
Send all messages in the buffer
|
|
"""
|
|
self.draining = True
|
|
while(self.pending_messages() > 0):
|
|
self.send_next()
|
|
|
|
def send_next(self):
|
|
"""
|
|
Send the (chronologically) next message in the buffer.
|
|
"""
|
|
if(not(self.is_full() or self.draining)):
|
|
return
|
|
|
|
event = self.next()
|
|
if(event != None):
|
|
self.feed_socket.send(json.dumps(event), self.zmq.NOBLOCK)
|
|
self.sent_count += 1
|
|
|
|
def __len__(self):
|
|
"""
|
|
Buffer's length is same as internal map holding separate
|
|
sorted arrays of events keyed by source id.
|
|
"""
|
|
return len(self.data_buffer)
|
|
|
|
|
|
class MergedParallelBuffer(ParallelBuffer):
|
|
"""
|
|
Merges multiple streams of events into single messages.
|
|
"""
|
|
|
|
def __init__(self):
|
|
ParallelBuffer.__init__(self)
|
|
|
|
self.init()
|
|
|
|
def init(self):
|
|
pass
|
|
|
|
@property
|
|
def get_id(self):
|
|
return "MERGE"
|
|
|
|
def open(self):
|
|
self.pull_socket = self.bind_merge()
|
|
self.feed_socket = self.bind_result()
|
|
|
|
def next(self):
|
|
"""Get the next merged message from the feed buffer."""
|
|
if(not(self.is_full() or self.draining)):
|
|
return
|
|
|
|
#get the raw event from the passthrough transform.
|
|
result = self.data_buffer["PASSTHROUGH"].pop(0)['value']
|
|
for source, events in self.data_buffer.iteritems():
|
|
if source == "PASSTHROUGH":
|
|
continue
|
|
if len(events) > 0:
|
|
cur = events.pop(0)
|
|
result[source] = cur['value']
|
|
return result
|
|
|
|
|
|
class BaseTransform(Component):
|
|
"""
|
|
Top level execution entry point for the transform::
|
|
|
|
- connects to the feed socket to subscribe to events
|
|
- connets to the result socket (most oftened bound by a TransformsMerge) to PUSH transforms
|
|
- processes all messages received from feed, until DONE message received
|
|
- pushes all transforms
|
|
- sends DONE to result socket, closes all sockets and context
|
|
|
|
Parent class for feed transforms. Subclass and override transform
|
|
method to create a new derived value from the combined feed.
|
|
"""
|
|
|
|
def __init__(self, name):
|
|
Component.__init__(self)
|
|
self.state = {}
|
|
self.state['name'] = name
|
|
|
|
self.init()
|
|
|
|
def init(self):
|
|
pass
|
|
|
|
@property
|
|
def get_id(self):
|
|
return self.state['name']
|
|
|
|
def open(self):
|
|
"""
|
|
Establishes zmq connections.
|
|
"""
|
|
#create the feed.
|
|
self.feed_socket = self.connect_feed()
|
|
#create the result PUSH
|
|
self.result_socket = self.connect_merge()
|
|
|
|
def do_work(self):
|
|
"""
|
|
Loops until feed's DONE message is received:
|
|
- receive an event from the data feed
|
|
- call transform (subclass' method) on event
|
|
- send the transformed event
|
|
"""
|
|
socks = dict(self.poll.poll(self.heartbeat_timeout))
|
|
|
|
if self.control_in in socks and socks[self.control_in] == self.zmq.POLLIN:
|
|
msg = self.control_in.recv()
|
|
|
|
if self.feed_socket in socks and socks[self.feed_socket] == self.zmq.POLLIN:
|
|
message = self.feed_socket.recv()
|
|
if message == str(CONTROL_PROTOCOL.DONE):
|
|
self.signal_done()
|
|
return
|
|
|
|
event = json.loads(message)
|
|
cur_state = self.transform(event)
|
|
cur_state['dt'] = event['dt']
|
|
cur_state['id'] = self.state['name']
|
|
|
|
self.result_socket.send(json.dumps(cur_state), self.zmq.NOBLOCK)
|
|
|
|
def transform(self, event):
|
|
"""
|
|
Must return the transformed value as a map with::
|
|
|
|
{name:"name of new transform", value: "value of new field"}
|
|
|
|
Transforms run in parallel and results are merged into a single map, so
|
|
transform names must be unique. Best practice is to use the self.state
|
|
object initialized from the transform configuration, and only set the
|
|
transformed value::
|
|
|
|
self.state['value'] = transformed_value
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class PassthroughTransform(BaseTransform):
|
|
|
|
def __init__(self):
|
|
BaseTransform.__init__(self, "PASSTHROUGH")
|
|
|
|
def init(self):
|
|
pass
|
|
|
|
def transform(self, event):
|
|
return {'value':event}
|
|
|
|
|
|
class DataSource(Component):
|
|
"""
|
|
Baseclass for data sources. Subclass and implement send_all - usually this
|
|
means looping through all records in a store, converting to a dict, and
|
|
calling send(map).
|
|
"""
|
|
def __init__(self, source_id):
|
|
Component.__init__(self)
|
|
self.id = source_id
|
|
self.init()
|
|
|
|
def init(self):
|
|
self.cur_event = None
|
|
|
|
@property
|
|
def get_id(self):
|
|
return self.id
|
|
|
|
def open(self):
|
|
#create the data sink. Based on http://zguide.zeromq.org/py:tasksink2
|
|
self.data_socket = self.connect_data()
|
|
|
|
def send(self, event):
|
|
"""
|
|
event is expected to be a dict
|
|
sets id and type properties in the dict
|
|
sends to the data_socket.
|
|
"""
|
|
event['id'] = self.id
|
|
event['type'] = self.get_type()
|
|
self.data_socket.send(json.dumps(event))
|
|
|
|
def get_type(self):
|
|
raise NotImplementedError
|