fix: avoid quadratic buffering of long SSE lines

ld_eventsource/async_reader.py

@@ -12,7 +12,8 @@ class _AsyncBufferedLineReader:
     @staticmethod
     async def lines_from(chunks: AsyncIterator[bytes]) -> AsyncIterator[str]:
         last_char_was_cr = False
-        partial_line = None
+        # Join fragments only when a line ends to avoid repeatedly copying the partial line.
+        pending_fragments: list = []
 
         async for chunk in chunks:
             if len(chunk) == 0:
@@ -25,14 +26,22 @@ async def lines_from(chunks: AsyncIterator[bytes]) -> AsyncIterator[str]:
                     lines.pop(0)
                     if len(lines) == 0:
                         continue
-            if partial_line is not None:
-                lines[0] = partial_line + lines[0]
-                partial_line = None
+
             last_char = chunk[-1]
+            terminated = last_char in (10, 13)
+
+            if pending_fragments:
+                pending_fragments.append(lines[0])
+                if len(lines) == 1 and not terminated:
+                    continue  # this chunk continues the pending line, but does not end it
+                # This chunk ends the pending line, so join the fragments one time.
+                lines[0] = b"".join(pending_fragments)
+                pending_fragments = []
+
             if last_char == 13:
                 last_char_was_cr = True
             elif last_char != 10:
-                partial_line = lines.pop()
+                pending_fragments = [lines.pop()]
             for line in lines:
                 yield line.decode()
 

ld_eventsource/reader.py

@@ -16,7 +16,8 @@ def lines_from(chunks):
         series of strings, each of which is one line of text. The line does not include the terminator.
         """
         last_char_was_cr = False
-        partial_line = None
+        # Join fragments only when a line ends to avoid repeatedly copying the partial line.
+        pending_fragments = []
 
         for chunk in chunks:
             if len(chunk) == 0:
@@ -35,18 +36,24 @@ def lines_from(chunks):
                     lines.pop(0)
                     if len(lines) == 0:
                         continue  # ran out of data, continue to get next chunk
-            if partial_line is not None:
-                # On our last time through the loop, we ended up with an unterminated line, so we should
-                # treat our first parsed line here as a continuation of that.
-                lines[0] = partial_line + lines[0]
-                partial_line = None
+
+            last_char = chunk[-1]
+            terminated = last_char in (10, 13)
+
+            if pending_fragments:
+                pending_fragments.append(lines[0])
+                if len(lines) == 1 and not terminated:
+                    continue  # this chunk continues the pending line, but does not end it
+                # This chunk ends the pending line, so join the fragments one time.
+                lines[0] = b"".join(pending_fragments)
+                pending_fragments = []
+
             # Check whether the buffer really ended in a terminator. If it did not, then the last line in
             # lines is a partial line and should not be emitted yet.
-            last_char = chunk[-1]
             if last_char == 13:
                 last_char_was_cr = True  # remember this in case the next chunk starts with \n
             elif last_char != 10:
-                partial_line = lines.pop()  # remove last element which is the partial line
+                pending_fragments = [lines.pop()]  # remove last element which is the partial line
             for line in lines:
                 yield line.decode()
 

ld_eventsource/testing/test_async_reader.py

@@ -66,6 +66,25 @@ async def test_line_reader_empty_chunk():
     assert lines == ["hello", "world"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("terminator", ["\r", "\n", "\r\n"])
+async def test_line_reader_fragments_with_empty_chunks(terminator):
+    chunks = [s.replace("*", terminator).encode()
+              for s in ["first", "", " line", "*", "second*", "unterminated"]]
+    lines = await lines_from_bytes(*chunks)
+    assert lines == ["first line", "second"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("terminator", ["\r", "\n", "\r\n"])
+async def test_line_reader_utf8_line_spanning_many_chunks(terminator):
+    line = "caf\u00e9 " * 1000
+    data = (line + terminator + "next line" + terminator).encode()
+    chunks = [data[i:i + 7] for i in range(0, len(data), 7)]
+    lines = await lines_from_bytes(*chunks)
+    assert lines == [line, "next line"]
+
+
 @pytest.mark.asyncio
 async def test_sse_reader_simple_event():
     items = await events_from_lines("data: hello", "")

ld_eventsource/testing/test_reader.py

@@ -25,6 +25,10 @@ def terminator(self, request):
                 ["first line*second line*third", " line*fourth line*"],
                 ["first line", "second line", "third line", "fourth line"],
             ],
+            [
+                ["first", "", " line", "*", "second*", "unterminated"],
+                ["first line", "second"],
+            ],
         ]
     )
     def inputs_outputs(self, terminator, request):
@@ -38,7 +42,9 @@ def test_parsing(self, inputs_outputs):
 
     def test_mixed_terminators(self):
         chunks = [
-            b"first line\nsecond line\r\nthird line\r",
+            b"first line\nsecond line\r\nthird ",
+            b"line\r",
+            b"",
             b"\nfourth line\r",
             b"\r\nlast\r\n",
         ]
@@ -52,6 +58,12 @@ def test_mixed_terminators(self):
         ]
         assert list(_BufferedLineReader.lines_from(chunks)) == expected
 
+    def test_utf8_line_spanning_many_chunks(self, terminator):
+        line = "caf\u00e9 " * 1000
+        data = (line + terminator + "next line" + terminator).encode()
+        chunks = [data[i:i + 7] for i in range(0, len(data), 7)]
+        assert list(_BufferedLineReader.lines_from(chunks)) == [line, "next line"]
+
 
 class TestSSEReader:
     def expect_output(self, lines, expected):