Permalink
Browse files

big rewrite, NEXT: handle heartbeat correctly.

  • Loading branch information...
1 parent d350ebf commit 4228521d7dbafd27b732ccfde39b6c238b8ec7e6 @yihuang committed Dec 10, 2011
View
144 Apps.hs
@@ -1,89 +1,85 @@
-{-# LANGUAGE OverloadedStrings, PatternGuards #-}
+{-# LANGUAGE OverloadedStrings, TemplateHaskell #-}
module Apps where
-import Data.Char
-import Data.Monoid
-import Data.ByteString (ByteString)
-import qualified Data.ByteString.Char8 as S
-import Data.Map (Map)
+import Control.Exception (fromException)
+import Control.Monad (forever, when, forM_)
+import Control.Monad.IO.Class (liftIO)
+import Control.Applicative
+import Control.Concurrent.MVar
+
+import Data.Map (Map)
import qualified Data.Map as M
+import Data.Monoid (mappend, mconcat)
+import Data.Attoparsec
+import Data.Attoparsec.Char8 (skipSpace)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString.Char8 as S
+
+import Network.WebSockets.Lite
+
+echo :: WSLite ()
+echo = forever $ recvBS >>= send
-import Control.Exception
-import Control.Monad
-import Control.Monad.IO.Class
-import Control.Concurrent
+close' :: WSLite ()
+close' = return ()
-import Network.WebSockets
+data ChatMessage = ChatJoin ByteString
+ | ChatData ByteString
+ | ChatError ByteString
-import Network.Sockjs.Types
-import Network.Sockjs
+chatParser :: Parser ChatMessage
+chatParser = ChatJoin <$> (string "join" *> skipSpace *> takeByteString)
+ <|> ChatData <$> takeByteString
-echo :: TextProtocol p => Request -> WebSockets p ()
-echo req = do
- acceptRequest req
- _ <- startHBThread
- sendSockjs SockjsOpen
- forever $ do
- msg <- receiveSockjs
- sendSockjsData msg
+instance UpProtocol ChatMessage where
+ decode s = parseOnly chatParser s
-close :: TextProtocol p => Request -> WebSockets p ()
-close req = do
- acceptRequest req
- sendSockjs SockjsOpen
+instance DownProtocol ChatMessage where
+ encode (ChatData s) = s
+ encode (ChatError e) = "error: " `mappend` e
+ encode (ChatJoin name) = name `mappend` " joined"
-type ServerState p = Map ByteString (Sink p)
+type ChatState = MVar (Map ByteString Sink)
-clientExists :: Protocol p => ByteString -> ServerState p -> Bool
-clientExists name = maybe False (const True) . M.lookup name
+newChatState :: IO ChatState
+newChatState = newMVar M.empty
-chat :: TextProtocol p => MVar (ServerState p) -> Request -> WebSockets p ()
-chat state req = do
- acceptRequest req
- sendSockjs SockjsOpen
+chat :: ChatState -> WSLite ()
+chat clients = do
+ name <- recvJoin
sink <- getSink
- msg <- receiveSockjs
- clients <- liftIO $ readMVar state
- case msg of
- _ | not (prefix `S.isPrefixOf` msg) ->
- sendSockjsData "Wrong Annoucement!"
- | any ($ name)
- [S.null, S.any isPunctuation, S.any isSpace] ->
- sendSockjsData $
- "Name cannot " `mappend`
- "contain punctuation or whitespace, and " `mappend`
- "cannot be empty"
- | clientExists name clients ->
- sendSockjsData "User already exists"
- | otherwise -> do
- liftIO $ modifyMVar_ state $ \s -> do
- let s' = M.insert name sink s
- sendSink sink $ sockjsData $
- "Welcome! Users: " `mappend`
- S.intercalate ", " (M.keys s)
- broadcast (name `mappend` " joined") s'
- return s'
- _ <- startHBThread
- talk state name
- where
- prefix = "Hi! I am "
- name = S.drop (S.length prefix) msg
-
-broadcast :: TextProtocol p => ByteString -> ServerState p -> IO ()
-broadcast message clients =
- mapM_ (flip sendSink (sockjsData message)) $ M.elems clients
-
-talk :: TextProtocol p => MVar (ServerState p) -> ByteString -> WebSockets p ()
-talk state user = flip catchWsError catchDisconnect $ do
- msg <- receiveSockjs
- liftIO $ readMVar state >>= broadcast
- (user `mappend` ": " `mappend` msg)
- talk state user
+ exists <- liftIO $ modifyMVar clients $ \cs -> do
+ case M.lookup name cs of
+ Nothing -> return (M.insert name sink cs, False)
+ Just _ -> return (cs, True)
+ when exists $ fail' "User already exists."
+
+ flip catchError (handleDisconnect name) $ do
+ welcome name
+ broadcast $ ChatJoin name
+ forever $ do
+ msg <- recv
+ case msg of
+ ChatData s -> broadcast $ ChatData $ mconcat [name, ": ", s]
+ _ -> fail' "invalid message."
where
- catchDisconnect e = case fromException e of
- Just ConnectionClosed -> liftIO $ modifyMVar_ state $ \s -> do
- let s' = M.delete user s
- broadcast (user `mappend` " disconnected") s'
- return s'
+ fail' s = send (ChatError s) >> close
+ recvJoin = do msg <- recv
+ case msg of
+ ChatJoin name -> return name
+ _ -> fail' "invalid message."
+
+ broadcast msg = do
+ sinks <- M.elems <$> liftIO (readMVar clients)
+ forM_ sinks (flip sendSink msg)
+
+ welcome name = do
+ users <- filter (/=name) . M.keys <$> liftIO (readMVar clients)
+ send $ ChatData $ "Welcome! Users: " `mappend` S.intercalate ", " users
+
+ handleDisconnect name e = case fromException e of
+ Just ConnectionClosed -> do
+ liftIO $ modifyMVar_ clients $ return . M.delete name
+ broadcast $ ChatData $ mconcat [name, " disconnected."]
_ -> return ()
@@ -1,4 +1,5 @@
-module Network.Sockjs.Timer where
+module Control.Concurrent.Timer where
+
import Control.Concurrent
type TimerId = ThreadId
View
@@ -0,0 +1,7 @@
+module Control.Monad.Utils where
+
+ifM :: Monad m => m Bool -> m a -> m a -> m a
+ifM mb m1 m2 = do
+ b <- mb
+ if b then m1 else m2
+
View
@@ -0,0 +1,12 @@
+module Data.ByteString.Utils where
+
+import Data.ByteString (ByteString)
+import qualified Data.ByteString.Lazy as L
+import qualified Data.ByteString as S
+
+toLazy :: ByteString -> L.ByteString
+toLazy = L.fromChunks . (:[])
+
+toStrict :: L.ByteString -> ByteString
+toStrict = S.concat . L.toChunks
+
View
@@ -0,0 +1,78 @@
+module Data.Enumerator.Utils where
+
+import Control.Applicative
+import Control.Monad.IO.Class (liftIO)
+
+import Data.ByteString (ByteString)
+import qualified Data.ByteString.Lazy as L
+import Data.Int (Int64)
+import Data.Enumerator
+import qualified Data.Enumerator.List as EL
+import Blaze.ByteString.Builder (Builder)
+import qualified Blaze.ByteString.Builder as B
+
+import Control.Concurrent.STM
+import Network.WebSockets.Lite.Emulate
+
+-- | An `Enumerator' only do IO.
+ioEnum :: IO () -> Enumerator a IO b
+ioEnum io step = liftIO io >> returnI step
+
+enumSingle :: Monad m => a -> Enumerator a m b
+enumSingle = enumChunks . (:[])
+
+enumChunks :: Monad m => [a] -> Enumerator a m b
+enumChunks xs = checkContinue0 $ \_ f -> f (Chunks xs) >>== returnI
+
+-- | set a limit to the stream size, but don't break chunk.
+limit :: Monad m => Int64 -> Enumeratee ByteString ByteString m b
+limit n step | n <= 0 = return step
+limit n (Continue k) = continue loop where
+ loop (Chunks []) = continue loop
+ loop (Chunks xs) = iter where
+ len = L.length (L.fromChunks xs)
+ iter = if len <= n
+ then k (Chunks xs) >>== limit (n - len)
+ else k (Chunks xs) >>== (`yield` Chunks [])
+ loop EOF = k EOF >>== (`yield` EOF)
+limit _ step = return step
+
+-- | fetch multiple (at least one) items from TChan at a time, if TChan is empty, block on it.
+readTChan' :: TChan a -> STM [a]
+readTChan' ch = (:) <$> readTChan ch <*> readRest ch
+ where
+ readRest ch = do
+ empty <- isEmptyTChan ch
+ if empty
+ then return []
+ else (:) <$> readTChan ch <*> readRest ch
+
+-- | fetch multiple (at least one) chunks from TChan at a time, and combine them into one.
+enumStreamChanContents :: StreamChan a -> Enumerator [a] IO b
+enumStreamChanContents ch = checkContinue0 $ \loop f -> do
+ streams <- liftIO $ atomically $ readTChan' ch
+ let chunks = takeWhile isChunk streams
+ datas = concat [xs | (Chunks xs) <- chunks]
+ if null chunks
+ then f EOF >>== returnI
+ else f (Chunks [datas]) >>== loop
+ where isChunk (Chunks _) = True
+ isChunk _ = False
+
+-- | like `Enumerator.List.concatMap' , but terminate when return Nothing
+concatMapMaybe :: Monad m => (ao -> Maybe [ai])
+ -> Enumeratee ao ai m b
+concatMapMaybe f = checkDone (continue . step) where
+ step k EOF = yield (Continue k) EOF
+ step k (Chunks xs) = loop k xs
+
+ loop k [] = continue (step k)
+ loop k (x:xs) = do
+ case f x of
+ Nothing -> k EOF >>==
+ (`yield` (Chunks xs))
+ Just fx -> k (Chunks fx) >>==
+ checkDoneEx (Chunks xs) (`loop` xs)
+
+chunking :: Monad m => Enumeratee ByteString Builder m a
+chunking = EL.concatMap ((:[B.flush]) . B.fromByteString)
View
31 Main.hs
@@ -1,46 +1,39 @@
{-# LANGUAGE OverloadedStrings, TemplateHaskell #-}
-import System.IO.Unsafe (unsafePerformIO)
import System.Environment (getArgs)
-
import Data.Maybe
-import qualified Data.Map as M
-
import Control.Applicative
-import Control.Concurrent
-import Network.Wai
+import Network.Wai (Application)
import Network.Wai.Handler.Warp
-import Network.WebSockets (TextProtocol)
import Data.FileEmbed (embedDir)
import qualified Network.Wai.Handler.WebSockets as WaiWS
import qualified Network.Wai.Application.Static as Static
import Network.Wai.Application.Sockjs
-import Apps (echo, chat, close, ServerState, chat)
-
-serverState :: TextProtocol p => MVar (ServerState p)
-serverState = unsafePerformIO $ newMVar M.empty
+import Apps (echo, chat, close', newChatState, ChatState, chat)
staticApp :: Application
staticApp = Static.staticApp Static.defaultFileServerSettings
-- { Static.ssFolder = Static.embeddedLookup $ Static.toEmbedded $(embedDir "static") }
{ Static.ssFolder = Static.fileSystemLookup "static" }
-wsApps :: TextProtocol p => AppRoute p
-wsApps = [ ( ["echo"], (echo, Nothing) )
- , ( ["chat"], (chat serverState, Just ["websocket"]) )
- , ( ["close"], (close, Nothing) )
- , ( ["disabled_websocket_echo"], (echo, Just ["websocket"]) )
- ]
+mkApps :: ChatState -> WSLiteRoute
+mkApps st = [ ( ["echo"], (echo, Nothing) )
+ , ( ["chat"], (chat st, Just ["websocket"]) )
+ , ( ["close"], (close', Nothing) )
+ , ( ["disabled_websocket_echo"], (echo, Just ["websocket"]) )
+ ]
main :: IO ()
main = do
port <- read . fromMaybe "8080" . listToMaybe <$> getArgs
sockjsState <- newSockjsState
+ chatState <- newChatState
putStrLn $ "http://localhost:"++show port++"/static/client.html"
+ let apps = mkApps chatState
runSettings defaultSettings
{ settingsPort = port
- , settingsIntercept = WaiWS.intercept (websocketApp wsApps)
- } $ httpRoute [(["static"], staticApp)] (sockjsApp sockjsState wsApps)
+ , settingsIntercept = WaiWS.intercept (wsApps apps)
+ } $ waiRoute [(["static"], staticApp)] (waiApps sockjsState apps)
Oops, something went wrong.

0 comments on commit 4228521

Please sign in to comment.