fix(libp2p): registrar topology notification fixes

Author: joeltg

packages/libp2p/src/registrar.ts

@@ -9,6 +9,7 @@ import type { PeerStore } from '@libp2p/interface/peer-store'
 import type { Topology } from '@libp2p/interface/topology'
 import type { ConnectionManager } from '@libp2p/interface-internal/connection-manager'
 import type { StreamHandlerOptions, StreamHandlerRecord, Registrar, StreamHandler } from '@libp2p/interface-internal/registrar'
+import { PeerSet } from '@libp2p/peer-collections'
 
 const log = logger('libp2p:registrar')
 
@@ -26,11 +27,13 @@ export interface RegistrarComponents {
  * Responsible for notifying registered protocols of events in the network.
  */
 export class DefaultRegistrar implements Registrar {
+  private readonly topologyPeers: Map<string, PeerSet>
   private readonly topologies: Map<string, Map<string, Topology>>
   private readonly handlers: Map<string, StreamHandlerRecord>
   private readonly components: RegistrarComponents
 
   constructor (components: RegistrarComponents) {
+    this.topologyPeers = new Map()
     this.topologies = new Map()
     this.handlers = new Map()
     this.components = components
@@ -130,6 +133,7 @@ export class DefaultRegistrar implements Registrar {
     }
 
     topologies.set(id, topology)
+    this.topologyPeers.set(id, new PeerSet())
 
     return id
   }
@@ -138,6 +142,7 @@ export class DefaultRegistrar implements Registrar {
    * Unregister topology
    */
   unregister (id: string): void {
+    this.topologyPeers.delete(id)
     for (const [protocol, topologies] of this.topologies.entries()) {
       if (topologies.has(id)) {
         topologies.delete(id)
@@ -185,9 +190,8 @@ export class DefaultRegistrar implements Registrar {
    */
   _onPeerUpdate (evt: CustomEvent<PeerUpdate>): void {
     const { peer, previous } = evt.detail
-    const removed = (previous?.protocols ?? []).filter(protocol => !peer.protocols.includes(protocol))
-    const added = peer.protocols.filter(protocol => !(previous?.protocols ?? []).includes(protocol))
 
+    const removed = (previous?.protocols ?? []).filter(protocol => !peer.protocols.includes(protocol))
     for (const protocol of removed) {
       const topologies = this.topologies.get(protocol)
 
@@ -196,26 +200,47 @@ export class DefaultRegistrar implements Registrar {
         continue
       }
 
-      for (const topology of topologies.values()) {
-        topology.onDisconnect?.(peer.id)
+      for (const [id, topology] of topologies.entries()) {
+        const peers = this.topologyPeers.get(id)
+        if (peers == null) {
+          continue
+        }
+
+        for (const peer of peers) {
+          topology.onDisconnect?.(peer)
+        }
+
+        peers.clear()
       }
     }
 
-    for (const protocol of added) {
+    const connections = this.components.connectionManager.getConnections(peer.id)
+    if (connections.length === 0) {
+      return
+    }
+
+    const nonTransientConnection = connections.find((connection) => connection.transient === false)
+    for (const protocol of peer.protocols) {
       const topologies = this.topologies.get(protocol)
 
       if (topologies == null) {
         // no topologies are interested in this protocol
         continue
       }
 
-      for (const connection of this.components.connectionManager.getConnections(peer.id)) {
-        for (const topology of topologies.values()) {
-          if (connection.transient && topology.notifyOnTransient !== true) {
-            continue
-          }
+      
+      for (const [id, topology] of topologies.entries()) {
+        const peers = this.topologyPeers.get(id)
+        if (peers?.has(peer.id)) {
+          continue
+        }
 
-          topology.onConnect?.(peer.id, connection)
+        if (topology.notifyOnTransient) {
+          topology.onConnect?.(peer.id, connections[0])
+          peers?.add(peer.id)
+        } else if (nonTransientConnection != null) {
+          topology.onConnect?.(peer.id, nonTransientConnection)
+          peers?.add(peer.id)
         }
       }
     }