mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 11:03:38 +08:00
fawce
145 lines
3.4 KiB
Python
145 lines
3.4 KiB
Python
"""
|
|
Abstract base class for Feed and Merge.
|
|
|
|
Component
|
|
|
|
|
Aggregate
|
|
|
|
|
/ \
|
|
Feed Merge
|
|
|
|
"""
|
|
import logbook
|
|
|
|
import zipline.protocol as zp
|
|
from zipline.core.component import Component
|
|
from zipline.protocol import CONTROL_PROTOCOL, COMPONENT_TYPE
|
|
from zipline.transitions import WorkflowMeta
|
|
from zipline.utils.protocol_utils import Enum
|
|
|
|
|
|
log = logbook.Logger('Aggregate')
|
|
|
|
# =================
|
|
# State Transitions
|
|
# =================
|
|
|
|
INIT, READY, DRAINING = AGGREGATE_STATES = \
|
|
Enum( 'INIT', 'READY', 'DRAINING')
|
|
|
|
AGGREGATE_TRANSITIONS = dict(
|
|
do_start = (-1 , INIT) ,
|
|
do_run = (INIT , READY) ,
|
|
do_drain = (READY , DRAINING) ,
|
|
)
|
|
|
|
# =========
|
|
# Component
|
|
# =========
|
|
|
|
class Aggregate(Component):
|
|
"""
|
|
Abstract superclass to Merge & Feed. Acts on two sockets
|
|
|
|
- pull_socket
|
|
- feed_socket
|
|
|
|
Both use ``sources`` for buffering.
|
|
|
|
Feed and Merge define these differently.
|
|
"""
|
|
|
|
abstract = True
|
|
__metaclass__ = WorkflowMeta
|
|
|
|
@property
|
|
def get_type(self):
|
|
return COMPONENT_TYPE.CONDUIT
|
|
|
|
def add_source(self, source_id):
|
|
self.sources[source_id] = []
|
|
|
|
# -------------
|
|
# Core Methods
|
|
# -------------
|
|
|
|
def do_work(self):
|
|
|
|
# -------------
|
|
# Work Dispatch
|
|
# -------------
|
|
if self.socks.get(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.sources) == self.ds_finished_counter:
|
|
# Drain any remaining messages in the buffer
|
|
log.debug("Draining Feed")
|
|
|
|
self.state = DRAINING
|
|
|
|
self.drain()
|
|
self.signal_done()
|
|
else:
|
|
event = self.unframe(message)
|
|
self.append(event)
|
|
|
|
if self.is_full():
|
|
event = self.next()
|
|
|
|
if event:
|
|
self.send(event)
|
|
else:
|
|
pass
|
|
|
|
# -------------
|
|
# Flow Control
|
|
# -------------
|
|
|
|
def drain(self):
|
|
"""
|
|
Send all messages in the buffer.
|
|
"""
|
|
while self.pending_messages() > 0:
|
|
event = self.next()
|
|
self.heartbeat()
|
|
if event:
|
|
self.send(event)
|
|
|
|
def send(self, event):
|
|
"""
|
|
Send the (chronologically) next message in the buffer.
|
|
"""
|
|
self.feed_socket.send(self.frame(event), self.zmq.NOBLOCK)
|
|
self.sent_counters[event.source_id] += 1
|
|
self.sent_count += 1
|
|
|
|
def is_full(self):
|
|
"""
|
|
Indicates whether the buffer has messages in buffer for all
|
|
un-DONE, blocking sources.
|
|
"""
|
|
for source_id, events in self.sources.iteritems():
|
|
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.sources.itervalues():
|
|
total += len(events)
|
|
return total
|
|
|
|
def __len__(self):
|
|
"""
|
|
Buffer's length is same as internal map holding separate
|
|
sorted arrays of events keyed by source id.
|
|
"""
|
|
return len(self.sources)
|