Make TextLogStream handle its own buffering rather than relying on Te…

…xtIOWrapper
python · freakboy3742 · Aug 6, 2024 · Aug 1, 2024 · Aug 2, 2024 · Aug 5, 2024
commit f8ff4421e4a267aa87a7103f8740a84c66a0e3c6
@@ -1,5 +1,6 @@
 import io
 import sys
+from threading import RLock
 
 
 # The maximum length of a log message in bytes, including the level marker and
@@ -11,9 +12,9 @@
 MAX_BYTES_PER_WRITE = 4000
 
 # UTF-8 uses a maximum of 4 bytes per character, so limiting text writes to this
-# size ensures that TextIOWrapper can always avoid exceeding MAX_BYTES_PER_WRITE.
+# size ensures that we can always avoid exceeding MAX_BYTES_PER_WRITE.
 # However, if the actual number of bytes per character is smaller than that,
-# then TextIOWrapper may still join multiple consecutive text writes into binary
+# then we may still join multiple consecutive text writes into binary
 # writes containing a larger number of characters.
 MAX_CHARS_PER_WRITE = MAX_BYTES_PER_WRITE // 4
 
@@ -35,9 +36,10 @@ def init_streams(android_log_write, stdout_prio, stderr_prio):
 class TextLogStream(io.TextIOWrapper):
     def __init__(self, android_log_write, prio, tag, **kwargs):
         kwargs.setdefault("encoding", "UTF-8")
-        kwargs.setdefault("line_buffering", True)
         super().__init__(BinaryLogStream(android_log_write, prio, tag), **kwargs)
-        self._CHUNK_SIZE = MAX_BYTES_PER_WRITE
+        self._lock = RLock()
+        self._pending_bytes = []
+        self._pending_bytes_count = 0
 
     def __repr__(self):
         return f"<TextLogStream {self.buffer.tag!r}>"
@@ -52,15 +54,44 @@ def write(self, s):
         s = str.__str__(s)
 
         # We want to emit one log message per line wherever possible, so split
-        # the string before sending it to the superclass. Note that
-        # "".splitlines() == [], so nothing will be logged for an empty string.
-        for line in s.splitlines(keepends=True):
-            while line:
-                super().write(line[:MAX_CHARS_PER_WRITE])
-                line = line[MAX_CHARS_PER_WRITE:]
+        # the string into lines first. Note that "".splitlines() == [], so
+        # nothing will be logged for an empty string.
+        with self._lock:
+            for line in s.splitlines(keepends=True):
+                while line:
+                    chunk = line[:MAX_CHARS_PER_WRITE]
+                    line = line[MAX_CHARS_PER_WRITE:]
+                    self._write_chunk(chunk)
 
         return len(s)
 
+    # The size and behavior of TextIOWrapper's buffer is not part of its public
+    # API, so we handle buffering ourselves to avoid truncation.
+    def _write_chunk(self, s):
+        b = s.encode(self.encoding, self.errors)
+        if self._pending_bytes_count + len(b) > MAX_BYTES_PER_WRITE:
+            self.flush()
+
+        self._pending_bytes.append(b)
+        self._pending_bytes_count += len(b)
+        if (
+            self.write_through
+            or b.endswith(b"\n")
+            or self._pending_bytes_count > MAX_BYTES_PER_WRITE
+        ):
+            self.flush()
+
+    def flush(self):
+        with self._lock:
+            self.buffer.write(b"".join(self._pending_bytes))
+            self._pending_bytes.clear()
+            self._pending_bytes_count = 0
+
+    # Since this is a line-based logging system, line buffering cannot be turned
+    # off, i.e. a newline always causes a flush.
+    def line_buffering(self):
+        return True
+
 
 class BinaryLogStream(io.RawIOBase):
     def __init__(self, android_log_write, prio, tag):

@@ -147,6 +147,13 @@ def write(s, lines=None, *, write_len=None):
                 write("f\n\ng", ["exxf", ""])
                 write("\n", ["g"])
 
+                # Since this is a line-based logging system, line buffering
+                # cannot be turned off, i.e. a newline always causes a flush.
+                stream.reconfigure(line_buffering=False)
+                self.assertTrue(stream.line_buffering)
+
+                # However, buffering can be turned off completely if you want a
+                # flush after every write.
                 with self.unbuffered(stream):
                     write("\nx", ["", "x"])
                     write("\na\n", ["", "a"])
@@ -209,21 +216,21 @@ def __str__(self):
                 # (MAX_BYTES_PER_WRITE).
                 #
                 # ASCII (1 byte per character)
-                write(("foobar" * 700) + "\n",
-                      [("foobar" * 666) + "foob",  # 4000 bytes
-                       "ar" + ("foobar" * 33)])  # 200 bytes
+                write(("foobar" * 700) + "\n",  # 4200 bytes in
+                      [("foobar" * 666) + "foob",  # 4000 bytes out
+                       "ar" + ("foobar" * 33)])  # 200 bytes out
 
                 # "Full-width" digits 0-9 (3 bytes per character)
                 s = "\uff10\uff11\uff12\uff13\uff14\uff15\uff16\uff17\uff18\uff19"
-                write((s * 150) + "\n",
-                      [s * 100,  # 3000 bytes
-                       s * 50])  # 1500 bytes
+                write((s * 150) + "\n",  # 4500 bytes in
+                      [s * 100,  # 3000 bytes out
+                       s * 50])  # 1500 bytes out
 
                 s = "0123456789"
-                write(s * 200, [])
-                write(s * 150, [])
-                write(s * 51, [s * 350])  # 3500 bytes
-                write("\n", [s * 51])  # 510 bytes
+                write(s * 200, [])  # 2000 bytes in
+                write(s * 150, [])  # 1500 bytes in
+                write(s * 51, [s * 350])  # 510 bytes in, 3500 bytes out
+                write("\n", [s * 51])  # 0 bytes in, 510 bytes out
 
     def test_bytes(self):
         for stream_name, level in [("stdout", "I"), ("stderr", "W")]: