Improve SSE chunk buffering performance

src/httpx2/httpx2/_sse.py

@@ -112,12 +112,12 @@ def decode(self, line: str) -> ServerSentEvent | None:
 
         return None
 
-    def check_pending(self, pending: str) -> None:
+    def check_pending(self, pending_size: int) -> None:
         """
         Bound the total bytes buffered for the in-progress event, including
         a trailing line that has not been terminated by a newline yet.
         """
-        self._check_size(len(pending.encode("utf-8")))
+        self._check_size(pending_size)
 
     def _check_size(self, pending_size: int = 0) -> None:
         if self._max_event_size is not None and self._event_size + pending_size > self._max_event_size:
@@ -126,12 +126,13 @@ def _check_size(self, pending_size: int = 0) -> None:
 
 class _SSELineDecoder:
     def __init__(self) -> None:
-        self._buffer = ""
+        self._parts: list[str] = []
+        self._pending_size = 0
         self._trailing_cr = False
 
     @property
-    def pending(self) -> str:
-        return self._buffer
+    def pending_size(self) -> int:
+        return self._pending_size
 
     def decode(self, text: str) -> list[str]:
         if self._trailing_cr:
@@ -141,20 +142,36 @@ def decode(self, text: str) -> list[str]:
             self._trailing_cr = True
             text = text[:-1]
 
-        text = self._buffer + text.replace("\r\n", "\n").replace("\r", "\n")
+        text = text.replace("\r\n", "\n").replace("\r", "\n")
+        if "\n" not in text:
+            self._append(text)
+            return []
+
         lines = text.split("\n")
-        self._buffer = lines.pop()
+        self._append(lines[0])
+        lines[0] = self._consume_pending()
+        self._append(lines.pop())
         return lines
 
     def flush(self) -> list[str]:
         if self._trailing_cr:
-            self._buffer += "\n"
+            self._append("\n")
             self._trailing_cr = False
-        if not self._buffer:
+        buffer = self._consume_pending()
+        if not buffer:
             return []
-        lines = self._buffer.split("\n")
-        self._buffer = ""
-        return lines
+        return buffer.split("\n")
+
+    def _append(self, text: str) -> None:
+        if text:
+            self._parts.append(text)
+            self._pending_size += len(text.encode("utf-8"))
+
+    def _consume_pending(self) -> str:
+        pending = "".join(self._parts)
+        self._parts = []
+        self._pending_size = 0
+        return pending
 
 
 class _SSEParser:
@@ -164,7 +181,7 @@ def __init__(self, max_event_size: int | None = None) -> None:
 
     def decode(self, text: str) -> Iterator[ServerSentEvent]:
         yield from self._decode_lines(self._line_decoder.decode(text))
-        self._event_decoder.check_pending(self._line_decoder.pending)
+        self._event_decoder.check_pending(self._line_decoder.pending_size)
 
     def flush(self) -> Iterator[ServerSentEvent]:
         yield from self._decode_lines(self._line_decoder.flush())

tests/httpx2/test_sse.py

@@ -497,3 +497,21 @@ def handler(request: httpx2.Request) -> httpx2.Response:
             events = [event async for event in source]
 
     assert [event.data for event in events] == ["hi"]
+
+
+def test_many_chunks_without_line_separator() -> None:
+    parts = [f"{index:04d}".encode() for index in range(1_000)]
+
+    def chunks() -> Iterator[bytes]:
+        yield b"data: "
+        yield from parts
+        yield b"\n\n"
+
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, content=chunks(), headers={"Content-Type": "text/event-stream"})
+
+    with httpx2.Client(transport=httpx2.MockTransport(handler)) as client:
+        with client.sse("http://testserver/sse") as source:
+            (event,) = list(source)
+
+    assert event.data == b"".join(parts).decode()