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