git.iko.soy/typeable/wai
on master help sign in
Merge pull request #549 from asariley/cleanup-stream
2f8abffMichael Snoyman
wai-websockets/Network/Wai/Handler/WebSockets.hs+17−16
compare with master
unifiedsplit
@@ -7,6 +7,7 @@ module Network.Wai.Handler.WebSockets
7 , runWebSockets
7 , runWebSockets
8 ) where
8 ) where
9
9
10+import Control.Exception (bracket)
10import Data.ByteString (ByteString)
11import Data.ByteString (ByteString)
11import qualified Data.ByteString.Char8 as BC
12import qualified Data.ByteString.Char8 as BC
12import qualified Data.ByteString.Lazy as BL
13import qualified Data.ByteString.Lazy as BL
@@ -94,20 +95,20 @@ runWebSockets :: WS.ConnectionOptions
94 -> IO ByteString
95 -> IO ByteString
95 -> (ByteString -> IO ())
96 -> (ByteString -> IO ())
96 -> IO a
97 -> IO a
97−runWebSockets opts req app src sink = do
98+runWebSockets opts req app src sink = bracket mkStream WS.close (app . pc)
99+ where
100+ mkStream =
98− stream <- WS.makeStream
101+ WS.makeStream
99− (do
102+ (do
100− bs <- src
103+ bs <- src
101− return $ if BC.null bs then Nothing else Just bs)
104+ return $ if BC.null bs then Nothing else Just bs)
102− (\mbBl -> case mbBl of
105+ (\mbBl -> case mbBl of
103− Nothing -> return ()
106+ Nothing -> return ()
104− Just bl -> mapM_ sink (BL.toChunks bl))
107+ Just bl -> mapM_ sink (BL.toChunks bl))
105−
106− let pc = WS.PendingConnection
107− { WS.pendingOptions = opts
108− , WS.pendingRequest = req
109− , WS.pendingOnAccept = \_ -> return ()
110− , WS.pendingStream = stream
111− }
112
108
113− app pc
109+ pc stream = WS.PendingConnection
110+ { WS.pendingOptions = opts
111+ , WS.pendingRequest = req
112+ , WS.pendingOnAccept = \_ -> return ()
113+ , WS.pendingStream = stream
114+ }