accelerated-dht: cleanup peer from message sender on disconnection

fullrt/dht.go

@@ -15,12 +15,14 @@
 	"github.com/multiformats/go-multihash"
 
 	"github.com/libp2p/go-libp2p-routing-helpers/tracing"
+	"github.com/libp2p/go-libp2p/core/event"
 	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/core/peerstore"
 	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/libp2p/go-libp2p/core/routing"
+	"github.com/libp2p/go-libp2p/p2p/host/eventbus"
 	swarm "github.com/libp2p/go-libp2p/p2p/net/swarm"
 
 	"github.com/gogo/protobuf/proto"
@@ -98,6 +100,8 @@
 	bulkSendParallelism int
 
 	self peer.ID
+
+	peerConnectednessSubscriber event.Subscription
 }
 
 // NewFullRT creates a DHT client that tracks the full network. It takes a protocol prefix for the given network,
@@ -151,6 +155,11 @@
 		}
 	}
 
+	sub, err := h.EventBus().Subscribe(new(event.EvtPeerConnectednessChanged), eventbus.Name("fullrt-dht"))
+	if err != nil {
+		return nil, fmt.Errorf("peer connectedness subscription failed: %w", err)
+	}
+
 	ctx, cancel := context.WithCancel(context.Background())
 
 	self := h.ID()
@@ -195,14 +204,14 @@
 
 		crawlerInterval: fullrtcfg.crawlInterval,
 
-		bulkSendParallelism: fullrtcfg.bulkSendParallelism,
-
-		self: self,
+		bulkSendParallelism:         fullrtcfg.bulkSendParallelism,
+		self:                        self,
+		peerConnectednessSubscriber: sub,
 	}
 
-	rt.wg.Add(1)
+	rt.wg.Add(2)
 	go rt.runCrawler(ctx)
-
+	go rt.runSubscriber()
 	return rt, nil
 }
 
@@ -211,6 +220,29 @@
 	key   kadkey.Key
 }
 
+func (dht *FullRT) runSubscriber() {
+	defer dht.wg.Done()
+	ms, ok := dht.messageSender.(dht_pb.MessageSenderWithDisconnect)
+	if !ok {
+		return
+	}
+	defer dht.peerConnectednessSubscriber.Close()
+	for {
+		select {
+		case e := <-dht.peerConnectednessSubscriber.Out():
+			pc, ok := e.(event.EvtPeerConnectednessChanged)
+			if !ok {
+				logger.Errorf("invalid event message type: %T", e)
+			}
+			if pc.Connectedness != network.Connected {
+				ms.OnDisconnect(dht.ctx, pc.Peer)
+			}
+		case <-dht.ctx.Done():
+			return
+		}
+	}
+}
+
 func (dht *FullRT) TriggerRefresh(ctx context.Context) error {
 	select {
 	case <-ctx.Done():