allow reconnecting with to REPL with token

1 files changed, 43 insertions, 35 deletions

diff --git a/lambda-calcul/haskell/src/Minilang/REPL/Net.hs b/lambda-calcul/haskell/src/Minilang/REPL/Net.hs
index a067a15..af04901 100644
--- a/lambda-calcul/haskell/src/Minilang/REPL/Net.hs
+++ b/lambda-calcul/haskell/src/Minilang/REPL/Net.hs
@@ -3,17 +3,22 @@
 -- | A version of REPL that interacts through a websocket.
 module Minilang.REPL.Net where
 
-import Control.Concurrent.STM (atomically)
+import Control.Concurrent.STM (atomically, newTVar, readTVar)
 import Control.Concurrent.STM.TVar
-    ( newTVarIO, readTVarIO, writeTVar, TVar )
-import Control.Monad.Catch ( MonadThrow(..), MonadCatch(..) )
+  ( TVar,
+    newTVarIO,
+    readTVarIO,
+    writeTVar,
+  )
+import Control.Monad.Catch (MonadCatch (..), MonadThrow (..))
 import Control.Monad.State
-    ( MonadIO(..),
-      MonadTrans(..),
-      MonadState,
-      StateT(StateT),
-      evalStateT,
-      gets )
+  ( MonadIO (..),
+    MonadState,
+    MonadTrans (..),
+    StateT (StateT),
+    evalStateT,
+    gets,
+  )
 import Data.Aeson
   ( ToJSON,
     eitherDecode,
@@ -24,23 +29,29 @@ import Data.Aeson
 import Data.ByteString (ByteString)
 import qualified Data.ByteString.Lazy as LBS
 import Data.Either (fromRight)
-import Data.Map as Map ( Map )
+import Data.Map as Map (Map, insert, lookup)
 import Data.Serialize
-    ( getLazyByteString,
-      getWord64be,
-      putLazyByteString,
-      putWord64be,
-      Serialize(..) )
+  ( Serialize (..),
+    getLazyByteString,
+    getWord64be,
+    putLazyByteString,
+    putWord64be,
+  )
 import Data.Text (Text)
 import Data.Text.Encoding (decodeUtf8With)
 import Data.Text.Encoding.Error (lenientDecode)
-import Minilang.Log ( LoggerEnv(logInfo) )
+import Minilang.Log (LoggerEnv (logInfo))
 import Minilang.REPL.Pure (render)
 import Minilang.REPL.Run (runREPL)
 import Minilang.REPL.Types
-    ( initialREPL, In(In), MonadREPL(..), Out, REPLEnv )
+  ( In (In),
+    MonadREPL (..),
+    Out,
+    REPLEnv,
+    initialREPL,
+  )
 import Network.Wai (Application)
-import Network.Wai.Handler.WebSockets ( websocketsOr )
+import Network.Wai.Handler.WebSockets (websocketsOr)
 import qualified Network.WebSockets as WS
 
 -- | Configuration of the WS server
@@ -142,29 +153,26 @@ safeDecodeUtf8 = decodeUtf8With lenientDecode
 clientHandler :: LoggerEnv IO -> TVar (Map ByteString (TVar REPLEnv)) -> WS.PendingConnection -> IO ()
 clientHandler loggerEnv envs cnx = do
   let path = WS.requestPath $ WS.pendingRequest cnx
-  logInfo loggerEnv $ object ["action" .= ("StartedREPL" :: Text), "path" .= safeDecodeUtf8 path]
   conn <- WS.acceptRequest cnx
   -- storageFile <- mkStorageFile path
   -- withStorage storageFile $ \fileStorage -> do
-  replEnv <- newREPLEnv
+  replEnv <- findOrCreateREPL path
   WS.withPingThread conn 30 (pure ()) $
     evalStateT (runNet runREPL) (NetEnv conn loggerEnv replEnv)
   where
-    newREPLEnv = newTVarIO initialREPL
-
--- findOrCreateREPL :: (H.Store IO store) => store -> ByteString -> IO (TVar REPLEnv)
--- findOrCreateREPL storageFile path = do
---   (env, shouldLoad) <- atomically $ do
---     envMaps <- readTVar envs
---     (newEnv, shouldLoad) <- case Map.lookup path envMaps of
---       Nothing -> newTVar initialREPL >>= \v -> pure (v, True)
---       Just e -> pure (e, False)
---     writeTVar envs (Map.insert path newEnv envMaps)
---     pure (newEnv, shouldLoad)
---   -- FIXME there is a race condition here
---   if shouldLoad
---     then replay storageFile env
---     else pure env
+    findOrCreateREPL :: ByteString -> IO (TVar REPLEnv)
+    findOrCreateREPL path = do
+      (newEnv, reconnect) <- atomically $ do
+        envMaps <- readTVar envs
+        (newEnv, reconnect) <- case Map.lookup path envMaps of
+          Nothing -> newTVar initialREPL >>= (\e -> pure (e, True))
+          Just e -> pure (e, False)
+        writeTVar envs (Map.insert path newEnv envMaps)
+        pure (newEnv, reconnect)
+      if reconnect
+        then logInfo loggerEnv $ object ["action" .= ("ReconnectingREPL" :: Text), "path" .= safeDecodeUtf8 path]
+        else logInfo loggerEnv $ object ["action" .= ("StartingREPL" :: Text), "path" .= safeDecodeUtf8 path]
+      pure newEnv
 
 -- mkStorageFile path = do
 --   let f = takeFileName (unpack (safeDecodeUtf8 path))