@@ -17,12 +17,15 @@ import qualified Control.Concurrent as Conc (yield)
1717import Control.Exception as E
1818import qualified Data.ByteString as S
1919import Data.Char (chr)
20−import Data.IORef (IORef, newIORef, readIORef, writeIORef, atomicModifyIORef')
20+import Data.IORef (IORef, newIORef, readIORef, writeIORef)
2121import Data.Streaming.Network (bindPortTCP)
2222import Foreign.C.Error (Errno(..), eCONNABORTED)
2323import GHC.IO.Exception (IOException(..))
2424import qualified Network.HTTP2 as H2
2525import Network.Socket (Socket, close, accept, withSocketsDo, SockAddr(SockAddrInet, SockAddrInet6), setSocketOption, SocketOption(..))
26+#if MIN_VERSION_network(3,1,1)
27+import Network.Socket (gracefulClose)
28+#endif
2629import qualified Network.Socket.ByteString as Sock
2730import Network.Wai
2831import Network.Wai.Internal (ResponseReceived (ResponseReceived))
@@ -55,21 +58,38 @@ import Network.Socket (fdSocket)
5558#endif
5659
5760-- | Creating 'Connection' for plain HTTP based on a given socket.
58−socketConnection :: Socket -> IO Connection
59−socketConnection s = do
61+socketConnection :: Settings -> Socket -> IO Connection
62+#if MIN_VERSION_network(3,1,1)
63+socketConnection set s = do
64+#else
65+socketConnection _ s = do
66+#endif
6067 bufferPool <- newBufferPool
6168 writeBuf <- allocateBuffer bufferSize
6269 let sendall = Sock.sendAll s
70+ isH2 <- newIORef False -- HTTP/1.x
6371 return Connection {
6472 connSendMany = Sock.sendMany s
6573 , connSendAll = sendall
6674 , connSendFile = sendFile s writeBuf bufferSize sendall
75+#if MIN_VERSION_network(3,1,1)
76+ , connClose = do
77+ h2 <- readIORef isH2
78+ let tm = if h2 then settingsGracefulCloseTimeout2 set
79+ else settingsGracefulCloseTimeout1 set
80+ if tm == 0 then
81+ close s
82+ else
83+ gracefulClose s tm `E.catch` \(E.SomeException _) -> return ()
84+#else
6785 , connClose = close s
86+#endif
6887 , connFree = freeBuffer writeBuf
6988 , connRecv = receive s bufferPool
7089 , connRecvBuf = receiveBuf s
7190 , connWriteBuffer = writeBuf
7291 , connBufferSize = bufferSize
92+ , connHTTP2 = isH2
7393 }
7494
7595-- | Run an 'Application' on the given port.
@@ -131,7 +151,7 @@ runSettingsSocket set socket app = do
131151 setSocketCloseOnExec s
132152 -- NoDelay causes an error for AF_UNIX.
133153 setSocketOption s NoDelay 1 `E.catch` \(E.SomeException _) -> return ()
134− conn <- socketConnection s
154+ conn <- socketConnection set s
135155 return (conn, sa)
136156
137157 closeListenSocket = close socket
@@ -266,10 +286,7 @@ fork :: Settings
266286fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
267287 -- Call the user-supplied on exception code if any
268288 -- exceptions are thrown.
269− handle (settingsOnException set Nothing) .
270− -- Allocate a new IORef indicating whether the connection has been
271− -- closed, to avoid double-freeing a connection
272− withClosedRef $ \ref ->
289+ handle (settingsOnException set Nothing) $
273290 -- Run the connection maker to get a new connection, and ensure
274291 -- that the connection is closed. If the mkConn call throws an
275292 -- exception, we will leak the connection. If the mkConn call is
@@ -280,21 +297,13 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
280297 -- We grab the connection before registering timeouts since the
281298 -- timeouts will be useless during connection creation, due to the
282299 -- fact that async exceptions are still masked.
283− bracket mkConn (cleanUp ref) (serve unmask ref)
300+ bracket mkConn cleanUp (serve unmask)
284301 where
285− withClosedRef inner = newIORef False >>= inner
286−
287− closeConn ref conn = do
288− isClosed <- atomicModifyIORef' ref $ \x -> (True, x)
289− unless isClosed $ connClose conn
290−
291− cleanUp ref (conn, _) = closeConn ref conn `finally` connFree conn
302+ cleanUp (conn, _) = connClose conn `finally` connFree conn
292303
293304 -- We need to register a timeout handler for this thread, and
294− -- cancel that handler as soon as we exit. We additionally close
295− -- the connection immediately in case the child thread catches the
296− -- async exception or performs some long-running cleanup action.
297− serve unmask ref (conn, transport) = bracket register cancel $ \th -> do
305+ -- cancel that handler as soon as we exit.
306+ serve unmask (conn, transport) = bracket register cancel $ \th -> do
298307 -- We now have fully registered a connection close handler in
299308 -- the case of all exceptions, so it is safe to once again
300309 -- allow async exceptions.
@@ -306,8 +315,7 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
306315 -- above ensures the connection is closed.
307316 when goingon $ serveConnection conn ii th addr transport set app
308317 where
309− register = T.registerKillThread (timeoutManager ii)
310− (closeConn ref conn)
318+ register = T.registerKillThread (timeoutManager ii) (connClose conn)
311319 cancel = T.cancel
312320
313321 onOpen adr = increase counter >> settingsOnOpen set adr
@@ -334,10 +342,19 @@ serveConnection conn ii th origAddr transport settings app = do
334342 istatus <- newIORef False
335343 if settingsHTTP2Enabled settings && h2 then do
336344 rawRecvN <- makeReceiveN bs (connRecv conn) (connRecvBuf conn)
345+ -- This thread becomes the sender in http2 library.
346+ -- In the case of event source, one request comes and one
347+ -- worker gets busy. But it is likely that the receiver does
348+ -- not receive any data at all while the sender is sending
349+ -- output data from the worker. It's not good enough to tickle
350+ -- the time handler in the receiver only. So, we should tickle
351+ -- the time handler in both the receiver and the sender.
337352 let recvN = wrappedRecvN th istatus (settingsSlowlorisSize settings) rawRecvN
353+ sendBS x = connSendAll conn x >> T.tickle th
338354 -- fixme: origAddr
339355 checkTLS
340− http2 conn ii origAddr settings recvN app
356+ setConnHTTP2 conn True
357+ http2 conn transport ii origAddr settings recvN sendBS app
341358 else do
342359 src <- mkSource (wrappedRecv conn th istatus (settingsSlowlorisSize settings))
343360 writeIORef istatus True
@@ -412,8 +429,7 @@ serveConnection conn ii th origAddr transport settings app = do
412429 errorResponse e = settingsOnExceptionResponse settings e
413430
414431 http1 firstRequest addr istatus src = do
415− (req', mremainingRef, idxhdr, nextBodyFlush) <- recvRequest firstRequest settings conn ii th addr src
416− let req = req' { isSecure = isTransportSecure transport }
432+ (req, mremainingRef, idxhdr, nextBodyFlush) <- recvRequest firstRequest settings conn ii th addr src transport
417433 keepAlive <- processRequest istatus src req mremainingRef idxhdr nextBodyFlush
418434 `E.catch` \e -> do
419435 settingsOnException settings (Just req) e