Issue #26721: Change StreamRequestHandler.wfile to BufferedIOBase

7 files changed, 131 insertions, 15 deletions

diff --git a/Doc/library/http.server.rst b/Doc/library/http.server.rst
index c3584f5..c1ea873 100644
--- a/Doc/library/http.server.rst
+++ b/Doc/library/http.server.rst
@@ -98,8 +98,8 @@ of which this module provides three different variants:
 
    .. attribute:: rfile
 
-      Contains an input stream, positioned at the start of the optional input
-      data.
+      An :class:`io.BufferedIOBase` input stream, ready to read from
+      the start of the optional input data.
 
    .. attribute:: wfile
 
@@ -107,6 +107,9 @@ of which this module provides three different variants:
       client. Proper adherence to the HTTP protocol must be used when writing to
       this stream.
 
+      .. versionchanged:: 3.6
+         This is an :class:`io.BufferedIOBase` stream.
+
    :class:`BaseHTTPRequestHandler` has the following attributes:
 
    .. attribute:: server_version
diff --git a/Doc/library/socketserver.rst b/Doc/library/socketserver.rst
index dac9281..3eb27e3 100644
--- a/Doc/library/socketserver.rst
+++ b/Doc/library/socketserver.rst
@@ -409,6 +409,15 @@ Request Handler Objects
    read or written, respectively, to get the request data or return data
    to the client.
 
+   The :attr:`rfile` attributes of both classes support the
+   :class:`io.BufferedIOBase` readable interface, and
+   :attr:`DatagramRequestHandler.wfile` supports the
+   :class:`io.BufferedIOBase` writable interface.
+
+   .. versionchanged:: 3.6
+      :attr:`StreamRequestHandler.wfile` also supports the
+      :class:`io.BufferedIOBase` writable interface.
+
 
 Examples
 --------
diff --git a/Doc/whatsnew/3.6.rst b/Doc/whatsnew/3.6.rst
index d3a588b..a8802a4 100644
--- a/Doc/whatsnew/3.6.rst
+++ b/Doc/whatsnew/3.6.rst
@@ -373,6 +373,12 @@ defined in :mod:`http.server`, :mod:`xmlrpc.server` and
 protocol.
 (Contributed by Aviv Palivoda in :issue:`26404`.)
 
+The :attr:`~socketserver.StreamRequestHandler.wfile` attribute of
+:class:`~socketserver.StreamRequestHandler` classes now implements
+the :class:`io.BufferedIOBase` writable interface.  In particular,
+calling :meth:`~io.BufferedIOBase.write` is now guaranteed to send the
+data in full.  (Contributed by Martin Panter in :issue:`26721`.)
+
 
 subprocess
 ----------
diff --git a/Lib/socketserver.py b/Lib/socketserver.py
index c6d38c7..41a3766 100644
--- a/Lib/socketserver.py
+++ b/Lib/socketserver.py
@@ -132,6 +132,7 @@ try:
     import threading
 except ImportError:
     import dummy_threading as threading
+from io import BufferedIOBase
 from time import monotonic as time
 
 __all__ = ["BaseServer", "TCPServer", "UDPServer",
@@ -743,7 +744,10 @@ class StreamRequestHandler(BaseRequestHandler):
             self.connection.setsockopt(socket.IPPROTO_TCP,
                                        socket.TCP_NODELAY, True)
         self.rfile = self.connection.makefile('rb', self.rbufsize)
-        self.wfile = self.connection.makefile('wb', self.wbufsize)
+        if self.wbufsize == 0:
+            self.wfile = _SocketWriter(self.connection)
+        else:
+            self.wfile = self.connection.makefile('wb', self.wbufsize)
 
     def finish(self):
         if not self.wfile.closed:
@@ -756,6 +760,24 @@ class StreamRequestHandler(BaseRequestHandler):
         self.wfile.close()
         self.rfile.close()
 
+class _SocketWriter(BufferedIOBase):
+    """Simple writable BufferedIOBase implementation for a socket
+
+    Does not hold data in a buffer, avoiding any need to call flush()."""
+
+    def __init__(self, sock):
+        self._sock = sock
+
+    def writable(self):
+        return True
+
+    def write(self, b):
+        self._sock.sendall(b)
+        with memoryview(b) as view:
+            return view.nbytes
+
+    def fileno(self):
+        return self._sock.fileno()
 
 class DatagramRequestHandler(BaseRequestHandler):
 
diff --git a/Lib/test/test_socketserver.py b/Lib/test/test_socketserver.py
index 9a90729..3f4dfa1 100644
--- a/Lib/test/test_socketserver.py
+++ b/Lib/test/test_socketserver.py
@@ -3,6 +3,7 @@ Test suite for socketserver.
 """
 
 import contextlib
+import io
 import os
 import select
 import signal
@@ -376,6 +377,84 @@ if HAVE_FORKING:
             self.active_children.clear()
 
 
+class SocketWriterTest(unittest.TestCase):
+    def test_basics(self):
+        class Handler(socketserver.StreamRequestHandler):
+            def handle(self):
+                self.server.wfile = self.wfile
+                self.server.wfile_fileno = self.wfile.fileno()
+                self.server.request_fileno = self.request.fileno()
+
+        server = socketserver.TCPServer((HOST, 0), Handler)
+        self.addCleanup(server.server_close)
+        s = socket.socket(
+            server.address_family, socket.SOCK_STREAM, socket.IPPROTO_TCP)
+        with s:
+            s.connect(server.server_address)
+        server.handle_request()
+        self.assertIsInstance(server.wfile, io.BufferedIOBase)
+        self.assertEqual(server.wfile_fileno, server.request_fileno)
+
+    @unittest.skipUnless(threading, 'Threading required for this test.')
+    def test_write(self):
+        # Test that wfile.write() sends data immediately, and that it does
+        # not truncate sends when interrupted by a Unix signal
+        pthread_kill = test.support.get_attribute(signal, 'pthread_kill')
+
+        class Handler(socketserver.StreamRequestHandler):
+            def handle(self):
+                self.server.sent1 = self.wfile.write(b'write data\n')
+                # Should be sent immediately, without requiring flush()
+                self.server.received = self.rfile.readline()
+                big_chunk = bytes(test.support.SOCK_MAX_SIZE)
+                self.server.sent2 = self.wfile.write(big_chunk)
+
+        server = socketserver.TCPServer((HOST, 0), Handler)
+        self.addCleanup(server.server_close)
+        interrupted = threading.Event()
+
+        def signal_handler(signum, frame):
+            interrupted.set()
+
+        original = signal.signal(signal.SIGUSR1, signal_handler)
+        self.addCleanup(signal.signal, signal.SIGUSR1, original)
+        response1 = None
+        received2 = None
+        main_thread = threading.get_ident()
+
+        def run_client():
+            s = socket.socket(server.address_family, socket.SOCK_STREAM,
+                socket.IPPROTO_TCP)
+            with s, s.makefile('rb') as reader:
+                s.connect(server.server_address)
+                nonlocal response1
+                response1 = reader.readline()
+                s.sendall(b'client response\n')
+
+                reader.read(100)
+                # The main thread should now be blocking in a send() syscall.
+                # But in theory, it could get interrupted by other signals,
+                # and then retried. So keep sending the signal in a loop, in
+                # case an earlier signal happens to be delivered at an
+                # inconvenient moment.
+                while True:
+                    pthread_kill(main_thread, signal.SIGUSR1)
+                    if interrupted.wait(timeout=float(1)):
+                        break
+                nonlocal received2
+                received2 = len(reader.read())
+
+        background = threading.Thread(target=run_client)
+        background.start()
+        server.handle_request()
+        background.join()
+        self.assertEqual(server.sent1, len(response1))
+        self.assertEqual(response1, b'write data\n')
+        self.assertEqual(server.received, b'client response\n')
+        self.assertEqual(server.sent2, test.support.SOCK_MAX_SIZE)
+        self.assertEqual(received2, test.support.SOCK_MAX_SIZE - 100)
+
+
 class MiscTestCase(unittest.TestCase):
 
     def test_all(self):
diff --git a/Lib/wsgiref/simple_server.py b/Lib/wsgiref/simple_server.py
index da74d7b..f71563a 100644
--- a/Lib/wsgiref/simple_server.py
+++ b/Lib/wsgiref/simple_server.py
@@ -11,7 +11,6 @@ module.  See also the BaseHTTPServer module docs for other API information.
 """
 
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from io import BufferedWriter
 import sys
 import urllib.parse
 from wsgiref.handlers import SimpleHandler
@@ -127,17 +126,11 @@ class WSGIRequestHandler(BaseHTTPRequestHandler):
         if not self.parse_request(): # An error code has been sent, just exit
             return
 
-        # Avoid passing the raw file object wfile, which can do partial
-        # writes (Issue 24291)
-        stdout = BufferedWriter(self.wfile)
-        try:
-            handler = ServerHandler(
-                self.rfile, stdout, self.get_stderr(), self.get_environ()
-            )
-            handler.request_handler = self      # backpointer for logging
-            handler.run(self.server.get_app())
-        finally:
-            stdout.detach()
+        handler = ServerHandler(
+            self.rfile, self.wfile, self.get_stderr(), self.get_environ()
+        )
+        handler.request_handler = self      # backpointer for logging
+        handler.run(self.server.get_app())
 
 
 
diff --git a/Misc/NEWS b/Misc/NEWS
index e0af6a8..3fcdf69 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,10 @@ What's New in Python 3.6.0 alpha 3
 Library
 -------
 
+- Issue #26721: Change the socketserver.StreamRequestHandler.wfile attribute
+  to implement BufferedIOBase. In particular, the write() method no longer
+  does partial writes.
+
 - Issue #22115: Added methods trace_add, trace_remove and trace_info in the
   tkinter.Variable class.  They replace old methods trace_variable, trace,
   trace_vdelete and trace_vinfo that use obsolete Tcl commands and might