Make upload-file.py work better on CDC-ACM console.

Author: jmattsson

scripts/upload-file.py

@@ -15,12 +15,15 @@
 
 # The loader we send to NodeMCU so that we may upload a (binary) file safely.
 # Uses STX/ETX/DLE framing and escaping.
+# The CDC-ACM console gets overwhelmed unless we throttle the send by using
+# an ack scheme. We use a fake prompt for simplicity's sake for that.
 loader = b'''
 (function()
   local function transmission_receiver(chunk_cb)
     local inframe = false
     local escaped = false
     local done = false
+    local len = 0
     local STX = 2
     local ETX = 3
     local DLE = 16
@@ -30,6 +33,11 @@
     end
     return function(data)
       if done then return end
+      len = len + #data
+      while len >= @BLOCKSIZE@ do
+        len = len - @BLOCKSIZE@
+        console.write("> ")
+      end
       local from
       local to
       for i = 1, #data
@@ -38,10 +46,7 @@
         if inframe
         then
           if not from then from = i end -- first valid byte
-          if escaped
-          then
-            escaped = false
-          else
+          if escaped then escaped = false else
             if b == DLE
             then
               escaped = true
@@ -57,8 +62,7 @@
         else -- look for an (unescaped) STX to sync frame start
           if b == DLE then escaped = true
           elseif b == STX and not escaped then inframe = true
-          else escaped = false
-          end
+          else escaped = false end
         end
         -- else ignore byte outside of framing
       end
@@ -70,17 +74,19 @@
   local function file_saver(name)
     local f = io.open(name, "w")
     return function(chunk)
-      if chunk then f:write(chunk)
-      else
+      if chunk then f:write(chunk) else
         f:close()
         console.on("data", 0, nil)
         console.mode(console.INTERACTIVE)
+        console.write("done")
       end
     end
   end
 
-  console.on("data", 0, transmission_receiver(file_saver("@FILENAME@")))
+  console.on("data", 0, transmission_receiver(file_saver(
+    "@FILENAME@")))
   console.mode(console.NONINTERACTIVE)
+  console.write("ready")
 end)()
 '''
 
@@ -91,6 +97,7 @@ def parse_args():
     parser.add_argument("name", nargs="?", help="Name to upload file as.")
     parser.add_argument("-p", "--port", default="/dev/ttyUSB0", help="Serial port (default: /dev/ttyUSB0).")
     parser.add_argument("-b", "--bitrate", type=int, default=115200, help="Bitrate (default: 115200).")
+    parser.add_argument("-s", "--blocksize", type=int, default=80, help="Block size of file data, tweak for speed/reliability of upload (default: 80)")
     return parser.parse_args()
 
 def load_file(filename):
@@ -103,26 +110,51 @@ def load_file(filename):
         print(f"Error reading file {filename}: {e}")
         sys.exit(1)
 
-def wait_prompt(ser):
+def xprint(msg):
+    print(msg, end='', flush=True)
+
+def wait_prompt(ser, ignore):
     """Wait until we see the '> ' prompt, or the serial times out"""
     buf = bytearray()
     b = ser.read()
-    while b != b'':
-        buf.extend(b)
+    timeout = 5
+    while timeout > 0:
+        if b == b'':
+            timeout -= 1
+            xprint('!')
+        else:
+            buf.extend(b)
+        if not ignore and buf.find(b'Lua error:') != -1:
+            xprint(buf.decode())
+            line = ser.readline()
+            while line != b'':
+                xprint(line.decode())
+                line = ser.readline()
+            sys.exit(1)
         if buf.find(b'> ') != -1:
             return True
         b = ser.read()
+    xprint(buf.decode())
+    return False
+
+def wait_line_match(ser, match, timeout):
+    """Wait until the 'match' string is found within a line, or times out"""
+    line = ser.readline()
+    while timeout > 0:
+        if line.find(match) != -1:
+            return True
+        elif line == b'':
+            timeout -= 1
+            xprint('!')
     return False
 
 def sync(ser):
     """Get ourselves to a clean prompt so we can understand the output"""
     ser.write(b'\x03\x03\n')
-    wait_prompt(ser)
+    if not wait_prompt(ser, True):
+        return False
     ser.write(b"print('sync')\n")
-    line = ser.readline()
-    while line != b"sync\n":
-       line = ser.readline()
-    return wait_prompt(ser)
+    return wait_line_match(ser, b'sync', 5) and wait_prompt(ser, True)
 
 def cleanup():
     """Cleanup function to send final data and close the serial port."""
@@ -140,7 +172,27 @@ def line_interactive_send(ser, data):
     for line in data.split(b'\n'):
         ser.write(line)
         ser.write(b'\n')
-        wait_prompt(ser)
+        if not wait_prompt(ser, False):
+            return False
+        xprint('.')
+    return True
+
+def chunk_data(data, size):
+    """Split a data block into chunks"""
+    return (data[0+i:size+i] for i in range(0, len(data), size))
+
+def chunk_interactive_send(ser, data, size):
+    """Send the data chunked into blocks, waiting for an ack in between"""
+    n=0
+    for chunk in chunk_data(data, size):
+        ser.write(chunk)
+        if len(chunk) == size and not wait_prompt(ser, False):
+            print(f"failed after sending {n} blocks")
+            return False
+        xprint('.')
+        n += 1
+    print(f" ok, sent {n} blocks")
+    return True
 
 def transmission(data):
     """Perform STX/ETX/DLE framing and escaping of the data"""
@@ -159,30 +211,48 @@ def transmission(data):
     upload_name = args.name if args.name else args.file
 
     file_data = load_file(args.file)
+    print(f"Loaded {len(file_data)} bytes of file contents")
+
+    blocksize = bytes(str(args.blocksize).encode())
 
     try:
-        ser = serial.Serial(args.port, args.bitrate, timeout=1)
+        ser = serial.Serial(port=args.port, baudrate=args.bitrate, timeout=1)
     except serial.SerialException as e:
         print(f"Error opening serial port {args.port}: {e}")
         sys.exit(1)
 
-    print("Synchronising serial...")
+    print("Synchronising serial...", end='')
     if not sync(ser):
-        print("NodeMCU not responding\n")
+        print("\nNodeMCU not responding\n")
         sys.exit(1)
 
-    print(f'Uploading "{args.file}" as "{upload_name}"')
+    print(f' ok\nUploading "{args.file}" as "{upload_name}"')
 
     atexit.register(cleanup)
 
-    print("Sending loader...")
-    line_interactive_send(
-        ser, loader.replace(b"@FILENAME@", upload_name.encode()))
+    xprint("Sending loader")
+    ok = line_interactive_send(
+        ser, loader.replace(
+            b"@FILENAME@", upload_name.encode()).replace(
+            b"@BLOCKSIZE@", blocksize))
+
+    if ok:
+        xprint(" ok\nWaiting for go-ahead...")
+        ok = wait_line_match(ser, b"ready", 5)
 
-    print("Sending file contents...")
-    ser.write(transmission(file_data))
-    wait_prompt(ser)
+    if ok:
+        xprint(f" ok\nSending file contents (using blocksize {args.blocksize})")
+        ok = chunk_interactive_send(
+            ser, transmission(file_data), int(blocksize))
+    if ok:
+        xprint("Waiting for final ack...")
+        ok = wait_line_match(ser, b"done", 5)
+        ser.write(b"\n")
+
+    if not ok or not wait_prompt(ser, False):
+        print("transmission timed out")
+        sys.exit(1)
 
     ser.close()
     ser = None
-    print("Done.")
+    print(" ok\nUpload complete.")