From a7ad72ea88a28330d5f52eb7ffbba098ef03086c Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 6 Jan 2023 08:48:23 +1300
Subject: [PATCH] identify: cache the snapshot

---
 p2p/host/basic/basic_host.go |  2 +-
 p2p/protocol/identify/id.go  | 35 +++++++++++++++++++++++++----------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/p2p/host/basic/basic_host.go b/p2p/host/basic/basic_host.go
index 2a500151fc..d4040f272f 100644
--- a/p2p/host/basic/basic_host.go
+++ b/p2p/host/basic/basic_host.go
@@ -181,7 +181,7 @@ func NewHost(n network.Network, opts *HostOpts) (*BasicHost, error) {
 
 	h.updateLocalIpAddr()
 
-	if h.emitters.evtLocalProtocolsUpdated, err = h.eventbus.Emitter(&event.EvtLocalProtocolsUpdated{}); err != nil {
+	if h.emitters.evtLocalProtocolsUpdated, err = h.eventbus.Emitter(&event.EvtLocalProtocolsUpdated{}, eventbus.Stateful); err != nil {
 		return nil, err
 	}
 	if h.emitters.evtLocalAddrsUpdated, err = h.eventbus.Emitter(&event.EvtLocalAddressesUpdated{}, eventbus.Stateful); err != nil {
diff --git a/p2p/protocol/identify/id.go b/p2p/protocol/identify/id.go
index 2a1af713c7..0acd2a8050 100644
--- a/p2p/protocol/identify/id.go
+++ b/p2p/protocol/identify/id.go
@@ -113,6 +113,11 @@ type idService struct {
 		evtPeerIdentificationFailed    event.Emitter
 	}
 
+	currentSnapshot struct {
+		sync.Mutex
+		snapshot *identifySnapshot
+	}
+
 	addPeerHandlerCh chan addPeerHandlerReq
 	rmPeerHandlerCh  chan rmPeerHandlerReq
 
@@ -180,6 +185,8 @@ func NewIDService(h host.Host, opts ...Option) (*idService, error) {
 	// register protocols that do not depend on peer records.
 	h.SetStreamHandler(ID, s.sendIdentifyResp)
 	h.SetStreamHandler(IDPush, s.pushHandler)
+	// do this after adding the stream handlers, so that these protocols will be included
+	s.updateSnapshot()
 
 	h.Network().Notify((*netNotifiee)(s))
 	return s, nil
@@ -260,6 +267,7 @@ func (ids *idService) loop() {
 			if !more {
 				return
 			}
+			ids.updateSnapshot()
 			switch e.(type) {
 			case event.EvtLocalAddressesUpdated:
 				for pid := range phs {
@@ -457,20 +465,29 @@ func readAllIDMessages(r pbio.Reader, finalMsg proto.Message) error {
 	return fmt.Errorf("too many parts")
 }
 
-func (ids *idService) getSnapshot() *identifySnapshot {
-	snapshot := new(identifySnapshot)
+func (ids *idService) updateSnapshot() {
+	log.Debug("updating Identify snapshot")
+	snapshot := &identifySnapshot{
+		addrs:     ids.Host.Addrs(),
+		protocols: ids.Host.Mux().Protocols(),
+	}
 	if !ids.disableSignedPeerRecord {
 		if cab, ok := peerstore.GetCertifiedAddrBook(ids.Host.Peerstore()); ok {
 			snapshot.record = cab.GetPeerRecord(ids.Host.ID())
 		}
 	}
-	snapshot.addrs = ids.Host.Addrs()
-	snapshot.protocols = ids.Host.Mux().Protocols()
-	return snapshot
+
+	ids.currentSnapshot.Lock()
+	defer ids.currentSnapshot.Unlock()
+	ids.currentSnapshot.snapshot = snapshot
 }
 
 func (ids *idService) writeChunkedIdentifyMsg(c network.Conn, s network.Stream) error {
-	snapshot := ids.getSnapshot()
+	ids.currentSnapshot.Lock()
+	snapshot := ids.currentSnapshot.snapshot
+	ids.currentSnapshot.Unlock()
+	log.Debugw("sending snapshot with protocols", "protos", snapshot.protocols)
+
 	mes := ids.createBaseIdentifyResponse(c, snapshot)
 	sr := ids.getSignedRecord(snapshot)
 	mes.SignedPeerRecord = sr
@@ -479,15 +496,13 @@ func (ids *idService) writeChunkedIdentifyMsg(c network.Conn, s network.Stream)
 	if sr == nil || proto.Size(mes) <= legacyIDSize {
 		return writer.WriteMsg(mes)
 	}
+
 	mes.SignedPeerRecord = nil
 	if err := writer.WriteMsg(mes); err != nil {
 		return err
 	}
-
 	// then write just the signed record
-	m := &pb.Identify{SignedPeerRecord: sr}
-	err := writer.WriteMsg(m)
-	return err
+	return writer.WriteMsg(&pb.Identify{SignedPeerRecord: sr})
 }
 
 func (ids *idService) createBaseIdentifyResponse(