Support zerocopy send

docs/server-behavior.md

@@ -56,6 +56,10 @@ Applications should generally treat `HEAD` requests in the same manner as `GET`
 
 One exception to this might be if your application serves large file downloads, in which case you might wish to only generate the response headers.
 
+### Zero copy send
+
+When you install httptools, Uvicorn supports [ASGI Zero Copy Send Extension](https://asgi.readthedocs.io/en/latest/extensions.html#zero-copy-send). Using the method specified by the extension, you can call the zero-copy interface of the operating system to send files.
+
 ---
 
 ## Timeouts

tests/protocols/for_test_sendfile.py

@@ -0,0 +1,20 @@
+async def app(scope, receive, send):
+    with open("./README.md", "rb") as file:
+        content_length = len(file.read())
+        file.seek(0, 0)
+        await send(
+            {
+                "type": "http.response.start",
+                "status": 200,
+                "headers": [
+                    (b"Content-Length", str(content_length).encode("ascii")),
+                    (b"Content-Type", b"text/plain; charset=utf8"),
+                ],
+            }
+        )
+        await send(
+            {
+                "type": "http.response.zerocopysend",
+                "file": file.fileno(),
+            }
+        )

tests/protocols/test_http.py

@@ -1,7 +1,10 @@
 import asyncio
 import contextlib
 import logging
+import os
+import subprocess
 
+import httpx
 import pytest
 
 from tests.response import Response
@@ -123,6 +126,9 @@ def clear_buffer(self):
     def set_protocol(self, protocol):
         pass
 
+    def set_write_buffer_limits(self, high=None, low=None):
+        pass
+
 
 class MockLoop(asyncio.AbstractEventLoop):
     def __init__(self, event_loop):
@@ -744,3 +750,26 @@ def test_invalid_http_request(request_line, protocol_cls, caplog, event_loop):
         protocol.data_received(request)
         assert not protocol.transport.buffer
         assert "Invalid HTTP request received." in caplog.messages
+
+
+@pytest.mark.skipif(
+    not hasattr(os, "sendfile"), reason="Only test in platform that has sendfile"
+)
+def test_sendfile():
+    with subprocess.Popen(
+        "python -m uvicorn tests.protocols.for_test_sendfile:app".split(" ")
+    ) as process:
+        with open("./README.md") as file:
+            file_content = file.read()
+
+        while True:
+            try:
+                httpx.get("http://127.0.0.1:8000")
+                break
+            except httpx.ConnectError:
+                continue
+
+        response = httpx.get("http://127.0.0.1:8000")
+        assert response.text == file_content
+
+        process.terminate()

uvicorn/protocols/http/httptools_impl.py

@@ -1,6 +1,7 @@
 import asyncio
 import http
 import logging
+import os
 import re
 import urllib
 from typing import Callable
@@ -38,6 +39,46 @@ def _get_status_line(status_code):
     status_code: _get_status_line(status_code) for status_code in range(100, 600)
 }
 
+try:
+    os.sendfile
+    HAS_SENDFILE = True
+except AttributeError:
+    HAS_SENDFILE = False
+else:
+    # Note: because uvloop don't implements loop.sendfile, so use os.sendfile as here
+    # Related link: https://github.com/MagicStack/uvloop/issues/228
+    async def sendfile(socket_fd: int, file_fd: int, offset: int, count: int) -> None:
+        loop = asyncio.get_event_loop()
+        future = loop.create_future()
+
+        def call_sendfile(
+            socket_fd: int, file_fd: int, offset: int, count: int, registered: bool
+        ) -> None:
+            if registered:
+                loop.remove_writer(socket_fd)
+            try:
+                sent_count = os.sendfile(socket_fd, file_fd, offset, count)
+            except BaseException as exc:
+                future.set_exception(exc)
+            else:
+                if count - sent_count > 0:
+                    new_offset = offset + sent_count
+                    new_count = count - sent_count
+                    loop.add_writer(
+                        socket_fd,
+                        call_sendfile,
+                        socket_fd,
+                        file_fd,
+                        new_offset,
+                        new_count,
+                        True,
+                    )
+                else:
+                    future.set_result(None)
+
+        call_sendfile(socket_fd, file_fd, offset, count, False)
+        return await future
+
 
 class HttpToolsProtocol(asyncio.Protocol):
     def __init__(
@@ -212,6 +253,9 @@ def on_url(self, url):
             "query_string": parsed_url.query if parsed_url.query else b"",
             "headers": self.headers,
         }
+        if HAS_SENDFILE and not is_ssl(self.transport):
+            extensions = self.scope.setdefault("extensions", {})
+            extensions["http.response.zerocopysend"] = {}
 
     def on_header(self, name: bytes, value: bytes):
         name = name.lower()
@@ -239,6 +283,7 @@ def on_headers_complete(self):
 
         existing_cycle = self.cycle
         self.cycle = RequestResponseCycle(
+            loop=self.loop,
             scope=self.scope,
             transport=self.transport,
             flow=self.flow,
@@ -332,6 +377,7 @@ def timeout_keep_alive_handler(self):
 class RequestResponseCycle:
     def __init__(
         self,
+        loop,
         scope,
         transport,
         flow,
@@ -344,6 +390,7 @@ def __init__(
         keep_alive,
         on_response,
     ):
+        self.loop = loop
         self.scope = scope
         self.transport = transport
         self.flow = flow
@@ -369,6 +416,12 @@ def __init__(
         self.chunked_encoding = None
         self.expected_content_length = 0
 
+        # Sendfile
+        self.allow_sendfile = HAS_SENDFILE and not is_ssl(transport)
+        if self.allow_sendfile:
+            # Set the buffer to 0 to avoid the problem of sending file before headers.
+            transport.set_write_buffer_limits(0)
+
     # ASGI exception wrapper
     async def run_asgi(self, app):
         try:
@@ -480,31 +533,65 @@ async def send(self, message):
 
         elif not self.response_complete:
             # Sending response body
-            if message_type != "http.response.body":
-                msg = "Expected ASGI message 'http.response.body', but got '%s'."
-                raise RuntimeError(msg % message_type)
-
-            body = message.get("body", b"")
-            more_body = message.get("more_body", False)
+            use_sendfile = False
+            if message_type == "http.response.body":
+                body = message.get("body", b"")
+                more_body = message.get("more_body", False)
+            elif self.allow_sendfile and message_type == "http.response.zerocopysend":
+                file_fd = message["file"]
+                socket_fd = self.transport.get_extra_info("socket").fileno()
+                sendfile_offset = message.get("offset", None)
+                if sendfile_offset is None:
+                    sendfile_offset = os.lseek(file_fd, 0, os.SEEK_CUR)
+                sendfile_count = message.get("count", None)
+                if sendfile_count is None:
+                    sendfile_count = os.stat(file_fd).st_size - sendfile_offset
+                more_body = message.get("more_body", False)
+                use_sendfile = True
+            else:
+                if self.allow_sendfile:
+                    expect_message_types = (
+                        "http.response.body",
+                        "http.response.zerocopysend",
+                    )
+                else:
+                    expect_message_types = ("http.response.body",)
+                msg = "Expected ASGI message %s, but got '%s'."
+                raise RuntimeError(msg % expect_message_types, message_type)
 
             # Write response body
             if self.scope["method"] == "HEAD":
                 self.expected_content_length = 0
             elif self.chunked_encoding:
-                if body:
-                    content = [b"%x\r\n" % len(body), body, b"\r\n"]
+                if not use_sendfile:
+                    if body:
+                        content = [b"%x\r\n" % len(body), body, b"\r\n"]
+                    else:
+                        content = []
+                    if not more_body:
+                        content.append(b"0\r\n\r\n")
+                    self.transport.write(b"".join(content))
                 else:
-                    content = []
-                if not more_body:
-                    content.append(b"0\r\n\r\n")
-                self.transport.write(b"".join(content))
+                    self.transport.write(b"%x\r\n" % sendfile_count)
+                    await self.flow.drain()
+                    await sendfile(socket_fd, file_fd, sendfile_offset, sendfile_count)
+                    if more_body:
+                        self.transport.write(b"\r\n")
+                    else:
+                        self.transport.write(b"\r\n0\r\n\r\n")
             else:
-                num_bytes = len(body)
+                if not use_sendfile:
+                    num_bytes = len(body)
+                    self.transport.write(body)
+                else:
+                    num_bytes = sendfile_count
+                    await self.flow.drain()
+                    await sendfile(socket_fd, file_fd, sendfile_offset, sendfile_count)
+
                 if num_bytes > self.expected_content_length:
                     raise RuntimeError("Response content longer than Content-Length")
                 else:
                     self.expected_content_length -= num_bytes
-                self.transport.write(body)
 
             # Handle response completion
             if not more_body: