Implement loop.sendfile()

asvetlov · Jan 27, 2018 · Jan 22, 2018 · Jan 22, 2018 · Jan 22, 2018 · Jan 22, 2018
commit fa07b3547a05f0872d483d490159596a395ee39c
diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
@@ -38,6 +38,7 @@
 from . import coroutines
 from . import events
 from . import futures
+from . import protocols
 from . import sslproto
 from . import tasks
 from .log import logger
@@ -155,6 +156,60 @@ def _run_until_complete_cb(fut):
     futures._get_loop(fut).stop()
 
 
+
+class _SendfileProtocol(protocols.Protocol):
+    def __init__(self, transp):
+        # transport should be _FlowControlMixin instance
+        self._transport = transp
+        self._proto = transp.get_protocol()
+        self._resume_reading = transp.is_reading()
+        self._write_paused = transp._protocol_paused
+        transp.pause_reading()
+        transp.set_protocol(self)
+        if self._write_paused:
+            self._paused = self._transport._loop.create_future()
+        else:
+            self._paused = None
+
+    def connection_made(self, transport):
+        raise RuntimeError("Broken state, "
+                           "connection should be established already.")
+
+    def connection_lost(self, exc):
+        if self._paused is not None:
+            if exc is None:
+                self._paused.set_result(True)
+            else:
+                self._paused.set_exception(exc)
+        self._proto.connection_lost(exc)
+
+    def pause_writing(self):
+        if self._paused is not None:
+            return
+        self._paused = self._transport._loop.create_future()
+
+    def resume_writing(self):
+        if self._paused is None:
+            return
+        self._paused.set_result(False)
+        self._paused = None
+
+    def data_received(self, data):
+        raise RuntimeError("Broken state, reading should be paused")
+
+    def eof_received(self):
+        raise RuntimeError("Broken state, reading should be paused")
+
+    async def restore(self):
+        self._transport.set_protocol(self._proto)
+        if self._resume_reading:
+            self._transport.resume_reading()
+        if self._paused is not None:
+            await self._paused
+        if self._write_paused:
+            self._proto.resume_writing()
+
+
 class Server(events.AbstractServer):
 
     def __init__(self, loop, sockets):
@@ -866,6 +921,62 @@ async def _create_connection_transport(
 
         return transport, protocol
 
+    async def sendfile(self, transport, file, offset=0, count=None,
+                       *, fallback=True):
+        """Send a file through a transport.
+
+        Return amount of sent bytes.
+        """
+        mode = getattr(transport, '_sendfile_compatible',
+                       constants._SendfileMode.UNSUPPORTED)
+        if mode is constants._SendfileMode.UNSUPPORTED:
+            raise RuntimeError(
+                f"sendfile is not supported for transport {transport!r}")
+        if mode is constants._SendfileMode.NATIVE:
+            try:
+                await self._sendfile_native(transport, file,
+                                            offset, count)
+                return
+            except events.SendfileNotAvailableError as exc:
+                if not fallback:
+                    raise
+        # the mode is FALLBACK or fallback is True
+        await self._sendfile_fallback(transport, file,
+                                      offset, count)
+
+    async def _sendfile_native(self, transp, file, offset, count):
+        raise events.SendfileNotAvailableError(
+            "sendfile syscall is not supported")
+
+    async def _sendfile_fallback(self, transp, file, offset, count):
+        if offset:
+            file.seek(offset)
+        blocksize = min(count, 16384) if count else 16384
+        buf = bytearray(blocksize)
+        total_sent = 0
+        proto = _SendfileProtocol(transp)
+        try:
+            while True:
+                if count:
+                    blocksize = min(count - total_sent, blocksize)
+                    if blocksize <= 0:
+                        break
+                fut = proto._paused
+                if fut is not None:
+                    if await fut:
+                        # eof received
+                        return
+                view = memoryview(buf)[:blocksize]
+                read = file.readinto(view)
+                if not read:
+                    break  # EOF
+                transp.write(view)
+                total_sent += read
+        finally:
+            if total_sent > 0 and hasattr(file, 'seek'):
+                file.seek(offset + total_sent)
+            await proto.restore()
+
     async def start_tls(self, transport, protocol, sslcontext, *,
                         server_side=False,
                         server_hostname=None,

diff --git a/Lib/asyncio/constants.py b/Lib/asyncio/constants.py
@@ -1,3 +1,5 @@
+import enum
+
 # After the connection is lost, log warnings after this many write()s.
 LOG_THRESHOLD_FOR_CONNLOST_WRITES = 5
 
@@ -11,3 +13,10 @@
 
 # Number of seconds to wait for SSL handshake to complete
 SSL_HANDSHAKE_TIMEOUT = 10.0
+
+# The enum should be here to break circular dependencies between
+# base_events and sslproto
+class _SendfileMode(enum.Enum):
+    UNSUPPORTED = enum.auto()
+    NATIVE = enum.auto()
+    FALLBACK = enum.auto()
diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
@@ -317,6 +317,14 @@ async def create_server(
         """
         raise NotImplementedError
 
+    async def sendfile(self, transport, file, offset=0, count=None,
+                       *, fallback=True):
+        """Send a file through a transport.
+
+        Return an amount of sent bytes.
+        """
+        raise NotImplementedError
+
     async def start_tls(self, transport, protocol, sslcontext, *,
                         server_side=False,
                         server_hostname=None,

diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
@@ -540,6 +540,22 @@ def _sock_accept(self, fut, registered, sock):
         else:
             fut.set_result((conn, address))
 
+    async def _sendfile_native(self, transp, file, offset, count):
+        del self._transports[transp._sock_fd]
+        resume_reading = transp.is_reading()
+        transp.pause_reading()
+        fut = transp._make_empty_waiter()
+        await fut
+        sock = transp._sock
+        try:
+            return await self.sock_sendfile(sock, file, offset, count,
+                                            fallback=False)
+        finally:
+            transp._reset_empty_waiter()
+            if resume_reading:
+                transp.resume_reading()
+            self._transports[transp._sock_fd] = transp
+
     def _process_events(self, event_list):
         for key, mask in event_list:
             fileobj, (reader, writer) = key.fileobj, key.data
@@ -695,12 +711,14 @@ def get_write_buffer_size(self):
 class _SelectorSocketTransport(_SelectorTransport):
 
     _start_tls_compatible = True
+    _sendfile_compatible = constants._SendfileMode.NATIVE
 
     def __init__(self, loop, sock, protocol, waiter=None,
                  extra=None, server=None):
         super().__init__(loop, sock, protocol, extra, server)
         self._eof = False
         self._paused = False
+        self._empty_waiter = None
 
         # Disable the Nagle algorithm -- small writes will be
         # sent without waiting for the TCP ACK.  This generally
@@ -765,6 +783,9 @@ def write(self, data):
                             f'not {type(data).__name__!r}')
         if self._eof:
             raise RuntimeError('Cannot call write() after write_eof()')
+        if self._empty_waiter is not None:
+            raise RuntimeError('Cannot call write() when loop.sendfile() '
+                               'is not finished')
         if not data:
             return
 
@@ -798,6 +819,8 @@ def _write_ready(self):
         assert self._buffer, 'Data should not be empty'
 
         if self._conn_lost:
+            if self._empty_waiter is not None:
+                self._empty_waiter.set_result(True)
             return
         try:
             n = self._sock.send(self._buffer)
@@ -807,12 +830,16 @@ def _write_ready(self):
             self._loop._remove_writer(self._sock_fd)
             self._buffer.clear()
             self._fatal_error(exc, 'Fatal write error on socket transport')
+            if self._empty_waiter is not None:
+                self._empty_waiter.set_exception(exc)
         else:
             if n:
                 del self._buffer[:n]
             self._maybe_resume_protocol()  # May append to buffer.
             if not self._buffer:
                 self._loop._remove_writer(self._sock_fd)
+                if self._empty_waiter is not None:
+                    self._empty_waiter.set_result(False)
                 if self._closing:
                     self._call_connection_lost(None)
                 elif self._eof:
@@ -828,6 +855,17 @@ def write_eof(self):
     def can_write_eof(self):
         return True
 
+    def _make_empty_waiter(self):
+        if self._empty_waiter is not None:
+            raise RuntimeError("Empty waiter is already set")
+        self._empty_waiter = self._loop.create_future()
+        if not self._buffer:
+            self._empty_waiter.set_result(None)
+        return self._empty_waiter
+
+    def _reset_empty_waiter(self):
+        self._empty_waiter = None
+
 
 class _SelectorDatagramTransport(_SelectorTransport):
 

diff --git a/Lib/asyncio/sslproto.py b/Lib/asyncio/sslproto.py
@@ -282,6 +282,8 @@ def feed_appdata(self, data, offset=0):
 class _SSLProtocolTransport(transports._FlowControlMixin,
                             transports.Transport):
 
+    _sendfile_compatible = constants._SendfileMode.FALLBACK
+
     def __init__(self, loop, ssl_protocol):
         self._loop = loop
         # SSLProtocol instance
@@ -365,6 +367,11 @@ def get_write_buffer_size(self):
         """Return the current size of the write buffer."""
         return self._ssl_protocol._transport.get_write_buffer_size()
 
+    @property
+    def _protocol_paused(self):
+        # Required for sendfile fallback pause_writing/resume_writing logic
+        return self._ssl_protocol._transport._protocol_paused
+
     def write(self, data):
         """Write some data bytes to the transport.
 

diff --git a/Lib/test/test_asyncio/test_base_events.py b/Lib/test/test_asyncio/test_base_events.py
@@ -1788,7 +1788,7 @@ def runner(loop):
             outer_loop.close()
 
 
-class BaseLoopSendfileTests(test_utils.TestCase):
+class BaseLoopSockSendfileTests(test_utils.TestCase):
 
     DATA = b"12345abcde" * 16 * 1024  # 160 KiB