Fixing a performance issue in SSE client (#221)

* Fixing a performance issue in SSE client

* Merged with master

* Using a more efficient mechanism to buffer and parse incoming SSE data

* Minor improvement to how tail string is read

* Using string slicing to reduce the window

* Updated comment

Author: hiranya911

CHANGELOG.md

@@ -4,6 +4,8 @@
   type.googleapis.com/google.firebase.fcm.v1.FcmError to set error code.
 - [fixed] Ensuring that `UserRecord.tokens_valid_after_time` always
   returns an integer, and never returns `None`.
+- [fixed] Fixing a performance issue in the `db.listen()` API
+  where it was taking a long time to process large RTDB nodes.
 
 # v2.13.0
 

firebase_admin/_sseclient.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""SSEClient module to stream realtime updates in the Firebase Database."""
+"""SSEClient module to stream realtime updates from the Firebase Database.
+
+Based on a similar implementation from Pyrebase.
+"""
 
 import re
 import time
@@ -37,6 +40,34 @@ def rebuild_auth(self, prepared_request, response):
         pass
 
 
+class _EventBuffer(object):
+    """A helper class for buffering and parsing raw SSE data."""
+
+    def __init__(self):
+        self._buffer = []
+        self._tail = ''
+
+    def append(self, char):
+        self._buffer.append(char)
+        self._tail += char
+        self._tail = self._tail[-4:]
+
+    def truncate(self):
+        head, sep, _ = self.buffer_string.rpartition('\n')
+        rem = head + sep
+        self._buffer = list(rem)
+        self._tail = rem[-4:]
+
+    @property
+    def is_end_of_field(self):
+        last_two_chars = self._tail[-2:]
+        return last_two_chars == '\n\n' or last_two_chars == '\r\r' or self._tail == '\r\n\r\n'
+
+    @property
+    def buffer_string(self):
+        return ''.join(self._buffer)
+
+
 class SSEClient(object):
     """SSE client implementation."""
 
@@ -58,7 +89,7 @@ def __init__(self, url, session, retry=3000, **kwargs):
         self.buf = u'' # Keep data here as it streams in
 
         headers = self.requests_kwargs.get('headers', {})
-        # The SSE spec requires making requests with Cache-Control: nocache
+        # The SSE spec requires making requests with Cache-Control: no-cache
         headers['Cache-Control'] = 'no-cache'
         # The 'Accept' header is not required, but explicit > implicit
         headers['Accept'] = 'text/event-stream'
@@ -82,32 +113,28 @@ def _connect(self):
         else:
             raise StopIteration()
 
-    def _event_complete(self):
-        """Checks if the event is completed by matching regular expression."""
-        return re.search(end_of_field, self.buf) is not None
-
     def __iter__(self):
         return self
 
     def __next__(self):
-        while not self._event_complete():
-            try:
-                nextchar = next(self.resp_iterator)
-                self.buf += nextchar
-            except (StopIteration, requests.RequestException):
-                time.sleep(self.retry / 1000.0)
-                self._connect()
-                # The SSE spec only supports resuming from a whole message, so
-                # if we have half a message we should throw it out.
-                head, sep, tail = self.buf.rpartition('\n')
-                self.buf = head + sep
-                continue
+        if not re.search(end_of_field, self.buf):
+            temp_buffer = _EventBuffer()
+            while not temp_buffer.is_end_of_field:
+                try:
+                    nextchar = next(self.resp_iterator)
+                    temp_buffer.append(nextchar)
+                except (StopIteration, requests.RequestException):
+                    time.sleep(self.retry / 1000.0)
+                    self._connect()
+                    # The SSE spec only supports resuming from a whole message, so
+                    # if we have half a message we should throw it out.
+                    temp_buffer.truncate()
+                    continue
+            self.buf = temp_buffer.buffer_string
 
         split = re.split(end_of_field, self.buf)
         head = split[0]
-        tail = ''.join(split[1:])
-
-        self.buf = tail
+        self.buf = '\n\n'.join(split[1:])
         event = Event.parse(head)
 
         if event.data == 'credential is no longer valid':
@@ -150,7 +177,7 @@ def parse(cls, raw):
           raw: the raw data to parse.
 
         Returns:
-          Event: newly intialized ``Event`` object with the parameters  initialized.
+          Event: A new ``Event`` with the parameters initialized.
         """
         event = cls()
         for line in raw.split('\n'):

tests/test_sseclient.py

@@ -19,10 +19,10 @@
 import six
 
 from firebase_admin import _sseclient
-from tests.testutils import MockAdapter
+from tests import testutils
 
 
-class MockSSEClientAdapter(MockAdapter):
+class MockSSEClientAdapter(testutils.MockAdapter):
 
     def __init__(self, payload, recorder):
         super(MockSSEClientAdapter, self).__init__(payload, 200, recorder)
@@ -73,6 +73,17 @@ def test_single_event(self):
         assert event_payload["path"] == "/"
         assert len(recorder) == 2
 
+    def test_large_event(self):
+        data = 'a' * int(0.1 * 1024 * 1024)
+        payload = 'event: put\ndata: {"path":"/","data":"' + data + '"}\n\n'
+        recorder = []
+        sseclient = self.init_sse(payload, recorder)
+        event = next(sseclient)
+        event_payload = json.loads(event.data)
+        assert event_payload["data"] == data
+        assert event_payload["path"] == "/"
+        assert len(recorder) == 1
+
     def test_multiple_events(self):
         payload = 'event: put\ndata: {"path":"/foo","data":"testevent1"}\n\n'
         payload += 'event: put\ndata: {"path":"/bar","data":"testevent2"}\n\n'
@@ -88,6 +99,26 @@ def test_multiple_events(self):
         assert event_payload["path"] == "/bar"
         assert len(recorder) == 1
 
+    def test_event_separators(self):
+        payload = 'event: put\ndata: {"path":"/foo","data":"testevent1"}\n\n'
+        payload += 'event: put\ndata: {"path":"/bar","data":"testevent2"}\r\r'
+        payload += 'event: put\ndata: {"path":"/baz","data":"testevent3"}\r\n\r\n'
+        recorder = []
+        sseclient = self.init_sse(payload, recorder)
+        event = next(sseclient)
+        event_payload = json.loads(event.data)
+        assert event_payload["data"] == "testevent1"
+        assert event_payload["path"] == "/foo"
+        event = next(sseclient)
+        event_payload = json.loads(event.data)
+        assert event_payload["data"] == "testevent2"
+        assert event_payload["path"] == "/bar"
+        event = next(sseclient)
+        event_payload = json.loads(event.data)
+        assert event_payload["data"] == "testevent3"
+        assert event_payload["path"] == "/baz"
+        assert len(recorder) == 1
+
 
 class TestEvent(object):
     """Test cases for server-side events"""