Merge ~blake-rouse/maas:faster-tftp into maas:master

[Blake Rouse]

Fix time tracking with TFTP requests. No longer use monkey patching for the TFTP service as its not required. Use sendfile to send data instead of reading the data into userspace.

be76a50... by Blake Rouse on 2019-03-06

Try using a threadpool to send the binary data.

f0bd30e... by Blake Rouse on 2019-03-07

Use a new eventloop on its own thread to handle read requests.

Unmerged commits

f0bd30e... by Blake Rouse on 2019-03-07

Use a new eventloop on its own thread to handle read requests.

be76a50... by Blake Rouse on 2019-03-06

Try using a threadpool to send the binary data.

af10523... by Blake Rouse on 2019-03-06

Fix time tracking with TFTP requests. Use sendfile to send data instead of reading the data into userspace.

229af11... by Blake Rouse on 2019-03-05

Try to read tftp request from the filesystem before processing the request with the boot methods.

diff --git a/required-packages/dev b/required-packages/dev
index 0fc3b6d..d799095 100644
--- a/required-packages/dev
+++ b/required-packages/dev
@@ -36,4 +36,5 @@ python-tempita
 python-twisted
 python-yaml
 socat
+tftp-hpa
 xvfb
diff --git a/src/provisioningserver/monkey.py b/src/provisioningserver/monkey.py
index f69da32..0f23b75 100644
--- a/src/provisioningserver/monkey.py
+++ b/src/provisioningserver/monkey.py
@@ -24,99 +24,6 @@ def add_term_error_code_to_tftp():
             "Terminate transfer due to option negotiation")
 
 
-def fix_tftp_requests():
-    """Use intelligence in determining IPv4 vs IPv6 when creatinging a session.
-
-       Specifically, look at addr[0] and pass iface to listenUDP based on that.
-
-       See https://bugs.launchpad.net/ubuntu/+source/python-tx-tftp/1614581
-    """
-    import tftp.protocol
-
-    from tftp.datagram import (
-        OP_WRQ,
-        ERRORDatagram,
-        ERR_NOT_DEFINED,
-        ERR_ACCESS_VIOLATION,
-        ERR_FILE_EXISTS,
-        ERR_ILLEGAL_OP,
-        OP_RRQ,
-        ERR_FILE_NOT_FOUND
-    )
-    from tftp.bootstrap import (
-        RemoteOriginWriteSession,
-        RemoteOriginReadSession,
-    )
-    from tftp.netascii import NetasciiReceiverProxy, NetasciiSenderProxy
-    from twisted.internet import reactor
-    from twisted.internet.defer import inlineCallbacks, returnValue
-    from twisted.python.context import call
-    from tftp.errors import (
-        FileExists,
-        Unsupported,
-        AccessViolation,
-        BackendError,
-        FileNotFound,
-    )
-    from netaddr import IPAddress
-
-    @inlineCallbacks
-    def new_startSession(self, datagram, addr, mode):
-        # Set up a call context so that we can pass extra arbitrary
-        # information to interested backends without adding extra call
-        # arguments, or switching to using a request object, for example.
-        context = {}
-        if self.transport is not None:
-            # Add the local and remote addresses to the call context.
-            local = self.transport.getHost()
-            context["local"] = local.host, local.port
-            context["remote"] = addr
-        try:
-            if datagram.opcode == OP_WRQ:
-                fs_interface = yield call(
-                    context, self.backend.get_writer, datagram.filename)
-            elif datagram.opcode == OP_RRQ:
-                fs_interface = yield call(
-                    context, self.backend.get_reader, datagram.filename)
-        except Unsupported as e:
-            self.transport.write(ERRORDatagram.from_code(
-                ERR_ILLEGAL_OP,
-                u"{}".format(e).encode("ascii", "replace")).to_wire(), addr)
-        except AccessViolation:
-            self.transport.write(
-                ERRORDatagram.from_code(ERR_ACCESS_VIOLATION).to_wire(), addr)
-        except FileExists:
-            self.transport.write(
-                ERRORDatagram.from_code(ERR_FILE_EXISTS).to_wire(), addr)
-        except FileNotFound:
-            self.transport.write(
-                ERRORDatagram.from_code(ERR_FILE_NOT_FOUND).to_wire(), addr)
-        except BackendError as e:
-            self.transport.write(ERRORDatagram.from_code(
-                ERR_NOT_DEFINED,
-                u"{}".format(e).encode("ascii", "replace")).to_wire(), addr)
-        else:
-            if IPAddress(addr[0]).version == 6:
-                iface = '::'
-            else:
-                iface = ''
-            if datagram.opcode == OP_WRQ:
-                if mode == b'netascii':
-                    fs_interface = NetasciiReceiverProxy(fs_interface)
-                session = RemoteOriginWriteSession(
-                    addr, fs_interface, datagram.options, _clock=self._clock)
-                reactor.listenUDP(0, session, iface)
-                returnValue(session)
-            elif datagram.opcode == OP_RRQ:
-                if mode == b'netascii':
-                    fs_interface = NetasciiSenderProxy(fs_interface)
-                session = RemoteOriginReadSession(
-                    addr, fs_interface, datagram.options, _clock=self._clock)
-                reactor.listenUDP(0, session, iface)
-                returnValue(session)
-    tftp.protocol.TFTP._startSession = new_startSession
-
-
 def get_patched_URI():
     """Create the patched `twisted.web.client.URI` to handle IPv6."""
     import re
@@ -343,7 +250,6 @@ def augment_twisted_deferToThreadPool():
 
 def add_patches_to_txtftp():
     add_term_error_code_to_tftp()
-    fix_tftp_requests()
 
 
 def add_patches_to_twisted():
diff --git a/src/provisioningserver/rackdservices/tests/test_tftp.py b/src/provisioningserver/rackdservices/tests/test_tftp.py
index 64a89e7..26d21af 100644
--- a/src/provisioningserver/rackdservices/tests/test_tftp.py
+++ b/src/provisioningserver/rackdservices/tests/test_tftp.py
@@ -50,8 +50,8 @@ from provisioningserver.rackdservices.tftp import (
     Port,
     TFTPBackend,
     TFTPService,
+    TransferTimeTrackingIPv6TFTP,
     TransferTimeTrackingSession,
-    TransferTimeTrackingTFTP,
     UDPServer,
 )
 from provisioningserver.rpc.exceptions import BootConfigNoResponse
@@ -93,6 +93,7 @@ from twisted.internet.defer import (
 )
 from twisted.internet.protocol import Protocol
 from twisted.internet.task import Clock
+from twisted.internet.utils import getProcessOutputAndValue
 from twisted.python import context
 from zope.interface.verify import verifyObject
 
@@ -973,6 +974,59 @@ class TestTFTPService(MAASTestCase):
         })
 
 
+class TestTFTPService_Sendfile(MAASTestCase):
+
+    run_tests_with = MAASTwistedRunTest.make_factory(timeout=30)
+
+    scenarios = (
+        ('div_block_size', {
+            'data': factory.make_string(size=1024 * 1024).encode('utf-8')
+        }),
+        ('non_div_block_size', {
+            'data': factory.make_string(size=(1024 * 1024) + 1).encode('utf-8')
+        }),
+    )
+
+    def test_tftp_service_with_sendfile(self):
+        # Make the fake file to read.
+        example_root = self.make_dir()
+        filename = factory.make_name('file')
+        factory.make_file(example_root, filename, self.data)
+
+        # Create the TFTP service.
+        example_client_service = Mock()
+        example_port = factory.pick_port()
+        tftp_service = TFTPService(
+            resource_root=example_root, client_service=example_client_service,
+            port=example_port)
+        tftp_service.updateServers()
+        tftp_service.startService()
+
+        # Run subprocess to make request for file over TFTP.
+        output_dir = self.make_dir()
+        cmd = [
+            'tftp', '-m', 'binary',
+            'localhost', str(example_port),
+            '-c', 'get', filename,
+        ]
+
+        def tftp_cb(result):
+            tftp_service.stopService()
+
+            out, err, code = result
+            self.assertEquals(0, code, "\nstdout: %s\nstrerr: %s" % (out, err))
+
+            with open(os.path.join(output_dir, filename), 'rb') as fp:
+                content = fp.read()
+            self.assertEqual(self.data, content)
+
+        self.assertFalse(os.path.exists(os.path.join(output_dir, filename)))
+        d = getProcessOutputAndValue(
+            cmd[0], cmd[1:], path=output_dir, reactor=reactor)
+        d.addCallbacks(tftp_cb, tftp_cb)
+        return d
+
+
 class TestTransferTimeTrackingSession(MAASTestCase):
 
     def test_track_time(self):
@@ -991,12 +1045,12 @@ class TestTransferTimeTrackingSession(MAASTestCase):
             metrics)
 
 
-class TestTransferTimeTrackingTFTP(MAASTestCase):
-    """Tests for `TransferTimeTrackingTFTP`."""
+class TestTransferTimeTrackingIPv6TFTP(MAASTestCase):
+    """Tests for `TransferTimeTrackingIPv6TFTP`."""
 
     def clean_filename(self, path):
         datagram = RQDatagram(path, b'octet', {})
-        tftp = TransferTimeTrackingTFTP(sentinel.backend)
+        tftp = TransferTimeTrackingIPv6TFTP(sentinel.backend)
         return tftp._clean_filename(datagram)
 
     def test_clean_filename(self):
diff --git a/src/provisioningserver/rackdservices/tftp.py b/src/provisioningserver/rackdservices/tftp.py
index 6254283..d86289c 100644
--- a/src/provisioningserver/rackdservices/tftp.py
+++ b/src/provisioningserver/rackdservices/tftp.py
@@ -8,11 +8,16 @@ __all__ = [
     "TFTPService",
     ]
 
+import asyncio
+import threading
 from functools import partial
+import os
+import socket
 from socket import (
     AF_INET,
     AF_INET6,
 )
+import struct
 from time import time
 
 from netaddr import IPAddress
@@ -46,12 +51,36 @@ from provisioningserver.utils.twisted import (
     RPCFetcher,
 )
 from tftp.backend import FilesystemSynchronousBackend
+from tftp.bootstrap import (
+    RemoteOriginReadSession,
+    RemoteOriginWriteSession,
+    TFTPBootstrap,
+)
+from tftp.datagram import (
+    ERR_ACCESS_VIOLATION,
+    ERR_FILE_EXISTS,
+    ERR_FILE_NOT_FOUND,
+    ERR_ILLEGAL_OP,
+    ERR_NOT_DEFINED,
+    ERRORDatagram,
+    OP_DATA,
+    OP_RRQ,
+    OP_WRQ,
+)
 from tftp.errors import (
+    AccessViolation,
     BackendError,
+    FileExists,
     FileNotFound,
+    Unsupported,
+)
+from tftp.netascii import (
+    NetasciiReceiverProxy,
+    NetasciiSenderProxy,
 )
 from tftp.protocol import TFTP
 from tftp.session import ReadSession
+from tftp.util import timedCaller
 from twisted.application import internet
 from twisted.application.service import MultiService
 from twisted.internet import (
@@ -69,7 +98,9 @@ from twisted.internet.defer import (
     returnValue,
     succeed,
 )
+from twisted.internet.asyncioreactor import AsyncioSelectorReactor
 from twisted.internet.task import deferLater
+from twisted.python.context import call
 from twisted.python.filepath import FilePath
 
 
@@ -372,10 +403,20 @@ class TFTPBackend(FilesystemSynchronousBackend):
         # unix compatiable.
         file_name = file_name.replace(b'\\', b'/')
         log_request(file_name)
-        d = self.get_boot_method(file_name)
-        d.addCallback(partial(self.handle_boot_method, file_name))
-        d.addErrback(self.no_response_errback, file_name)
-        d.addErrback(self.all_is_lost_errback)
+
+        def boot_methods(_):
+            # Failure most likely because file doesn't exist. So send the
+            # request to the boot methods.
+            d = self.get_boot_method(file_name)
+            d.addCallback(partial(self.handle_boot_method, file_name))
+            d.addErrback(self.no_response_errback, file_name)
+            d.addErrback(self.all_is_lost_errback)
+            return d
+
+        # Try to load the file directly from the filesystem first, failure will
+        # pass the request onto the boot methods for processing.
+        d = super(TFTPBackend, self).get_reader(file_name)
+        d.addErrback(boot_methods)
         return d
 
 
@@ -412,7 +453,66 @@ class UDPServer(internet.UDPServer):
         return p
 
 
-class TransferTimeTrackingSession(ReadSession):
+class ReadSendfileSession(ReadSession):
+    """A `ReadSession` that can use `os.sendfile`."""
+
+    def __init__(self, reader, _clock=None):
+        super().__init__(reader, _clock)
+
+    def startProtocol(self):
+        self.blocknum_noroll = 0
+        self.use_sendfile = getattr(self.reader, 'file_obj', None) is not None
+        if self.use_sendfile:
+            self.size = self.reader.size
+        super().startProtocol()
+
+    def nextBlock(self):
+        """ACK datagram for the previous block has been received. Attempt to
+        read the next block, that will be sent."""
+        self.blocknum += 1
+        self.blocknum_noroll += 1
+        if self.use_sendfile:
+            self.timeout_watchdog = timedCaller(
+                (0,) + self.timeout, self.sendWithSendfile,
+                self.timedOut, clock=self._clock)
+            return succeed(None)
+        else:
+            d = maybeDeferred(self.reader.read, self.block_size)
+            d.addCallbacks(
+                callback=self.dataFromReader, errback=self.readFailed)
+            return d
+
+    def sendWithSendfile(self):
+        d = maybeDeferred(self._sendWithSendfile)
+        d.addErrback(self.readFailed)
+        return d
+
+    def _sendWithSendfile(self):
+        # Reached maximum number of blocks. Rolling over
+        if self.blocknum == 65536:
+            self.blocknum = 0
+
+        offset = (self.blocknum_noroll - 1) * self.block_size
+        count = self.size - offset
+        if count > self.block_size:
+            count = self.block_size
+        if count == 0:
+            self.transport.socket.send(
+                struct.pack(b'!HH', OP_DATA, self.blocknum))
+            self.completed = True
+        else:
+            self.transport.socket.send(
+                struct.pack(b'!HH', OP_DATA, self.blocknum),
+                socket.MSG_MORE)
+            os.sendfile(
+                self.transport.socket.fileno(), self.reader.file_obj.fileno(),
+                offset, count)
+            if count < self.block_size:
+                self.completed = True
+
+
+class TransferTimeTrackingSession(ReadSendfileSession):
+    """A `ReadSendfileSession` that tracks the latency of sending."""
 
     def __init__(
             self, filename, reader, _clock=None,
@@ -435,18 +535,95 @@ class TransferTimeTrackingSession(ReadSession):
         super().cancel()
 
 
-class TransferTimeTrackingTFTP(TFTP):
+class RemoteOriginReadTimeTrackingSession(RemoteOriginReadSession):
+    """A `RemoteOriginReadSession` that uses a
+    `TransferTimeTrackingSession`."""
+
+    def __init__(
+            self, filename, remote, reader, options=None, _clock=None):
+        TFTPBootstrap.__init__(self, remote, reader, options, _clock)
+        self.session = TransferTimeTrackingSession(
+            filename, reader, self._clock)
+
+    def stopProtocol(self):
+        super().stopProtocol()
+        self._clock.stop()
+
+
+def asyncio_worker(session, iface):
+    loop = asyncio.new_event_loop()
+    read_reactor = AsyncioSelectorReactor(eventloop=loop)
+    session._clock = read_reactor
+    session.session._clock = read_reactor
+    read_reactor.listenUDP(0, session, iface)
+    read_reactor.run(installSignalHandlers=False)
+
+
+class TransferTimeTrackingIPv6TFTP(TFTP):
+    """A `TFTP` that tracks the TFTP transfer times and supports IPv6."""
+
+    def __init__(self, backend, _clock=None):
+        super().__init__(backend, _clock)
 
     @inlineCallbacks
     def _startSession(self, datagram, addr, mode):
-        session = yield super()._startSession(datagram, addr, mode)
-        stream_session = getattr(session, 'session', None)
-        # replace the standard ReadSession with one that tracks transfer time
-        if stream_session is not None:
-            filename = self._clean_filename(datagram)
-            session.session = TransferTimeTrackingSession(
-                filename, stream_session.reader, _clock=stream_session._clock)
-        returnValue(session)
+        # Set up a call context so that we can pass extra arbitrary
+        # information to interested backends without adding extra call
+        # arguments, or switching to using a request object, for example.
+        context = {}
+        if self.transport is not None:
+            # Add the local and remote addresses to the call context.
+            local = self.transport.getHost()
+            context["local"] = local.host, local.port
+            context["remote"] = addr
+        try:
+            if datagram.opcode == OP_WRQ:
+                fs_interface = yield call(
+                    context, self.backend.get_writer, datagram.filename)
+            elif datagram.opcode == OP_RRQ:
+                fs_interface = yield call(
+                    context, self.backend.get_reader, datagram.filename)
+        except Unsupported as e:
+            self.transport.write(ERRORDatagram.from_code(
+                ERR_ILLEGAL_OP,
+                u"{}".format(e).encode("ascii", "replace")).to_wire(), addr)
+        except AccessViolation:
+            self.transport.write(
+                ERRORDatagram.from_code(ERR_ACCESS_VIOLATION).to_wire(), addr)
+        except FileExists:
+            self.transport.write(
+                ERRORDatagram.from_code(ERR_FILE_EXISTS).to_wire(), addr)
+        except FileNotFound:
+            self.transport.write(
+                ERRORDatagram.from_code(ERR_FILE_NOT_FOUND).to_wire(), addr)
+        except BackendError as e:
+            self.transport.write(ERRORDatagram.from_code(
+                ERR_NOT_DEFINED,
+                u"{}".format(e).encode("ascii", "replace")).to_wire(), addr)
+        else:
+            if IPAddress(addr[0]).version == 6:
+                iface = '::'
+            else:
+                iface = ''
+            if datagram.opcode == OP_WRQ:
+                if mode == b'netascii':
+                    fs_interface = NetasciiReceiverProxy(fs_interface)
+                session = RemoteOriginWriteSession(
+                    addr, fs_interface, datagram.options, _clock=self._clock)
+                reactor.listenUDP(0, session, iface)
+                returnValue(session)
+            elif datagram.opcode == OP_RRQ:
+                if mode == b'netascii':
+                    fs_interface = NetasciiSenderProxy(fs_interface)
+                # Spawn a new reactor in its own thread to handle the read
+                # request. The session will stop the reactor once complete.
+                session = RemoteOriginReadTimeTrackingSession(
+                    self._clean_filename(datagram), addr, fs_interface,
+                    datagram.options)
+                t = threading.Thread(
+                    target=asyncio_worker, args=(session, iface))
+                t.start()
+                returnValue(session)
 
     def _clean_filename(self, datagram):
         filename = datagram.filename.decode('ascii')
@@ -522,7 +699,7 @@ class TFTPService(MultiService, object):
         for address in addrs_desired - addrs_established:
             if not IPAddress(address).is_link_local():
                 tftp_service = UDPServer(
-                    self.port, TransferTimeTrackingTFTP(self.backend),
+                    self.port, TransferTimeTrackingIPv6TFTP(self.backend),
                     interface=address)
                 tftp_service.setName(address)
                 tftp_service.setServiceParent(self)