{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} module Reanimate.Driver.Server ( daemon ) where import Control.Concurrent import Control.Exception (finally) import qualified Control.Exception as E import Control.Monad (forM_, forever, unless, void, when) import qualified Data.ByteString.Char8 as BS import qualified Data.Foldable as F import qualified Data.Map as Map import qualified Data.Text as T import Network.Socket (AddrInfo (..), AddrInfoFlag (..), SocketOption (..), SocketType (Stream), accept, bind, close, defaultHints, getAddrInfo, listen, socket, setCloseOnExecIfNeeded, setSocketOption, withFdSocket, withSocketsDo) import Network.Socket.ByteString (recv) import Network.WebSockets import Paths_reanimate (getDataFileName) import System.IO (hPutStrLn, stderr) import Web.Browser (openBrowser) opts :: ConnectionOptions opts = defaultConnectionOptions { connectionCompressionOptions = PermessageDeflateCompression defaultPermessageDeflate } daemon :: IO () daemon = do state <- newMVar (0, Map.empty) connsRef <- newMVar Map.empty self <- myThreadId dTid <- daemonReceive self $ \msg -> case msg of WebStatus _status -> return () WebError _err -> return () WebFrameCount count -> do void $ swapMVar state (count, Map.empty) conns <- readMVar connsRef F.forM_ conns $ \(conn) -> do sendWebMessage conn (WebFrameCount count) when (Map.null conns) openViewer WebFrame nth path -> do modifyMVar_ state $ \(count, frames) -> pure (count, Map.insert nth path frames) conns <- readMVar connsRef F.forM_ conns $ \conn -> do sendWebMessage conn (WebFrame nth path) openViewer let options = ServerOptions { serverHost = "127.0.0.1" , serverPort = 9161 , serverConnectionOptions = opts , serverRequirePong = Nothing } runServerWithOptions options (\pending -> do tid <- myThreadId conn <- acceptRequest pending modifyMVar_ connsRef $ pure . Map.insert tid conn (count, frames) <- readMVar state when (count > 0) $ do sendWebMessage conn (WebFrameCount count) forM_ (Map.toList frames) $ \(nth, path) -> sendWebMessage conn (WebFrame nth path) let loop = do -- FIXME: We don't use msg here. _msg <- receiveData conn :: IO T.Text loop cleanup = do modifyMVar_ connsRef $ pure . Map.delete tid nConns <- Map.size <$> readMVar connsRef when (nConns == 0) $ do threadDelay (second * 5) nConns' <- Map.size <$> readMVar connsRef when (nConns'==0) $ killThread self loop `finally` cleanup) `finally` (killThread dTid) second :: Int second = 10^(6::Int) daemonReceive :: ThreadId -> (WebMessage -> IO ()) -> IO ThreadId daemonReceive parent cb = withSocketsDo $ do addr <- resolve sock <- open addr forkIO $ handler sock `finally` close sock where handler sock = forever $ E.bracketOnError (accept sock) (close . fst) $ \(conn, _peer) -> do inp <- BS.unpack <$> recv conn 4096 case words inp of ["frame_count", n] -> cb $ WebFrameCount (read n) ["frame", nth, path] -> cb $ WebFrame (read nth) path ["stop"] -> killThread parent [] -> return () _ -> error $ "Bad message: " ++ inp close conn resolve = do let hints = defaultHints { addrFlags = [AI_PASSIVE] , addrSocketType = Stream } head <$> getAddrInfo (Just hints) (Just "127.0.0.1") (Just "9162") oSocket addr = socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr) open addr = E.bracketOnError (oSocket addr) close $ \sock -> do setSocketOption sock ReuseAddr 1 withFdSocket sock setCloseOnExecIfNeeded bind sock $ addrAddress addr listen sock 1024 return sock openViewer :: IO () openViewer = do url <- getDataFileName "viewer-elm/dist/index.html" bSucc <- openBrowser url unless bSucc $ hPutStrLn stderr $ "Failed to open browser. Manually visit: " ++ url ------------------------------------------------------------------------------- -- Websocket API data WebMessage = WebStatus String | WebError String | WebFrameCount Int | WebFrame Int FilePath sendWebMessage :: Connection -> WebMessage -> IO () sendWebMessage conn msg = sendTextData conn $ case msg of WebStatus txt -> T.pack "status\n" <> T.pack txt WebError txt -> T.pack "error\n" <> T.pack txt WebFrameCount n -> T.pack $ "frame_count\n" ++ show n WebFrame n path -> T.pack $ "frame\n" ++ show n ++ "\n" ++ path