From b77e105d5ce4af6b56769cc14c983da060a9fc90 Mon Sep 17 00:00:00 2001
From: Morgan McMillian <gilag@monkeystew.com>
Date: Mon, 27 Jan 2020 22:09:39 -0800
Subject: [PATCH] reinitiate connection on websocket error resolves issue #15

---
 partybot-pnut.py | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/partybot-pnut.py b/partybot-pnut.py
index 3e65fda..6fe4fa3 100644
--- a/partybot-pnut.py
+++ b/partybot-pnut.py
@@ -13,8 +13,10 @@ from database import db_session, init_db
 from sqlalchemy import and_
 from models import Karma, Optout, Queue, Preferences, MdnpRequests
 
+_startup = threading.Event()
 _shutdown = threading.Event()
 _connected = threading.Event()
+_error = threading.Event()
 
 def subscribe(connection_id):
     url = f"https://api.pnut.io/v0/channels/{config['CHANNEL']}/messages"
@@ -23,7 +25,9 @@ def subscribe(connection_id):
     r = requests.get(url, headers=headers)
     if r.status_code == 200:
         _connected.set()
-        send(config['CHANNEL'], "Partybot online and ready! \o/")
+        if _startup.isSet():
+            send(config['CHANNEL'], "Partybot online and ready! \o/")
+            _startup.clear()
     else:
         logger.error(r)
 
@@ -282,7 +286,8 @@ def on_message(ws, message):
     msg = json.loads(message)
 
     if not _connected.isSet() and 'connection_id' in msg['meta']:
-        send(config['CHANNEL'], "...connecting circuits...")
+        if _startup.isSet():
+            send(config['CHANNEL'], "...connecting circuits...")
         logger.debug("connection_id: " + msg['meta']['connection_id'])
         subscribe(msg['meta']['connection_id'])
         return
@@ -317,26 +322,24 @@ def on_message(ws, message):
 def on_error(ws, error):
     logger.error("on_error: !!! ERROR !!!")
     logger.error(error)
-    # _shutdown.set()
+    _error.set()
 
 def on_close(ws):
-    send(config['CHANNEL'], "...shutdown initiated...")
+    # send(config['CHANNEL'], "...shutdown initiated...")
     logger.debug("on_close: ### CLOSED ###")
-    # _shutdown.set()
+    _connected.clear()
 
 def on_open(ws):
 
     def run(*args):
-        while not _shutdown.isSet():
+        while not _error.isSet():
             qmsg = Queue.query.one_or_none()
             if qmsg:
                 send(config['CHANNEL'], qmsg.msg)
                 db_session.delete(qmsg)
                 db_session.commit()
-            time.sleep(3)
             ws.send(".")
-        time.sleep(1)
-        ws.close()
+            time.sleep(3)
         logger.debug("*** terminate ***")
     
     t = threading.Thread(target=run)
@@ -361,4 +364,8 @@ if __name__ == "__main__":
     ws = websocket.WebSocketApp(ws_url, on_message=on_message, 
         on_error=on_error, on_close=on_close)
     ws.on_open = on_open
-    ws.run_forever()
+    r = True
+    _startup.set()
+    while r:
+        _error.clear()
+        r = ws.run_forever()