warp/recv: optimistic non-blocking recv before parking the thread

recv/Network/Socket/BufferPool.hs

@@ -24,9 +24,11 @@ module Network.Socket.BufferPool (
     -- * Recv
     Recv,
     receive,
+    receiveNoWait,
     BufferPool,
     newBufferPool,
     withBufferPool,
+    tryWithBufferPool,
 
     -- * RecvN
     RecvN,

recv/Network/Socket/BufferPool/Buffer.hs

@@ -1,6 +1,7 @@
 module Network.Socket.BufferPool.Buffer (
     newBufferPool,
     withBufferPool,
+    tryWithBufferPool,
     mallocBS,
     copy,
 ) where
@@ -44,6 +45,26 @@ withBufferPool (BufferPool l h ref) f = do
     writeIORef ref $ unsafeDrop consumed buf
     return $ unsafeTake consumed buf
 
+-- | Like 'withBufferPool' for fillers that can decline to fill:
+--   a negative return value from the filler leaves the pool untouched
+--   and produces 'Nothing'.
+tryWithBufferPool
+    :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO (Maybe ByteString)
+tryWithBufferPool (BufferPool l h ref) f = do
+    buf0 <- readIORef ref
+    buf <-
+        if BS.length buf0 >= l
+            then return buf0
+            else mallocBS h
+    consumed <- withForeignBuffer buf f
+    if consumed < 0
+        then do
+            writeIORef ref buf
+            return Nothing
+        else do
+            writeIORef ref $ unsafeDrop consumed buf
+            return $ Just $ unsafeTake consumed buf
+
 withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO Int) -> IO Int
 withForeignBuffer (PS ps s l) f = withForeignPtr ps $ \p -> f (castPtr p `plusPtr` s) l
 {-# INLINE withForeignBuffer #-}

recv/Network/Socket/BufferPool/Recv.hs

@@ -1,14 +1,22 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE OverloadedStrings #-}
 
 module Network.Socket.BufferPool.Recv (
     receive,
+    receiveNoWait,
     makeRecvN,
 ) where
 
 import qualified Data.ByteString as BS
 import Data.ByteString.Internal (ByteString (..), unsafeCreate)
 import Data.IORef
 import Network.Socket (Socket, recvBuf)
+#ifndef mingw32_HOST_OS
+import Foreign.C.Types (CInt (..), CSize (..))
+import Foreign.Ptr (Ptr, castPtr)
+import Network.Socket (withFdSocket)
+import System.Posix.Types (CSsize (..))
+#endif
 
 import Network.Socket.BufferPool.Buffer
 import Network.Socket.BufferPool.Types
@@ -20,6 +28,27 @@ import Network.Socket.BufferPool.Types
 receive :: Socket -> BufferPool -> Recv
 receive sock pool = withBufferPool pool $ \ptr size -> recvBuf sock ptr size
 
+-- | Like 'receive' but never blocks and never involves the IO manager:
+--   'Nothing' means no data was available (or an error occurred, which a
+--   subsequent blocking 'receive' will report properly). @Just \"\"@ is EOF.
+--   On Windows this always returns 'Nothing'.
+receiveNoWait :: Socket -> BufferPool -> IO (Maybe ByteString)
+#ifndef mingw32_HOST_OS
+receiveNoWait sock pool = tryWithBufferPool pool $ \ptr size ->
+    withFdSocket sock $ \fd -> do
+        -- The socket is non-blocking, so an unsafe foreign call is fine:
+        -- recv(2) returns immediately either way. Both EAGAIN and real
+        -- errors map to a negative result, deferring to the blocking path
+        -- so errors surface there with their errno intact.
+        n <- c_recv fd (castPtr ptr) (fromIntegral size) 0
+        return $ fromIntegral n
+
+foreign import ccall unsafe "recv"
+    c_recv :: CInt -> Ptr CSsize -> CSize -> CInt -> IO CSsize
+#else
+receiveNoWait _ _ = return Nothing
+#endif
+
 ----------------------------------------------------------------
 
 -- | This function returns a receiving function

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

@@ -214,13 +214,10 @@ processRequest settings ii conn app th istatus src req mremainingRef idxhdr next
     keepAlive <- readIORef keepAliveRef
 
     -- We just send a Response and it takes a time to
-    -- receive a Request again. If we immediately call recv,
-    -- it is likely to fail and cause the IO manager to do some work.
-    -- It is very costly, so we yield to another Haskell
-    -- thread hoping that the next Request will arrive
-    -- when this Haskell thread will be re-scheduled.
-    -- This improves performance at least when
-    -- the number of cores is small.
+    -- receive a Request again. Yield to other Haskell threads so the
+    -- next Request has a chance to arrive before we call recv: the
+    -- non-blocking fast path in 'makeGracefulRecv' then reads it without
+    -- involving the IO manager at all.
     Conc.yield
 
     if keepAlive

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

@@ -144,22 +144,33 @@ socketConnection set s = do
 -- actively using this 'Socket'.
 makeGracefulRecv :: Socket -> BufferPool -> ServerState -> TVar Int -> Recv
 makeGracefulRecv sock pool ss appsInProgress = do
-    sockWait <-
+    -- Fast path: under load the next request has usually already arrived,
+    -- so read it without registering with the IO manager or entering STM.
+    -- One epoll_ctl + one futex sleep/wake pair saved per request.
+    -- A request already in the kernel buffer is served even if shutdown
+    -- began meanwhile, which merely restores pre-graceful-shutdown behavior
+    -- for bytes the client already sent.
+    mbs <- receiveNoWait sock pool
+    case mbs of
+        Just bs -> return bs
+        Nothing -> slowPath
+  where
+    slowPath = do
+        sockWait <-
 #if !WINDOWS && MIN_VERSION_network(3,2,2)
-        waitReadSocketSTM sock
+            waitReadSocketSTM sock
 #else
-        -- FIXME: 'waitReadSocketSTM' doesn't work on WINDOWS, and actually
-        -- blocks indefinitely, so we fall back to going straight to 'recv'.
-        pure (pure ())
+            -- FIXME: 'waitReadSocketSTM' doesn't work on WINDOWS, and actually
+            -- blocks indefinitely, so we fall back to going straight to 'recv'.
+            pure (pure ())
 #endif
-    isShuttingDown <- atomically $
-        -- when shutting down
-        (checkShutdown $> True)
-        <|>
-        -- else wait for socket readiness and do non-blocking read
-        (sockWait $> False)
-    if isShuttingDown then pure "" else recv
-  where
+        isShuttingDown <- atomically $
+            -- when shutting down
+            (checkShutdown $> True)
+            <|>
+            -- else wait for socket readiness and do non-blocking read
+            (sockWait $> False)
+        if isShuttingDown then pure "" else recv
     recv = receive sock pool
     checkShutdown = do
        check =<< currentShuttingDownStateSTM ss