Iteratively encode JSON responses

changelog.d/8013.feature

@@ -0,0 +1 @@
+Iteratively encode JSON to avoid blocking the reactor.

synapse/http/server.py

@@ -22,10 +22,10 @@
 import urllib
 from http import HTTPStatus
 from io import BytesIO
-from typing import Any, Callable, Dict, Tuple, Union
+from typing import Any, Callable, Dict, List, Tuple, Union
 
 import jinja2
-from canonicaljson import encode_canonical_json, encode_pretty_printed_json, json
+from canonicaljson import _canonical_encoder, _pretty_encoder, json
 
 from twisted.internet import defer
 from twisted.python import failure
@@ -492,6 +492,59 @@ class RootOptionsRedirectResource(OptionsResource, RootRedirect):
     pass
 
 
+class _JsonProducer:
+    """
+    Iteratively write JSON to the request.
+    """
+    # The minimum number of bytes for each chunk. Note that the last chunk will
+    # usually be smaller than this.
+    min_chunk_size = 1024
+
+    def __init__(self, request, json_encoder, json_object):
+        self.request = request
+        self._generator = json_encoder.iterencode(json_object)
+
+    def start(self):
+        self.request.registerProducer(self, False)
+
+    def _send_data(self, data: List[str]):
+        """
+        Send a list of strings as a response to the request.
+        """
+        if not data:
+            return
+        self.request.write("".join(data).encode("utf-8"))
+
+    def resumeProducing(self):
+        # We've stopped producing in the meantime.
+        if not self.request:
+            return
+
+        # Get the next chunk and write it to the request. Calling write will
+        # spin the reactor (and might be re-entrant).
+        buffer = []
+        buffered_bytes = 0
+        while buffered_bytes < self.min_chunk_size:
+            try:
+                data = next(self._generator)
+                buffer.append(data)
+                buffered_bytes += len(data)
+            except StopIteration:
+                # Everything is serialized, write any data, then finalize the
+                # producer.
+                self._send_data(buffer)
+                self.request.unregisterProducer()
+                self.request.finish()
+                self.stopProducing()
+                return
+
+        self._send_data(data)
+
+    def stopProducing(self):
+        self._generator = None
+        self.request = None
+
+
 def respond_with_json(
     request: Request,
     code: int,
@@ -526,15 +579,24 @@ def respond_with_json(
         return None
 
     if pretty_print:
-        json_bytes = encode_pretty_printed_json(json_object) + b"\n"
+        encoder = _pretty_encoder
     else:
         if canonical_json or synapse.events.USE_FROZEN_DICTS:
-            # canonicaljson already encodes to bytes
-            json_bytes = encode_canonical_json(json_object)
+            encoder = _canonical_encoder
         else:
-            json_bytes = json.dumps(json_object).encode("utf-8")
+            # TODO Re-use this.
+            encoder = json.JSONEncoder(separators=(",", ":"))
+
+    request.setResponseCode(code)
+    request.setHeader(b"Content-Type", b"application/json")
+    request.setHeader(b"Cache-Control", b"no-cache, no-store, must-revalidate")
 
-    return respond_with_json_bytes(request, code, json_bytes, send_cors=send_cors)
+    if send_cors:
+        set_cors_headers(request)
+
+    producer = _JsonProducer(request, encoder, json_object)
+    producer.start()
+    return NOT_DONE_YET
 
 
 def respond_with_json_bytes(

synapse/rest/key/v2/remote_key_resource.py

@@ -15,12 +15,12 @@
 import logging
 from typing import Dict, Set
 
-from canonicaljson import encode_canonical_json, json
+from canonicaljson import json
 from signedjson.sign import sign_json
 
 from synapse.api.errors import Codes, SynapseError
 from synapse.crypto.keyring import ServerKeyFetcher
-from synapse.http.server import DirectServeJsonResource, respond_with_json_bytes
+from synapse.http.server import DirectServeJsonResource, respond_with_json
 from synapse.http.servlet import parse_integer, parse_json_object_from_request
 
 logger = logging.getLogger(__name__)
@@ -223,4 +223,4 @@ async def query_keys(self, request, query, query_remote_on_cache_miss=False):
 
             results = {"server_keys": signed_keys}
 
-            respond_with_json_bytes(request, 200, encode_canonical_json(results))
+            respond_with_json(request, 200, results, canonical_json=True)