feat: federated learning demo on the actor framework

- supervisor: ignore stale death messages (OneForAll restart storm),
  drop normally-exited children instead of respawning, add a
  restart-rate limit
- register actors before their thread starts so UUID dispatch and
  getSelf cannot race the spawn
- killActor now kills the thread so links see Killed
- fulfill connection promises with the pool winner on simultaneous
  connects
- clamp the PaySim test split on small datasets and split shared
  transaction types across trainers in non-IID partitioning

.gitignore

@@ -1 +1,3 @@
 dist-newstyle
+demo-out/
+fl_model_*.bin

actors.cabal

@@ -70,6 +70,20 @@ library
       Control.Actor.Core
       Control.Actor.Supervision
       Control.Actor.Network
+      Control.Actor.Registry
+      FL.Types
+      FL.CRDT
+      FL.NN
+      FL.Retry
+      FL.Data.PaySim
+      FL.Actors.Logger
+      FL.Actors.Evaluator
+      FL.Actors.Trainer
+      FL.Actors.Aggregator
+      FL.Actors.Coordinator
+      FL.Actors.PeerCoordinator
+      FL.Provider
+      FL.P2P
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -78,10 +92,34 @@ library
     -- other-extensions:
 
     -- Other library packages from which modules are imported.
-    build-depends:    base ^>=4.18.3.0, stm, uuid, bytestring ^>=0.12.0.0, containers ^>=0.8, mtl ^>=2.3.0, binary ^>=0.8.9, network ^>=3.2
+    build-depends:
+        base ^>=4.18.3.0,
+        stm,
+        uuid,
+        bytestring ^>=0.12.0.0,
+        containers ^>=0.8,
+        mtl ^>=2.3.0,
+        binary ^>=0.8.9,
+        network ^>=3.2,
+        hmatrix,
+        cassava,
+        vector,
+        random,
+        text,
+        time
 
     -- Directories containing source files.
     hs-source-dirs:   src
 
     -- Base language which the package is written in.
     default-language: GHC2021
+
+executable fl-actors
+    import:           warnings
+    main-is:          Main.hs
+    hs-source-dirs:   app
+    build-depends:
+        base ^>=4.18.3.0,
+        actors,
+        optparse-applicative
+    default-language: GHC2021

app/Main.hs

@@ -0,0 +1,150 @@
+module Main where
+
+import Control.Actor.Types (NodeAddr (..))
+import FL.P2P (runPeer)
+import FL.Provider (runCoordinator, runTrainer)
+import FL.Types
+import Options.Applicative
+
+-- ---------------------------------------------------------------------------
+-- CLI
+-- ---------------------------------------------------------------------------
+
+data Mode = Provider | P2P deriving (Show, Eq)
+
+data Role = Coordinator | Trainer | Peer deriving (Show, Eq)
+
+data CLIArgs = CLIArgs
+  { argMode         :: Mode
+  , argRole         :: Role
+  , argPort         :: Int
+  , argBasePort     :: Int
+  , argRounds       :: Int
+  , argEpochs       :: Int
+  , argLR           :: Double
+  , argBatch        :: Int
+  , argNumTrainers  :: Int
+  , argTrainerId    :: Int
+  , argCoordHost    :: String
+  , argCoordPort    :: Int
+  , argKnownPeers   :: [String]
+  , argDataPath     :: Maybe FilePath
+  , argPartition    :: PartitionMode
+  , argNorm         :: NormMode
+  }
+
+parseMode :: String -> Either String Mode
+parseMode "provider" = Right Provider
+parseMode "p2p"      = Right P2P
+parseMode s          = Left ("unknown mode: " <> s)
+
+parsePartition :: String -> Either String PartitionMode
+parsePartition "iid"    = Right IID
+parsePartition "noniid" = Right NonIID
+parsePartition s        = Left ("unknown partition: " <> s)
+
+parseNorm :: String -> Either String NormMode
+parseNorm "minmax" = Right MinMax
+parseNorm "zscore" = Right ZScore
+parseNorm s        = Left ("unknown norm: " <> s)
+
+parseRole :: String -> Either String Role
+parseRole "coordinator" = Right Coordinator
+parseRole "trainer"     = Right Trainer
+parseRole "peer"        = Right Peer
+parseRole s             = Left ("unknown role: " <> s)
+
+parsePeer :: String -> NodeAddr
+parsePeer s =
+  let (host, rest) = break (== ':') s
+  in NodeAddr host (read (drop 1 rest))
+
+argsParser :: Parser CLIArgs
+argsParser = CLIArgs
+  <$> option (eitherReader parseMode)
+        (long "mode" <> short 'm' <> metavar "provider|p2p"
+         <> help "Topology mode" <> value Provider <> showDefaultWith show)
+  <*> option (eitherReader parseRole)
+        (long "role" <> short 'r' <> metavar "coordinator|trainer|peer"
+         <> help "Node role" <> value Coordinator <> showDefaultWith show)
+  <*> option auto
+        (long "port" <> metavar "PORT"
+         <> help "This node's port" <> value 9100 <> showDefault)
+  <*> option auto
+        (long "base-port" <> metavar "PORT"
+         <> help "Base port (coordinator uses this)" <> value 9100 <> showDefault)
+  <*> option auto
+        (long "rounds" <> metavar "N"
+         <> help "Number of federated rounds" <> value 10 <> showDefault)
+  <*> option auto
+        (long "epochs" <> short 'e' <> metavar "N"
+         <> help "Local training epochs per round" <> value 5 <> showDefault)
+  <*> option auto
+        (long "lr" <> metavar "LR"
+         <> help "Learning rate" <> value 0.01 <> showDefault)
+  <*> option auto
+        (long "batch" <> metavar "N"
+         <> help "Mini-batch size" <> value 32 <> showDefault)
+  <*> option auto
+        (long "trainers" <> short 't' <> metavar "N"
+         <> help "Expected number of trainers/peers" <> value 3 <> showDefault)
+  <*> option auto
+        (long "id" <> metavar "N"
+         <> help "Trainer/peer index (0-based)" <> value 0 <> showDefault)
+  <*> strOption
+        (long "coord-host" <> metavar "HOST"
+         <> help "Coordinator hostname (trainer mode)" <> value "localhost")
+  <*> option auto
+        (long "coord-port" <> metavar "PORT"
+         <> help "Coordinator port (trainer mode)" <> value 9100 <> showDefault)
+  <*> many (strOption
+        (long "peer" <> metavar "HOST:PORT"
+         <> help "Known peer address (p2p mode, repeat for each peer)"))
+  <*> optional (strOption
+        (long "data" <> short 'd' <> metavar "PATH"
+         <> help "Path to PaySim CSV (uses synthetic data if absent)"))
+  <*> option (eitherReader parsePartition)
+        (long "partition" <> metavar "iid|noniid"
+         <> help "Data partition mode" <> value IID <> showDefaultWith show)
+  <*> option (eitherReader parseNorm)
+        (long "norm" <> metavar "minmax|zscore"
+         <> help "Feature normalization mode" <> value MinMax <> showDefaultWith show)
+
+main :: IO ()
+main = do
+  args <- execParser $
+    info (argsParser <**> helper)
+      ( fullDesc
+        <> progDesc "Federated Learning on the Haskell Actor Framework"
+        <> header "fl-actors — FedAvg fraud detection"
+      )
+
+  let cfg = FLConfig
+        { cfgRounds       = argRounds args
+        , cfgEpochs       = argEpochs args
+        , cfgLearningRate = argLR args
+        , cfgBatchSize    = argBatch args
+        , cfgNumTrainers  = argNumTrainers args
+        , cfgDataPath     = argDataPath args
+        , cfgPartition    = argPartition args
+        , cfgNorm         = argNorm args
+        , cfgBasePort     = argBasePort args
+        }
+
+  case (argMode args, argRole args) of
+    (Provider, Coordinator) ->
+      runCoordinator cfg { cfgBasePort = argPort args }
+
+    (Provider, Trainer) -> do
+      let coordAddr = NodeAddr (argCoordHost args) (fromIntegral (argCoordPort args))
+      runTrainer cfg (argTrainerId args) coordAddr
+
+    (P2P, Peer) -> do
+      let peers = map parsePeer (argKnownPeers args)
+          cfg'  = cfg { cfgBasePort = argPort args }
+      runPeer cfg' (argTrainerId args) peers
+
+    (mode, role) ->
+      putStrLn $ "Invalid combination: mode=" <> show mode <> " role=" <> show role
+        <> "\n  provider mode uses: --role=coordinator or --role=trainer"
+        <> "\n  p2p mode uses: --role=peer"

scripts/demo-p2p.sh

@@ -0,0 +1,85 @@
+#!/usr/bin/env bash
+#
+# Demo: federated learning in P2P mode (full mesh, no central coordinator).
+# Starts N peer nodes as separate processes; each trains locally, gossips
+# updates via PeerSync (CRDT-gated rounds), and aggregates independently.
+# At the end every peer must hold the SAME model — the script verifies the
+# saved model files are byte-identical.
+#
+# Usage:
+#   scripts/demo-p2p.sh [ROUNDS] [PEERS] [EPOCHS] [/path/to/paysim.csv]
+#
+# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
+# meaningless, but the gossip protocol and aggregation are fully exercised).
+set -euo pipefail
+
+ROUNDS="${1:-3}"
+PEERS="${2:-3}"
+EPOCHS="${3:-2}"
+DATA="${4:-}"
+
+BASE_PORT=9600
+OUT="demo-out/p2p"
+
+cd "$(dirname "$0")/.."
+echo "==> Building..."
+cabal build exe:fl-actors >/dev/null
+BIN="$(cabal list-bin fl-actors)"
+
+mkdir -p "$OUT"
+rm -f "$OUT"/*.log fl_model_peer-*.bin
+
+DATA_ARGS=()
+[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
+
+PIDS=()
+cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
+trap cleanup EXIT INT TERM
+
+for i in $(seq 0 $((PEERS - 1))); do
+  PORT=$((BASE_PORT + i))
+  ARGS=(--mode p2p --role peer --id "$i" --port "$PORT"
+        --rounds "$ROUNDS" --trainers "$PEERS" --epochs "$EPOCHS")
+  for j in $(seq 0 $((PEERS - 1))); do
+    [ "$j" = "$i" ] && continue
+    ARGS+=(--peer "localhost:$((BASE_PORT + j))")
+  done
+  echo "==> Starting peer-$i on port $PORT"
+  "$BIN" "${ARGS[@]}" "${DATA_ARGS[@]}" > "$OUT/peer$i.log" 2>&1 &
+  PIDS+=($!)
+done
+
+echo "==> Federation running; waiting for all $PEERS peers to finish..."
+echo "    (live log: tail -f $OUT/peer0.log)"
+FAIL=0
+for pid in "${PIDS[@]}"; do
+  wait "$pid" || FAIL=1
+done
+PIDS=()
+if [ "$FAIL" -ne 0 ]; then
+  echo "!!  A peer exited abnormally; check $OUT/ logs"
+  exit 1
+fi
+
+echo
+echo "==> Evaluation results (peer-0's view):"
+grep "Eval round=" "$OUT/peer0.log" || true
+echo
+
+MODELS=(fl_model_peer-*.bin)
+if [ "${#MODELS[@]}" -ne "$PEERS" ]; then
+  echo "!!  Expected $PEERS model files, found ${#MODELS[@]}; check $OUT/ logs"
+  exit 1
+fi
+
+FIRST="${MODELS[0]}"
+for m in "${MODELS[@]:1}"; do
+  if ! cmp -s "$FIRST" "$m"; then
+    echo "!!  Divergence: $m differs from $FIRST"
+    md5sum "${MODELS[@]}"
+    exit 1
+  fi
+done
+echo "==> Success: all $PEERS peers converged to a byte-identical model:"
+ls -la "${MODELS[@]}"
+echo "    Full logs in $OUT/"

scripts/demo-provider.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+#
+# Demo: federated learning in PROVIDER mode (star topology).
+# Starts one coordinator and N trainer nodes as separate processes, runs the
+# federation, then prints the evaluation log and the saved model file.
+#
+# Usage:
+#   scripts/demo-provider.sh [ROUNDS] [TRAINERS] [EPOCHS] [/path/to/paysim.csv]
+#
+# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
+# meaningless, but the message flow and aggregation are fully exercised).
+set -euo pipefail
+
+ROUNDS="${1:-3}"
+TRAINERS="${2:-3}"
+EPOCHS="${3:-2}"
+DATA="${4:-}"
+
+BASE_PORT=9500
+OUT="demo-out/provider"
+
+cd "$(dirname "$0")/.."
+echo "==> Building..."
+cabal build exe:fl-actors >/dev/null
+BIN="$(cabal list-bin fl-actors)"
+
+mkdir -p "$OUT"
+rm -f "$OUT"/*.log fl_model_coordinator.bin
+
+DATA_ARGS=()
+[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
+
+PIDS=()
+cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
+trap cleanup EXIT INT TERM
+
+echo "==> Starting coordinator on port $BASE_PORT (rounds=$ROUNDS, trainers=$TRAINERS, epochs=$EPOCHS)"
+"$BIN" --mode provider --role coordinator \
+  --port "$BASE_PORT" --rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
+  "${DATA_ARGS[@]}" > "$OUT/coordinator.log" 2>&1 &
+COORD=$!
+PIDS+=("$COORD")
+sleep 1
+
+for i in $(seq 0 $((TRAINERS - 1))); do
+  echo "==> Starting trainer-$i on port $((BASE_PORT + i + 1))"
+  "$BIN" --mode provider --role trainer --id "$i" \
+    --base-port "$BASE_PORT" --coord-host localhost --coord-port "$BASE_PORT" \
+    --rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
+    "${DATA_ARGS[@]}" > "$OUT/trainer$i.log" 2>&1 &
+  PIDS+=($!)
+done
+
+echo "==> Federation running; waiting for the coordinator to finish..."
+echo "    (live log: tail -f $OUT/coordinator.log)"
+if ! wait "$COORD"; then
+  echo "!!  Coordinator exited abnormally; last log lines:"
+  tail -20 "$OUT/coordinator.log"
+  exit 1
+fi
+
+echo
+echo "==> Evaluation results:"
+grep "Eval round=" "$OUT/coordinator.log" || true
+echo
+if [ -f fl_model_coordinator.bin ]; then
+  echo "==> Success: model saved to fl_model_coordinator.bin ($(stat -c %s fl_model_coordinator.bin) bytes)"
+  echo "    Full logs in $OUT/"
+else
+  echo "!!  No model file was produced; check $OUT/ logs"
+  exit 1
+fi

src/Control/Actor.hs

@@ -5,8 +5,8 @@ module Control.Actor
   , module Control.Actor.Core
   , module Control.Actor.Supervision
   , module Control.Actor.Network
-  , pass
+  -- Initialization
-  , continue
+  , initRuntime
   -- Demo
   , pingActor
   , forwardActorWithCell
@@ -27,41 +27,52 @@ import Control.Concurrent.STM
   ( TMVar, TVar
   , atomically, newEmptyTMVarIO, newTVarIO, readTMVar, readTVarIO, writeTQueue, writeTVar
   )
-import Control.Monad ((<=<))
+import Control.Monad ((<=<), void)
 import Control.Monad.Reader (MonadIO (..))
 import Unsafe.Coerce (unsafeCoerce)
+import Control.Exception (SomeException, try)
+import Control.Actor.Registry (createRegistry, registerActor, lookupRemoteActor)
 
-continue :: r -> Actor u r
+initRuntime :: NodeAddr -> (NodeAddr -> IO (Transport, NodeAddr)) -> IO Runtime
-continue x = (,) (Just x) <$> state
+initRuntime myAddr createTransport = do
+  (transport, actualAddr) <- createTransport myAddr
+  rt0 <- newRuntime actualAddr transport
+  let rt = rt0 { rtSendRemote = \addr nm -> getOrCreateConn addr >>= cast' nm }
+  withRuntime rt $ void $ spawnActor sysHandlerFn stopOnDeath ()
+  withRuntime rt $ void createRegistry
+  tListen transport $ \ch -> do
+    result <- try @SomeException $ withRuntime rt $ handleNewConn ch
+    case result of
+      Left _   -> chClose ch
+      Right () -> return ()
+  return rt
 
-pass :: Actor u r
-pass = (,) Nothing <$> state
 
 -- Demo
 
-pingActor :: String -> Actor () String
+pingActor :: Handler String () String
-pingActor msg = return (Just ("Hello, " <> msg <> "!"), ())
+pingActor msg = continue ("Hello, " <> msg <> "!")
 
-forwardActorWithCell :: TMVar (ActorRef String String) -> String -> Actor () String
+forwardActorWithCell :: TMVar (ActorRef String String) -> Handler String () String
 forwardActorWithCell cell msg = do
   pingRef <- liftIO $ atomically $ readTMVar cell
   reply   <- call msg pingRef
   case reply of
     Nothing -> liftIO $ putStrLn "forwardActorWithCell: received empty reply!"
     Just x  -> liftIO $ putStrLn ("forwardActorWithCell: received reply - " <> x)
-  return (reply, ())
+  maybeContinue reply
 
-repeatActor :: String -> TMVar (ActorRef String String) -> () -> Actor () String
+repeatActor :: String -> TMVar (ActorRef String String) -> Handler () () String
 repeatActor r cell () = do
   ref             <- liftIO $ atomically $ readTMVar cell
   (SomeActorRef self) <- getSelf
   cast r ref
   castIn 1000 () (unsafeCoerce self)
-  return (Nothing, ())
+  pass
 
 system :: IO ()
 system = do
-  rt <- initRuntime (NodeAddr "localhost" 9000)
+  rt <- initRuntime (NodeAddr "localhost" 9000) createTCPTransport
   withRuntime rt $ do
     pingCell    <- liftIO newEmptyTMVarIO
     forwardCell <- liftIO newEmptyTMVarIO
@@ -81,25 +92,25 @@ system = do
 
 -- Network demo
 
-newNodeActor :: NodeAddr -> Actor () NodeId
+newNodeActor :: Handler NodeAddr () NodeId
 newNodeActor = continue <=< liftRuntime . connect Nothing
 
 -- | Actor on node 2: echo — returns whatever it received.
-echoActor :: String -> Actor () String
+echoActor :: Handler String () String
 echoActor = continue
 
 -- | Actor on node 2: printer — side-effects only.
-printerActor :: String -> Actor () ()
+printerActor :: Handler String () ()
 printerActor msg = liftIO (putStrLn $ "  [remote/print] " <> msg) >> pass
 
 -- | Trivial actor used as a stub when we only care about death notifications.
-noopActor :: () -> Actor () ()
+noopActor :: Handler () () ()
-noopActor _ = return (Nothing, ())
+noopActor _ = pass
 
 -- | Captures the received string into a TVar (used as a node-1 receiver actor).
-recvActorFn :: TVar (Maybe String) -> String -> Actor () ()
+recvActorFn :: TVar (Maybe String) -> Handler String () ()
 recvActorFn var msg =
-  liftIO (atomically $ writeTVar var (Just msg)) >> return (Nothing, ())
+  liftIO (atomically $ writeTVar var (Just msg)) >> pass
 
 -- | End-to-end networking demo.
 --
@@ -114,8 +125,8 @@ networkDemo = do
 
   -- Port 0 lets the OS pick a free port each run, avoiding stale-socket
   -- conflicts when re-running in the same GHCi session.
-  rt1 <- initRuntime (NodeAddr "localhost" 0)
+  rt1 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
-  rt2 <- initRuntime (NodeAddr "localhost" 0)
+  rt2 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
   let addr1 = rtNodeId rt1
       addr2 = rtNodeId rt2
 
@@ -124,27 +135,27 @@ networkDemo = do
   node2 <- withRuntime rt1 $ connect Nothing addr2
   node1 <- withRuntime rt2 $ connect Nothing addr1
 
-  -- Helper: build a RemoteRef that is visible from one node for an actor
+  -- Spawn and register workers on node 2.
-  -- that was spawned on a different node. The actor's own id uses NodeId 0
-  -- (self), but from the caller's node it is addressed by the assigned id.
-  let toRemote nid ref =
-        let ActorId _ u = someActorId (SomeActorRef ref)
-        in  RemoteRef (ActorId nid u)
-
-  -- Spawn workers on node 2.
   echoRef  <- withRuntime rt2 $ spawnActor echoActor   stopOnDeath ()
   printRef <- withRuntime rt2 $ spawnActor printerActor stopOnDeath ()
+  withRuntime rt2 $ registerActor "echo"    echoRef
+  withRuntime rt2 $ registerActor "printer" printRef
 
-  let remoteEcho  = toRemote node2 echoRef  :: ActorRef String String
+  -- Small delay to let NMRegistryUpdate propagate to node 1.
-      remotePrint = toRemote node2 printRef :: ActorRef String ()
+  threadDelay 100_000
 
   -- 1. Remote cast: fire-and-forget from node 1 to node 2.
   putStrLn "\n[1] remote cast  node 1 -> node 2"
+  Just remotePrint <- withRuntime rt1 $ lookupRemoteActor "printer"
   withRuntime rt1 $ cast' "hello from node 1" remotePrint
   threadDelay 300_000
 
   -- 2. Remote call: request/reply across nodes.
   putStrLn "\n[2] remote call  node 1 <-> node 2"
+  Just remoteEcho' <- withRuntime rt1 $ lookupRemoteActor "echo"
+  let
+    remoteEcho :: ActorRef String String
+    remoteEcho = (unsafeCoerce remoteEcho')
   reply <- withRuntime rt1 $ call' "ping" remoteEcho
   putStrLn $ "    reply: " <> show reply
 
@@ -152,7 +163,9 @@ networkDemo = do
   putStrLn "\n[3] reverse cast  node 2 -> node 1"
   recvVar <- newTVarIO (Nothing :: Maybe String)
   recvRef <- withRuntime rt1 $ spawnActor (recvActorFn recvVar) stopOnDeath ()
-  let remoteRecv = toRemote node1 recvRef :: ActorRef String ()
+  withRuntime rt1 $ registerActor "receiver" recvRef
+  threadDelay 100_000
+  Just remoteRecv <- withRuntime rt2 $ lookupRemoteActor "receiver"
   withRuntime rt2 $ cast' "hello from node 2" remoteRecv
   threadDelay 300_000
   readTVarIO recvVar >>= \v -> putStrLn $ "    received: " <> show v
@@ -175,13 +188,11 @@ networkDemo = do
   let mortalId = someActorId (SomeActorRef mortalRef)
       watchId  = someActorId (SomeActorRef watchRef)
 
-  -- mortalRef's links: on death, send NMDeath to node 1's address.
+  -- mortalRef's links: on death, send NMDeath to node 1 (by NodeId).
-  withRuntime rt2 $
+  withRuntime rt2 $ linkActorTo (RemoteTarget watchId  node1) mortalRef
-    linkActorTo (RemoteTarget watchId  addr1) mortalRef
   -- watchRef's links: declare interest in mortalId so routeRemoteDeath
   -- on node 1 can find this actor when NMDeath arrives.
-  withRuntime rt1 $
+  withRuntime rt1 $ linkActorTo (RemoteTarget mortalId node2) watchRef
-    linkActorTo (RemoteTarget mortalId addr2) watchRef
 
   -- Kill the mortal actor by posting directly to its death queue.
   case mortalRef of

src/Control/Actor/Core.hs

@@ -6,9 +6,12 @@ module Control.Actor.Core
   , state
   , getSelf
   , Actor
+  , ActorResult (..)
+  , Handler
   , liftRuntime
   , notifyOfDeath
   , spawnActor
+  , spawnActorAs
   , linkActorTo
   , linkTo
   , killActor
@@ -18,11 +21,17 @@ module Control.Actor.Core
   , cast
   , call
   , castIn
+  , pass
+  , passWith
+  , continue
+  , maybeContinue
+  , become
+  , lastMessageFrom
   ) where
 
 import Control.Actor.Runtime
 import Control.Actor.Types
-import Control.Concurrent (forkIO, threadDelay)
+import Control.Concurrent (forkIO, killThread, threadDelay)
 import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
 import Control.Concurrent.STM
   ( atomically
@@ -43,7 +52,9 @@ import Control.Exception
   , throwIO
   , try
   )
-import Control.Monad (forM_, void)
+import Control.Monad (forM_, join, void)
+import Data.Maybe (fromMaybe)
+import System.Timeout (timeout)
 import Control.Monad.Reader
   ( MonadIO (..)
   , MonadReader (..)
@@ -53,6 +64,7 @@ import Control.Monad.Reader
   )
 import Data.Binary (Binary, decode, encode)
 import Data.Map qualified as Map
+import Data.UUID (UUID)
 import Data.UUID.V4 (nextRandom)
 
 newtype ActorM u r = ActorM
@@ -68,6 +80,9 @@ runActorM m rt s = runReaderT (unActorM m) (s, rt)
 state :: ActorM u u
 state = asks (asEnv . fst)
 
+lastMessageFrom :: ActorM u NodeId
+lastMessageFrom = asks (asLatestFrom . fst)
+
 getSelf :: ActorM u SomeActorRef
 getSelf = do
   (as, rt) <- ask
@@ -78,33 +93,46 @@ getSelf = do
     Just ref -> return ref
     Nothing  -> error "getSelf: actor not found in runtime"
 
-type Actor u r = ActorM u (Maybe r, u)
+data ActorResult msg u r = ActorResult
+  { actorReply  :: Maybe r
+  , actorState  :: u
+  , actorBecome :: Maybe (msg -> ActorM u (ActorResult msg u r))
+  }
+
+type Actor msg u r = ActorM u (ActorResult msg u r)
+
+type Handler msg u r = msg -> Actor msg u r
 
 liftRuntime :: RuntimeM a -> ActorM u a
 liftRuntime = ActorM . withReaderT snd . unRuntimeM
 
 notifyOfDeath :: Runtime -> DeathMessage -> DeathTarget -> IO ()
 notifyOfDeath _  dm (LocalTarget q)         = atomically $ writeTQueue q dm
-notifyOfDeath rt dm (RemoteTarget _ peerAddr) =
+notifyOfDeath rt dm (RemoteTarget _ nodeId) =
-  withRuntime rt $ rtSendRemote rt peerAddr (NMDeath (dmActorId dm) (toRemoteExitReason (dmReason dm)))
+  withRuntime rt $ do
+    maybeAddr <- lookupNode nodeId
+    case maybeAddr of
+      Nothing   -> return ()
+      Just addr -> rtSendRemote rt addr (NMDeath (dmActorId dm) (toRemoteExitReason (dmReason dm)))
 
-spawnActor ::
+spawnActorAs ::
   (Binary m, Binary r) =>
-  (m -> Actor u r) ->
+  UUID ->
+  Handler m u r ->
   (DeathMessage -> ActorM u (SupervisorAction u)) ->
   u ->
   RuntimeM (ActorRef m r)
-spawnActor actorFn deathFn initState = do
+spawnActorAs uuid actorFn deathFn initState = do
+  liftIO $ putStrLn $ "spawning actor with UUID " <> show uuid
   rt      <- ask
   mailbox <- liftIO newTQueueIO
   deathQ  <- liftIO newTQueueIO
   links   <- liftIO $ newTVarIO []
-  uuid    <- liftIO nextRandom
   let actorId    = ActorId thisNodeId uuid
-      actorState = ActorState actorId links initState
+      actorState = ActorState actorId links initState thisNodeId
       actorRef   = LocalRef mailbox deathQ actorState
 
-  let loop as = do
+  let loop fn as = do
         event <-
           atomically $
             (Left <$> readTQueue mailbox)
@@ -112,21 +140,27 @@ spawnActor actorFn deathFn initState = do
         case event of
           Left envelope ->
             case envelope of
-              Cast msg -> do
+              (from, Cast msg) -> do
-                (_, u') <- runActorM (actorFn msg) rt as
+                ActorResult _ u' mFn <- runActorM (fn msg) rt as {asLatestFrom = from}
-                loop as {asEnv = u'}
+                loop (fromMaybe fn mFn) as {asEnv = u'}
-              Call msg mv -> do
+              (from, Call msg mv) -> do
-                (reply, u') <- runActorM (actorFn msg) rt as
+                ActorResult reply u' mFn <- runActorM (fn msg) rt as {asLatestFrom = from}
                 putMVar mv reply
-                loop as {asEnv = u'}
+                loop (fromMaybe fn mFn) as {asEnv = u'}
           Right dm -> do
             action <- runActorM (deathFn dm) rt as
             case action of
               Stop       -> return ()
-              Continue u -> loop as {asEnv = u}
+              Continue u -> loop fn as {asEnv = u}
 
+  -- The thread blocks on `ready` until it is registered in rtActors:
+  -- otherwise messages dispatched by UUID in the window before the insert
+  -- would be dropped, and a death racing the insert would leave a stale
+  -- entry behind (the death path's delete running before the insert).
+  ready <- liftIO newEmptyMVar
   tid <- liftIO $ forkIO $ do
-    result <- try @SomeException (loop actorState)
+    takeMVar ready
+    result <- try @SomeException (loop actorFn actorState)
     let reason = case result of
           Right () -> Normal
           Left exc -> case fromException exc of
@@ -141,8 +175,19 @@ spawnActor actorFn deathFn initState = do
 
   liftIO $ atomically $
     modifyTVar (rtActors rt) (Map.insert actorId (tid, SomeActorRef actorRef))
+  liftIO $ putMVar ready ()
   return actorRef
 
+spawnActor ::
+  (Binary m, Binary r) =>
+  Handler m u r ->
+  (DeathMessage -> ActorM u (SupervisorAction u)) ->
+  u ->
+  RuntimeM (ActorRef m r)
+spawnActor actorFn deathFn initState = do
+  uuid <- liftIO nextRandom
+  spawnActorAs uuid actorFn deathFn initState
+
 linkActorTo :: DeathTarget -> ActorRef m r -> RuntimeM ()
 linkActorTo target (LocalRef {arState}) =
   liftIO $ atomically $ modifyTVar (asLinks arState) (target :)
@@ -153,9 +198,16 @@ linkTo target = do
   as <- asks fst
   liftIO $ atomically $ modifyTVar (asLinks as) (target :)
 
+-- | Forcibly terminate an actor's thread. Previously this only posted a
+-- 'DeathMessage' to the target's death queue, which a deathFn returning
+-- 'Continue' would silently ignore — and even on 'Stop' the links were told
+-- the exit was 'Normal'. Killing the thread makes links see 'Killed'.
 killActor :: ActorRef m r -> ActorM u ()
-killActor (LocalRef {arDeathQ, arState}) =
+killActor ref@(LocalRef {}) = do
-  liftIO $ atomically $ writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
+  rt <- asks snd
+  liftIO $ do
+    actors <- readTVarIO (rtActors rt)
+    forM_ (Map.lookup (actorRefId ref) actors) (killThread . fst)
 killActor (RemoteRef _) =
   liftIO $ putStrLn "killActor: remote kill not yet implemented"
 
@@ -164,19 +216,29 @@ stopOnDeath _ = return Stop
 
 cast' :: (Binary msg) => msg -> ActorRef msg reply -> RuntimeM ()
 cast' msg (LocalRef {arMsgQ}) =
-  liftIO $ atomically $ writeTQueue arMsgQ (Cast msg)
+  liftIO $ atomically $ writeTQueue arMsgQ (thisNodeId, Cast msg)
 cast' msg (RemoteRef (ActorId nodeId uuid)) = do
   rt <- ask
   maybeAddr <- lookupNode nodeId
   case maybeAddr of
     Nothing   -> liftIO $ putStrLn $ "cast: no node in lookup table with id " <> show nodeId
-    Just addr -> rtSendRemote rt addr (NMCast uuid (encode msg))
+    Just addr -> do
+      result <- liftIO $ try @SomeException $ withRuntime rt $ rtSendRemote rt addr (NMCast uuid (encode msg))
+      case result of
+        Left e  -> liftIO $ putStrLn $ "cast: send to " <> show addr <> " failed: " <> show e
+        Right _ -> return ()
+
+-- | How long a call waits for a reply before giving up with Nothing.
+-- Without this, a call to a dead/missing actor (or over a dead connection)
+-- would block the caller forever.
+callTimeoutMicros :: Int
+callTimeoutMicros = 10_000_000
 
 call' :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> RuntimeM (Maybe reply)
 call' msg (LocalRef {arMsgQ}) = liftIO $ do
   mv <- newEmptyMVar
-  atomically $ writeTQueue arMsgQ (Call msg mv)
+  atomically $ writeTQueue arMsgQ (thisNodeId, Call msg mv)
-  takeMVar mv
+  join <$> timeout callTimeoutMicros (takeMVar mv)
 call' msg (RemoteRef (ActorId nodeId uuid)) = do
   rt <- ask
   corrId <- liftIO $ atomically $ do
@@ -193,9 +255,9 @@ call' msg (RemoteRef (ActorId nodeId uuid)) = do
       return Nothing
     Just addr -> do
       rtSendRemote rt addr (NMCall uuid corrId (rtNodeId rt) (encode msg))
-      raw <- liftIO $ takeMVar replyVar
+      raw <- liftIO $ timeout callTimeoutMicros (takeMVar replyVar)
       liftIO $ atomically $ modifyTVar (rtPending rt) (Map.delete corrId)
-      return $ Just (decode raw)
+      return $ decode <$> raw
 
 cast :: (Binary msg) => msg -> ActorRef msg reply -> ActorM u ()
 cast = (liftRuntime .) . cast'
@@ -209,3 +271,19 @@ castIn ms msg ref = do
 
 call :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> ActorM u (Maybe reply)
 call = (liftRuntime .) . call'
+
+continue :: r -> Actor msg u r
+continue x = (\u -> ActorResult (Just x) u Nothing) <$> state
+
+maybeContinue :: Maybe r -> Actor msg u r
+maybeContinue Nothing = pass
+maybeContinue (Just r) = continue r
+
+pass :: Actor msg u r
+pass = (\u -> ActorResult Nothing u Nothing) <$> state
+
+passWith :: u -> Actor msg u r
+passWith u = return (ActorResult Nothing u Nothing)
+
+become :: Handler msg u r -> Actor msg u r
+become fn = (\u -> ActorResult Nothing u (Just fn)) <$> state

src/Control/Actor/Network.hs

@@ -10,83 +10,96 @@ module Control.Actor.Network
   , routeRemoteDeath
   , validateAndDispatch
   , sysHandlerFn
-  , initRuntime
+  , createTCPTransport
   ) where
 
 import Control.Actor.Core
-  ( Actor, ActorM, cast', liftRuntime, linkActorTo, spawnActor, state, stopOnDeath )
+  ( ActorM, ActorResult (..), Handler, cast', liftRuntime, linkActorTo
-import Control.Actor.Runtime (Runtime (..), RuntimeM, newRuntime, withRuntime)
+  , pass, spawnActor, state, stopOnDeath )
-import Control.Actor.Supervision (ChildSpec (..), RestartStrategy (..), childWithRef, supervise')
+import Control.Actor.Registry (deregisterNode)
+import Control.Actor.Runtime (Runtime (..), RuntimeM, addrToNodeId, withRuntime)
 import Control.Actor.Transport (ConnHandle (..), Transport (..), createTCPTransport)
 import Control.Actor.Types
-import Control.Concurrent (ThreadId, forkIO)
+import Control.Concurrent (forkIO)
 import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
 import Control.Concurrent.STM
   ( atomically
   , modifyTVar
   , newEmptyTMVar
-  , newEmptyTMVarIO
-  , putTMVar
   , readTMVar
   , readTVar
   , readTVarIO
   , tryPutTMVar
-  , tryTakeTMVar
   , writeTQueue
   , writeTVar
   )
-import Control.Exception (SomeException, try)
+import Control.Exception (SomeException, throwIO, toException, try)
-import Control.Monad (forM_, forever, unless, void)
+import Control.Monad (forM_, forever, unless, void, when)
-import Control.Monad.Reader (MonadIO (..), MonadReader (..))
+import Control.Monad.Reader (MonadIO (..), MonadReader (..), asks)
 import Data.Binary (decode, decodeOrFail, encode)
 import Data.ByteString.Lazy (ByteString)
 import Data.Map qualified as Map
-import Data.UUID (UUID)
+import Data.Maybe (fromMaybe)
+import System.Timeout (timeout)
 
 -- Connection actors
 
-connActorFn :: NetworkMessage -> Actor ConnHandle ()
+connActorFn :: Handler NetworkMessage ConnHandle ()
 connActorFn nm = do
   ch <- state
   liftIO $ chSend ch (encode nm)
-  return (Nothing, ch)
+  return $ ActorResult Nothing ch Nothing
 
 connDeathFn :: NodeAddr -> DeathMessage -> ActorM ConnHandle (SupervisorAction ConnHandle)
 connDeathFn peer _ = do
   ch  <- state
+  aid <- asks (asId . fst)
   rt  <- liftRuntime ask
-  liftIO $ do
+  -- Only clear the pool entry if it is ours: with simultaneous connects two
-    atomically $ modifyTVar (rtConnections rt) (Map.delete peer)
+  -- conn actors can exist for the same peer, and the loser must not evict
-    chClose ch
+  -- the active one.
+  wasActive <- liftIO $ atomically $ do
+    conns <- readTVar (rtConnections rt)
+    case Map.lookup peer conns of
+      Just ref | actorRefId ref == aid -> do
+        modifyTVar (rtConnections rt) (Map.delete peer)
+        return True
+      _else -> return False
+  liftIO $ chClose ch
+  when wasActive $ do
+    mNodeId <- liftRuntime $ addrToNodeId peer
+    forM_ mNodeId $ \nodeId -> liftRuntime $ deregisterNode nodeId
   return Stop
 
-routerActorFn :: NodeAddr -> ByteString -> Actor () ()
+routerActorFn :: NodeAddr -> Handler ByteString () ()
-routerActorFn senderAddr raw = do
+routerActorFn senderAddr raw =
-  let nm = decode raw :: NetworkMessage
+  case decodeOrFail raw of
+    Left _ -> do
+      liftIO $ putStrLn $ "router: dropping undecodable frame from " <> show senderAddr
+      pass
+    Right (_, _, nm) -> do
       valid <- liftRuntime $ validateAndDispatch senderAddr nm
-  unless valid $ liftIO $ putStrLn "router: dropping invalid message"
+      unless valid $ liftIO $
-  return (Nothing, ())
+        putStrLn $ "router: dropping invalid message from " <> show senderAddr
+      pass
 
 -- Connection supervision tree
 
 spawnConnTree :: NodeAddr -> ConnHandle -> RuntimeM (ActorRef NetworkMessage ())
 spawnConnTree peer ch = do
   rt        <- ask
-  routerCell <- liftIO newEmptyTMVarIO
+  routerRef <- spawnActor (routerActorFn peer) stopOnDeath ()
-  connCell   <- liftIO newEmptyTMVarIO
+  connRef   <- spawnActor connActorFn (connDeathFn peer) ch
-  supervise' OneForAll
+  -- When conn dies (peer disconnected), notify router so it shuts down too
-    [ childWithRef (routerActorFn peer) stopOnDeath () routerCell
+  case routerRef of
-    , ChildSpec
+    LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) connRef
-        { csRun = \target -> do
+    RemoteRef _         -> return ()
-            ref <- spawnActor connActorFn (connDeathFn peer) ch
+  -- And vice versa: if the router dies, close the connection down cleanly
-            linkActorTo target ref
+  case connRef of
-            return ref
+    LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) routerRef
-        , csOnSpawn = \ref -> case ref of
+    RemoteRef _         -> return ()
-            LocalRef { arDeathQ, arState } -> do
+  liftIO $ case connRef of
-              atomically $ do
+    LocalRef {arDeathQ, arState} ->
-                void $ tryTakeTMVar connCell
-                putTMVar connCell ref
-              routerRef <- atomically $ readTMVar routerCell
       void $ forkIO $ do
         result <- try @SomeException $ forever $
           chRecv ch >>= \raw -> withRuntime rt $ cast' raw routerRef
@@ -95,9 +108,7 @@ spawnConnTree peer ch = do
             writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
           Right () -> return ()
     RemoteRef _ -> return ()
-        }
+  return connRef
-    ]
-  liftIO $ atomically $ readTMVar connCell
 
 -- Incoming connection handler
 
@@ -109,13 +120,26 @@ handleNewConn ch = do
     NMHandshake peerAddr -> do
       ref <- spawnConnTree peerAddr ch
       liftIO $ atomically $ do
+        conns <- readTVar (rtConnections rt)
+        -- On a simultaneous connect two sockets (and two conn trees) exist
+        -- for the same peer. Only the pool entry is deduplicated; the extra
+        -- tree must stay alive because it is the read path for whatever the
+        -- peer sends on its socket. It cleans itself up when that socket
+        -- closes. Local waiters get the pool winner, never the extra ref.
+        let winner = Map.findWithDefault ref peerAddr conns
+        unless (Map.member peerAddr conns) $
           modifyTVar (rtConnections rt) (Map.insert peerAddr ref)
+        table <- readTVar (rtNodeTable rt)
+        unless (any (== peerAddr) (Map.elems table)) $ do
+          nid <- readTVar (rtNextNodeId rt)
+          writeTVar (rtNextNodeId rt) (nid + 1)
+          modifyTVar (rtNodeTable rt) (Map.insert nid peerAddr)
         promises <- readTVar (rtConnPromises rt)
         case Map.lookup peerAddr promises of
           Nothing -> return ()
           Just p  -> do
             modifyTVar (rtConnPromises rt) (Map.delete peerAddr)
-            void $ tryPutTMVar p ref
+            void $ tryPutTMVar p (Right winner)
     _else -> liftIO $ chClose ch
 
 -- Connection pool
@@ -137,9 +161,20 @@ getOrCreateConn peer = do
             return (Right (Right p))
   case action of
     Left ref              -> return ref
-    Right (Left promise)  -> liftIO $ atomically $ readTMVar promise
+    Right (Left promise)  -> do
+      -- Someone else is connecting; wait for their outcome so a failed
+      -- connect propagates instead of blocking us forever.
+      outcome <- liftIO $ atomically $ readTMVar promise
+      either (liftIO . throwIO) return outcome
     Right (Right promise) -> do
-      ch <- liftIO $ tConnect (rtTransport rt) peer
+      result <- liftIO $ try @SomeException $ tConnect (rtTransport rt) peer
+      case result of
+        Left e -> do
+          liftIO $ atomically $ do
+            modifyTVar (rtConnPromises rt) (Map.delete peer)
+            void $ tryPutTMVar promise (Left e)
+          liftIO $ throwIO e
+        Right ch -> do
           liftIO $ chSend ch (encode (NMHandshake (rtNodeId rt)))
           ref <- spawnConnTree peer ch
           liftIO $ atomically $ do
@@ -147,29 +182,21 @@ getOrCreateConn peer = do
             unless (Map.member peer conns) $
               modifyTVar (rtConnections rt) (Map.insert peer ref)
             modifyTVar (rtConnPromises rt) (Map.delete peer)
-        void $ tryPutTMVar promise ref
+            void $ tryPutTMVar promise (Right ref)
           return ref
 
 -- Message dispatch
 
-findByUUID :: UUID -> Map.Map ActorId (ThreadId, SomeActorRef) -> Maybe SomeActorRef
-findByUUID uuid actors =
-  case Map.toList (Map.filterWithKey (\(ActorId _ u) _ -> u == uuid) actors) of
-    []            -> Nothing
-    (_, (_, r)):_ -> Just r
-
 routeRemoteDeath :: NodeAddr -> ActorId -> RemoteExitReason -> RuntimeM ()
 routeRemoteDeath senderAddr (ActorId _ uuid) reason = do
   rt           <- ask
+  mSenderNid   <- addrToNodeId senderAddr
   liftIO $ do
-    localNodeId <- atomically $ do
+    let deadId     = ActorId (fromMaybe 0 mSenderNid) uuid
-      table <- readTVar (rtNodeTable rt)
-      return $ Map.foldrWithKey (\k v acc -> if v == senderAddr then Just k else acc) Nothing table
-    let deadId     = ActorId (case localNodeId of { Just n -> n; Nothing -> 0 }) uuid
         exitReason = case reason of
           RNormal      -> Normal
           RKilled      -> Killed
-          RException s -> Exception (error s)
+          RException s -> Exception (toException (userError s))
         dm = DeathMessage deadId exitReason
     actors <- readTVarIO (rtActors rt)
     forM_ (Map.elems actors) $ \(_, SomeActorRef ref) ->
@@ -177,8 +204,8 @@ routeRemoteDeath senderAddr (ActorId _ uuid) reason = do
         LocalRef {arDeathQ, arState} -> do
           links <- readTVarIO (asLinks arState)
           forM_ links $ \case
-            RemoteTarget (ActorId _ uid) peerAddr
+            RemoteTarget (ActorId _ uid) nodeId
-              | uid == uuid && peerAddr == senderAddr ->
+              | uid == uuid && Just nodeId == mSenderNid ->
                   atomically $ writeTQueue arDeathQ dm
             _else -> return ()
         RemoteRef _ -> return ()
@@ -198,22 +225,25 @@ validateAndDispatch senderAddr nm = case nm of
 
   NMCast uuid payload -> do
     rt <- ask
+    rid <- remoteNodeId
     liftIO $ do
       actors <- readTVarIO (rtActors rt)
       case findByUUID uuid actors of
-        Nothing ->
+        Nothing -> do
+          putStrLn $ "didn't find actor: " <> show uuid
           return False
         Just (SomeActorRef (LocalRef {arMsgQ})) ->
           case decodeOrFail payload of
             Left _            -> return False
             Right (_, _, msg) -> do
-              atomically $ writeTQueue arMsgQ (Cast msg)
+              atomically $ writeTQueue arMsgQ (rid, Cast msg)
               return True
         Just (SomeActorRef (RemoteRef _)) ->
           return False
 
   NMCall uuid corrId returnAddr payload -> do
     rt <- ask
+    rid <- remoteNodeId
     liftIO $ do
       actors <- readTVarIO (rtActors rt)
       case findByUUID uuid actors of
@@ -224,12 +254,16 @@ validateAndDispatch senderAddr nm = case nm of
             Left _            -> return False
             Right (_, _, msg) -> do
               mv <- newEmptyMVar
-              atomically $ writeTQueue arMsgQ (Call msg mv)
+              atomically $ writeTQueue arMsgQ (rid, Call msg mv)
               void $ forkIO $ do
-                reply <- takeMVar mv
+                -- Bounded wait: if the target actor dies before replying,
+                -- this thread must not linger forever.
+                reply <- timeout 30_000_000 (takeMVar mv)
                 case reply of
                   Nothing        -> return ()
-                  Just rv -> withRuntime rt $ do
+                  Just Nothing   -> return ()
+                  Just (Just rv) -> void $ try @SomeException $
+                    withRuntime rt $ do
                       connRef <- getOrCreateConn returnAddr
                       cast' (NMReply corrId (encode rv)) connRef
               return True
@@ -240,10 +274,13 @@ validateAndDispatch senderAddr nm = case nm of
     routeRemoteDeath senderAddr deadId reason
     return True
 
+  where
+    remoteNodeId = fromMaybe thisNodeId <$> addrToNodeId senderAddr
+
 -- System event handler
 
-sysHandlerFn :: NetworkMessage -> Actor () ()
+sysHandlerFn :: Handler NetworkMessage () ()
-sysHandlerFn _ = return (Nothing, ())
+sysHandlerFn _ = return $ ActorResult Nothing () Nothing
 
 -- Node connection
 
@@ -255,6 +292,9 @@ connect suggestedId peer = do
   rt     <- ask
   nodeId <- liftIO $ atomically $ do
     table <- readTVar (rtNodeTable rt)
+    case Map.foldrWithKey (\k v acc -> if v == peer then Just k else acc) Nothing table of
+      Just existing -> return existing
+      Nothing -> do
         nid <- readTVar (rtNextNodeId rt)
         let assigned = case suggestedId of
               Just n | n /= 0 && not (Map.member n table) -> n
@@ -264,18 +304,3 @@ connect suggestedId peer = do
         return assigned
   void $ getOrCreateConn peer
   return nodeId
-
--- Runtime initialization
-
-initRuntime :: NodeAddr -> IO Runtime
-initRuntime myAddr = do
-  (transport, actualAddr) <- createTCPTransport myAddr
-  rt0 <- newRuntime actualAddr transport
-  let rt = rt0 { rtSendRemote = \addr nm -> getOrCreateConn addr >>= cast' nm }
-  withRuntime rt $ void $ spawnActor sysHandlerFn stopOnDeath ()
-  tListen transport $ \ch -> do
-    result <- try @SomeException $ withRuntime rt $ handleNewConn ch
-    case result of
-      Left _   -> chClose ch
-      Right () -> return ()
-  return rt

src/Control/Actor/Registry.hs

@@ -0,0 +1,141 @@
+module Control.Actor.Registry
+  ( registerActor,
+    lookupActor,
+    lookupRemoteActor,
+    registryUUID,
+    registry,
+    registerDeath,
+    deregisterNode,
+    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), NodeId, 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 (RMDeregisterNode nodeId) = do
+  u <- state
+  passWith (Map.filter (\(ActorId nid _) -> nid /= nodeId) u)
+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
+  mRef <- getActorByUUID registryUUID
+  return $ case mRef of
+    Nothing                    -> Nothing
+    Just (SomeActorRef regRef) -> Just (unsafeCoerce regRef)
+
+registerDeath :: ActorId -> RuntimeM ()
+registerDeath (ActorId _ uuid) = do
+  r <- registry
+  mapM_ (cast' (RMDeath uuid)) r
+
+deregisterNode :: NodeId -> RuntimeM ()
+deregisterNode nodeId = do
+  r <- registry
+  mapM_ (cast' (RMDeregisterNode nodeId)) 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))
+        _else -> go rest
+
+createRegistry :: RuntimeM (ActorRef RegistryMsg (Maybe ActorId))
+createRegistry =
+  spawnActorAs registryUUID registryHandlerFn registryDeathFn Map.empty

src/Control/Actor/Runtime.hs

@@ -6,18 +6,24 @@ module Control.Actor.Runtime
   , newRuntime
   , withRuntime
   , lookupNode
+  , addrToNodeId
+  , getActorRef
+  , getActorByUUID
   ) where
 
 import Control.Actor.Transport (Transport)
 import Control.Actor.Types
 import Control.Concurrent (ThreadId)
+import Control.Exception (SomeException)
 import Control.Concurrent.MVar (MVar)
 import Control.Concurrent.STM
-  ( TMVar, TVar, atomically, newTVarIO, readTVar )
+  ( TMVar, TVar, atomically, newTVarIO, readTVar, readTVarIO )
 import Control.Monad.Reader
-  ( MonadIO (..), MonadReader (..), ReaderT (..), runReaderT )
+  ( MonadIO (..), MonadReader (..), ReaderT (..), runReaderT, asks )
 import Data.ByteString.Lazy (ByteString)
+import Data.Foldable (foldl')
 import Data.Map qualified as Map
+import Data.UUID (UUID)
 
 data Runtime = Runtime
   { rtNodeId       :: NodeAddr
@@ -28,7 +34,7 @@ data Runtime = Runtime
   , rtNodeTable    :: TVar (Map.Map NodeId NodeAddr)
   , rtTransport    :: Transport
   , rtConnections  :: TVar (Map.Map NodeAddr (ActorRef NetworkMessage ()))
-  , rtConnPromises :: TVar (Map.Map NodeAddr (TMVar (ActorRef NetworkMessage ())))
+  , rtConnPromises :: TVar (Map.Map NodeAddr (TMVar (Either SomeException (ActorRef NetworkMessage ()))))
   , rtSendRemote   :: NodeAddr -> NetworkMessage -> RuntimeM ()
   }
 
@@ -68,3 +74,30 @@ lookupNode nodeId = do
   liftIO $ atomically $ do
     table <- readTVar (rtNodeTable rt)
     return $ Map.lookup nodeId table
+
+addrToNodeId :: NodeAddr -> RuntimeM (Maybe NodeId)
+addrToNodeId addr = do
+  rt <- ask
+  liftIO $ atomically $ do
+    table <- readTVar (rtNodeTable rt)
+    return $ Map.foldrWithKey (\k v acc -> if v == addr then Just k else acc) Nothing table
+
+getActorRef :: ActorId -> RuntimeM (Maybe SomeActorRef)
+getActorRef aid = do
+  actVar <- asks rtActors
+  actors <- liftIO $ readTVarIO actVar
+  return $ snd <$> Map.lookup aid actors
+
+firstByKey' :: Ord k => (k -> Bool) -> Map.Map k v -> Maybe v
+firstByKey' p = snd . foldl' step (False, Nothing) . Map.toAscList
+  where
+    step (True, mv) _ = (True, mv)
+    step (False, _) (k,v)
+      | p k       = (True, Just v)
+      | otherwise = (False, Nothing)
+
+getActorByUUID :: UUID -> RuntimeM (Maybe SomeActorRef)
+getActorByUUID uuid = do
+  actVar <- asks rtActors
+  actors <- liftIO $ readTVarIO actVar
+  return $ snd <$> firstByKey' (\(ActorId _ x) -> x == uuid) actors

src/Control/Actor/Supervision.hs

@@ -16,17 +16,15 @@ module Control.Actor.Supervision
   , killSlot
   ) where
 
-import Control.Actor.Core (Actor, ActorM, liftRuntime, linkActorTo, spawnActor)
+import Control.Actor.Core (ActorM, Handler, liftRuntime, linkActorTo, spawnActor)
 import Control.Actor.Runtime (Runtime (..), RuntimeM, withRuntime)
 import Control.Actor.Types
 import Data.Binary (Binary)
 import Control.Concurrent (ThreadId, forkIO, killThread)
 import Control.Concurrent.STM
   ( TMVar
-  , TQueue
   , TVar
   , atomically
-  , flushTQueue
   , modifyTVar
   , newTQueueIO
   , newTVarIO
@@ -36,18 +34,21 @@ import Control.Concurrent.STM
   , tryTakeTMVar
   , writeTVar
   )
-import Control.Monad (forever, void)
+import Control.Monad (void)
+import GHC.Clock (getMonotonicTime)
 import Control.Monad.Reader (MonadIO (..), MonadReader (..))
 import Data.Map qualified as Map
 
 data ChildSpec = forall m r. (Binary m, Binary r) => ChildSpec
   { csRun     :: DeathTarget -> RuntimeM (ActorRef m r)
-  , csOnSpawn :: ActorRef m r -> IO ()
+  -- | Runs after every (re)spawn — including supervisor restarts — so it can
+  -- refresh shared ref cells and re-register names.
+  , csOnSpawn :: ActorRef m r -> RuntimeM ()
   }
 
 child ::
   (Binary m, Binary r) =>
-  (m -> Actor u r) ->
+  Handler m u r ->
   (DeathMessage -> ActorM u (SupervisorAction u)) ->
   u ->
   ChildSpec
@@ -62,7 +63,7 @@ child msgFn deathFn initState =
 
 childWithRef ::
   (Binary m, Binary r) =>
-  (m -> Actor u r) ->
+  Handler m u r ->
   (DeathMessage -> ActorM u (SupervisorAction u)) ->
   u ->
   TMVar (ActorRef m r) ->
@@ -73,7 +74,7 @@ childWithRef msgFn deathFn initState cell =
         ref <- spawnActor msgFn deathFn initState
         linkActorTo target ref
         return ref
-    , csOnSpawn = \ref -> atomically $ do
+    , csOnSpawn = \ref -> liftIO $ atomically $ do
         void $ tryTakeTMVar cell
         putTMVar cell ref
     }
@@ -91,7 +92,7 @@ spawnSlot :: DeathTarget -> ChildSpec -> RuntimeM ChildSlot
 spawnSlot target spec@ChildSpec {csRun, csOnSpawn} = do
   rt  <- ask
   ref <- csRun target
-  _   <- liftIO $ csOnSpawn ref
+  csOnSpawn ref
   let aid = someActorId (SomeActorRef ref)
   tid <- liftIO $ do
     actors <- readTVarIO (rtActors rt)
@@ -103,6 +104,15 @@ spawnSlot target spec@ChildSpec {csRun, csOnSpawn} = do
 killSlot :: ChildSlot -> IO ()
 killSlot ChildSlot {slotTid} = killThread slotTid
 
+-- | Max child restarts within 'restartWindow' seconds before the supervisor
+-- gives up and stops all children — otherwise a child that crashes right
+-- after spawning would respawn in a hot loop forever.
+restartLimit :: Int
+restartLimit = 5
+
+restartWindow :: Double
+restartWindow = 10
+
 supervise' :: RestartStrategy -> [ChildSpec] -> RuntimeM ()
 supervise' strategy specs = do
   rt        <- ask
@@ -110,13 +120,35 @@ supervise' strategy specs = do
   let target = LocalTarget supDeathQ
   slots    <- mapM (spawnSlot target) specs
   slotsVar <- liftIO $ newTVarIO slots
-  liftIO $ void $ forkIO $ forever $ do
+  let loop restarts = do
-    DeathMessage deadId _ <- atomically $ readTQueue supDeathQ
+        DeathMessage deadId reason <- atomically $ readTQueue supDeathQ
         slots' <- readTVarIO slotsVar
+        if all ((/= deadId) . slotId) slots'
+          -- Stale death: killSlot is asynchronous, so children replaced by a
+          -- OneForAll/RestForOne restart report their deaths after the fact.
+          -- Acting on one would kill the fresh children and cascade forever.
+          then loop restarts
+          else case reason of
+            -- A normal exit is completion, not a crash: drop the slot.
+            Normal -> do
+              atomically $ modifyTVar slotsVar (filter ((/= deadId) . slotId))
+              loop restarts
+            _crash -> do
+              now <- getMonotonicTime
+              let recent = filter (> now - restartWindow) restarts
+              if length recent >= restartLimit
+                then do
+                  putStrLn $ "supervisor: " <> show restartLimit
+                    <> " restarts within " <> show restartWindow
+                    <> "s, giving up"
+                  mapM_ killSlot slots'
+                else do
                   withRuntime rt $ case strategy of
                     OneForOne  -> doOneForOne target slotsVar slots' deadId
-      OneForAll  -> doOneForAll target slotsVar supDeathQ slots'
+                    OneForAll  -> doOneForAll target slotsVar slots'
-      RestForOne -> doRestForOne target slotsVar supDeathQ slots' deadId
+                    RestForOne -> doRestForOne target slotsVar slots' deadId
+                  loop (now : recent)
+  liftIO $ void $ forkIO $ loop []
 
 supervise :: RestartStrategy -> [ChildSpec] -> ActorM u ()
 supervise = (liftRuntime .) . supervise'
@@ -135,30 +167,24 @@ doOneForOne target slotsVar slots deadId =
 doOneForAll ::
   DeathTarget ->
   TVar [ChildSlot] ->
-  TQueue DeathMessage ->
   [ChildSlot] ->
   RuntimeM ()
-doOneForAll target slotsVar supDeathQ slots = do
+doOneForAll target slotsVar slots = do
-  liftIO $ do
+  liftIO $ mapM_ killSlot slots
-    mapM_ killSlot slots
-    atomically $ void $ flushTQueue supDeathQ
   newSlots <- mapM (spawnSlot target . slotSpec) slots
   liftIO $ atomically $ writeTVar slotsVar newSlots
 
 doRestForOne ::
   DeathTarget ->
   TVar [ChildSlot] ->
-  TQueue DeathMessage ->
   [ChildSlot] ->
   ActorId ->
   RuntimeM ()
-doRestForOne target slotsVar supDeathQ slots deadId = do
+doRestForOne target slotsVar slots deadId = do
   let (before, fromDead) = break (\s -> slotId s == deadId) slots
   case fromDead of
     []        -> return ()
     _nonempty -> do
-      liftIO $ do
+      liftIO $ mapM_ killSlot (drop 1 fromDead)
-        mapM_ killSlot (drop 1 fromDead)
-        atomically $ void $ flushTQueue supDeathQ
       newSlots <- mapM (spawnSlot target . slotSpec) fromDead
       liftIO $ atomically $ writeTVar slotsVar (before ++ newSlots)

src/Control/Actor/Types.hs

@@ -16,9 +16,12 @@ module Control.Actor.Types
   , ActorRef (..)
   , SomeActorRef (..)
   , someActorId
+  , actorRefId
   , SupervisorAction (..)
   , CorrelationId
   , Envelope (..)
+  , RegistryMsg (..)
+  , findByUUID
   ) where
 
 import Control.Concurrent.MVar (MVar)
@@ -28,6 +31,8 @@ import Data.Binary (Binary)
 import Data.ByteString.Lazy (ByteString)
 import Data.UUID (UUID)
 import GHC.Generics (Generic)
+import Data.Map qualified as Map
+import Control.Concurrent (ThreadId)
 
 data NodeAddr = NodeAddr
   { nodeHost :: String
@@ -59,7 +64,7 @@ data DeathMessage = DeathMessage
 
 data DeathTarget
   = LocalTarget (TQueue DeathMessage)
-  | RemoteTarget ActorId NodeAddr
+  | RemoteTarget ActorId NodeId
 
 data RemoteExitReason
   = RNormal
@@ -80,7 +85,7 @@ data NetworkMessage
   | NMCall UUID CorrelationId NodeAddr ByteString
   | NMReply CorrelationId ByteString
   | NMDeath ActorId RemoteExitReason
-  deriving (Generic)
+  deriving (Generic, Show)
 
 instance Binary NetworkMessage
 
@@ -88,6 +93,7 @@ data ActorState u = ActorState
   { asId    :: ActorId
   , asLinks :: TVar [DeathTarget]
   , asEnv   :: u
+  , asLatestFrom :: NodeId
   }
 
 type CorrelationId = Integer
@@ -98,7 +104,7 @@ data Envelope msg reply
 
 data ActorRef msg reply
   = forall u. LocalRef
-      { arMsgQ   :: TQueue (Envelope msg reply)
+      { arMsgQ   :: TQueue (NodeId, Envelope msg reply)
       , arDeathQ :: TQueue DeathMessage
       , arState  :: ActorState u
       }
@@ -110,6 +116,26 @@ someActorId :: SomeActorRef -> ActorId
 someActorId (SomeActorRef (LocalRef {arState})) = asId arState
 someActorId (SomeActorRef (RemoteRef aid))      = aid
 
+actorRefId :: ActorRef msg r -> ActorId
+actorRefId (LocalRef {arState}) = asId arState
+actorRefId (RemoteRef aid)      = aid
+
 data SupervisorAction u
   = Stop
   | Continue u
+
+data RegistryMsg
+  = RMRegister        String UUID
+  | RMLookup          String
+  | RMDeregister      String
+  | RMDeath           UUID
+  | RMDeregisterNode  NodeId
+  deriving (Generic)
+
+instance Binary RegistryMsg
+
+findByUUID :: UUID -> Map.Map ActorId (ThreadId, SomeActorRef) -> Maybe SomeActorRef
+findByUUID uuid actors =
+  case Map.toList (Map.filterWithKey (\(ActorId _ u) _ -> u == uuid) actors) of
+    []            -> Nothing
+    (_, (_, r)):_ -> Just r

src/FL/Actors/Aggregator.hs

@@ -0,0 +1,43 @@
+module FL.Actors.Aggregator
+  ( AggregatorState (..)
+  , aggregatorHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Concurrent.STM (TMVar, atomically, readTMVar)
+import Control.Monad.IO.Class (liftIO)
+import Data.List (sortOn)
+import FL.Actors.Logger (logFn)
+import FL.NN (fedAvg)
+import FL.Types
+
+data AggregatorState = AggregatorState
+  { asExpected     :: Int
+  , asUpdates      :: [ModelUpdate]
+  , asCoordCell    :: TMVar (ActorRef CoordinatorMsg ())  -- filled after coordinator spawns
+  , asLogRef       :: ActorRef LogMsg ()
+  , asCurrentRound :: Int
+  }
+
+aggregatorHandler :: Handler AggregatorMsg AggregatorState ()
+aggregatorHandler (SubmitUpdate upd) = do
+  st <- state
+  if muRound upd /= asCurrentRound st
+    then pass
+    else do
+      let updates' = upd : filter (\u -> muTrainerId u /= muTrainerId upd) (asUpdates st)
+      if length updates' >= asExpected st
+        then do
+          let pairs    = [(muSampleSize u, wData (muWeights u)) | u <- sortOn muTrainerId updates']
+              avgW     = Weights (fedAvg pairs)
+              nextRound = asCurrentRound st + 1
+              newModel  = GlobalModel nextRound avgW
+          logFn (asLogRef st) "Aggregator" INFO $
+            "Round " <> show (asCurrentRound st)
+            <> ": aggregated " <> show (length updates') <> " updates"
+          coordRef <- liftIO $ atomically $ readTMVar (asCoordCell st)
+          cast (AggregationComplete newModel) coordRef
+          passWith st { asUpdates = [], asCurrentRound = nextRound }
+        else
+          passWith st { asUpdates = updates' }

src/FL/Actors/Coordinator.hs

@@ -0,0 +1,103 @@
+module FL.Actors.Coordinator
+  ( CoordinatorState (..)
+  , coordinatorHandler
+  , coordinatorDeathHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Data.Binary (encodeFile)
+import Control.Concurrent.MVar (MVar, putMVar)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (liftIO)
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as Map
+import Data.Set qualified as Set
+import Data.UUID.V4 (nextRandom)
+import FL.Actors.Logger (logFn)
+import FL.CRDT (ORSet, orAdd, orMembers)
+import FL.Types
+
+data CoordinatorState = CoordinatorState
+  { csTrainers         :: ORSet String
+  , csTrainerRefs      :: Map String (ActorRef TrainerMsg ())
+  , csExpectedTrainers :: Int
+  , csCurrentRound     :: Int
+  , csTotalRounds      :: Int
+  , csGlobalModel      :: GlobalModel
+  , csAggRef           :: ActorRef AggregatorMsg ()
+  , csEvalRef          :: ActorRef EvaluatorMsg ()
+  , csLogRef           :: ActorRef LogMsg ()
+  , csDoneVar          :: MVar ()
+  }
+
+coordinatorHandler :: Handler CoordinatorMsg CoordinatorState ()
+coordinatorHandler (RegisterTrainer tid aid) = do
+  st <- state
+  -- The ActorId in the message carries the *sender's* node numbering (its own
+  -- node id is always 0 to itself); rebase it onto our node id for the sender,
+  -- otherwise every cast to the trainer is dropped ("no node with id 0").
+  from <- lastMessageFrom
+  uuid <- liftIO nextRandom
+  let ActorId _ trainerUUID = aid
+      ref       = RemoteRef (ActorId from trainerUUID) :: ActorRef TrainerMsg ()
+      trainers' = orAdd uuid tid (csTrainers st)
+      refs'     = Map.insert tid ref (csTrainerRefs st)
+      st'       = st { csTrainers = trainers', csTrainerRefs = refs' }
+      numRegistered = Set.size (orMembers trainers')
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Trainer registered: " <> tid
+    <> " (" <> show numRegistered <> "/" <> show (csExpectedTrainers st) <> ")"
+  -- Auto-start once all expected trainers have registered (per the OR-Set)
+  if numRegistered >= csExpectedTrainers st
+    then startRound st'
+    else passWith st'
+
+coordinatorHandler StartFederation = do
+  st <- state
+  startRound st
+
+coordinatorHandler (AggregationComplete newModel) = do
+  st <- state
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Round " <> show (gmRound newModel - 1) <> " aggregation complete"
+  cast (EvaluateModel newModel) (csEvalRef st)
+  let st' = st { csGlobalModel = newModel, csCurrentRound = gmRound newModel }
+  if csCurrentRound st' >= csTotalRounds st'
+    then do
+      liftIO $ encodeFile "fl_model_coordinator.bin" (gmWeights newModel)
+      logFn (csLogRef st') "Coordinator" INFO
+        "Model saved to fl_model_coordinator.bin"
+      passWith st'
+    else
+      startRound st'
+
+coordinatorHandler (EvaluationDone metrics) = do
+  st <- state
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Eval round=" <> show (mRound metrics)
+    <> " acc="  <> show4 (mAccuracy  metrics)
+    <> " prec=" <> show4 (mPrecision metrics)
+    <> " rec="  <> show4 (mRecall    metrics)
+    <> " f1="   <> show4 (mF1        metrics)
+    <> " auc="  <> show4 (mAUCROC    metrics)
+  liftIO $ when (mRound metrics >= csTotalRounds st) $
+    putMVar (csDoneVar st) ()
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+startRound :: CoordinatorState -> Actor CoordinatorMsg CoordinatorState ()
+startRound st = do
+  let model = csGlobalModel st
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Starting round " <> show (gmRound model + 1)
+    <> " with " <> show (Map.size (csTrainerRefs st)) <> " trainers"
+  forM_ (Map.elems (csTrainerRefs st)) $ \ref ->
+    cast (StartTraining model) ref
+  passWith st
+
+coordinatorDeathHandler :: DeathMessage -> ActorM CoordinatorState (SupervisorAction CoordinatorState)
+coordinatorDeathHandler _ = do
+  st <- state
+  return (Continue st)

src/FL/Actors/Evaluator.hs

@@ -0,0 +1,36 @@
+module FL.Actors.Evaluator
+  ( EvaluatorState (..)
+  , evaluatorHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Concurrent.STM (TMVar, atomically, readTMVar)
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn)
+import FL.NN (LocalData, evaluate, unflattenWeights)
+import FL.Types
+
+data EvaluatorState = EvaluatorState
+  { esTestData  :: LocalData
+  , esCoordCell :: TMVar (ActorRef CoordinatorMsg ())  -- filled after coordinator spawns
+  , esLogRef    :: ActorRef LogMsg ()
+  }
+
+evaluatorHandler :: Handler EvaluatorMsg EvaluatorState ()
+evaluatorHandler (EvaluateModel gm) = do
+  st <- state
+  let ws      = unflattenWeights (wData (gmWeights gm))
+      metrics = evaluate ws (esTestData st) (gmRound gm)
+  logFn (esLogRef st) "Evaluator" INFO $
+    "Round " <> show (gmRound gm)
+    <> "  acc="  <> show4 (mAccuracy  metrics)
+    <> "  prec=" <> show4 (mPrecision metrics)
+    <> "  rec="  <> show4 (mRecall    metrics)
+    <> "  f1="   <> show4 (mF1        metrics)
+    <> "  auc="  <> show4 (mAUCROC    metrics)
+  coordRef <- liftIO $ atomically $ readTMVar (esCoordCell st)
+  cast (EvaluationDone metrics) coordRef
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)

src/FL/Actors/Logger.hs

@@ -0,0 +1,37 @@
+module FL.Actors.Logger
+  ( loggerHandler
+  , logFn
+  , logFn'
+  ) where
+
+import Control.Actor.Core (ActorM, Handler, cast, cast', pass)
+import Control.Actor.Runtime (RuntimeM)
+import Control.Actor.Types
+import Control.Monad.IO.Class (liftIO)
+import Data.Time.Clock (getCurrentTime)
+import FL.Types
+
+type LogState = ()
+
+loggerHandler :: Handler LogMsg LogState ()
+loggerHandler (Log LogEntry{leSource, leLevel, leMsg}) = do
+  t <- liftIO getCurrentTime
+  liftIO $ putStrLn $
+    "[" <> show t <> "] [" <> showLevel leLevel <> "] "
+    <> leSource <> ": " <> leMsg
+  pass
+  where
+    showLevel DEBUG = "DEBUG"
+    showLevel INFO  = "INFO "
+    showLevel WARN  = "WARN "
+    showLevel ERROR = "ERROR"
+
+-- | Log via the logger actor. Goes through the normal cast path so it works
+-- for both local and remote logger refs (a remote ref used to be a silent
+-- no-op, losing all trainer-node logs in provider mode).
+logFn :: ActorRef LogMsg () -> String -> LogLevel -> String -> ActorM u ()
+logFn ref src lvl msg = cast (Log (LogEntry src lvl msg)) ref
+
+-- | Same as 'logFn' but usable outside an actor (node startup code).
+logFn' :: ActorRef LogMsg () -> String -> LogLevel -> String -> RuntimeM ()
+logFn' ref src lvl msg = cast' (Log (LogEntry src lvl msg)) ref

src/FL/Actors/PeerCoordinator.hs

@@ -0,0 +1,166 @@
+{-# LANGUAGE StrictData #-}
+
+module FL.Actors.PeerCoordinator
+  ( PeerCoordinatorState (..)
+  , peerCoordinatorHandler
+  , peerCoordinatorDeathHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Data.Binary (encodeFile)
+import Unsafe.Coerce (unsafeCoerce)
+import Control.Concurrent.MVar (MVar, putMVar)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (liftIO)
+import Data.List (sortOn)
+import Data.Set qualified as Set
+import Data.UUID.V4 (nextRandom)
+import FL.Actors.Logger (logFn)
+import FL.CRDT
+import FL.NN (fedAvg)
+import FL.Types
+
+data PeerCoordinatorState = PeerCoordinatorState
+  { pcsNodeId       :: String
+  , pcsRoundCounter :: GCounter        -- completed training rounds per peer (monotone, never reset)
+  , pcsParticipants :: ORSet String    -- active peers ever seen (monotone, never reset)
+  , pcsCurrentRound :: Int
+  , pcsTotalRounds  :: Int
+  , pcsGlobalModel  :: GlobalModel
+  , pcsPeerRefs     :: [ActorRef PeerCoordinatorMsg ()]
+  , pcsSelfRef      :: ActorRef PeerCoordinatorMsg ()
+  , pcsTrainerRef   :: ActorRef TrainerMsg ()
+  , pcsEvalRef      :: ActorRef EvaluatorMsg ()
+  , pcsLogRef       :: ActorRef LogMsg ()
+  , pcsCollected    :: [ModelUpdate]   -- updates for the current round + buffered future rounds
+  , pcsNumPeers     :: Int  -- total number of peers (including self)
+  , pcsDoneVar      :: MVar ()
+  }
+
+-- | Insert an update, deduplicating on (trainer, round). Stale (past-round)
+-- updates are dropped; future-round updates from fast peers are buffered so
+-- they are not lost — without the round key a fast peer's round-(r+1) update
+-- would overwrite its round-r update in a slower peer and rounds would mix.
+insertUpdate :: Int -> ModelUpdate -> [ModelUpdate] -> [ModelUpdate]
+insertUpdate currentRound upd collected
+  | muRound upd < currentRound = collected
+  | otherwise =
+      upd : filter (\u -> (muTrainerId u, muRound u) /= (muTrainerId upd, muRound upd)) collected
+
+currentRoundUpdates :: PeerCoordinatorState -> [ModelUpdate]
+currentRoundUpdates st = [u | u <- pcsCollected st, muRound u == pcsCurrentRound st]
+
+-- | A round is complete when the CRDTs agree with the collected updates:
+-- the OR-Set has seen all peers, the G-Counter shows every one of them has
+-- finished training the current round, and we hold all their updates.
+roundComplete :: PeerCoordinatorState -> Bool
+roundComplete st =
+  let r       = pcsCurrentRound st
+      members = orMembers (pcsParticipants st)
+      trained = Set.filter (\p -> gcGet p (pcsRoundCounter st) > r) members
+  in Set.size members >= pcsNumPeers st
+     && Set.size trained >= pcsNumPeers st
+     && length (currentRoundUpdates st) >= pcsNumPeers st
+
+peerCoordinatorHandler :: Handler PeerCoordinatorMsg PeerCoordinatorState ()
+peerCoordinatorHandler (SetPeers peerIds) = do
+  st <- state
+  (SomeActorRef self) <- getSelf
+  let peers   = map (\aid -> RemoteRef aid :: ActorRef PeerCoordinatorMsg ()) peerIds
+      selfRef = unsafeCoerce self :: ActorRef PeerCoordinatorMsg ()
+  passWith st { pcsPeerRefs = peers, pcsSelfRef = selfRef }
+
+peerCoordinatorHandler PeerStartRound = do
+  st <- state
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "Starting local training for P2P round " <> show (pcsCurrentRound st + 1)
+  cast (StartTraining (pcsGlobalModel st)) (pcsTrainerRef st)
+  pass
+
+peerCoordinatorHandler (PeerTrainingDone upd) = do
+  st <- state
+  uuid <- liftIO nextRandom
+  let counter'      = gcIncrement (pcsNodeId st) (pcsRoundCounter st)
+      participants' = orAdd uuid (pcsNodeId st) (pcsParticipants st)
+      collected'    = insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
+      payload       = PeerSyncPayload
+                        { pspCounter      = counter'
+                        , pspParticipants = participants'
+                        , pspUpdate       = Just upd
+                        }
+  logFn (pcsLogRef st) (pcsNodeId st) INFO
+    "Training done, broadcasting PeerSync"
+  forM_ (pcsPeerRefs st) $ \ref ->
+    cast (PeerSync payload) ref
+  let st' = st { pcsRoundCounter = counter'
+               , pcsParticipants = participants'
+               , pcsCollected    = collected'
+               }
+  if roundComplete st'
+    then finishRound st'
+    else passWith st'
+
+peerCoordinatorHandler (PeerSync payload) = do
+  st <- state
+  -- Merge CRDTs
+  let counter'      = gcMerge (pcsRoundCounter st) (pspCounter payload)
+      participants' = orMerge (pcsParticipants st) (pspParticipants payload)
+      collected'    = case pspUpdate payload of
+        Nothing  -> pcsCollected st
+        Just upd -> insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
+      st' = st { pcsRoundCounter = counter'
+               , pcsParticipants = participants'
+               , pcsCollected    = collected'
+               }
+  if roundComplete st'
+    then finishRound st'
+    else passWith st'
+
+peerCoordinatorHandler (PeerEvaluationDone metrics) = do
+  st <- state
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "Eval round=" <> show (mRound metrics)
+    <> " acc="  <> show4 (mAccuracy  metrics)
+    <> " auc="  <> show4 (mAUCROC    metrics)
+  liftIO $ when (mRound metrics >= pcsTotalRounds st) $
+    putMVar (pcsDoneVar st) ()
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+finishRound :: PeerCoordinatorState -> Actor PeerCoordinatorMsg PeerCoordinatorState ()
+finishRound st = do
+  -- Sort so every peer folds the weighted average in the same order and
+  -- arrives at bit-identical models regardless of message arrival order.
+  let updates   = sortOn muTrainerId (currentRoundUpdates st)
+      pairs     = [(muSampleSize u, wData (muWeights u)) | u <- updates]
+      avgW      = Weights (fedAvg pairs)
+      nextRound = pcsCurrentRound st + 1
+      newModel  = GlobalModel nextRound avgW
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "P2P round " <> show (pcsCurrentRound st) <> " aggregated ("
+    <> show (length updates) <> " updates)"
+  cast (EvaluateModel newModel) (pcsEvalRef st)
+  let st' = st
+        { pcsGlobalModel  = newModel
+        , pcsCurrentRound = nextRound
+        -- Keep buffered updates from peers already in later rounds
+        , pcsCollected    = [u | u <- pcsCollected st, muRound u >= nextRound]
+        }
+  if nextRound >= pcsTotalRounds st
+    then do
+      let filename = "fl_model_" <> pcsNodeId st' <> ".bin"
+      liftIO $ encodeFile filename (gmWeights newModel)
+      logFn (pcsLogRef st') (pcsNodeId st') INFO $
+        "Model saved to " <> filename
+      passWith st'
+    else do
+      -- Schedule next round after a short delay to let PeerSync propagate
+      castIn 500 PeerStartRound (pcsSelfRef st')
+      passWith st'
+
+peerCoordinatorDeathHandler :: DeathMessage -> ActorM PeerCoordinatorState (SupervisorAction PeerCoordinatorState)
+peerCoordinatorDeathHandler _ = do
+  st <- state
+  return (Continue st)

src/FL/Actors/Trainer.hs

@@ -0,0 +1,45 @@
+module FL.Actors.Trainer
+  ( TrainerState (..)
+  , trainerHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn)
+import FL.NN (LocalData, flattenWeights, localDataSize, trainLocal, unflattenWeights)
+import FL.Types
+
+data TrainerState = TrainerState
+  { tsId        :: String
+  , tsData      :: LocalData
+  , tsLastRound :: Int
+  , tsAggRef    :: ActorRef AggregatorMsg ()
+  , tsLogRef    :: ActorRef LogMsg ()
+  , tsConfig    :: FLConfig
+  }
+
+trainerHandler :: Handler TrainerMsg TrainerState ()
+trainerHandler (StartTraining gm) = do
+  st <- state
+  let round_ = gmRound gm
+  if round_ <= tsLastRound st
+    then pass  -- idempotent: already processed this round
+    else do
+      logFn (tsLogRef st) (tsId st) INFO $
+        "Starting local training for round " <> show round_
+      let initW = unflattenWeights (wData (gmWeights gm))
+      updatedW <- liftIO $ trainLocal (tsConfig st) (tsId st) round_ initW (tsData st)
+      let n      = localDataSize (tsData st)
+          update = ModelUpdate
+            { muTrainerId  = tsId st
+            , muRound      = round_
+            , muSampleSize = n
+            , muWeights    = Weights (flattenWeights updatedW)
+            }
+      logFn (tsLogRef st) (tsId st) INFO $
+        "Round " <> show round_ <> " training complete, samples=" <> show n
+      cast (SubmitUpdate update) (tsAggRef st)
+      passWith st { tsLastRound = round_ }
+
+trainerHandler TrainerStop = pass

src/FL/CRDT.hs

@@ -0,0 +1,82 @@
+module FL.CRDT
+  ( GCounter
+  , gcEmpty
+  , gcIncrement
+  , gcMerge
+  , gcValue
+  , gcGet
+  , ORSet
+  , orEmpty
+  , orAdd
+  , orRemove
+  , orMerge
+  , orMembers
+  ) where
+
+import Data.Binary (Binary (..))
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as Map
+import Data.Set (Set)
+import Data.Set qualified as Set
+import Data.UUID (UUID)
+import GHC.Generics (Generic)
+
+-- ---------------------------------------------------------------------------
+-- G-Counter: each node has its own non-decreasing counter.
+-- Merge takes element-wise max. Value is the sum.
+-- ---------------------------------------------------------------------------
+
+newtype GCounter = GCounter (Map String Int)
+  deriving (Show, Eq, Generic)
+
+instance Binary GCounter where
+  put (GCounter m) = put (Map.toAscList m)
+  get = GCounter . Map.fromAscList <$> get
+
+gcEmpty :: GCounter
+gcEmpty = GCounter Map.empty
+
+gcIncrement :: String -> GCounter -> GCounter
+gcIncrement nodeId (GCounter m) =
+  GCounter (Map.insertWith (+) nodeId 1 m)
+
+gcMerge :: GCounter -> GCounter -> GCounter
+gcMerge (GCounter a) (GCounter b) =
+  GCounter (Map.unionWith max a b)
+
+gcValue :: GCounter -> Int
+gcValue (GCounter m) = sum (Map.elems m)
+
+gcGet :: String -> GCounter -> Int
+gcGet nodeId (GCounter m) = Map.findWithDefault 0 nodeId m
+
+-- ---------------------------------------------------------------------------
+-- OR-Set: observe-remove set. Each element tagged with a set of UUIDs.
+-- Add: insert a fresh UUID tag. Remove: delete all tags for that element.
+-- Merge: union of tag sets per element. An element is present iff it has
+-- at least one tag.
+-- ---------------------------------------------------------------------------
+
+newtype ORSet a = ORSet (Map a (Set UUID))
+  deriving (Show, Eq, Generic)
+
+instance (Ord a, Binary a) => Binary (ORSet a) where
+  put (ORSet m) = put (Map.toAscList (fmap Set.toAscList m))
+  get = ORSet . Map.fromAscList . map (fmap Set.fromAscList) <$> get
+
+orEmpty :: ORSet a
+orEmpty = ORSet Map.empty
+
+orAdd :: Ord a => UUID -> a -> ORSet a -> ORSet a
+orAdd tag x (ORSet m) =
+  ORSet (Map.insertWith Set.union x (Set.singleton tag) m)
+
+orRemove :: Ord a => a -> ORSet a -> ORSet a
+orRemove x (ORSet m) = ORSet (Map.delete x m)
+
+orMerge :: Ord a => ORSet a -> ORSet a -> ORSet a
+orMerge (ORSet a) (ORSet b) =
+  ORSet (Map.unionWith Set.union a b)
+
+orMembers :: Ord a => ORSet a -> Set a
+orMembers (ORSet m) = Map.keysSet (Map.filter (not . Set.null) m)

src/FL/Data/PaySim.hs

@@ -0,0 +1,231 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+module FL.Data.PaySim
+  ( loadAndPartition
+  -- , syntheticPartition
+  ) where
+
+import Data.ByteString.Lazy qualified as BL
+import Data.Csv qualified as Csv
+import Data.List (sortBy)
+import Data.Ord (comparing)
+import Data.Vector qualified as V
+import FL.NN (LocalData (..))
+import FL.Types (FLConfig (..), NormMode (..), PartitionMode (..))
+import GHC.Generics (Generic)
+import Numeric.LinearAlgebra hiding ((<>))
+import Numeric.LinearAlgebra qualified as LA
+import System.Random (mkStdGen, randomRs)
+
+-- ---------------------------------------------------------------------------
+-- PaySim CSV record
+-- ---------------------------------------------------------------------------
+
+data PaySimRow = PaySimRow
+  { rowStep           :: !Double
+  , rowType           :: !String
+  , rowAmount         :: !Double
+  , rowOldBalOrig     :: !Double
+  , rowNewBalOrig     :: !Double
+  , rowOldBalDest     :: !Double
+  , rowNewBalDest     :: !Double
+  , rowIsFraud        :: !Double
+  } deriving (Show, Generic)
+
+instance Csv.FromNamedRecord PaySimRow where
+  parseNamedRecord r =
+    PaySimRow
+      <$> r Csv..: "step"
+      <*> r Csv..: "type"
+      <*> r Csv..: "amount"
+      <*> r Csv..: "oldbalanceOrg"
+      <*> r Csv..: "newbalanceOrig"
+      <*> r Csv..: "oldbalanceDest"
+      <*> r Csv..: "newbalanceDest"
+      <*> r Csv..: "isFraud"
+
+rowToFeatures :: PaySimRow -> Vector Double
+rowToFeatures PaySimRow{..} =
+  let typeVec = case rowType of
+        "CASH_IN"   -> [1,0,0,0,0]
+        "CASH_OUT"  -> [0,1,0,0,0]
+        "DEBIT"     -> [0,0,1,0,0]
+        "PAYMENT"   -> [0,0,0,1,0]
+        "TRANSFER"  -> [0,0,0,0,1]
+        _           -> [0,0,0,0,0]
+      numeric = [ rowAmount
+                , rowOldBalOrig
+                , rowNewBalOrig
+                , rowOldBalDest
+                , rowNewBalDest
+                , rowNewBalOrig - rowOldBalOrig   -- balance change orig
+                , rowNewBalDest - rowOldBalDest   -- balance change dest
+                , rowAmount / (rowOldBalOrig + 1) -- amount ratio orig
+                , rowAmount / (rowOldBalDest + 1) -- amount ratio dest
+                , rowStep / 744                   -- normalized step (max 744 in PaySim)
+                ]
+      -- pad to 29 features (5 one-hot + 10 numeric + 14 zeros)
+      padded = typeVec ++ numeric ++ replicate 14 0.0
+  in fromList padded
+
+-- ---------------------------------------------------------------------------
+-- Load PaySim CSV
+-- ---------------------------------------------------------------------------
+
+loadPaySimCSV :: FilePath -> IO (Either String (Matrix Double, Vector Double))
+loadPaySimCSV path = do
+  raw <- BL.readFile path
+  case Csv.decodeByName raw of
+    Left err -> return (Left err)
+    Right (_, rows) -> do
+      let featureList = V.toList (V.map rowToFeatures rows)
+          labelList   = V.toList (V.map rowIsFraud rows)
+      return (Right (fromRows featureList, fromList labelList))
+
+-- ---------------------------------------------------------------------------
+-- Normalization
+-- ---------------------------------------------------------------------------
+
+normalizeMinMax :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
+normalizeMinMax m =
+  let mins  = fromList [minElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
+      maxs  = fromList [maxElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
+      range = cmap (\x -> if x == 0 then 1 else x) (maxs - mins)
+      norm  = (m - LA.repmat (asRow mins) (rows m) 1)
+              / LA.repmat (asRow range)  (rows m) 1
+  in (norm, mins, maxs)
+
+normalizeZScore :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
+normalizeZScore m =
+  let n     = fromIntegral (rows m) :: Double
+      means = fromList [sumElements (m ¿ [j]) / n | j <- [0 .. cols m - 1]]
+      vars  = fromList [sumElements (cmap (^(2::Int)) ((m ¿ [j]) - scalar (means ! j))) / n
+                       | j <- [0 .. cols m - 1]]
+      stds  = cmap (\v -> if v == 0 then 1 else sqrt v) vars
+      norm  = (m - LA.repmat (asRow means) (rows m) 1)
+              / LA.repmat (asRow stds)  (rows m) 1
+  in (norm, means, stds)
+
+applyNorm :: NormMode -> (Matrix Double, Vector Double, Vector Double) -> Matrix Double -> Matrix Double
+applyNorm mode (_, p1, p2) m =
+  case mode of
+    MinMax ->
+      let range = cmap (\x -> if x == 0 then 1 else x) (p2 - p1)
+      in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow range) (rows m) 1
+    ZScore ->
+      let stds = cmap (\x -> if x == 0 then 1 else x) p2
+      in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow stds)  (rows m) 1
+
+-- ---------------------------------------------------------------------------
+-- IID partition: deterministic shard by trainer index
+-- ---------------------------------------------------------------------------
+
+partitionIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
+partitionIID trainerIdx numTrainers feats labels =
+  let n       = rows feats
+      -- Shuffle with fixed seed for reproducibility across processes
+      seed    = 2024
+      rng     = mkStdGen seed
+      sorted  = sortBy (comparing fst)
+                  (zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1])
+      idxs    = map snd sorted
+      shard   = chunkedShard trainerIdx numTrainers idxs
+      f       = feats ? shard
+      l       = fromList [labels ! i | i <- shard]
+  in LocalData f l
+
+chunkedShard :: Int -> Int -> [a] -> [a]
+chunkedShard idx total xs =
+  let n       = length xs
+      size    = n `div` total
+      start   = idx * size
+      end     = if idx == total - 1 then n else start + size
+  in take (end - start) (drop start xs)
+
+-- ---------------------------------------------------------------------------
+-- Non-IID partition: each trainer gets dominant transaction type
+-- (approximated via deterministic row assignment)
+-- ---------------------------------------------------------------------------
+
+partitionNonIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
+partitionNonIID trainerIdx numTrainers feats labels =
+  let n         = rows feats
+      -- Assign rows to trainers based on dominant type one-hot (cols 0–4)
+      -- Trainer i "owns" type (i mod 5); remaining trainers share type 0
+      ownType   = trainerIdx `mod` 5
+      isOwn i   = let row = toList (flatten (feats ? [i]))
+                  in length row > ownType && row !! ownType > 0.5
+      -- With more than 5 trainers several trainers own the same type;
+      -- split that type's rows among them instead of giving each trainer
+      -- all of them (which duplicated data across shards).
+      sharers   = [i | i <- [0 .. numTrainers - 1], i `mod` 5 == ownType]
+      shareIdx  = length (takeWhile (/= trainerIdx) sharers)
+      -- Gather own-type rows + 10% of other rows for minority
+      ownIdxs   = chunkedShard shareIdx (length sharers) [i | i <- [0 .. n-1], isOwn i]
+      otherIdxs = [i | i <- [0 .. n-1], not (isOwn i)]
+      minority  = chunkedShard trainerIdx numTrainers otherIdxs
+      allIdxs   = ownIdxs ++ minority
+      f         = feats ? allIdxs
+      l         = fromList [labels ! i | i <- allIdxs]
+  in LocalData f l
+
+-- ---------------------------------------------------------------------------
+-- Synthetic data fallback
+-- ---------------------------------------------------------------------------
+
+syntheticData :: Int -> Int -> Int -> LocalData
+syntheticData trainerIdx numTrainers nSamples =
+  let seed    = trainerIdx * 13337 + numTrainers
+      rng     = mkStdGen seed
+      -- 29 features from normal distribution
+      featVals = take (nSamples * 29) (randomRs (-3.0 :: Double, 3.0) rng)
+      feats   = (nSamples >< 29) featVals
+      -- Base 2% fraud rate, +1% per trainer index for variety (real PaySim
+      -- is ~0.1%, but synthetic runs need enough positives to learn from)
+      fraudRate = 0.02 + fromIntegral trainerIdx * 0.01
+      rng2    = mkStdGen (seed + 1)
+      labelVals = take nSamples (randomRs (0.0 :: Double, 1.0) rng2)
+      labels  = fromList [if v < fraudRate then 1.0 else 0.0 | v <- labelVals]
+  in LocalData feats labels
+
+-- ---------------------------------------------------------------------------
+-- Top-level entry point used by Provider and P2P startup
+-- ---------------------------------------------------------------------------
+
+-- Returns (train partition, test data).
+-- Test data is the first 10% of the full dataset (at least 1000 rows),
+-- taken before partitioning so every node evaluates on the same split.
+loadAndPartition
+  :: FLConfig
+  -> Int    -- this node's trainer index (0-based)
+  -> IO (LocalData, LocalData)
+loadAndPartition FLConfig{..} trainerIdx = do
+  result <- case cfgDataPath of
+    Nothing   -> return (Left "no data path")
+    Just path -> loadPaySimCSV path
+  case result of
+    Left _ -> do
+      -- Synthetic fallback
+      let testData  = syntheticData (cfgNumTrainers + 1) cfgNumTrainers 2000
+          trainData = syntheticData trainerIdx cfgNumTrainers 10000
+      return (trainData, testData)
+    Right (feats, labels) -> do
+      let n         = rows feats
+          -- Aim for 1000 test rows for stable metrics, but never take more
+          -- than half the data — max alone made testN exceed n on small
+          -- CSVs and crashed the train subMatrix with a negative size.
+          testN     = min (max 1000 (n `div` 10)) (n `div` 2)
+          testFeats = subMatrix (0, 0) (testN, cols feats) feats
+          testLabs  = subVector 0 testN labels
+          trainF    = subMatrix (testN, 0) (n - testN, cols feats) feats
+          trainL    = subVector testN (n - testN) labels
+          (normF, p1, p2) = case cfgNorm of
+            MinMax -> normalizeMinMax trainF
+            ZScore -> normalizeZScore trainF
+          normTest  = applyNorm cfgNorm (normF, p1, p2) testFeats
+          trainData = case cfgPartition of
+            IID    -> partitionIID    trainerIdx cfgNumTrainers normF trainL
+            NonIID -> partitionNonIID trainerIdx cfgNumTrainers normF trainL
+          testData  = LocalData normTest testLabs
+      return (trainData, testData)

src/FL/NN.hs

@@ -0,0 +1,290 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE StrictData #-}
+
+module FL.NN
+  ( MLPWeights (..)
+  , LocalData (..)
+  , initWeights
+  , trainLocal
+  , evaluate
+  , flattenWeights
+  , unflattenWeights
+  , fedAvg
+  , localDataSize
+  ) where
+
+import Data.List (foldl', sortBy)
+import Data.Ord (Down (..), comparing)
+import FL.Types (FLConfig (..), Metrics (..))
+import Numeric.LinearAlgebra hiding ((<>))
+import Numeric.LinearAlgebra qualified as LA
+import System.Random (mkStdGen, randomRs)
+
+-- ---------------------------------------------------------------------------
+-- Types
+-- ---------------------------------------------------------------------------
+
+data MLPWeights = MLPWeights
+  { w1 :: Matrix Double  -- 64 × 29
+  , b1 :: Vector Double  -- 64
+  , w2 :: Matrix Double  -- 32 × 64
+  , b2 :: Vector Double  -- 32
+  , w3 :: Matrix Double  --  2 × 32
+  , b3 :: Vector Double  --  2
+  }
+
+data LocalData = LocalData
+  { ldFeatures :: Matrix Double  -- N × 29
+  , ldLabels   :: Vector Double  -- N  (0.0 or 1.0)
+  }
+
+localDataSize :: LocalData -> Int
+localDataSize = rows . ldFeatures
+
+-- ---------------------------------------------------------------------------
+-- Weight initialization (Xavier uniform)
+-- ---------------------------------------------------------------------------
+
+xavierMatrix :: Int -> Int -> Int -> Matrix Double
+xavierMatrix seed fanIn fanOut =
+  let bound = sqrt (6.0 / fromIntegral (fanIn + fanOut))
+      n     = fanIn * fanOut
+      vals  = take n (randomRs (-bound, bound) (mkStdGen seed))
+  in (fanOut >< fanIn) vals
+
+initWeights :: IO MLPWeights
+initWeights = do
+  -- Use fixed seeds so multiple processes with same call order get different
+  -- weight initializations only if we pass different seeds; for now deterministic.
+  let w1' = xavierMatrix 42  29 64
+      b1' = konst 0 64
+      w2' = xavierMatrix 137 64 32
+      b2' = konst 0 32
+      w3' = xavierMatrix 999 32  2
+      b3' = konst 0 2
+  return MLPWeights{ w1=w1', b1=b1', w2=w2', b2=b2', w3=w3', b3=b3' }
+
+-- ---------------------------------------------------------------------------
+-- Activation functions
+-- ---------------------------------------------------------------------------
+
+relu :: Matrix Double -> Matrix Double
+relu = cmap (max 0)
+
+reluGrad :: Matrix Double -> Matrix Double
+reluGrad = cmap (\x -> if x > 0 then 1 else 0)
+
+softmaxRows :: Matrix Double -> Matrix Double
+softmaxRows m = fromRows (map softmaxVec (toRows m))
+  where
+    -- Subtract the row max before exponentiating so large logits
+    -- can't overflow exp and poison the weights with NaN.
+    softmaxVec v =
+      let mx = maxElement v
+          e  = cmap (\x -> exp (x - mx)) v
+          s  = sumElements e
+      in cmap (/ s) e
+
+-- ---------------------------------------------------------------------------
+-- Forward pass
+-- ---------------------------------------------------------------------------
+
+-- Returns (output probs N×2, z1, a1, z2, a2) for backprop.
+forwardWithCache
+  :: MLPWeights
+  -> Matrix Double  -- N × 29
+  -> (Matrix Double, Matrix Double, Matrix Double, Matrix Double, Matrix Double)
+forwardWithCache MLPWeights{..} x =
+  let z1  = x LA.<> tr w1 + repRows (rows x) b1
+      a1  = relu z1
+      z2  = a1 LA.<> tr w2 + repRows (rows x) b2
+      a2  = relu z2
+      z3  = a2 LA.<> tr w3 + repRows (rows x) b3
+      out = softmaxRows z3
+  in (out, z1, a1, z2, a2)
+
+repRows :: Int -> Vector Double -> Matrix Double
+repRows n v = fromRows (replicate n v)
+
+colSums :: Matrix Double -> Vector Double
+colSums m = fromList [sumElements (m ¿ [j]) | j <- [0 .. cols m - 1]]
+
+-- ---------------------------------------------------------------------------
+-- Loss
+-- ---------------------------------------------------------------------------
+
+crossEntropyLoss :: Matrix Double -> Vector Double -> Double
+crossEntropyLoss probs labels =
+  let n = fromIntegral (rows probs)
+      logProbs = cmap (log . max 1e-15) probs
+      oneHot   = buildOneHot (rows probs) (map round (toList labels))
+      loss     = negate (sumElements (oneHot * logProbs)) / n
+  in loss
+
+buildOneHot :: Int -> [Int] -> Matrix Double
+buildOneHot n ys =
+  fromRows [fromList [if i == y then 1.0 else 0.0 | i <- [0, 1]] | y <- ys]
+
+-- ---------------------------------------------------------------------------
+-- Backward pass (mini-batch SGD)
+-- ---------------------------------------------------------------------------
+
+stepBatch
+  :: MLPWeights
+  -> Matrix Double  -- batch features N × 29
+  -> Vector Double  -- batch labels N
+  -> Double         -- learning rate
+  -> (MLPWeights, Double)
+stepBatch ws@MLPWeights{..} x labels lr =
+  let n              = fromIntegral (rows x) :: Double
+      (out, z1, a1, z2, a2) = forwardWithCache ws x
+      oneHot         = buildOneHot (rows x) (map round (toList labels))
+      -- Output delta: dL/dZ3 = (probs - oneHot) / N
+      dZ3            = cmap (/ n) (out - oneHot)           -- N × 2
+      dW3            = tr dZ3 LA.<> a2                     -- 2 × 32
+      db3            = colSums dZ3                         -- 2
+      dA2            = dZ3 LA.<> w3                        -- N × 32
+      dZ2            = dA2 * reluGrad z2                   -- N × 32
+      dW2            = tr dZ2 LA.<> a1                     -- 32 × 64
+      db2            = colSums dZ2                         -- 32
+      dA1            = dZ2 LA.<> w2                        -- N × 64
+      dZ1            = dA1 * reluGrad z1                   -- N × 64
+      dW1            = tr dZ1 LA.<> x                      -- 64 × 29
+      db1            = colSums dZ1                         -- 64
+      ws'            = MLPWeights
+        { w1 = w1 - cmap (* lr) dW1
+        , b1 = b1 - cmap (* lr) db1
+        , w2 = w2 - cmap (* lr) dW2
+        , b2 = b2 - cmap (* lr) db2
+        , w3 = w3 - cmap (* lr) dW3
+        , b3 = b3 - cmap (* lr) db3
+        }
+      loss = crossEntropyLoss out labels
+  in (ws', loss)
+
+-- ---------------------------------------------------------------------------
+-- Training
+-- ---------------------------------------------------------------------------
+
+-- Shuffle row indices with a seed, then split into mini-batches.
+shuffleIndices :: Int -> Int -> [Int]
+shuffleIndices seed n =
+  let rng = mkStdGen seed
+      tagged = zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1]
+  in map snd (sortBy (comparing fst) tagged)
+
+trainEpoch
+  :: MLPWeights
+  -> Matrix Double  -- features N × 29
+  -> Vector Double  -- labels N
+  -> Int            -- batch size
+  -> Double         -- learning rate
+  -> Int            -- seed for shuffle
+  -> (MLPWeights, Double)
+trainEpoch ws feats labels batchSize lr seed =
+  let n       = rows feats
+      idxs    = shuffleIndices seed n
+      batches = chunksOf batchSize idxs
+      step (w, totalLoss) batch =
+        let bFeats  = feats ? batch
+            bLabels = fromList [labels ! i | i <- batch]
+            (w', l) = stepBatch w bFeats bLabels lr
+        in (w', totalLoss + l)
+      (wFinal, lossSum) = foldl' step (ws, 0.0) batches
+  in (wFinal, lossSum / fromIntegral (length batches))
+
+chunksOf :: Int -> [a] -> [[a]]
+chunksOf _ [] = []
+chunksOf n xs = let (h, t) = splitAt n xs in h : chunksOf n t
+
+-- | Train for the configured number of epochs. The shuffle seed is derived
+-- from the trainer id and round so batches differ across rounds and trainers
+-- (a bare epoch counter would replay the identical shuffle every round).
+trainLocal :: FLConfig -> String -> Int -> MLPWeights -> LocalData -> IO MLPWeights
+trainLocal FLConfig{..} trainerId roundNum initW LocalData{..} = do
+  let seedBase = foldl' (\h c -> 31 * h + fromEnum c) (roundNum * 7919) trainerId
+      go ws epoch =
+        let (ws', _loss) = trainEpoch ws ldFeatures ldLabels cfgBatchSize cfgLearningRate (seedBase + epoch)
+        in if epoch >= cfgEpochs then ws'
+           else go ws' (epoch + 1)
+  return (go initW 1)
+
+-- ---------------------------------------------------------------------------
+-- Evaluation
+-- ---------------------------------------------------------------------------
+
+evaluate :: MLPWeights -> LocalData -> Int -> Metrics
+evaluate ws LocalData{..} round_ =
+  let n         = rows ldFeatures
+      (probs, _, _, _, _) = forwardWithCache ws ldFeatures
+      predLabels = fromList [if (probs ! i ! 1) >= 0.5 then 1.0 else 0.0 | i <- [0 .. n - 1]]
+      scores     = fromList [probs ! i ! 1 | i <- [0 .. n - 1]]
+      tp = sumElements (predLabels * ldLabels)
+      fp = sumElements (predLabels * (1 - ldLabels))
+      fn = sumElements ((1 - predLabels) * ldLabels)
+      tn = sumElements ((1 - predLabels) * (1 - ldLabels))
+      acc  = (tp + tn) / fromIntegral n
+      prec = if tp + fp == 0 then 0 else tp / (tp + fp)
+      rec  = if tp + fn == 0 then 0 else tp / (tp + fn)
+      f1   = if prec + rec == 0 then 0 else 2 * prec * rec / (prec + rec)
+      auc  = computeAUCROC (toList scores) (toList ldLabels)
+  in Metrics acc prec rec f1 auc round_
+
+computeAUCROC :: [Double] -> [Double] -> Double
+computeAUCROC scores labels =
+  let nPos = fromIntegral (length (filter (== 1) labels))
+      nNeg = fromIntegral (length (filter (== 0) labels))
+  in if nPos == 0 || nNeg == 0
+     then 0.5
+     else
+       let sorted = sortBy (comparing (Down . fst)) (zip scores labels)
+           points = scanl (\(tp, fp) (_, y) -> (tp + y, fp + (1 - y))) (0, 0) sorted
+           trapz (tp0, fp0) (tp1, fp1) =
+             let fpr0 = fp0 / nNeg; tpr0 = tp0 / nPos
+                 fpr1 = fp1 / nNeg; tpr1 = tp1 / nPos
+             in (fpr1 - fpr0) * (tpr0 + tpr1) / 2.0
+       in sum (zipWith trapz points (tail points))
+
+-- ---------------------------------------------------------------------------
+-- Serialization bridge
+-- ---------------------------------------------------------------------------
+
+flattenWeights :: MLPWeights -> [Double]
+flattenWeights MLPWeights{..} =
+  toList (flatten w1)
+  ++ toList b1
+  ++ toList (flatten w2)
+  ++ toList b2
+  ++ toList (flatten w3)
+  ++ toList b3
+
+-- Fixed architecture: 29→64→32→2
+-- w1: 64*29=1856, b1: 64, w2: 32*64=2048, b2: 32, w3: 2*32=64, b3: 2
+unflattenWeights :: [Double] -> MLPWeights
+unflattenWeights xs =
+  let (xw1, r1) = splitAt (64*29) xs
+      (xb1, r2) = splitAt 64     r1
+      (xw2, r3) = splitAt (32*64) r2
+      (xb2, r4) = splitAt 32     r3
+      (xw3, r5) = splitAt (2*32)  r4
+      xb3        = take 2 r5
+  in MLPWeights
+       { w1 = (64 >< 29) xw1
+       , b1 = fromList xb1
+       , w2 = (32 >< 64) xw2
+       , b2 = fromList xb2
+       , w3 = ( 2 >< 32) xw3
+       , b3 = fromList xb3
+       }
+
+-- ---------------------------------------------------------------------------
+-- FedAvg
+-- ---------------------------------------------------------------------------
+
+-- Weighted average of flat weight vectors.
+fedAvg :: [(Int, [Double])] -> [Double]
+fedAvg [] = []
+fedAvg updates =
+  let totalN   = fromIntegral (sum (map fst updates)) :: Double
+      weighted (n, ws) = map (* (fromIntegral n / totalN)) ws
+  in foldl1 (zipWith (+)) (map weighted updates)

src/FL/P2P.hs

@@ -0,0 +1,182 @@
+{-# LANGUAGE RecordWildCards #-}
+module FL.P2P
+  ( runPeer
+  ) where
+
+import Control.Actor
+import Control.Actor.Registry (registerActor, lookupRemoteActor)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
+import Control.Concurrent.STM
+  ( TMVar, atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
+import Control.Monad (forM, forM_, void)
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn, logFn', loggerHandler)
+import FL.Actors.PeerCoordinator
+import FL.CRDT (gcEmpty, orEmpty)
+import FL.Data.PaySim (loadAndPartition)
+import FL.NN
+  ( LocalData, evaluate, flattenWeights, initWeights
+  , localDataSize, trainLocal, unflattenWeights
+  )
+import FL.Retry (retryConnect, retryLookup)
+import FL.Types
+
+-- ---------------------------------------------------------------------------
+-- P2P-specific actor state types
+-- ---------------------------------------------------------------------------
+
+-- Evaluator reads the PeerCoordinator ref lazily from a TMVar.
+-- This breaks the circular dependency: evaluator is spawned before PC,
+-- but PC's ref is filled into the TMVar right after PC spawns.
+data P2PEvalState = P2PEvalState
+  { p2eTestData :: LocalData
+  , p2ePcCell   :: TMVar (ActorRef PeerCoordinatorMsg ())
+  , p2eLogRef   :: ActorRef LogMsg ()
+  }
+
+p2pEvalHandler :: Handler EvaluatorMsg P2PEvalState ()
+p2pEvalHandler (EvaluateModel gm) = do
+  st <- state
+  let ws      = unflattenWeights (wData (gmWeights gm))
+      metrics = evaluate ws (p2eTestData st) (gmRound gm)
+  logFn (p2eLogRef st) "Evaluator" INFO $
+    "P2P round " <> show (gmRound gm)
+    <> " acc=" <> show4 (mAccuracy metrics)
+    <> " auc=" <> show4 (mAUCROC  metrics)
+  pcRef <- liftIO $ atomically $ readTMVar (p2ePcCell st)
+  cast (PeerEvaluationDone metrics) pcRef
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+-- Trainer likewise reads the PeerCoordinator ref lazily.
+data P2PTrainerState = P2PTrainerState
+  { p2tId        :: String
+  , p2tData      :: LocalData
+  , p2tLastRound :: Int
+  , p2tPcCell    :: TMVar (ActorRef PeerCoordinatorMsg ())
+  , p2tLogRef    :: ActorRef LogMsg ()
+  , p2tConfig    :: FLConfig
+  }
+
+p2pTrainerHandler :: Handler TrainerMsg P2PTrainerState ()
+p2pTrainerHandler (StartTraining gm) = do
+  st <- state
+  let round_ = gmRound gm
+  if round_ <= p2tLastRound st
+    then pass
+    else do
+      logFn (p2tLogRef st) (p2tId st) INFO $
+        "P2P training round " <> show round_
+      let initW = unflattenWeights (wData (gmWeights gm))
+      updatedW <- liftIO $ trainLocal (p2tConfig st) (p2tId st) round_ initW (p2tData st)
+      let n   = localDataSize (p2tData st)
+          upd = ModelUpdate (p2tId st) round_ n (Weights (flattenWeights updatedW))
+      pcRef <- liftIO $ atomically $ readTMVar (p2tPcCell st)
+      cast (PeerTrainingDone upd) pcRef
+      passWith st { p2tLastRound = round_ }
+p2pTrainerHandler TrainerStop = pass
+
+-- ---------------------------------------------------------------------------
+-- Peer node startup
+-- ---------------------------------------------------------------------------
+
+runPeer :: FLConfig -> Int -> [NodeAddr] -> IO ()
+runPeer cfg peerIdx knownPeers = do
+  let FLConfig{..} = cfg
+      peerPort = cfgBasePort
+      peerId   = "peer-" <> show peerIdx
+
+  rt      <- initRuntime (NodeAddr "localhost" (fromIntegral peerPort)) createTCPTransport
+  doneVar <- newEmptyMVar
+
+  withRuntime rt $ do
+    forM_ knownPeers $ \addr -> retryConnect 20 addr
+    liftIO $ threadDelay 300_000
+
+    logRef <- spawnActor loggerHandler stopOnDeath ()
+    registerActor (peerId <> "-logger") logRef
+
+    initW <- liftIO initWeights
+    let initModel = GlobalModel 0 (Weights (flattenWeights initW))
+
+    (trainData, testData) <- liftIO $ loadAndPartition cfg peerIdx
+
+    -- Shared TMVar: filled with the PeerCoordinator ref after it spawns.
+    -- Evaluator and Trainer read from it lazily on their first message.
+    pcCell <- liftIO newEmptyTMVarIO
+
+    evalRef <- spawnActor p2pEvalHandler stopOnDeath
+      (P2PEvalState testData pcCell logRef)
+
+    trainerRef <- spawnActor p2pTrainerHandler stopOnDeath
+      P2PTrainerState
+        { p2tId        = peerId
+        , p2tData      = trainData
+        , p2tLastRound = -1
+        , p2tPcCell    = pcCell
+        , p2tLogRef    = logRef
+        , p2tConfig    = cfg
+        }
+
+    -- PeerCoordinator — placeholder self-ref (replaced by SetPeers)
+    let dummyId  = ActorId 0 (read "00000000-0000-0000-0000-000000000002")
+        dummySelf = RemoteRef dummyId :: ActorRef PeerCoordinatorMsg ()
+
+    let pcSt = PeerCoordinatorState
+          { pcsNodeId       = peerId
+          , pcsRoundCounter = gcEmpty
+          , pcsParticipants = orEmpty
+          , pcsCurrentRound = 0
+          , pcsTotalRounds  = cfgRounds
+          , pcsGlobalModel  = initModel
+          , pcsPeerRefs     = []
+          , pcsSelfRef      = dummySelf
+          , pcsTrainerRef   = trainerRef
+          , pcsEvalRef      = evalRef
+          , pcsLogRef       = logRef
+          , pcsCollected    = []
+          , pcsNumPeers     = cfgNumTrainers
+          , pcsDoneVar      = doneVar
+          }
+
+    -- Supervised PeerCoordinator. On every (re)spawn the shared cell is
+    -- refreshed and the registry name re-registered, so trainer/evaluator
+    -- keep reaching the live incarnation. (State restarts fresh.)
+    supervise' OneForOne
+      [ ChildSpec
+          { csRun = \target -> do
+              ref <- spawnActor peerCoordinatorHandler peerCoordinatorDeathHandler pcSt
+              linkActorTo target ref
+              return ref
+          , csOnSpawn = \ref -> do
+              liftIO $ atomically $ do
+                void $ tryTakeTMVar pcCell
+                putTMVar pcCell ref
+              registerActor ("peercoord-" <> peerId) ref
+          }
+      ]
+    pcRef <- liftIO $ atomically $ readTMVar pcCell
+
+    logFn' logRef peerId INFO $
+      "Peer ready, waiting for " <> show (cfgNumTrainers - 1) <> " other peers"
+
+    -- Discover all other peer coordinators
+    let otherIds = ["peer-" <> show i | i <- [0 .. cfgNumTrainers - 1], i /= peerIdx]
+    peerPcRefs <- forM otherIds $ \pid ->
+      retryLookup 40
+        (lookupRemoteActor ("peercoord-" <> pid)
+          :: RuntimeM (Maybe (ActorRef PeerCoordinatorMsg ())))
+
+    -- Wire up PeerCoordinator with peers
+    let peerIds = map actorRefId peerPcRefs
+    cast' (SetPeers peerIds) pcRef
+
+    liftIO $ threadDelay 200_000
+    cast' PeerStartRound pcRef
+    logFn' logRef peerId INFO "P2P round 1 started"
+
+  takeMVar doneVar
+  threadDelay 300_000
+  putStrLn $ "\n=== Peer " <> show peerIdx <> " federation complete ==="

src/FL/Provider.hs

@@ -0,0 +1,149 @@
+{-# LANGUAGE RecordWildCards #-}
+module FL.Provider
+  ( runCoordinator
+  , runTrainer
+  ) where
+
+import Control.Actor
+import Control.Actor.Registry (lookupRemoteActor, registerActor)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
+import Control.Concurrent.STM
+  ( atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
+import Control.Monad (void)
+import Control.Monad.IO.Class (liftIO)
+import Data.Map.Strict qualified as Map
+import FL.Actors.Aggregator
+import FL.Actors.Coordinator
+import FL.Actors.Evaluator
+import FL.Actors.Logger
+import FL.Actors.Trainer
+import FL.CRDT (orEmpty)
+import FL.Data.PaySim (loadAndPartition)
+import FL.NN (flattenWeights, initWeights)
+import FL.Retry (retryConnect, retryLookup)
+import FL.Types
+
+-- ---------------------------------------------------------------------------
+-- Coordinator node
+-- ---------------------------------------------------------------------------
+
+runCoordinator :: FLConfig -> IO ()
+runCoordinator cfg = do
+  let FLConfig{..} = cfg
+  rt      <- initRuntime (NodeAddr "localhost" (fromIntegral cfgBasePort)) createTCPTransport
+  doneVar <- newEmptyMVar
+
+  withRuntime rt $ do
+    -- Logger
+    logRef <- spawnActor loggerHandler stopOnDeath ()
+    registerActor "logger" logRef
+
+    -- Initial global model
+    initW <- liftIO initWeights
+    let initModel = GlobalModel 0 (Weights (flattenWeights initW))
+
+    -- Shared TMVar for coordinator ref: filled after coordinator spawns.
+    -- Aggregator and Evaluator read from it lazily.
+    coordCell <- liftIO newEmptyTMVarIO
+
+    -- Load test data (coordinator holds the shared test split)
+    (_, testData) <- liftIO $ loadAndPartition cfg 0
+
+    -- Evaluator
+    evalRef <- spawnActor evaluatorHandler stopOnDeath
+      (EvaluatorState testData coordCell logRef)
+    registerActor "evaluator" evalRef
+
+    -- Aggregator
+    aggRef <- spawnActor aggregatorHandler stopOnDeath
+      (AggregatorState cfgNumTrainers [] coordCell logRef 0)
+    registerActor "aggregator" aggRef
+
+    -- Coordinator (knows aggRef and evalRef at spawn time)
+    let initCoordSt = CoordinatorState
+          { csTrainers         = orEmpty
+          , csTrainerRefs      = Map.empty
+          , csExpectedTrainers = cfgNumTrainers
+          , csCurrentRound     = 0
+          , csTotalRounds      = cfgRounds
+          , csGlobalModel      = initModel
+          , csAggRef           = aggRef
+          , csEvalRef          = evalRef
+          , csLogRef           = logRef
+          , csDoneVar          = doneVar
+          }
+    -- Supervised coordinator. On every (re)spawn — including restarts — the
+    -- shared cell is refreshed and the registry name re-registered, so
+    -- aggregator/evaluator/trainers keep reaching the live incarnation.
+    -- (Coordinator state itself restarts fresh: round progress is lost.)
+    supervise' OneForOne
+      [ ChildSpec
+          { csRun = \target -> do
+              ref <- spawnActor coordinatorHandler coordinatorDeathHandler initCoordSt
+              linkActorTo target ref
+              return ref
+          , csOnSpawn = \ref -> do
+              liftIO $ atomically $ do
+                void $ tryTakeTMVar coordCell
+                putTMVar coordCell ref
+              registerActor "coordinator" ref
+          }
+      ]
+    void $ liftIO $ atomically $ readTMVar coordCell
+
+    logFn' logRef "System" INFO $
+      "Coordinator ready, waiting for " <> show cfgNumTrainers <> " trainers"
+
+  takeMVar doneVar
+  threadDelay 300_000
+  putStrLn "\n=== Provider federation complete ==="
+
+-- ---------------------------------------------------------------------------
+-- Trainer node
+-- ---------------------------------------------------------------------------
+
+runTrainer :: FLConfig -> Int -> NodeAddr -> IO ()
+runTrainer cfg trainerIdx coordAddr = do
+  let FLConfig{..} = cfg
+  let trainerPort = cfgBasePort + trainerIdx + 1
+      trainerId   = "trainer-" <> show trainerIdx
+
+  rt <- initRuntime (NodeAddr "localhost" (fromIntegral trainerPort)) createTCPTransport
+
+  withRuntime rt $ do
+    retryConnect 20 coordAddr
+    liftIO $ threadDelay 500_000  -- wait for registry sync
+
+    -- Look up aggregator and logger on the coordinator node (retry until
+    -- available; the logger is registered first, so if the aggregator
+    -- resolves the logger will too)
+    aggRef <- retryLookup 20
+                (lookupRemoteActor "aggregator" :: RuntimeM (Maybe (ActorRef AggregatorMsg ())))
+    logRef <- retryLookup 20
+                (lookupRemoteActor "logger" :: RuntimeM (Maybe (ActorRef LogMsg ())))
+
+    (trainData, _) <- liftIO $ loadAndPartition cfg trainerIdx
+
+    let trainerSt = TrainerState
+          { tsId        = trainerId
+          , tsData      = trainData
+          , tsLastRound = -1
+          , tsAggRef    = aggRef
+          , tsLogRef    = logRef
+          , tsConfig    = cfg
+          }
+    trainerRef <- spawnActor trainerHandler stopOnDeath trainerSt
+    registerActor trainerId trainerRef
+
+    -- Register with coordinator
+    coordRef <- retryLookup 20
+                  (lookupRemoteActor "coordinator" :: RuntimeM (Maybe (ActorRef CoordinatorMsg ())))
+    cast' (RegisterTrainer trainerId (actorRefId trainerRef)) coordRef
+
+    logFn' logRef trainerId INFO "Registered with coordinator"
+
+  waitForever
+
+waitForever :: IO ()
+waitForever = threadDelay maxBound >> waitForever

src/FL/Retry.hs

@@ -0,0 +1,34 @@
+-- | Startup helpers shared by the provider and P2P entry points: retry a
+-- node connection or a registry lookup until it succeeds or attempts run out.
+module FL.Retry
+  ( retryConnect
+  , retryLookup
+  ) where
+
+import Control.Actor (ActorRef, NodeAddr, RuntimeM, connect, withRuntime)
+import Control.Concurrent (threadDelay)
+import Control.Exception (SomeException, try)
+import Control.Monad (void)
+import Control.Monad.IO.Class (liftIO)
+import Control.Monad.Reader (ask)
+
+retryConnect :: Int -> NodeAddr -> RuntimeM ()
+retryConnect 0 addr = error $ "retryConnect: could not connect to " <> show addr
+retryConnect n addr = do
+  rt     <- ask
+  result <- liftIO $ try @SomeException $ withRuntime rt $ void (connect Nothing addr)
+  case result of
+    Right () -> return ()
+    Left  _  -> do
+      liftIO $ threadDelay 1_000_000
+      retryConnect (n - 1) addr
+
+retryLookup :: Int -> RuntimeM (Maybe (ActorRef msg r)) -> RuntimeM (ActorRef msg r)
+retryLookup 0 _ = error "retryLookup: exhausted retries"
+retryLookup n action = do
+  result <- action
+  case result of
+    Just ref -> return ref
+    Nothing  -> do
+      liftIO $ threadDelay 500_000
+      retryLookup (n - 1) action

src/FL/Types.hs

@@ -0,0 +1,160 @@
+module FL.Types
+  ( -- Config
+    FLConfig (..)
+  , PartitionMode (..)
+  , NormMode (..)
+    -- Wire types
+  , Weights (..)
+  , GlobalModel (..)
+  , ModelUpdate (..)
+  , Metrics (..)
+    -- Logger messages
+  , LogLevel (..)
+  , LogEntry (..)
+  , LogMsg (..)
+    -- Actor messages
+  , EvaluatorMsg (..)
+  , TrainerMsg (..)
+  , AggregatorMsg (..)
+  , CoordinatorMsg (..)
+  , PeerSyncPayload (..)
+  , PeerCoordinatorMsg (..)
+  ) where
+
+import Control.Actor.Types (ActorId)
+import Data.Binary (Binary)
+import FL.CRDT (GCounter, ORSet)
+import GHC.Generics (Generic)
+
+-- ---------------------------------------------------------------------------
+-- Config (not a wire type — parsed from CLI, passed around locally)
+-- ---------------------------------------------------------------------------
+
+data PartitionMode = IID | NonIID
+  deriving (Show, Eq, Generic)
+instance Binary PartitionMode
+
+data NormMode = MinMax | ZScore
+  deriving (Show, Eq, Generic)
+instance Binary NormMode
+
+data FLConfig = FLConfig
+  { cfgRounds       :: Int
+  , cfgEpochs       :: Int
+  , cfgLearningRate :: Double
+  , cfgBatchSize    :: Int
+  , cfgNumTrainers  :: Int
+  , cfgDataPath     :: Maybe FilePath
+  , cfgPartition    :: PartitionMode
+  , cfgNorm         :: NormMode
+  , cfgBasePort     :: Int
+  } deriving (Show, Generic)
+
+-- ---------------------------------------------------------------------------
+-- Wire types (sent over the network, must be Binary)
+-- ---------------------------------------------------------------------------
+
+newtype Weights = Weights { wData :: [Double] }
+  deriving (Show, Eq, Generic)
+instance Binary Weights
+
+data GlobalModel = GlobalModel
+  { gmRound   :: Int
+  , gmWeights :: Weights
+  } deriving (Show, Eq, Generic)
+instance Binary GlobalModel
+
+data ModelUpdate = ModelUpdate
+  { muTrainerId  :: String
+  , muRound      :: Int
+  , muSampleSize :: Int
+  , muWeights    :: Weights
+  } deriving (Show, Eq, Generic)
+instance Binary ModelUpdate
+
+data Metrics = Metrics
+  { mAccuracy  :: Double
+  , mPrecision :: Double
+  , mRecall    :: Double
+  , mF1        :: Double
+  , mAUCROC    :: Double
+  , mRound     :: Int
+  } deriving (Show, Eq, Generic)
+instance Binary Metrics
+
+-- ---------------------------------------------------------------------------
+-- Logger
+-- ---------------------------------------------------------------------------
+
+data LogLevel = DEBUG | INFO | WARN | ERROR
+  deriving (Show, Eq, Ord, Generic)
+instance Binary LogLevel
+
+data LogEntry = LogEntry
+  { leSource :: String
+  , leLevel  :: LogLevel
+  , leMsg    :: String
+  } deriving (Show, Generic)
+instance Binary LogEntry
+
+data LogMsg = Log LogEntry
+  deriving (Show, Generic)
+instance Binary LogMsg
+
+-- ---------------------------------------------------------------------------
+-- Evaluator
+-- ---------------------------------------------------------------------------
+
+data EvaluatorMsg = EvaluateModel GlobalModel
+  deriving (Show, Generic)
+instance Binary EvaluatorMsg
+
+-- ---------------------------------------------------------------------------
+-- Trainer
+-- ---------------------------------------------------------------------------
+
+data TrainerMsg
+  = StartTraining GlobalModel
+  | TrainerStop
+  deriving (Show, Generic)
+instance Binary TrainerMsg
+
+-- ---------------------------------------------------------------------------
+-- Aggregator (provider mode)
+-- ---------------------------------------------------------------------------
+
+data AggregatorMsg = SubmitUpdate ModelUpdate
+  deriving (Show, Generic)
+instance Binary AggregatorMsg
+
+-- ---------------------------------------------------------------------------
+-- Coordinator (provider mode)
+-- ---------------------------------------------------------------------------
+
+data CoordinatorMsg
+  = RegisterTrainer String ActorId
+  | AggregationComplete GlobalModel
+  | EvaluationDone Metrics
+  | StartFederation
+  deriving (Show, Generic)
+instance Binary CoordinatorMsg
+
+-- ---------------------------------------------------------------------------
+-- PeerCoordinator (P2P mode)
+-- ---------------------------------------------------------------------------
+
+data PeerSyncPayload = PeerSyncPayload
+  { pspCounter      :: !GCounter
+  , pspParticipants :: !(ORSet String)
+  , pspUpdate       :: !(Maybe ModelUpdate)
+  } deriving (Show, Generic)
+instance Binary PeerSyncPayload
+
+data PeerCoordinatorMsg
+  = SetPeers ![ActorId]
+  | PeerStartRound
+  | PeerTrainingDone !ModelUpdate
+  | PeerSync !PeerSyncPayload
+  | PeerEvaluationDone !Metrics
+  deriving (Show, Generic)
+instance Binary PeerCoordinatorMsg