Merge pull request #605 from yesodweb/602-better-async-timeoutthread

Better async exception behavior for timeouts #602

warp-tls/Network/Wai/Handler/WarpTLS.hs

@@ -345,6 +345,7 @@ httpOverTls TLSSettings{..} s bs0 params = do
       , connSendAll          = sendall
       , connSendFile         = sendfile
       , connClose            = close'
+      , connFree             = freeBuffer writeBuf
       , connRecv             = recv ref
       , connRecvBuf          = recvBuf ref
       , connWriteBuffer      = writeBuf
@@ -355,8 +356,7 @@ httpOverTls TLSSettings{..} s bs0 params = do
         sendfile fid offset len hook headers =
             readSendFile writeBuf bufferSize sendall fid offset len hook headers
 
-        close' = freeBuffer writeBuf `finally`
-                 void (tryIO $ TLS.bye ctx) `finally`
+        close' = void (tryIO $ TLS.bye ctx) `finally`
                  TLS.contextClose ctx
 
         -- TLS version of recv with a cache for leftover input data.

warp-tls/warp-tls.cabal

@@ -21,7 +21,7 @@ Library
   Build-Depends:     base                          >= 4        && < 5
                    , bytestring                    >= 0.9
                    , wai                           >= 3.2      && < 3.3
-                   , warp                          >= 3.2      && < 3.3
+                   , warp                          >= 3.2.10   && < 3.3
                    , data-default-class            >= 0.0.1
                    , tls                           >= 1.3.5
                    , cryptonite                    >= 0.12

warp/ChangeLog.md

@@ -1,3 +1,9 @@
+## 3.2.10
+
+* Add `connFree` to `Connection`
+* Close socket connections on timeout triggered
+* Timeout exceptions extend from `SomeAsyncException`
+
 ## 3.2.9
 
 * Fixing a space leak. [#586] https://github.com/yesodweb/wai/pull/586

warp/Network/Wai/Handler/Warp/HTTP2/Worker.hs

@@ -236,7 +236,8 @@ worker :: Context -> S.Settings -> Application -> Responder -> T.Manager -> IO (
 worker ctx@Context{inputQ,controlQ} set app responder tm = do
     sinfo <- newStreamInfo
     tcont <- newThreadContinue
-    E.bracket (T.registerKillThread tm) T.cancel $ go sinfo tcont
+    let timeoutAction = return () -- cannot close the shared connection
+    E.bracket (T.registerKillThread tm timeoutAction) T.cancel $ go sinfo tcont
   where
     go sinfo tcont th = do
         setThreadContinue tcont True

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

@@ -35,6 +35,7 @@ import qualified Network.Wai.Handler.Warp.FdCache as F
 import qualified Network.Wai.Handler.Warp.FileInfoCache as I
 import Network.Wai.Handler.Warp.HTTP2 (http2, isHTTP2)
 import Network.Wai.Handler.Warp.Header
+import Network.Wai.Handler.Warp.IORef
 import Network.Wai.Handler.Warp.ReadInt
 import Network.Wai.Handler.Warp.Recv
 import Network.Wai.Handler.Warp.Request
@@ -63,7 +64,8 @@ socketConnection s = do
         connSendMany = Sock.sendMany s
       , connSendAll = sendall
       , connSendFile = sendFile s writeBuf bufferSize sendall
-      , connClose = close s >> freeBuffer writeBuf
+      , connClose = close s
+      , connFree = freeBuffer writeBuf
       , connRecv = receive s bufferPool
       , connRecvBuf = receiveBuf s
       , connWriteBuffer = writeBuf
@@ -262,6 +264,10 @@ fork :: Settings
      -> InternalInfo0
      -> IO ()
 fork set mkConn addr app counter ii0 = settingsFork set $ \ unmask ->
+    -- Allocate a new IORef indicating whether the connection has been
+    -- closed, to avoid double-freeing a connection
+    withClosedRef $ \ref ->
+
     -- Run the connection maker to get a new connection, and ensure
     -- that the connection is closed. If the mkConn call throws an
     -- exception, we will leak the connection. If the mkConn call is
@@ -272,11 +278,14 @@ fork set mkConn addr app counter ii0 = settingsFork set $ \ unmask ->
     -- We grab the connection before registering timeouts since the
     -- timeouts will be useless during connection creation, due to the
     -- fact that async exceptions are still masked.
-    bracket mkConn closeConn $ \(conn, transport) ->
+    bracket mkConn (\(conn, _) -> closeConn ref conn `finally` connFree conn)
+    $ \(conn, transport) ->
 
     -- We need to register a timeout handler for this thread, and
-    -- cancel that handler as soon as we exit.
-    bracket (T.registerKillThread (timeoutManager0 ii0)) T.cancel $ \th ->
+    -- cancel that handler as soon as we exit. We additionally close
+    -- the connection immediately in case the child thread catches the
+    -- async exception or performs some long-running cleanup action.
+    bracket (T.registerKillThread (timeoutManager0 ii0) (closeConn ref conn)) T.cancel $ \th ->
 
     let ii1 = toInternalInfo1 ii0 th
         -- We now have fully registered a connection close handler
@@ -294,7 +303,11 @@ fork set mkConn addr app counter ii0 = settingsFork set $ \ unmask ->
        -- bracket with closeConn above ensures the connection is closed.
        when goingon $ serveConnection conn ii1 addr transport set app
   where
-    closeConn (conn, _transport) = connClose conn
+    withClosedRef inner = newIORef False >>= inner
+
+    closeConn ref conn = do
+        isClosed <- atomicModifyIORef' ref $ \x -> (True, x)
+        unless isClosed $ connClose conn
 
     onOpen adr    = increase counter >> settingsOnOpen  set adr
     onClose adr _ = decrease counter >> settingsOnClose set adr

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

@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 
 module Network.Wai.Handler.Warp.Timeout (
@@ -97,19 +98,24 @@ register mgr onTimeout = do
     return h
 
 -- | Registering a timeout action of killing this thread.
-registerKillThread :: Manager -> IO Handle
-registerKillThread m = do
+registerKillThread :: Manager -> TimeoutAction -> IO Handle
+registerKillThread m onTimeout = do
     -- If we hold ThreadId, the stack and data of the thread is leaked.
     -- If we hold Weak ThreadId, the stack is released. However, its
     -- data is still leaked probably because of a bug of GHC.
     -- So, let's just use ThreadId and release ThreadId by
     -- overriding the timeout action by "cancel".
     tid <- myThreadId
-    register m $ E.throwTo tid TimeoutThread
+    -- First run the timeout action in case the child thread is masked.
+    register m $ onTimeout `E.finally` E.throwTo tid TimeoutThread
 
 data TimeoutThread = TimeoutThread
     deriving Typeable
-instance E.Exception TimeoutThread
+instance E.Exception TimeoutThread where
+#if MIN_VERSION_base(4,7,0)
+    toException = E.asyncExceptionToException
+    fromException = E.asyncExceptionFromException
+#endif
 instance Show TimeoutThread where
     show TimeoutThread = "Thread killed by Warp's timeout reaper"
 

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

@@ -92,8 +92,13 @@ data Connection = Connection {
     , connSendAll     :: ByteString -> IO ()
     -- | The sending function for files in HTTP/1.1.
     , connSendFile    :: SendFile
-    -- | The connection closing function.
+    -- | The connection closing function. Warp guarantees it will only be
+    -- called once. Other functions (like 'connRecv') may be called after
+    -- 'connClose' is called.
     , connClose       :: IO ()
+    -- | Free any buffers allocated. Warp guarantees it will only be
+    -- called once, and no other functions will be called after it.
+    , connFree        :: IO ()
     -- | The connection receiving function. This returns "" for EOF.
     , connRecv        :: Recv
     -- | The connection receiving function. This tries to fill the buffer.

warp/warp.cabal

@@ -1,5 +1,5 @@
 Name:                warp
-Version:             3.2.9
+Version:             3.2.10
 Synopsis:            A fast, light-weight web server for WAI applications.
 License:             MIT
 License-file:        LICENSE