feat: add global registry actor to system

src/Control/Actor/Registry.hs

@@ -0,0 +1,131 @@
+module Control.Actor.Registry
+  ( registerActor,
+    lookupActor,
+    lookupRemoteActor,
+    registryUUID,
+    registry,
+    registerDeath,
+    createRegistry,
+  )
+where
+
+import Control.Actor.Core (ActorM, ActorResult (..), Handler, call, call', cast', liftRuntime, linkActorTo, pass, spawnActorAs, state, lastMessageFrom, cast, passWith)
+import Control.Actor.Runtime (Runtime (..), RuntimeM, getActorByUUID, withRuntime)
+import Control.Actor.Types (ActorId (..), ActorRef (..), DeathMessage (..), DeathTarget (LocalTarget), RegistryMsg (..), SomeActorRef (..), SupervisorAction (Continue), actorRefId, thisNodeId, findByUUID)
+import Control.Concurrent.STM (readTVarIO)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (MonadIO (liftIO))
+import Control.Monad.Reader (ask)
+import Data.Map qualified as Map
+import Data.UUID (UUID, fromWords)
+import Unsafe.Coerce (unsafeCoerce)
+
+registryUUID :: UUID
+registryUUID = fromWords 0 0 0 1
+
+type RegistryState = Map.Map String ActorId
+
+registryHandlerFn :: Handler RegistryMsg RegistryState (Maybe ActorId)
+registryHandlerFn msg@(RMRegister name uuid) = do
+  u <- state
+  from <- lastMessageFrom
+  when (from == thisNodeId) $ do
+    rt <- liftRuntime ask
+    nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+    forM_ (Map.keys nodeTable) $ \nodeId -> do
+      cast msg (RemoteRef (ActorId nodeId registryUUID))
+  passWith (Map.insert name (ActorId from uuid) u)
+registryHandlerFn (RMLookup name) = do
+  u <- state
+  return $ ActorResult (Just (Map.lookup name u)) u Nothing
+registryHandlerFn (RMDeregister name) = do
+  u <- state
+  return $ ActorResult Nothing (Map.delete name u) Nothing
+registryHandlerFn msg@(RMDeath deadUUID) = do
+  u <- state
+  from <- lastMessageFrom
+
+  let (_, kept) = Map.partition (\(ActorId _ uuid) -> uuid == deadUUID) u
+
+  when (from == thisNodeId) $ do
+    rt <- liftRuntime ask
+    nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+    forM_ (Map.keys nodeTable) $ \nodeId -> do
+      cast msg (RemoteRef (ActorId nodeId registryUUID))
+
+  passWith kept
+
+registryDeathFn :: DeathMessage -> ActorM RegistryState (SupervisorAction RegistryState)
+registryDeathFn (DeathMessage (ActorId _ deadUUID) _) = do
+  u <- state
+  let (_, kept) = Map.partition (\(ActorId _ uuid) -> uuid == deadUUID) u
+  rt <- liftRuntime ask
+  nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+  forM_ (Map.keys nodeTable) $ \nodeId -> do
+    cast (RMDeath deadUUID) (RemoteRef (ActorId nodeId registryUUID))
+
+  return $ Continue kept
+
+registry :: RuntimeM (Maybe (ActorRef RegistryMsg (Maybe ActorId)))
+registry = do
+  unsafeCoerce $ getActorByUUID registryUUID
+
+registerDeath :: ActorId -> RuntimeM ()
+registerDeath (ActorId _ uuid) = do
+  r <- registry
+  mapM_ (cast' (RMDeath uuid)) r
+
+registerActor :: String -> ActorRef msg r -> RuntimeM ()
+registerActor name ref = do
+  rt <- ask
+  let ActorId _ uuid = actorRefId ref
+  actors <- liftIO $ readTVarIO (rtActors rt)
+  case findByUUID registryUUID actors of
+    Just (SomeActorRef regRef) -> do
+      let reg = unsafeCoerce regRef :: ActorRef RegistryMsg (Maybe ActorId)
+      cast' (RMRegister name uuid) reg
+      case reg of
+        LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) ref
+        RemoteRef _ -> return ()
+    Nothing -> liftIO $ putStrLn "registerActor: registry not found"
+
+lookupActor :: String -> RuntimeM (Maybe (ActorRef msg r))
+lookupActor name = do
+  rt <- ask
+  actors <- liftIO $ readTVarIO (rtActors rt)
+  case findByUUID registryUUID actors of
+    Nothing -> return Nothing
+    Just (SomeActorRef regRef) -> do
+      result <-
+        call'
+          (RMLookup name)
+          (unsafeCoerce regRef :: ActorRef RegistryMsg (Maybe ActorId))
+      case result of
+        Just (Just (ActorId 0 uuid)) -> do
+          actors' <- liftIO $ readTVarIO (rtActors rt)
+          return $ fmap (\(SomeActorRef r) -> unsafeCoerce r) (findByUUID uuid actors')
+        Just (Just aid) -> return $ Just (RemoteRef aid)
+        _else -> return Nothing
+
+lookupRemoteActor :: String -> RuntimeM (Maybe (ActorRef msg r))
+lookupRemoteActor name = do
+  rt <- ask
+  table <- liftIO $ readTVarIO (rtNodeTable rt)
+  go (Map.keys table)
+  where
+    go [] = return Nothing
+    go (peerNodeId : rest) = do
+      let remoteReg =
+            RemoteRef (ActorId peerNodeId registryUUID) ::
+              ActorRef RegistryMsg (Maybe ActorId)
+      result <- call' (RMLookup name) remoteReg
+      case result of
+        Just (Just (ActorId 0 uuid)) ->
+          return $ Just (RemoteRef (ActorId peerNodeId uuid))
+        Just (Just aid) ->
+          return $ Just (RemoteRef aid)
+        _else -> go rest
+
+createRegistry :: RuntimeM (ActorRef RegistryMsg (Maybe ActorId))
+createRegistry =
+  spawnActorAs registryUUID registryHandlerFn registryDeathFn Map.empty