From 124c4e9ea07d203393cdc602b0a48179cf9b5280 Mon Sep 17 00:00:00 2001 From: Chris Wendt Date: Mon, 13 Mar 2023 16:07:54 -0600 Subject: [PATCH] refactor --- src/Network/WebSockets/Server.hs | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/Network/WebSockets/Server.hs b/src/Network/WebSockets/Server.hs index 4293e3d..827ed37 100644 --- a/src/Network/WebSockets/Server.hs +++ b/src/Network/WebSockets/Server.hs @@ -19,7 +19,7 @@ module Network.WebSockets.Server -------------------------------------------------------------------------------- -import Control.Concurrent (MVar, takeMVar, tryPutMVar, +import Control.Concurrent (takeMVar, tryPutMVar, newEmptyMVar) import qualified Control.Concurrent.Async as Async import Control.Exception (Exception, bracket, @@ -121,19 +121,23 @@ runServerWithOptions opts app = S.withSocketsDo $ -- whenever a pong is received. connOpts' = connOpts { connectionOnPong = do - tryPutMVar heartbeat () + _ <- tryPutMVar heartbeat () connectionOnPong connOpts } - -- Kills the thread if a pong was not received within - -- the grace period. - reaper grace appAsync = do - result <- timeout (grace * 1000000) (takeMVar heartbeat) + whileJust io = do + result <- io case result of - Nothing -> appAsync `Async.cancelWith` PongTimeout - Just _ -> reaper grace appAsync + Nothing -> return () + Just _ -> whileJust io - runApp conn connOpts' app `Async.withAsync` reaper grace + -- Runs until a pong was not received within the grace + -- period. + heart = whileJust $ timeout (grace * 1000000) (takeMVar heartbeat) + + Async.race_ + (runApp conn connOpts' app) + (heart >> throwIO PongTimeout) mainThread = do (conn, _) <- S.accept sock