Keep updated LocalState head in Server object

2 files changed, 27 insertions, 17 deletions

diff --git a/src/Erebos/Network.hs b/src/Erebos/Network.hs
index dc24e93..64c2c44 100644
--- a/src/Erebos/Network.hs
+++ b/src/Erebos/Network.hs
@@ -94,6 +94,7 @@ data Server = Server
     { serverStorage :: Storage
     , serverOptions :: ServerOptions
     , serverOrigHead :: Head LocalState
+    , serverCurrentHead :: MVar (Head LocalState)
     , serverIdentity_ :: MVar UnifiedIdentity
     , serverThreads :: MVar [ThreadId]
     , serverSocket :: MVar Socket
@@ -271,6 +272,7 @@ forkServerThread server label act = do
 startServer :: ServerOptions -> Head LocalState -> (String -> IO ()) -> [SomeService] -> IO Server
 startServer serverOptions serverOrigHead logd' serverServices = do
     let serverStorage = headStorage serverOrigHead
+    serverCurrentHead <- newMVar serverOrigHead
     serverIdentity_ <- newMVar $ headLocalIdentity serverOrigHead
     serverThreads <- newMVar []
     serverSocket <- newEmptyMVar
@@ -1077,19 +1079,22 @@ runPeerServiceOn mbservice newStreams paddr peer handler = liftIO $ do
                                     , svcPrintOp = atomically . logd
                                     , svcNewStreams = newStreams
                                     }
-                            reloadHead (serverOrigHead server) >>= \case
-                                Nothing -> atomically $ do
-                                    logd $ "current head deleted"
-                                    putTMVar (peerServiceState peer) svcs
-                                    putTMVar (serverServiceStates server) global
-                                Just h -> do
-                                    (rsp, (s', gs')) <- runServiceHandler h inp ps gs handler
-                                    moveKeys (peerStorage peer) (serverStorage server)
-                                    when (not (null rsp)) $ do
-                                        sendToPeerList peer rsp
-                                    atomically $ do
-                                        putTMVar (peerServiceState peer) $ M.insert svc (SomeServiceState proxy s') svcs
-                                        putTMVar (serverServiceStates server) $ M.insert svc (SomeServiceGlobalState proxy gs') global
+                            modifyMVar_ (serverCurrentHead server) $ \ph -> do
+                                reloadHead ph >>= \case
+                                    Nothing -> atomically $ do
+                                        logd $ "current head deleted"
+                                        putTMVar (peerServiceState peer) svcs
+                                        putTMVar (serverServiceStates server) global
+                                        return ph
+                                    Just h -> do
+                                        ( rsp, ( s', gs', h' ) ) <- runServiceHandler h inp ps gs handler
+                                        moveKeys (peerStorage peer) (serverStorage server)
+                                        when (not (null rsp)) $ do
+                                            sendToPeerList peer rsp
+                                        atomically $ do
+                                            putTMVar (peerServiceState peer) $ M.insert svc (SomeServiceState proxy s') svcs
+                                            putTMVar (serverServiceStates server) $ M.insert svc (SomeServiceGlobalState proxy gs') global
+                                        return h'
                 _ -> do
                     atomically $ logd $ "can't run service handler on peer with incomplete identity " ++ show paddr
 
diff --git a/src/Erebos/Service.hs b/src/Erebos/Service.hs
index 75551eb..dcc3df5 100644
--- a/src/Erebos/Service.hs
+++ b/src/Erebos/Service.hs
@@ -155,19 +155,24 @@ instance MonadHead LocalState (ServiceHandler s) where
         return x
     getLocalHeadCache _ = gets svcLocalCache
 
-runServiceHandler :: Service s => Head LocalState -> ServiceInput s -> ServiceState s -> ServiceGlobalState s -> ServiceHandler s () -> IO ([ServiceReply s], (ServiceState s, ServiceGlobalState s))
+runServiceHandler
+    :: Service s
+    => Head LocalState -> ServiceInput s -> ServiceState s -> ServiceGlobalState s
+    -> ServiceHandler s ()
+    -> IO ( [ ServiceReply s], ( ServiceState s, ServiceGlobalState s, Head LocalState ) )
 runServiceHandler h input svc global shandler = do
     let sstate = ServiceHandlerState { svcValue = svc, svcGlobal = global, svcLocal = headStoredObject h, svcLocalCache = headCache h }
         ServiceHandler handler = shandler
     (runExceptT $ flip runStateT sstate $ execWriterT $ flip runReaderT input $ handler) >>= \case
         Left err -> do
             svcPrintOp input $ "service failed: " ++ showErebosError err
-            return ([], (svc, global))
+            return ( [], ( svc, global, h ) )
         Right (rsp, sstate')
-            | svcLocal sstate' == svcLocal sstate -> return (rsp, (svcValue sstate', svcGlobal sstate'))
+            | svcLocal sstate' == svcLocal sstate -> return ( rsp, ( svcValue sstate', svcGlobal sstate', h ))
             | otherwise -> replaceHead h (svcLocal sstate') >>= \case
                 Left (Just h') -> runServiceHandler h' input svc global shandler
-                _              -> return (rsp, (svcValue sstate', svcGlobal sstate'))
+                Left Nothing   -> return ( rsp, ( svcValue sstate', svcGlobal sstate', h ) )
+                Right h'       -> return ( rsp, ( svcValue sstate', svcGlobal sstate', h' ) )
 
 svcGet :: ServiceHandler s (ServiceState s)
 svcGet = gets svcValue