From ae935b224244100f1d1ae0f09457dbf0ec4a1add Mon Sep 17 00:00:00 2001 From: fawce Date: Tue, 17 Jan 2012 15:11:11 -0500 Subject: [PATCH] reliably passing the test for merging datasources in the feed class --- data/feed.py | 76 +++++++++++++++++++++++++++++++++------------------ qbt_client.py | 22 +++++++++++---- test.py | 3 ++ 3 files changed, 70 insertions(+), 31 deletions(-) diff --git a/data/feed.py b/data/feed.py index d8ec5afe..7e0865dc 100644 --- a/data/feed.py +++ b/data/feed.py @@ -15,6 +15,9 @@ class DataFeed(object): self.data_buffer = {} #source_id -> [] self.subscriber_count = subscriber_count + self.received_count = 0 + self.sent_count = 0 + def start_data_workers(self): """Start a sub-process for each datasource.""" @@ -51,9 +54,8 @@ class DataFeed(object): # Prepare our context and sockets self.context = zmq.Context() - counter = 0 - self.start_data_workers() - + ds_finished_counter = 0 + #create the data sink. Based on http://zguide.zeromq.org/py:tasksink2 #see: http://zguide.zeromq.org/py:taskwork2 self.data_socket = self.context.socket(zmq.PULL) @@ -63,6 +65,9 @@ class DataFeed(object): self.feed_socket = self.context.socket(zmq.PUSH) self.feed_socket.bind(self.feed_address) + #start the data source workers + self.start_data_workers() + #wait for all feed subscribers self.sync_clients() @@ -71,40 +76,59 @@ class DataFeed(object): while True: message = self.data_socket.recv() event = json.loads(message) - #self.logger.info(" count " + str(counter) + " - " + str(event['dt'])) if(event["type"] == "DONE"): - self.logger.info("DONE") - source = event[u's'] - if(self.data_workers.has_key(source)): - del(self.data_workers[source]) - if(len(self.data_workers) == 0): + ds_finished_counter += 1 + if(len(self.data_workers) == ds_finished_counter): break else: self.data_buffer[event[u's']].append(event) - counter = counter + 1 - self.send_earliest_event() + self.received_count = self.received_count + 1 + self.send_next() #drain any remaining messages in the buffer - self.send_earliest_event(drain=True) + while(self.pending_messages() > 0): + self.send_next(drain=True) - self.logger.info("Collected {n} messages".format(n=counter)) + #send the DONE message + self.feed_socket.send("DONE") + + self.logger.info("received {n} messages, sent {m} messages".format(n=self.received_count, m=self.sent_count)) self.data_socket.close() self.feed_socket.close() self.context.term() - def send_earliest_event(self, drain=False): + + def send_next(self, drain=False): + if(not(self.buffers_full() or drain)): + return + + cur = None earliest = None - next_source = None - while True: #send messages as long as we have >0 messages from each source - for source, events in self.data_buffer.iteritems(): - if(not drain and len(events) == 0 and self.data_workers.has_key(source)): - #there's no way to know that we have the next message - return - if(len(events) > 0 and (earliest == None or earliest > events[0])): - earliest = events[0]['dt'] - next_source = source - + for source, events in self.data_buffer.iteritems(): + if len(events) == 0: + continue + cur = events + if(earliest == None) or (cur[0]['dt'] <= earliest[0]['dt']): + earliest = cur - event = self.data_buffer[next_source].pop(0) - self.feed_socket.send(json.dumps(event)) \ No newline at end of file + if(earliest != None): + event = earliest.pop(0) + self.feed_socket.send(json.dumps(event)) + self.sent_count += 1 + + + def buffers_full(self): + for source, events in self.data_buffer.iteritems(): + if (len(events) == 0): + return False + return True + + def pending_messages(self): + total = 0 + for source, events in self.data_buffer.iteritems(): + total += len(events) + return total + + + \ No newline at end of file diff --git a/qbt_client.py b/qbt_client.py index d7e5abf5..ad6e440a 100644 --- a/qbt_client.py +++ b/qbt_client.py @@ -9,13 +9,15 @@ from backtest.util import * class TestClient(object): def __init__(self,feed_address, sync_address): - self.context = zmq.Context() self.logger = logging.getLogger() self.feed_address = feed_address self.sync_address = sync_address def run(self): + self.logger.info("running the client") + self.context = zmq.Context() + self.data_feed = self.context.socket(zmq.PULL) self.data_feed.connect(self.feed_address) @@ -31,10 +33,20 @@ class TestClient(object): counter = 0 prev_dt = None while True: - counter += 1 msg = self.data_feed.recv() + counter += 1 + if(msg == "DONE"): + self.logger.info("DONE!") + break event = json.loads(msg) if(prev_dt != None): - assert(event['dt'] >= prev_dt) - self.logger.info("received {n} messages".format(n=counter)) - \ No newline at end of file + if(not event['dt'] >= prev_dt): + raise Exception("message arrived out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt)) + + prev_dt = event['dt'] + if(counter % 100 == 0): + self.logger.info("received {n} messages".format(n=counter)) + + self.logger.info("received {n} messages".format(n=counter)) + self.data_feed.close() + self.context.term() \ No newline at end of file diff --git a/test.py b/test.py index 8263c2bd..a574c2d1 100644 --- a/test.py +++ b/test.py @@ -31,6 +31,9 @@ def datafeed(): client = TestClient(feed.feed_address, feed.sync_address) client.run() + logger.info("feed has {pending} messages".format(pending=feed.pending_messages())) + assert(feed.pending_messages() == 0) + def pubsub(): proc1 = multiprocessing.Process(target=sub) proc2 = multiprocessing.Process(target=pub)