Fix Custody Columns (#14021)

prysmaticlabs · Jul 18, 2024 · 83e1c5b · 83e1c5b
1 parent 085fd62
commit 83e1c5b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/beacon-chain/p2p/custody.go b/beacon-chain/p2p/custody.go
@@ -63,8 +63,7 @@ func (s *Service) CustodyCountFromRemotePeer(pid peer.ID) (uint64, error) {
 		if err := peerRecord.Load(&custodyObj); err != nil {
 			return 0, errors.Wrap(err, "load custody_subnet_count")
 		}
-		actualCustodyCount := ssz.UnmarshallUint64(custodyBytes)
-
+		actualCustodyCount := ssz.UnmarshallUint64(custodyObj)
 		if actualCustodyCount > peerCustodiedSubnetCount {
 			peerCustodiedSubnetCount = actualCustodyCount
 		}

diff --git a/beacon-chain/p2p/discovery.go b/beacon-chain/p2p/discovery.go
@@ -371,6 +371,8 @@ func (s *Service) filterPeer(node *enode.Node) bool {
 
 	// Ignore nodes that are already active.
 	if s.peers.IsActive(peerData.ID) {
+		// Constantly update enr for known peers
+		s.peers.UpdateENR(node.Record(), peerData.ID)
 		return false
 	}
 

diff --git a/beacon-chain/p2p/peers/status.go b/beacon-chain/p2p/peers/status.go
@@ -159,6 +159,14 @@ func (p *Status) Add(record *enr.Record, pid peer.ID, address ma.Multiaddr, dire
 	p.addIpToTracker(pid)
 }
 
+func (p *Status) UpdateENR(record *enr.Record, pid peer.ID) {
+	p.store.Lock()
+	defer p.store.Unlock()
+	if peerData, ok := p.store.PeerData(pid); ok {
+		peerData.Enr = record
+	}
+}
+
 // Address returns the multiaddress of the given remote peer.
 // This will error if the peer does not exist.
 func (p *Status) Address(pid peer.ID) (ma.Multiaddr, error) {