reliably passing the test for merging datasources in the feed class

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))
+        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
+        
+    
+                

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))
-            
+                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()

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)