feat: cache NodeJS proxied requests (#12558)

* NodeProxyCache

* add changeset

* Last touch-ups

* add changeset

* Move NodeProxyCache to route_utils + last touch-ups

* Fix race condition

* Actually safer

* Immutable: no need dataclass

---------

Co-authored-by: gradio-pr-bot <[email protected]>

Author: cbensimon

.changeset/floppy-dingos-act.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+fix:feat: cache NodeJS proxied requests

gradio/route_utils.py

@@ -23,6 +23,7 @@
     TYPE_CHECKING,
     Any,
     BinaryIO,
+    NamedTuple,
     Union,
 )
 from urllib.parse import urlparse
@@ -1057,3 +1058,71 @@ def slugify(value):
     )
     value = re.sub(r"[^\w\s-]", "", value.lower())
     return re.sub(r"[-\s]+", "-", value).strip("-_")
+
+
+class NodeProxyCache:
+    """
+    Fan-out streaming cache for NodeJS requests proxying
+    """
+
+    class CacheEntry(NamedTuple):
+        head: bytearray
+        subs: list[asyncio.Queue[bytes | None]]
+        resp: asyncio.Future[httpx.Response | None]
+
+    class ProxyReq(NamedTuple):
+        method: str
+        url: str
+        headers: dict[str, str]
+
+    def __init__(self, client: httpx.AsyncClient):
+        self.client = client
+        self.cache: dict[str, NodeProxyCache.CacheEntry | None] = {}
+
+    async def get(self, req: ProxyReq):
+        key = f"{req.method} {req.url}"
+        key += "::".join(map(":".join, req.headers.items()))
+        res: asyncio.Queue[bytes | None] = asyncio.Queue()
+        if (entry := self.cache.get(key, None)) is None:
+            loop = asyncio.get_running_loop()
+            entry = NodeProxyCache.CacheEntry(bytearray(), [], loop.create_future())
+            asyncio.create_task(self.fetch(key, entry, req))
+            self.cache[key] = entry
+        entry.subs.append(res)
+        head = bytes(entry.head)
+        if (resp := await entry.resp) is None:
+            raise Error("Error while proxying request to Node server")
+        return resp.status_code, resp.headers, NodeProxyCache.iter_body(head, res)
+
+    async def fetch(self, key: str, entry: CacheEntry, req: ProxyReq):
+        try:
+            response = await self.client.send(
+                self.client.build_request(
+                    method=req.method,
+                    url=httpx.URL(req.url),
+                    headers=req.headers,
+                ),
+                stream=True,
+            )
+        except Exception:
+            entry.resp.set_result(None)
+            del self.cache[key]
+            raise
+        entry.resp.set_result(response)
+        try:
+            async for bytes_chunk in response.aiter_raw():
+                entry.head.extend(bytes_chunk)
+                for sub in entry.subs:
+                    sub.put_nowait(bytes_chunk)
+        finally:
+            for sub in entry.subs:
+                sub.put_nowait(None)
+            del self.cache[key]
+            await response.aclose()
+
+    @staticmethod
+    async def iter_body(head: bytes, queue: asyncio.Queue[bytes | None]):
+        if len(head) > 0:
+            yield head
+        while (chunk := await queue.get()) is not None:
+            yield chunk

gradio/routes.py

@@ -102,6 +102,7 @@
     GradioMultiPartParser,
     GradioUploadFile,
     MultiPartException,
+    NodeProxyCache,
     Request,
     compare_passwords_securely,
     create_lifespan_handler,
@@ -281,6 +282,7 @@ def __init__(
     # We're not overriding any defaults here
 
     client = httpx.AsyncClient()
+    proxy_cache = NodeProxyCache(client)
 
     @staticmethod
     async def proxy_to_node(
@@ -305,23 +307,23 @@ async def proxy_to_node(
         if mounted_path:
             server_url += mounted_path
 
-        headers = dict(request.headers)
+        headers = {}  # Do not include arbitrary headers from original request so NodeProxyCache can be effective
         headers["x-gradio-server"] = server_url
         headers["x-gradio-port"] = str(python_port)
 
         if os.getenv("GRADIO_LOCAL_DEV_MODE"):
             headers["x-gradio-local-dev-mode"] = "1"
 
-        new_request = App.client.build_request(
-            request.method, httpx.URL(url), headers=headers
-        )
-        node_response = await App.client.send(new_request, stream=True)
+        if (accept_language := request.headers.get("accept-language")) is not None:
+            headers["accept-language"] = accept_language
+
+        proxy_req = App.proxy_cache.ProxyReq(request.method, url, headers)
+        status, response_headers, aiter_raw = await App.proxy_cache.get(proxy_req)
 
         return StreamingResponse(
-            node_response.aiter_raw(),
-            status_code=node_response.status_code,
-            headers=node_response.headers,
-            background=BackgroundTask(node_response.aclose),
+            aiter_raw,
+            status_code=status,
+            headers=response_headers,
         )
 
     def configure_app(self, blocks: gradio.Blocks) -> None: