Put sinks in rooms not top-level

1 files changed, 32 insertions, 12 deletions

diff --git a/server/lib/Cafp/Main/Server.hs b/server/lib/Cafp/Main/Server.hs
index ac0c536..b965f20 100644
--- a/server/lib/Cafp/Main/Server.hs
+++ b/server/lib/Cafp/Main/Server.hs
@@ -28,17 +28,20 @@ type PlayerId = Int
 
 type Sink = B.ByteString -> IO ()
 
+data Room = Room
+    { roomSinks :: TVar (HMS.HashMap PlayerId Sink)
+    }
+
 data Server = Server
-    { serverRooms        :: TVar (HMS.HashMap RoomId ())
-    , serverSinks        :: TVar (HMS.HashMap PlayerId Sink)
+    { serverRooms        :: TVar (HMS.HashMap RoomId Room)
     , serverNextPlayerId :: TVar Int
     }
 
 newServer :: STM Server
-newServer = Server
-    <$> STM.newTVar HMS.empty
-    <*> STM.newTVar HMS.empty
-    <*> STM.newTVar 0
+newServer = Server <$> STM.newTVar HMS.empty <*> STM.newTVar 0
+
+newRoom :: STM Room
+newRoom = Room <$> STM.newTVar HMS.empty
 
 scottyApp :: IO Wai.Application
 scottyApp = Scotty.scottyApp $ do
@@ -60,20 +63,37 @@ routePendingConnection pending =
         [_, "rooms", roomId, "events"] -> Just roomId
         _                              -> Nothing
 
+newPlayerId :: Server -> STM PlayerId
+newPlayerId s = STM.stateTVar (serverNextPlayerId s) $ \x -> (x, x + 1)
+
+getOrCreateRoom :: Server -> RoomId -> STM Room
+getOrCreateRoom server roomId = do
+    rooms <- STM.readTVar $ serverRooms server
+    case HMS.lookup roomId rooms of
+        Just room -> pure room
+        Nothing   -> do
+            room <- newRoom
+            STM.writeTVar (serverRooms server) $ HMS.insert roomId room rooms
+            pure room
+
+joinPlayer :: Room -> PlayerId -> Sink -> STM ()
+joinPlayer room pid sink = STM.modifyTVar (roomSinks room) $ HMS.insert pid sink
+
+leavePlayer :: Room -> PlayerId -> STM ()
+leavePlayer room pid = STM.modifyTVar (roomSinks room) $ HMS.delete pid
+
 wsApp :: Server -> WS.ServerApp
 wsApp server pc = case routePendingConnection pc of
     Nothing -> WS.rejectRequest pc "Invalid URL"
     Just roomId -> do
-        playerId <- atomically . STM.stateTVar (serverNextPlayerId server) $
-            \x -> (x, x + 1)
+        playerId <- atomically $ newPlayerId server
+        room <- atomically $ getOrCreateRoom server roomId
         conn <- WS.acceptRequest pc
         WS.withPingThread conn 30 (pure ()) $ do
             WS.sendTextData conn $ Aeson.encode $ Welcome playerId
             bracket
-                (atomically . STM.modifyTVar (serverSinks server) $
-                    HMS.insert playerId (\bs -> WS.sendTextData conn bs))
-                (\() -> atomically . STM.modifyTVar (serverSinks server) $
-                    HMS.delete playerId)
+                (atomically $ joinPlayer room playerId (WS.sendTextData conn))
+                (\() -> atomically $ leavePlayer room playerId)
                 (\() -> loop conn)
   where
     loop conn = forever $ do