.gitlab-ci/bare-metal/serial_buffer.py

   1 #!/usr/bin/env python3
   2 #
   3 # Copyright © 2020 Google LLC
   4 #
   5 # Permission is hereby granted, free of charge, to any person obtaining a
   6 # copy of this software and associated documentation files (the "Software"),
   7 # to deal in the Software without restriction, including without limitation
   8 # the rights to use, copy, modify, merge, publish, distribute, sublicense,
   9 # and/or sell copies of the Software, and to permit persons to whom the
  10 # Software is furnished to do so, subject to the following conditions:
  11 #
  12 # The above copyright notice and this permission notice (including the next
  13 # paragraph) shall be included in all copies or substantial portions of the
  14 # Software.
  15 #
  16 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  17 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  18 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
  19 # THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  20 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  21 # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
  22 # IN THE SOFTWARE.
  23
  24 import argparse
  25 from datetime import datetime,timezone
  26 import queue
  27 import serial
  28 import threading
  29 import time
  30
  31 class SerialBuffer:
  32     def __init__(self, dev, filename, prefix):
  33         self.filename = filename
  34         self.dev = dev
  35
  36         if dev:
  37             self.f = open(filename, "wb+")
  38             self.serial = serial.Serial(dev, 115200, timeout=10)
  39         else:
  40             self.f = open(filename, "rb")
  41
  42         self.byte_queue = queue.Queue()
  43         self.line_queue = queue.Queue()
  44         self.prefix = prefix
  45         self.sentinel = object()
  46
  47         if self.dev:
  48             self.read_thread = threading.Thread(target=self.serial_read_thread_loop, daemon=True)
  49         else:
  50             self.read_thread = threading.Thread(target=self.serial_file_read_thread_loop, daemon=True)
  51         self.read_thread.start()
  52
  53         self.lines_thread = threading.Thread(target=self.serial_lines_thread_loop, daemon=True)
  54         self.lines_thread.start()
  55
  56     # Thread that just reads the bytes from the serial device to try to keep from
  57     # buffer overflowing it.
  58     def serial_read_thread_loop(self):
  59         greet = "Serial thread reading from %s\n" % self.dev
  60         self.byte_queue.put(greet.encode())
  61
  62         while True:
  63             try:
  64                 self.byte_queue.put(self.serial.read())
  65             except Exception as err:
  66                 print(self.prefix + str(err))
  67                 self.byte_queue.put(self.sentinel)
  68                 break
  69
  70     # Thread that just reads the bytes from the file of serial output that some
  71     # other process is appending to.
  72     def serial_file_read_thread_loop(self):
  73         greet = "Serial thread reading from %s\n" % self.filename
  74         self.byte_queue.put(greet.encode())
  75
  76         while True:
  77             line = self.f.readline()
  78             if line:
  79                 self.byte_queue.put(line)
  80             else:
  81                 time.sleep(0.1)
  82
  83     # Thread that processes the stream of bytes to 1) log to stdout, 2) log to
  84     # file, 3) add to the queue of lines to be read by program logic
  85
  86     def serial_lines_thread_loop(self):
  87         line = bytearray()
  88         while True:
  89             bytes = self.byte_queue.get(block=True)
  90
  91             if bytes == self.sentinel:
  92                 self.read_thread.join()
  93                 self.line_queue.put(self.sentinel)
  94                 break;
  95
  96             # Write our data to the output file if we're the ones reading from
  97             # the serial device
  98             if self.dev:
  99                 self.f.write(bytes)
 100                 self.f.flush()
 101
 102             for b in bytes:
 103                 line.append(b)
 104                 if b == b'\n'[0]:
 105                     line = line.decode(errors="replace")
 106
 107                     time = datetime.now().strftime('%y-%m-%d %H:%M:%S')
 108                     print("{time} {prefix}{line}".format(time=time, prefix=self.prefix, line=line), flush=True, end='')
 109
 110                     self.line_queue.put(line)
 111                     line = bytearray()
 112
 113     def get_line(self):
 114         line = self.line_queue.get()
 115         if line == self.sentinel:
 116             self.lines_thread.join()
 117         return line
 118
 119     def lines(self):
 120         return iter(self.get_line, self.sentinel)
 121
 122 def main():
 123     parser = argparse.ArgumentParser()
 124
 125     parser.add_argument('--dev', type=str, help='Serial device')
 126     parser.add_argument('--file', type=str, help='Filename for serial output', required=True)
 127     parser.add_argument('--prefix', type=str, help='Prefix for logging serial to stdout', nargs='?')
 128
 129     args = parser.parse_args()
 130
 131     ser = SerialBuffer(args.dev, args.file, args.prefix or "")
 132     for line in ser.lines():
 133         # We're just using this as a logger, so eat the produced lines and drop
 134         # them
 135         pass
 136
 137 if __name__ == '__main__':
 138     main()