@@ -48,10 +48,12 @@ import qualified Data.Set as Set
4848 ( empty
4949 , insert
5050 , delete
51+ , map
5152 , member
5253 , toList
5354 )
5455import Data.Foldable (forM_)
56+import Data.List (foldl')
5557import Data.Maybe (isJust, fromJust, isNothing, catMaybes)
5658import Data.Typeable (Typeable)
5759import Control.Category ((>>>))
@@ -77,6 +79,8 @@ import qualified Control.Exception as Exception
7779 , finally
7880 )
7981import Control.Concurrent (forkIO, killThread)
82+import Control.Distributed.Process.Internal.BiMultiMap (BiMultiMap)
83+import qualified Control.Distributed.Process.Internal.BiMultiMap as BiMultiMap
8084import Control.Distributed.Process.Internal.StrictMVar
8185 ( newMVar
8286 , withMVar
@@ -656,9 +660,9 @@ runNodeController node =
656660
657661data NCState = NCState
658662 { -- Mapping from remote processes to linked local processes
659− _links :: !(Map Identifier (Set ProcessId))
663+ _links :: !(BiMultiMap Identifier ProcessId ())
660664 -- Mapping from remote processes to monitoring local processes
661− , _monitors :: !(Map Identifier (Set (ProcessId, MonitorRef)))
665+ , _monitors :: !(BiMultiMap Identifier ProcessId MonitorRef)
662666 -- Process registry: names and where they live, mapped to the PIDs
663667 , _registeredHere :: !(Map String ProcessId)
664668 , _registeredOnNodes :: !(Map ProcessId [(NodeId,Int)])
@@ -674,8 +678,8 @@ newtype NC a = NC { unNC :: StateT NCState (ReaderT LocalNode IO) a }
674678 )
675679
676680initNCState :: NCState
677−initNCState = NCState { _links = Map.empty
678− , _monitors = Map.empty
681+initNCState = NCState { _links = BiMultiMap.empty
682+ , _monitors = BiMultiMap.empty
679683 , _registeredHere = Map.empty
680684 , _registeredOnNodes = Map.empty
681685 }
@@ -818,8 +822,8 @@ ncEffectMonitor from them mRef = do
818822 case (shouldLink, isLocal node (ProcessIdentifier from)) of
819823 (True, _) -> -- [Unified: first rule]
820824 case mRef of
821− Just ref -> modify' $ monitorsFor them ^: Set.insert (from, ref)
822− Nothing -> modify' $ linksFor them ^: Set.insert from
825+ Just ref -> modify' $ monitors ^: BiMultiMap.insert them from ref
826+ Nothing -> modify' $ links ^: BiMultiMap.insert them from ()
823827 (False, True) -> -- [Unified: second rule]
824828 notifyDied from them DiedUnknownId mRef
825829 (False, False) -> -- [Unified: third rule]
@@ -843,7 +847,7 @@ ncEffectUnlink from them = do
843847 postAsMessage from $ DidUnlinkNode nid
844848 SendPortIdentifier cid ->
845849 postAsMessage from $ DidUnlinkPort cid
846− modify' $ linksFor them ^: Set.delete from
850+ modify' $ links ^: BiMultiMap.delete them from ()
847851
848852-- [Unified: Table 11]
849853ncEffectUnmonitor :: ProcessId -> MonitorRef -> NC ()
@@ -851,7 +855,7 @@ ncEffectUnmonitor from ref = do
851855 node <- ask
852856 when (isLocal node (ProcessIdentifier from)) $
853857 postAsMessage from $ DidUnmonitor ref
854− modify' $ monitorsFor (monitorRefIdent ref) ^: Set.delete (from, ref)
858+ modify' $ monitors ^: BiMultiMap.delete (monitorRefIdent ref) from ref
855859
856860-- [Unified: Table 12]
857861ncEffectDied :: Identifier -> DiedReason -> NC ()
@@ -866,7 +870,7 @@ ncEffectDied ident reason = do
866870 let localOnly = case ident of NodeIdentifier _ -> True ; _ -> False
867871
868872 forM_ (Map.toList affectedLinks) $ \(them, uss) ->
869− forM_ uss $ \us ->
873+ forM_ uss $ \(us, _) ->
870874 when (localOnly <= isLocal node (ProcessIdentifier us)) $
871875 notifyDied us them reason Nothing
872876
@@ -875,7 +879,15 @@ ncEffectDied ident reason = do
875879 when (localOnly <= isLocal node (ProcessIdentifier us)) $
876880 notifyDied us them reason (Just ref)
877881
878− modify' $ (links ^= unaffectedLinks) . (monitors ^= unaffectedMons)
882+ let (unaffectedLinks', unaffectedMons') =
883+ case ident of
884+ ProcessIdentifier pid ->
885+ ( BiMultiMap.deleteAllBy2nd pid unaffectedLinks
886+ , BiMultiMap.deleteAllBy2nd pid unaffectedMons
887+ )
888+ _ -> (unaffectedLinks, unaffectedMons)
889+
890+ modify' $ (links ^= unaffectedLinks') . (monitors ^= unaffectedMons')
879891
880892 modify' $ registeredHere ^: Map.filter (\pid -> not $ ident `impliesDeathOf` ProcessIdentifier pid)
881893
@@ -1045,8 +1057,9 @@ ncEffectGetInfo from pid =
10451057 Nothing -> dispatch (isLocal node (ProcessIdentifier from))
10461058 from (ProcessInfoNone DiedUnknownId)
10471059 Just proc -> do
1048− itsLinks <- gets (^. linksFor them)
1049− itsMons <- gets (^. monitorsFor them)
1060+ itsLinks <- Set.map fst . BiMultiMap.lookup them <$>
1061+ gets (^. links)
1062+ itsMons <- BiMultiMap.lookup them <$> gets (^. monitors)
10501063 registered <- gets (^. registeredHere)
10511064 size <- liftIO $ queueSize $ processQueue $ proc
10521065
@@ -1082,8 +1095,8 @@ ncEffectGetNodeStats from _nid = do
10821095 NodeStats {
10831096 nodeStatsNode = localNodeId node
10841097 , nodeStatsRegisteredNames = Map.size $ ncState ^. registeredHere
1085− , nodeStatsMonitors = Map.size $ ncState ^. monitors
1086− , nodeStatsLinks = Map.size $ ncState ^. links
1098+ , nodeStatsMonitors = BiMultiMap.size $ ncState ^. monitors
1099+ , nodeStatsLinks = BiMultiMap.size $ ncState ^. links
10871100 , nodeStatsProcesses = Map.size (nodeState ^. localProcesses)
10881101 }
10891102 postAsMessage from stats
@@ -1204,10 +1217,10 @@ withLocalProc node pid p =
12041217-- Accessors --
12051218--------------------------------------------------------------------------------
12061219
1207−links :: Accessor NCState (Map Identifier (Set ProcessId))
1220+links :: Accessor NCState (BiMultiMap Identifier ProcessId ())
12081221links = accessor _links (\ls st -> st { _links = ls })
12091222
1210−monitors :: Accessor NCState (Map Identifier (Set (ProcessId, MonitorRef)))
1223+monitors :: Accessor NCState (BiMultiMap Identifier ProcessId MonitorRef)
12111224monitors = accessor _monitors (\ms st -> st { _monitors = ms })
12121225
12131226registeredHere :: Accessor NCState (Map String ProcessId)
@@ -1216,12 +1229,6 @@ registeredHere = accessor _registeredHere (\ry st -> st { _registeredHere = ry }
12161229registeredOnNodes :: Accessor NCState (Map ProcessId [(NodeId, Int)])
12171230registeredOnNodes = accessor _registeredOnNodes (\ry st -> st { _registeredOnNodes = ry })
12181231
1219−linksFor :: Identifier -> Accessor NCState (Set ProcessId)
1220−linksFor ident = links >>> DAC.mapDefault Set.empty ident
1221−
1222−monitorsFor :: Identifier -> Accessor NCState (Set (ProcessId, MonitorRef))
1223−monitorsFor ident = monitors >>> DAC.mapDefault Set.empty ident
1224−
12251232registeredHereFor :: String -> Accessor NCState (Maybe ProcessId)
12261233registeredHereFor ident = registeredHere >>> DAC.mapMaybe ident
12271234
@@ -1248,10 +1255,12 @@ registeredOnNodesFor ident = registeredOnNodes >>> DAC.mapMaybe ident
12481255-- * the notifications for typed channels to that process.
12491256--
12501257-- See https://github.com/haskell/containers/issues/14 for the bang on _v.
1251−splitNotif :: Identifier
1252− -> Map Identifier a
1253− -> (Map Identifier a, Map Identifier a)
1254−splitNotif ident = Map.partitionWithKey (\k !_v -> ident `impliesDeathOf` k)
1258+splitNotif :: (Ord a, Ord v)
1259+ => Identifier
1260+ -> BiMultiMap Identifier a v
1261+ -> (Map Identifier (Set (a,v)), BiMultiMap Identifier a v)
1262+splitNotif ident =
1263+ BiMultiMap.partitionWithKey (\k !_v -> ident `impliesDeathOf` k)
12551264
12561265--------------------------------------------------------------------------------
12571266-- Auxiliary --