From adc7e4f17fa413370aacaedfef3130f74affb35b Mon Sep 17 00:00:00 2001
From: Ben Darnell <ben@bendarnell.com>
Date: Sun, 16 Mar 2014 18:43:49 -0400
Subject: [PATCH] Add chunk and callback arguments to write_headers.

This restores an optimization that allows small responses to be written out
in the same system call as the headers.
---
 tornado/http1connection.py | 21 +++++++++++++++------
 tornado/web.py             | 16 +++++++++-------
 tornado/wsgi.py            |  6 +++++-
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/tornado/http1connection.py b/tornado/http1connection.py
index 3a1284c0..ac759a4b 100644
--- a/tornado/http1connection.py
+++ b/tornado/http1connection.py
@@ -173,7 +173,7 @@ class HTTP1Connection(object):
         # cycle and delay garbage collection of this connection.
         self._clear_request_state()
 
-    def write_headers(self, start_line, headers):
+    def write_headers(self, start_line, headers, chunk=None, callback=None):
         self._chunking = (
             # TODO: should this use self._version or start_line.version?
             self._version == 'HTTP/1.1' and
@@ -194,17 +194,26 @@ class HTTP1Connection(object):
             if b'\n' in line:
                 raise ValueError('Newline in header: ' + repr(line))
         if not self.stream.closed():
-            self.stream.write(b"\r\n".join(lines) + b"\r\n\r\n")
+            self._write_callback = stack_context.wrap(callback)
+            data = b"\r\n".join(lines) + b"\r\n\r\n"
+            if chunk:
+                data += self._format_chunk(chunk)
+            self.stream.write(data, self._on_write_complete)
 
-    def write(self, chunk, callback=None):
-        """Writes a chunk of output to the stream."""
+    def _format_chunk(self, chunk):
         if self._chunking and chunk:
             # Don't write out empty chunks because that means END-OF-STREAM
             # with chunked encoding
-            chunk = utf8("%x" % len(chunk)) + b"\r\n" + chunk + b"\r\n"
+            return utf8("%x" % len(chunk)) + b"\r\n" + chunk + b"\r\n"
+        else:
+            return chunk
+
+    def write(self, chunk, callback=None):
+        """Writes a chunk of output to the stream."""
         if not self.stream.closed():
             self._write_callback = stack_context.wrap(callback)
-            self.stream.write(chunk, self._on_write_complete)
+            self.stream.write(self._format_chunk(chunk),
+                              self._on_write_complete)
 
     def finish(self):
         """Finishes the request."""
diff --git a/tornado/web.py b/tornado/web.py
index 8442cd18..3650180d 100644
--- a/tornado/web.py
+++ b/tornado/web.py
@@ -777,6 +777,10 @@ class RequestHandler(object):
                 self._status_code, self._headers, chunk = \
                     transform.transform_first_chunk(
                         self._status_code, self._headers, chunk, include_footers)
+            # Ignore the chunk and only write the headers for HEAD requests
+            if self.request.method == "HEAD":
+                chunk = None
+
             # Finalize the cookie headers (which have been stored in a side
             # object so an outgoing cookie could be overwritten before it
             # is sent).
@@ -787,16 +791,14 @@ class RequestHandler(object):
             start_line = httputil.ResponseStartLine(self.request.version,
                                                     self._status_code,
                                                     self._reason)
-            self.request.connection.write_headers(start_line, self._headers)
+            self.request.connection.write_headers(start_line, self._headers,
+                                                  chunk, callback=callback)
         else:
             for transform in self._transforms:
                 chunk = transform.transform_chunk(chunk, include_footers)
-
-        # Ignore the chunk and only write the headers for HEAD requests
-        if self.request.method == "HEAD":
-            return
-
-        self.request.write(chunk, callback=callback)
+            # Ignore the chunk and only write the headers for HEAD requests
+            if self.request.method != "HEAD":
+                self.request.connection.write(chunk, callback=callback)
 
     def finish(self, chunk=None):
         """Finishes this response, ending the HTTP request."""
diff --git a/tornado/wsgi.py b/tornado/wsgi.py
index 8ff8a76e..8a4889e5 100644
--- a/tornado/wsgi.py
+++ b/tornado/wsgi.py
@@ -95,10 +95,14 @@ class _WSGIConnection(object):
         # so we can simply ignore the callback.
         pass
 
-    def write_headers(self, start_line, headers):
+    def write_headers(self, start_line, headers, chunk=None, callback=None):
         self.start_response(
             '%s %s' % (start_line.code, start_line.reason),
             [(native_str(k), native_str(v)) for (k, v) in headers.get_all()])
+        if chunk is not None:
+            self.write(chunk, callback)
+        elif callback is not None:
+            callback()
 
     def write(self, chunk, callback=None):
         self._write_buffer.append(chunk)