Added a refactoring to separate sending file logic from StaticRoute dispatcher

aiohttp/helpers.py

@@ -6,6 +6,7 @@
 import datetime
 import functools
 import io
+import mimetypes
 import os
 import re
 from urllib.parse import quote, urlencode
@@ -16,7 +17,6 @@
 
 from . import hdrs
 from .errors import InvalidURL
-
 try:
     from asyncio import ensure_future
 except ImportError:
@@ -529,3 +529,131 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def _cancel_task(self):
         self._cancelled = self._task.cancel()
+
+
+class FileSender:
+    """"A helper that can be used to send files.
+    """
+
+    def __init__(self, resp_factory, chunk_size):
+        self._response_factory = resp_factory
+        self._chunk_size = chunk_size
+        if bool(os.environ.get("AIOHTTP_NOSENDFILE")):
+            self._sendfile = self._sendfile_fallback
+
+    def _sendfile_cb(self, fut, out_fd, in_fd, offset,
+                     count, loop, registered):
+        if registered:
+            loop.remove_writer(out_fd)
+        try:
+            n = os.sendfile(out_fd, in_fd, offset, count)
+            if n == 0:  # EOF reached
+                n = count
+        except (BlockingIOError, InterruptedError):
+            n = 0
+        except Exception as exc:
+            fut.set_exception(exc)
+            return
+
+        if n < count:
+            loop.add_writer(out_fd, self._sendfile_cb, fut, out_fd, in_fd,
+                            offset + n, count - n, loop, True)
+        else:
+            fut.set_result(None)
+
+    @asyncio.coroutine
+    def _sendfile_system(self, req, resp, fobj, count):
+        """
+        Write `count` bytes of `fobj` to `resp` starting from `offset` using
+        the ``sendfile`` system call.
+
+        `req` should be a :obj:`aiohttp.web.Request` instance.
+
+        `resp` should be a :obj:`aiohttp.web.StreamResponse` instance.
+
+        `fobj` should be an open file object.
+
+        `offset` should be an integer >= 0.
+
+        `count` should be an integer > 0.
+        """
+        transport = req.transport
+
+        if transport.get_extra_info("sslcontext"):
+            yield from self._sendfile_fallback(req, resp, fobj, count)
+            return
+
+        yield from resp.drain()
+
+        loop = req.app.loop
+        out_fd = transport.get_extra_info("socket").fileno()
+        in_fd = fobj.fileno()
+        fut = create_future(loop)
+
+        self._sendfile_cb(fut, out_fd, in_fd, 0, count, loop, False)
+
+        yield from fut
+
+    @asyncio.coroutine
+    def _sendfile_fallback(self, req, resp, fobj, count):
+        """
+        Mimic the :meth:`_sendfile_system` method, but without using the
+        ``sendfile`` system call. This should be used on systems that don't
+        support the ``sendfile`` system call.
+
+        To avoid blocking the event loop & to keep memory usage low, `fobj` is
+        transferred in chunks controlled by the `chunk_size` argument to
+        :class:`StaticRoute`.
+        """
+        chunk_size = self._chunk_size
+
+        chunk = fobj.read(chunk_size)
+        while chunk and count > chunk_size:
+            resp.write(chunk)
+            yield from resp.drain()
+            count = count - chunk_size
+            chunk = fobj.read(chunk_size)
+
+        if chunk:
+            resp.write(chunk[:count])
+            yield from resp.drain()
+
+    if hasattr(os, "sendfile"):  # pragma: no cover
+        _sendfile = _sendfile_system
+    else:  # pragma: no cover
+        _sendfile = _sendfile_fallback
+
+    @asyncio.coroutine
+    def send(self, req, filepath):
+        from .web_exceptions import HTTPNotModified
+
+        st = filepath.stat()
+
+        modsince = req.if_modified_since
+        if modsince is not None and st.st_mtime <= modsince.timestamp():
+            raise HTTPNotModified()
+
+        ct, encoding = mimetypes.guess_type(str(filepath))
+        if not ct:
+            ct = 'application/octet-stream'
+
+        resp = self._response_factory()
+        resp.content_type = ct
+        if encoding:
+            resp.headers[hdrs.CONTENT_ENCODING] = encoding
+        resp.last_modified = st.st_mtime
+
+        file_size = st.st_size
+
+        resp.content_length = file_size
+        resp.set_tcp_cork(True)
+        try:
+            yield from resp.prepare(req)
+
+            with filepath.open('rb') as f:
+                yield from self._sendfile(req, resp, f, file_size)
+
+        finally:
+            resp.set_tcp_nodelay(True)
+
+        return resp

aiohttp/web_urldispatcher.py

@@ -3,7 +3,6 @@
 
 import keyword
 import collections
-import mimetypes
 import re
 import os
 import sys
@@ -19,9 +18,10 @@
 
 from . import hdrs, helpers
 from .abc import AbstractRouter, AbstractMatchInfo, AbstractView
+from .helpers import FileSender
 from .protocol import HttpVersion11
 from .web_exceptions import (HTTPMethodNotAllowed, HTTPNotFound,
-                             HTTPNotModified, HTTPExpectationFailed)
+                             HTTPExpectationFailed)
 from .web_reqrep import StreamResponse
 
 
@@ -461,9 +461,6 @@ def __init__(self, name, prefix, directory, *,
         self._chunk_size = chunk_size
         self._response_factory = response_factory
 
-        if bool(os.environ.get("AIOHTTP_NOSENDFILE")):
-            self._sendfile = self._sendfile_fallback
-
     def match(self, path):
         if not path.startswith(self._prefix):
             return None
@@ -581,36 +578,9 @@ def handle(self, request):
         if not filepath.is_file():
             raise HTTPNotFound()
 
-        st = filepath.stat()
-
-        modsince = request.if_modified_since
-        if modsince is not None and st.st_mtime <= modsince.timestamp():
-            raise HTTPNotModified()
-
-        ct, encoding = mimetypes.guess_type(str(filepath))
-        if not ct:
-            ct = 'application/octet-stream'
-
-        resp = self._response_factory()
-        resp.content_type = ct
-        if encoding:
-            resp.headers[hdrs.CONTENT_ENCODING] = encoding
-        resp.last_modified = st.st_mtime
-
-        file_size = st.st_size
-
-        resp.content_length = file_size
-        resp.set_tcp_cork(True)
-        try:
-            yield from resp.prepare(request)
-
-            with filepath.open('rb') as f:
-                yield from self._sendfile(request, resp, f, file_size)
-
-        finally:
-            resp.set_tcp_nodelay(True)
-
-        return resp
+        file_sender = FileSender(resp_factory=self._response_factory,
+                                 chunk_size=self._chunk_size)
+        return file_sender.send(request, filepath)
 
     def __repr__(self):
         name = "'" + self.name + "' " if self.name is not None else ""

tests/test_urldispatch.py

@@ -8,7 +8,7 @@
 from urllib.parse import unquote
 from multidict import CIMultiDict
 import aiohttp.web
-from aiohttp import hdrs, helpers
+from aiohttp import hdrs
 from aiohttp.web import (UrlDispatcher, Request, Response,
                          HTTPMethodNotAllowed, HTTPNotFound,
                          HTTPCreated)
@@ -576,56 +576,6 @@ def test_add_route_invalid_method(self):
             handler = self.make_handler()
             self.router.add_route('INVALID_METHOD', '/path', handler)
 
-    def test_static_handle_eof(self):
-        loop = mock.Mock()
-        route = self.router.add_static('/st',
-                                       os.path.dirname(aiohttp.__file__))
-        with mock.patch('aiohttp.web_urldispatcher.os') as m_os:
-            out_fd = 30
-            in_fd = 31
-            fut = helpers.create_future(self.loop)
-            m_os.sendfile.return_value = 0
-            route._sendfile_cb(fut, out_fd, in_fd, 0, 100, loop, False)
-            m_os.sendfile.assert_called_with(out_fd, in_fd, 0, 100)
-            self.assertTrue(fut.done())
-            self.assertIsNone(fut.result())
-            self.assertFalse(loop.add_writer.called)
-            self.assertFalse(loop.remove_writer.called)
-
-    def test_static_handle_again(self):
-        loop = mock.Mock()
-        route = self.router.add_static('/st',
-                                       os.path.dirname(aiohttp.__file__))
-        with mock.patch('aiohttp.web_urldispatcher.os') as m_os:
-            out_fd = 30
-            in_fd = 31
-            fut = helpers.create_future(self.loop)
-            m_os.sendfile.side_effect = BlockingIOError()
-            route._sendfile_cb(fut, out_fd, in_fd, 0, 100, loop, False)
-            m_os.sendfile.assert_called_with(out_fd, in_fd, 0, 100)
-            self.assertFalse(fut.done())
-            loop.add_writer.assert_called_with(out_fd, route._sendfile_cb,
-                                               fut, out_fd, in_fd, 0, 100,
-                                               loop, True)
-            self.assertFalse(loop.remove_writer.called)
-
-    def test_static_handle_exception(self):
-        loop = mock.Mock()
-        route = self.router.add_static('/st',
-                                       os.path.dirname(aiohttp.__file__))
-        with mock.patch('aiohttp.web_urldispatcher.os') as m_os:
-            out_fd = 30
-            in_fd = 31
-            fut = helpers.create_future(self.loop)
-            exc = OSError()
-            m_os.sendfile.side_effect = exc
-            route._sendfile_cb(fut, out_fd, in_fd, 0, 100, loop, False)
-            m_os.sendfile.assert_called_with(out_fd, in_fd, 0, 100)
-            self.assertTrue(fut.done())
-            self.assertIs(exc, fut.exception())
-            self.assertFalse(loop.add_writer.called)
-            self.assertFalse(loop.remove_writer.called)
-
     def fill_routes(self):
         route1 = self.router.add_route('GET', '/plain', self.make_handler())
         route2 = self.router.add_route('GET', '/variable/{name}',

tests/test_web_functional.py

@@ -1108,34 +1108,3 @@ def go(dirname, filename):
         file_st = os.stat(fname)
 
         self.loop.run_until_complete(go(here, filename))
-
-    def test_env_nosendfile(self):
-        directory = os.path.dirname(__file__)
-
-        with mock.patch.dict(os.environ, {'AIOHTTP_NOSENDFILE': '1'}):
-            route = web.StaticRoute(None, "/", directory)
-            self.assertEqual(route._sendfile, route._sendfile_fallback)
-
-
-class TestStaticFileSendfileFallback(StaticFileMixin,
-                                     unittest.TestCase):
-
-    def patch_sendfile(self, add_static):
-        def f(*args, **kwargs):
-            route = add_static(*args, **kwargs)
-            route._sendfile = route._sendfile_fallback
-            return route
-        return f
-
-
-@unittest.skipUnless(hasattr(os, "sendfile"),
-                     "sendfile system call not supported")
-class TestStaticFileSendfile(StaticFileMixin,
-                             unittest.TestCase):
-
-    def patch_sendfile(self, add_static):
-        def f(*args, **kwargs):
-            route = add_static(*args, **kwargs)
-            route._sendfile = route._sendfile_system
-            return route
-        return f