warp: compose response header directly into the write buffer

warp/Network/Wai/Handler/Warp/IO.hs

@@ -4,23 +4,38 @@ import Control.Exception (mask_)
 import Data.ByteString.Builder (Builder)
 import Data.ByteString.Builder.Extra (Next (Chunk, Done, More), runBuilder)
 import Data.IORef (IORef, readIORef, writeIORef)
+import Foreign.Ptr (plusPtr)
 import Network.Wai.Handler.Warp.Buffer
 import Network.Wai.Handler.Warp.Imports
 import Network.Wai.Handler.Warp.Types
 
 toBufIOWith
     :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO Integer
-toBufIOWith maxRspBufSize writeBufferRef io builder = do
+toBufIOWith = toBufIOWithOffset 0
+
+-- | Like 'toBufIOWith' but the first @offset@ bytes of the write buffer
+-- are assumed to be already filled (e.g. with a response header composed
+-- directly into the buffer). They are flushed together with the first
+-- batch of builder output and included in the returned total.
+-- @offset@ must not exceed the current buffer size.
+toBufIOWithOffset
+    :: Int
+    -> Int
+    -> IORef WriteBuffer
+    -> (ByteString -> IO ())
+    -> Builder
+    -> IO Integer
+toBufIOWithOffset offset0 maxRspBufSize writeBufferRef io builder = do
     writeBuffer <- readIORef writeBufferRef
-    loop writeBuffer firstWriter 0
+    loop writeBuffer offset0 firstWriter 0
   where
     firstWriter = runBuilder builder
-    loop writeBuffer writer bytesSent = do
+    loop writeBuffer offset writer bytesSent = do
         let buf = bufBuffer writeBuffer
             size = bufSize writeBuffer
-        (len, signal) <- writer buf size
-        bufferIO buf len io
-        let totalBytesSent = toInteger len + bytesSent
+        (len, signal) <- writer (buf `plusPtr` offset) (size - offset)
+        bufferIO buf (offset + len) io
+        let totalBytesSent = toInteger (offset + len) + bytesSent
         case signal of
             Done -> return totalBytesSent
             More minSize next
@@ -43,8 +58,8 @@ toBufIOWith maxRspBufSize writeBufferRef io builder = do
                         biggerWriteBuffer <- createWriteBuffer minSize
                         writeIORef writeBufferRef biggerWriteBuffer
                         return biggerWriteBuffer
-                    loop biggerWriteBuffer next totalBytesSent
-                | otherwise -> loop writeBuffer next totalBytesSent
+                    loop biggerWriteBuffer 0 next totalBytesSent
+                | otherwise -> loop writeBuffer 0 next totalBytesSent
             Chunk bs next -> do
                 io bs
-                loop writeBuffer next totalBytesSent
+                loop writeBuffer 0 next totalBytesSent

warp/Network/Wai/Handler/Warp/Response.hs

@@ -26,6 +26,7 @@ import Data.ByteString.Builder.HTTP.Chunked (
 import qualified Data.ByteString.Char8 as C8
 import qualified Data.CaseInsensitive as CI
 import Data.Function (on)
+import Data.IORef (readIORef)
 import Data.List (deleteBy)
 import Data.Streaming.ByteString.Builder (
     newByteStringBuilderRecv,
@@ -42,7 +43,7 @@ import Network.Wai.Handler.Warp.Buffer (toBuilderBuffer)
 import qualified Network.Wai.Handler.Warp.Date as D
 import Network.Wai.Handler.Warp.File
 import Network.Wai.Handler.Warp.Header
-import Network.Wai.Handler.Warp.IO (toBufIOWith)
+import Network.Wai.Handler.Warp.IO (toBufIOWith, toBufIOWithOffset)
 import Network.Wai.Handler.Warp.Imports
 import Network.Wai.Handler.Warp.ResponseHeader
 import Network.Wai.Handler.Warp.Settings
@@ -237,21 +238,32 @@ sendRsp conn _ _ ver s hs _ _ _ RspNoBody = do
 ----------------------------------------------------------------
 
 sendRsp conn _ th ver s hs _ maxRspBufSize _ (RspBuilder body needsChunked) = do
-    header <- composeHeaderBuilder ver s hs needsChunked
-    let hdrBdy
-            | needsChunked =
-                header
-                    <> chunkedTransferEncoding body
-                    <> chunkedTransferTerminator
-            | otherwise = header <> body
-        writeBufferRef = connWriteBuffer conn
+    writeBuffer <- readIORef writeBufferRef
     len <-
-        toBufIOWith
-            maxRspBufSize
-            writeBufferRef
-            (\bs -> connSendAll conn bs >> T.tickle th)
-            hdrBdy
+        if hdrLen < bufSize writeBuffer
+            then do
+                -- Compose the header directly into the connection write
+                -- buffer and run the body builder right after it, saving
+                -- a copy of the header bytes through an intermediate
+                -- ByteString.
+                _ <- composeHeaderPtr (bufBuffer writeBuffer) ver s hs'
+                toBufIOWithOffset hdrLen maxRspBufSize writeBufferRef send bdy
+            else do
+                -- Huge headers: fall back to composing a separate header
+                -- ByteString and letting the builder machinery copy it.
+                header <- composeHeaderBuilder ver s hs needsChunked
+                toBufIOWith maxRspBufSize writeBufferRef send (header <> bdy)
     return (Just s, Just len)
+  where
+    hs'
+        | needsChunked = addTransferEncoding hs
+        | otherwise = hs
+    hdrLen = composeHeaderLength s hs'
+    bdy
+        | needsChunked = chunkedTransferEncoding body <> chunkedTransferTerminator
+        | otherwise = body
+    writeBufferRef = connWriteBuffer conn
+    send bs = connSendAll conn bs >> T.tickle th
 
 ----------------------------------------------------------------
 

warp/Network/Wai/Handler/Warp/ResponseHeader.hs

@@ -1,7 +1,11 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE OverloadedStrings #-}
 
-module Network.Wai.Handler.Warp.ResponseHeader (composeHeader) where
+module Network.Wai.Handler.Warp.ResponseHeader (
+    composeHeader,
+    composeHeaderPtr,
+    composeHeaderLength,
+) where
 
 import qualified Data.ByteString as S
 import Data.ByteString.Internal (create)
@@ -18,15 +22,32 @@ import Network.Wai.Handler.Warp.Imports
 ----------------------------------------------------------------
 
 composeHeader :: H.HttpVersion -> H.Status -> H.ResponseHeaders -> IO ByteString
-composeHeader !httpversion !status !responseHeaders = create len $ \ptr -> do
-    ptr1 <- copyStatus ptr httpversion status
-    ptr2 <- copyHeaders ptr1 responseHeaders
-    void $ copyCRLF ptr2
+composeHeader !httpversion !status !responseHeaders =
+    create len $ \ptr ->
+        void $ composeHeaderPtr ptr httpversion status responseHeaders
+  where
+    !len = composeHeaderLength status responseHeaders
+
+-- | The exact number of bytes 'composeHeaderPtr' writes for this
+-- status line and header list (including the final CRLF).
+composeHeaderLength :: H.Status -> H.ResponseHeaders -> Int
+composeHeaderLength !status !responseHeaders =
+    17 + slen + foldl' fieldLength 0 responseHeaders
   where
-    !len = 17 + slen + foldl' fieldLength 0 responseHeaders
     fieldLength !l (!k, !v) = l + S.length (CI.original k) + S.length v + 4
     !slen = S.length $ H.statusMessage status
 
+-- | Compose the response header directly into the given buffer,
+-- returning the number of bytes written. The buffer must have room
+-- for at least 'composeHeaderLength' bytes.
+composeHeaderPtr
+    :: Ptr Word8 -> H.HttpVersion -> H.Status -> H.ResponseHeaders -> IO Int
+composeHeaderPtr !ptr !httpversion !status !responseHeaders = do
+    ptr1 <- copyStatus ptr httpversion status
+    ptr2 <- copyHeaders ptr1 responseHeaders
+    ptr3 <- copyCRLF ptr2
+    return $! ptr3 `minusPtr` ptr
+
 httpVer11 :: ByteString
 httpVer11 = "HTTP/1.1 "