Update comments to reflect what we're doing

Author: hyperthunk

distributed-process-tests/src/Control/Distributed/Process/Tests/CH.hs

@@ -39,6 +39,7 @@ import Control.Distributed.Process.Internal.Types
 import Control.Distributed.Process.Node
 import Control.Distributed.Process.Debug
 import Control.Distributed.Process.Management.Internal.Types
+import Control.Distributed.Process.Tests.Internal.Utils (shouldBe)
 import Control.Distributed.Process.Serializable (Serializable)
 
 import Test.HUnit (Assertion, assertFailure, assertBool)
@@ -1404,6 +1405,9 @@ testRegistryMonitoring TestTransport{..} = do
         w <- whereis "test" :: Process (Maybe ProcessId)
         forM_ w (const waitpoll)
     in do register "test" remote1
+          -- [race-condition] invalid test if we force exit before waitpoll runs
+          registered <- whereis "test"
+          registered `shouldBe` (equalTo (Just remote1))
           send remote1 ()
           waitpoll
           return ()
@@ -1412,15 +1416,19 @@ testRegistryMonitoring TestTransport{..} = do
   -- Many labels. Test if all labels associated with process
   -- are removed from registry when it dies.
   remote2 <- testRemote remoteNode
+  remoteX <- testRemote remoteNode
   runProcess localNode $
     let waitpoll = do
         w1 <- whereis "test-3" :: Process (Maybe ProcessId)
         w2 <- whereis "test-4" :: Process (Maybe ProcessId)
         forM_ (w1 <|> w2) (const waitpoll)
     in do register "test-3" remote2
           register "test-4" remote2
+          register "test-X" remoteX
           send remote2 ()
           waitpoll
+          rX <- whereis "test-X"
+          rX `shouldBe` (equalTo (Just remoteX))
           return ()
 
 {- XXX: waiting including patch for nsend for remote process

distributed-process-tests/src/Control/Distributed/Process/Tests/Mx.hs

@@ -176,14 +176,14 @@ testAgentEventHandling result = do
 
   stash result $ seenAlive && seenDead
 
-testMxMonitorEvents :: Process ()
-testMxMonitorEvents = do
+testMxRegEvents :: Process ()
+testMxRegEvents = do
 
   {- This test only deals with the local case, to ensure that we are being
      notified in the expected order - the remote cases related to the
      behaviour of the node controller are contained in the CH test suite. -}
 
-  let label = "testMxMonitorEvents"
+  let label = "testMxRegEvents"
   let agentLabel = "listener-agent"
   let delay = 1000000
   (regChan, regSink) <- newChan
@@ -223,9 +223,65 @@ testMxMonitorEvents = do
 
   mapM_ (flip kill $ "test-complete") [agent, p1, p2]
 
+testMxRegMon :: LocalNode -> Process ()
+testMxRegMon remoteNode = do
+
+  -- ensure that when a registered process dies, we get a notification that
+  -- it has been unregistered as well as seeing the name get removed
+
+  let label1 = "aaaaa"
+  let label2 = "bbbbb"
+  let isValid l = l ==label1 || l == label2
+  let agentLabel = "listener-agent"
+  let delay = 1000000
+  (regChan, regSink) <- newChan
+  (unRegChan, unRegSink) <- newChan
+  agent <- mxAgent (MxAgentId agentLabel) () [
+      mxSink $ \ev -> do
+        case ev of
+          MxRegistered pid label
+            | isValid label -> liftMX $ sendChan regChan (label, pid)
+          MxUnRegistered pid label
+            | isValid label -> liftMX $ sendChan unRegChan (label, pid)
+          _                 -> return ()
+        mxReady
+    ]
+
+  (sp, rp) <- newChan
+  liftIO $ forkProcess remoteNode $ do
+    getSelfPid >>= sendChan sp
+    expect :: Process ()
+
+  p1 <- receiveChan rp
+
+  register label1 p1
+  reg1 <- receiveChanTimeout delay regSink
+  reg1 `shouldBe` equalTo (Just (label1, p1))
+
+  register label2 p1
+  reg2 <- receiveChanTimeout delay regSink
+  reg2 `shouldBe` equalTo (Just (label2, p1))
+
+  n1 <- whereis label1
+  n1 `shouldBe` equalTo (Just p1)
+
+  n2 <- whereis label2
+  n2 `shouldBe` equalTo (Just p1)
+
+  kill p1 "goodbye"
+
+  unreg1 <- receiveChanTimeout delay unRegSink
+  unreg2 <- receiveChanTimeout delay unRegSink
+
+  sort [unreg1, unreg2]
+    `shouldBe` equalTo [Just (label1, p1), Just (label2, p1)]
+
+  kill agent "test-complete"
+
 tests :: TestTransport -> IO [Test]
 tests TestTransport{..} = do
   node1 <- newLocalNode testTransport initRemoteTable
+  node2 <- newLocalNode testTransport initRemoteTable
   return [
       testGroup "Mx Agents" [
           testCase "Event Handling"
@@ -252,7 +308,9 @@ tests TestTransport{..} = do
                             "fifth"]) testAgentPrioritisation)
       ]
     , testGroup "Mx Events" [
-        testCase "Monitor Events"
-          (runProcess node1 testMxMonitorEvents)
+        testCase "Name Registration Events"
+          (runProcess node1 testMxRegEvents)
+      , testCase "Post Death Name UnRegistration Events"
+          (runProcess node1 (testMxRegMon node2))
       ]
     ]

src/Control/Distributed/Process/Node.hs

@@ -34,7 +34,7 @@ import qualified Data.Map as Map
   , filter
   , insert
   , lookup
-  , insertLookupWithKey
+  , partition
   , partitionWithKey
   , elems
   , size
@@ -328,28 +328,23 @@ startDefaultTracer node' = do
 registryMonitorAgentId :: MxAgentId
 registryMonitorAgentId = MxAgentId "service.registry.monitoring"
 
--- note [registry monitoring agent]
--- this agent listens for 'MxRegistered' and 'MxUnRegistered' events and tracks
--- all remote 'ProcessId's that are stored in the registry.
---
--- When a remote process is registered, the agent starts monitoring it until it
--- is unregistered or the monitor notification arrives.
---
--- The agent keeps the amount of labels associated with each registered remote
--- process. This is necessary so the process is unmonitored only when it has
--- been unregistered from all of the labels.
---
+{- note [registry monitoring agent]
+   This agent listens for 'MxRegistered' and 'MxUnRegistered' events and tracks
+   all labels for remote 'ProcessId's that are stored in the registry.
+
+   When a remote process is registered, the agent starts monitoring it until it
+   is unregistered. We can safely ignore ProcessMonitorNotification's, since the
+   node controller is obligated to issue an MxUnRegistered event for any labels
+   connected to a dead process. Thus, our only responsibility here is to ensure
+   that we set up monitors for locally registered processes, and tear them down
+   once unregistration occurs.
+ -}
 
 registryMonitorAgent :: Process ProcessId
 registryMonitorAgent = do
   nid <- getSelfNode
-  -- For each process the map associates the 'MonitorRef' used to monitor it and
-  -- the amount of labels associated with it.
   mxAgent registryMonitorAgentId (Map.empty :: Map String MonitorRef)
-    [ mxSink $ \(ProcessMonitorNotification mref _ _) -> do
-        mxUpdateLocal $ Map.filter (== mref)
-        mxReady
-    , mxSink $ \ev -> do
+    [ mxSink $ \ev -> do
         case ev of
           MxRegistered pid label
             | processNodeId pid /= nid -> do
@@ -363,9 +358,9 @@ registryMonitorAgent = do
                 mxUpdateLocal (Map.delete label)
           _ -> return ()
         mxReady
-      -- remove async answers from mailbox
-    , mxSink $ \RegisterReply{} -> mxReady
-    , mxSink $ \DidUnmonitor{} -> mxReady
+      -- [note: no need to remove async answers from mailbox]
+      -- The framework simply discards any input you don't have a handler for.
+      -- See Management.hs `runAgent` for details.
     ]
 
 -- | Start and register the service processes on a node
@@ -961,7 +956,11 @@ ncEffectDied ident reason = do
 
   modify' $ (links ^= unaffectedLinks') . (monitors ^= unaffectedMons')
 
-  modify' $ registeredHere ^: Map.filter (\pid -> not $ ident `impliesDeathOf` ProcessIdentifier pid)
+  -- we now consider all labels for this identifier unregistered
+  let toDrop pid = not $ ident `impliesDeathOf` ProcessIdentifier pid
+  (keepNames, dropNames) <- Map.partition toDrop <$> gets (^. registeredHere)
+  mapM_ (\(p, l) -> liftIO $ trace node (MxUnRegistered l p)) (Map.toList dropNames)
+  modify' $ registeredHere ^= keepNames
 
   remaining <- fmap Map.toList (gets (^. registeredOnNodes)) >>=
       mapM (\(pid,nidlist) ->