Merge pull request #123 from tfgm-bud/cleanup

Clean up to fix a thead joining issue

Author: danpaquin

gdax/order_book.py

@@ -36,38 +36,39 @@ def on_open(self):
     def on_close(self):
         print("\n-- OrderBook Socket Closed! --")
 
+    def reset_book(self):
+        self._asks = RBTree()
+        self._bids = RBTree()
+        res = self._client.get_product_order_book(product_id=self.product_id, level=3)
+        for bid in res['bids']:
+            self.add({
+                'id': bid[2],
+                'side': 'buy',
+                'price': Decimal(bid[0]),
+                'size': Decimal(bid[1])
+            })
+        for ask in res['asks']:
+            self.add({
+                'id': ask[2],
+                'side': 'sell',
+                'price': Decimal(ask[0]),
+                'size': Decimal(ask[1])
+            })
+        self._sequence = res['sequence']
+
     def on_message(self, message):
         if self._log_to:
             pickle.dump(message, self._log_to)
 
         sequence = message['sequence']
         if self._sequence == -1:
-            self._asks = RBTree()
-            self._bids = RBTree()
-            res = self._client.get_product_order_book(product_id=self.product_id, level=3)
-            for bid in res['bids']:
-                self.add({
-                    'id': bid[2],
-                    'side': 'buy',
-                    'price': Decimal(bid[0]),
-                    'size': Decimal(bid[1])
-                })
-            for ask in res['asks']:
-                self.add({
-                    'id': ask[2],
-                    'side': 'sell',
-                    'price': Decimal(ask[0]),
-                    'size': Decimal(ask[1])
-                })
-            self._sequence = res['sequence']
-
+            self.reset_book()
+            return
         if sequence <= self._sequence:
             # ignore older messages (e.g. before order book initialization from getProductOrderBook)
             return
         elif sequence > self._sequence + 1:
-            print('Error: messages missing ({} - {}). Re-initializing websocket.'.format(sequence, self._sequence))
-            self.close()
-            self.start()
+            self.on_sequence_gap(self._sequence, sequence)
             return
 
         msg_type = message['type']
@@ -83,18 +84,11 @@ def on_message(self, message):
 
         self._sequence = sequence
 
-        # bid = self.get_bid()
-        # bids = self.get_bids(bid)
-        # bid_depth = sum([b['size'] for b in bids])
-        # ask = self.get_ask()
-        # asks = self.get_asks(ask)
-        # ask_depth = sum([a['size'] for a in asks])
-        # print('bid: %f @ %f - ask: %f @ %f' % (bid_depth, bid, ask_depth, ask))
+    def on_sequence_gap(self, gap_start, gap_end):
+        self.reset_book()
+        print('Error: messages missing ({} - {}). Re-initializing  book at sequence.'.format(
+            gap_start, gap_end, self._sequence))
 
-    def on_error(self, e):
-        self._sequence = -1
-        self.close()
-        self.start()
 
     def add(self, order):
         order = {
@@ -250,6 +244,7 @@ def set_bids(self, price, bids):
 
 
 if __name__ == '__main__':
+    import sys
     import time
     import datetime as dt
 
@@ -286,10 +281,18 @@ def on_message(self, message):
                 self._ask = ask
                 self._bid_depth = bid_depth
                 self._ask_depth = ask_depth
-                print('{}\tbid: {:.3f} @ {:.2f}\task: {:.3f} @ {:.2f}'.format(dt.datetime.now(), bid_depth, bid,
-                                                                              ask_depth, ask))
+                print('{} {} bid: {:.3f} @ {:.2f}\task: {:.3f} @ {:.2f}'.format(
+                    dt.datetime.now(), self.product_id, bid_depth, bid, ask_depth, ask))
 
     order_book = OrderBookConsole()
     order_book.start()
-    time.sleep(10)
-    order_book.close()
+    try:
+        while True:
+            time.sleep(10)
+    except KeyboardInterrupt:
+        order_book.close()
+
+    if order_book.error:
+        sys.exit(1)
+    else:
+        sys.exit(0)

gdax/websocket_client.py

@@ -23,6 +23,7 @@ def __init__(self, url="wss://ws-feed.gdax.com", products=None,
         self.products = products
         self.type = message_type
         self.stop = False
+        self.error = None
         self.ws = None
         self.thread = None
         self.auth = auth
@@ -36,6 +37,7 @@ def start(self):
         def _go():
             self._connect()
             self._listen()
+            self._disconnect()
 
         self.stop = False
         self.on_open()
@@ -80,22 +82,29 @@ def _listen(self):
                 if int(time.time() % 30) == 0:
                     # Set a 30 second ping to keep connection alive
                     self.ws.ping("keepalive")
-                msg = json.loads(self.ws.recv())
+                data = self.ws.recv()
+                msg = json.loads(data)
+            except ValueError as e:
+                self.on_error(e)
             except Exception as e:
                 self.on_error(e)
             else:
                 self.on_message(msg)
 
+    def _disconnect(self):
+        if self.type == "heartbeat":
+            self.ws.send(json.dumps({"type": "heartbeat", "on": False}))
+        try:
+            if self.ws:
+                self.ws.close()
+        except WebSocketConnectionClosedException as e:
+            pass
+
+        self.on_close()
+
     def close(self):
-        if not self.stop:
-            self.on_close()
-            self.stop = True
-            self.thread.join()
-            try:
-                if self.ws:
-                    self.ws.close()
-            except WebSocketConnectionClosedException as e:
-                pass
+        self.stop = True
+        self.thread.join()
 
     def on_open(self):
         if self.should_print:
@@ -111,10 +120,13 @@ def on_message(self, msg):
         if self.mongo_collection: # dump JSON to given mongo collection
             self.mongo_collection.insert_one(msg)
 
-    def on_error(self, e):
-        print(e)
+    def on_error(self, e, data=None):
+        self.error = e
+        self.stop
+        print('{} - data: {}'.format(e, data))
 
 if __name__ == "__main__":
+    import sys
     import gdax
     import time
 
@@ -126,8 +138,7 @@ def on_open(self):
             print("Let's count the messages!")
 
         def on_message(self, msg):
-            if 'price' in msg and 'type' in msg:
-                print("Message type:", msg["type"], "\t@ %.3f" % float(msg["price"]))
+            print(json.dumps(msg, indent=4, sort_keys=True))
             self.message_count += 1
 
         def on_close(self):
@@ -136,9 +147,14 @@ def on_close(self):
     wsClient = MyWebsocketClient()
     wsClient.start()
     print(wsClient.url, wsClient.products)
-    # Do some logic with the data
-    while wsClient.message_count < 10000:
-        print("\nMessageCount =", "%i \n" % wsClient.message_count)
-        time.sleep(1)
-
-    wsClient.close()
+    try:
+        while True:
+            print("\nMessageCount =", "%i \n" % wsClient.message_count)
+            time.sleep(1)
+    except KeyboardInterrupt:
+        wsClient.close()
+
+    if wsClient.error:
+        sys.exit(1)
+    else:
+        sys.exit(0)