Add logging, retry decorator, and a few refactors

This seems to be working all the way through.

1 files changed, 78 insertions, 42 deletions

diff --git a/Que/client.py b/Que/client.py
index 6958576..f4ee601 100755
--- a/Que/client.py
+++ b/Que/client.py
@@ -5,7 +5,9 @@ simple client for que.run
 
 import argparse
 import configparser
+import functools
 import http.client
+import logging
 import os
 import subprocess
 import sys
@@ -18,19 +20,62 @@ MAX_TIMEOUT = 99999999  # basically never timeout
 
 def auth(args):
     "Returns the auth key for the given ns from ~/.config/que.conf"
+    logging.debug("auth")
     namespace = args.target.split("/")[0]
     if namespace == "pub":
         return None
     conf_file = os.path.expanduser("~/.config/que.conf")
     if not os.path.exists(conf_file):
         sys.exit("you need a ~/.config/que.conf")
-        cfg = configparser.ConfigParser()
-        cfg.read(conf_file)
+    cfg = configparser.ConfigParser()
+    cfg.read(conf_file)
     return cfg[namespace]["key"]
 
 
+def autodecode(bytestring):
+    """Attempt to decode bytes `bs` into common codecs, preferably utf-8. If
+    no decoding is available, just return the raw bytes.
+
+    For all available codecs, see:
+    <https://docs.python.org/3/library/codecs.html#standard-encodings>
+
+    """
+    logging.debug("autodecode")
+    codecs = ["utf-8", "ascii"]
+    for codec in codecs:
+        try:
+            return bytestring.decode(codec)
+        except UnicodeDecodeError:
+            pass
+    return bytestring
+
+
+def retry(exception, tries=4, delay=3, backoff=2):
+    "Decorator for retrying an action."
+
+    def decorator(func):
+        @functools.wraps(func)
+        def func_retry(*args, **kwargs):
+            mtries, mdelay = tries, delay
+            while mtries > 1:
+                try:
+                    return func(*args, **kwargs)
+                except exception as ex:
+                    logging.debug(ex)
+                    logging.debug("retrying...")
+                    time.sleep(mdelay)
+                    mtries -= 1
+                    mdelay *= backoff
+            return func(*args, **kwargs)
+
+        return func_retry
+
+    return decorator
+
+
 def send(args):
     "Send a message to the que."
+    logging.debug("send")
     key = auth(args)
     data = args.infile
     req = request.Request(f"{args.host}/{args.target}")
@@ -38,6 +83,7 @@ def send(args):
     if key:
         req.add_header("Authorization", key)
     if args.serve:
+        logging.debug("serve")
         while not time.sleep(1):
             request.urlopen(req, data=data, timeout=MAX_TIMEOUT)
 
@@ -45,19 +91,22 @@ def send(args):
         request.urlopen(req, data=data, timeout=MAX_TIMEOUT)
 
 
-def recv(args):
-    "Receive a message from the que."
+def then(args, msg):
+    "Perform an action when passed `--then`."
+    if args.then:
+        logging.debug("then")
+        subprocess.run(
+            args.then.replace(r"\msg", msg).replace(r"\que", args.target),
+            check=False,
+            shell=True,
+        )
 
-    def _recv(_req):
-        msg = autodecode(_req.read())
-        print(msg)
-        if args.then:
-            subprocess.run(
-                args.then.replace(r"\msg", msg).replace(r"que", args.target),
-                shell=True,
-                check=False,
-            )
 
+@retry(http.client.IncompleteRead, tries=10, delay=5, backoff=1)
+@retry(http.client.RemoteDisconnected, tries=10, delay=2, backoff=2)
+def recv(args):
+    "Receive a message from the que."
+    logging.debug("recv on: %s", args.target)
     params = urllib.parse.urlencode({"poll": args.poll})
     req = request.Request(f"{args.host}/{args.target}?{params}")
     req.add_header("User-Agent", "Que/Client")
@@ -66,32 +115,23 @@ def recv(args):
         req.add_header("Authorization", key)
     with request.urlopen(req) as _req:
         if args.poll:
+            logging.debug("poll")
             while not time.sleep(1):
-                _recv(_req)
+                logging.debug("reading")
+                msg = autodecode(_req.readline())
+                logging.debug("read")
+                print(msg, end="")
+                then(args, msg)
         else:
-            _recv(_req)
-
-
-def autodecode(bytestring):
-    """Attempt to decode bytes `bs` into common codecs, preferably utf-8. If
-    no decoding is available, just return the raw bytes.
-
-    For all available codecs, see:
-    <https://docs.python.org/3/library/codecs.html#standard-encodings>
-
-    """
-    codecs = ["utf-8", "ascii"]
-    for codec in codecs:
-        try:
-            return bytestring.decode(codec)
-        except UnicodeDecodeError:
-            pass
-    return bytestring
+            msg = autodecode(_req.read())
+            print(msg)
+            then(args, msg)
 
 
 def get_args():
     "Command line parser"
     cli = argparse.ArgumentParser(description=__doc__)
+    cli.add_argument("--debug", action="store_true", help="log to stderr")
     cli.add_argument(
         "--host", default="http://que.run", help="where que-server is running"
     )
@@ -133,6 +173,12 @@ def get_args():
 
 if __name__ == "__main__":
     ARGV = get_args()
+    if ARGV.debug:
+        logging.basicConfig(
+            format="%(asctime)s %(message)s",
+            level=logging.DEBUG,
+            datefmt="%Y.%m.%d..%H.%M.%S",
+        )
     try:
         if ARGV.infile:
             send(ARGV)
@@ -140,13 +186,3 @@ if __name__ == "__main__":
             recv(ARGV)
     except KeyboardInterrupt:
         sys.exit(0)
-    except urllib.error.HTTPError as err:
-        print(err)
-        sys.exit(1)
-    except http.client.RemoteDisconnected as err:
-        print("disconnected... retrying in 5 seconds")
-        time.sleep(5)
-        if ARGV.infile:
-            send(ARGV)
-        else:
-            recv(ARGV)