From ee8ec92c6729ff3b20241ed8f7d25846bf645171 Mon Sep 17 00:00:00 2001
From: haskal <haskal@awoo.systems>
Date: Wed, 12 May 2021 16:32:02 -0400
Subject: [PATCH] handle errors and retry logic better

---
 megacom/__init__.py | 74 +++++++++++++++++++++++++++++++++++++++------
 1 file changed, 65 insertions(+), 9 deletions(-)

diff --git a/megacom/__init__.py b/megacom/__init__.py
index f2da0a2..8dbacd3 100644
--- a/megacom/__init__.py
+++ b/megacom/__init__.py
@@ -1,9 +1,11 @@
 import argparse
 import asyncio
 import contextlib
+import errno
 import re
 import sys
 import termios
+import traceback
 import tty
 from types import TracebackType
 from typing import Any, List, Optional, Tuple, Type
@@ -76,9 +78,33 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
     bytesize = MODE_LOOKUP["bytesize"][m.group(1)]
     parity = MODE_LOOKUP["parity"][m.group(2)]
     stopbits = MODE_LOOKUP["stopbits"][m.group(3)]
-    (serialin, serialout) = await serial_asyncio.open_serial_connection(
-        loop=asyncio.get_event_loop(), url=tty, baudrate=baud, bytesize=bytesize, parity=parity,
-        stopbits=stopbits)
+
+    printed_fnf = False
+
+    while True:
+        try:
+            (serialin, serialout) = await serial_asyncio.open_serial_connection(
+                loop=asyncio.get_event_loop(), url=tty, baudrate=baud, bytesize=bytesize,
+                parity=parity, stopbits=stopbits)
+            break
+        except serial.SerialException as e:
+            if e.errno == errno.ENOENT:
+                # the device could not be plugged in yet.. just wait
+                if not printed_fnf:
+                    printed_fnf = True
+                    stdout.write(f"waiting for {tty} to become available...\r\n".encode())
+                    await stdout.drain()
+            else:
+                # permanent failure
+                stdout.write(f"failed to open port because: {e}\r\n".encode())
+                await stdout.drain()
+                return
+        except Exception as e:
+            # permanant failure
+            stdout.write(f"failed to open port because: {e}\r\n".encode())
+            await stdout.drain()
+            return
+        await asyncio.sleep(2)
 
     stdout.write(f"megacom connected to {tty}:{baud}:{mode}\r\n".encode())
     await stdout.drain()
@@ -113,21 +139,51 @@ async def megacom(tty: str, baud: int, mode: str) -> None:
 
     stdin_to_serial = asyncio.create_task(connect_pipe(stdin, serialout, True))
     serial_to_stdout = asyncio.create_task(connect_pipe(serialin, stdout))
-    await exit_flag.wait()
+    time_to_exit = asyncio.create_task(exit_flag.wait())
+
+    do_retry = False
+
+    def handle_done(task):
+        nonlocal do_retry
+        if task.done():
+            exc = task.exception()
+            if exc is not None:
+                stdout.write(f"\r\n\r\nmegacom encountered error: {exc}\r\n".encode())
+                if isinstance(exc, serial.SerialException):
+                    do_retry = True
+            else:
+                task.result()
+        return task.done()
+
+    await asyncio.wait([time_to_exit, stdin_to_serial, serial_to_stdout], return_when=asyncio.FIRST_COMPLETED)
+    if handle_done(time_to_exit):
+        pass
+    elif handle_done(stdin_to_serial):
+        pass
+    elif handle_done(serial_to_stdout):
+        pass
 
     stdin_to_serial.cancel()
     serial_to_stdout.cancel()
     with contextlib.suppress(asyncio.CancelledError):
-        await stdin_to_serial
+        with contextlib.suppress(serial.SerialException):
+            await stdin_to_serial
     with contextlib.suppress(asyncio.CancelledError):
-        await serial_to_stdout
+        with contextlib.suppress(serial.SerialException):
+            await serial_to_stdout
 
-    stdout.write(b"\r\n\r\nmegacom is exiting\r\n")
+    if not do_retry:
+        stdout.write(b"\r\n\r\nmegacom is exiting\r\n")
 
-    await serialout.drain()
-    serialout.close()
+    with contextlib.suppress(serial.SerialException):
+        await serialout.drain()
+    with contextlib.suppress(serial.SerialException):
+        serialout.close()
     await stdout.drain()
 
+    if do_retry:
+        return await megacom(tty, baud, mode)
+
 
 def main() -> None:
     parser = argparse.ArgumentParser(prog="megacom",