network: fix peerstore Get/Put races

network/p2p/peerstore/peerstore.go

@@ -42,6 +42,7 @@ type PeerStore struct {
 	peerStoreCAB
 	connectionsRateLimitingCount  uint
 	connectionsRateLimitingWindow time.Duration
+	lock                          deadlock.Mutex
 }
 
 // addressData: holds the information associated with each phonebook address.
@@ -55,7 +56,6 @@ type addressData struct {
 
 	// networkNames: lists the networks to which the given address belongs.
 	networkNames map[string]bool
-	mu           *deadlock.RWMutex
 
 	// roles is the roles that this address serves.
 	roles phonebook.RoleSet
@@ -123,8 +123,10 @@ func (ps *PeerStore) UpdateRetryAfter(addr string, retryAfter time.Time) {
 // The provisional time should be updated after the connection with UpdateConnectionTime
 func (ps *PeerStore) GetConnectionWaitTime(addrOrPeerID string) (bool, time.Duration, time.Time) {
 	curTime := time.Now()
-	var timeSince time.Duration
-	var numElmtsToRemove int
+
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
 	peerID := peer.ID(addrOrPeerID)
 	metadata, err := ps.Get(peerID, addressDataKey)
 	if err != nil {
@@ -134,7 +136,10 @@ func (ps *PeerStore) GetConnectionWaitTime(addrOrPeerID string) (bool, time.Dura
 	if !ok {
 		return false, 0 /* not used */, curTime /* not used */
 	}
+
 	// Remove from recentConnectionTimes the times later than ConnectionsRateLimitingWindowSeconds
+	numElmtsToRemove := 0
+	timeSince := time.Duration(0)
 	for numElmtsToRemove < len(ad.recentConnectionTimes) {
 		timeSince = curTime.Sub(ad.recentConnectionTimes[numElmtsToRemove])
 		if timeSince >= ps.connectionsRateLimitingWindow {
@@ -146,6 +151,7 @@ func (ps *PeerStore) GetConnectionWaitTime(addrOrPeerID string) (bool, time.Dura
 
 	// Remove the expired elements from e.data[addr].recentConnectionTimes
 	ps.popNElements(numElmtsToRemove, peerID)
+
 	// If there are max number of connections within the time window, wait
 	metadata, _ = ps.Get(peerID, addressDataKey)
 	ad, ok = metadata.(addressData)
@@ -157,7 +163,6 @@ func (ps *PeerStore) GetConnectionWaitTime(addrOrPeerID string) (bool, time.Dura
 		return true, /* true */
 			ps.connectionsRateLimitingWindow - timeSince, curTime /* not used */
 	}
-
 	// Else, there is space in connectionsRateLimitingCount. The
 	// connection request of the caller will proceed
 	// Update curTime, since it may have significantly changed if waited
@@ -169,6 +174,9 @@ func (ps *PeerStore) GetConnectionWaitTime(addrOrPeerID string) (bool, time.Dura
 
 // UpdateConnectionTime updates the connection time for the given address.
 func (ps *PeerStore) UpdateConnectionTime(addrOrPeerID string, provisionalTime time.Time) bool {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
 	peerID := peer.ID(addrOrPeerID)
 	metadata, err := ps.Get(peerID, addressDataKey)
 	if err != nil {
@@ -206,6 +214,9 @@ func (ps *PeerStore) UpdateConnectionTime(addrOrPeerID string, provisionalTime t
 // existing items that aren't included in addressesThey are being removed
 // matching entries roles gets updated as needed and persistent peers are not touched
 func (ps *PeerStore) ReplacePeerList(addressesThey []*peer.AddrInfo, networkName string, role phonebook.Role) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
 	// prepare a map of items we'd like to remove.
 	removeItems := make(map[peer.ID]bool, 0)
 	peerIDs := ps.Peers()
@@ -214,7 +225,6 @@ func (ps *PeerStore) ReplacePeerList(addressesThey []*peer.AddrInfo, networkName
 		if data != nil {
 			ad := data.(addressData)
 			updated := false
-			ad.mu.RLock()
 			if ad.networkNames[networkName] && !ad.roles.IsPersistent(role) {
 				if ad.roles.Is(role) {
 					removeItems[pid] = true
@@ -223,8 +233,6 @@ func (ps *PeerStore) ReplacePeerList(addressesThey []*peer.AddrInfo, networkName
 					updated = true
 				}
 			}
-			ad.mu.RUnlock()
-
 			if updated {
 				_ = ps.Put(pid, addressDataKey, ad)
 			}
@@ -237,10 +245,8 @@ func (ps *PeerStore) ReplacePeerList(addressesThey []*peer.AddrInfo, networkName
 			// we already have this
 			// update the networkName and role
 			ad := data.(addressData)
-			ad.mu.Lock()
 			ad.networkNames[networkName] = true
 			ad.roles.Add(role)
-			ad.mu.Unlock()
 			_ = ps.Put(info.ID, addressDataKey, ad)
 
 			// do not remove this entry
@@ -263,6 +269,9 @@ func (ps *PeerStore) ReplacePeerList(addressesThey []*peer.AddrInfo, networkName
 // i.e. they won't be replaced by ReplacePeerList calls.
 // If a peer is already in the peerstore, its role will be updated.
 func (ps *PeerStore) AddPersistentPeers(addrInfo []*peer.AddrInfo, networkName string, role phonebook.Role) {
+	ps.lock.Lock()
+	defer ps.lock.Unlock()
+
 	for _, info := range addrInfo {
 		data, _ := ps.Get(info.ID, addressDataKey)
 		if data != nil {
@@ -289,7 +298,6 @@ func (ps *PeerStore) Length() int {
 func makePhonebookEntryData(networkName string, role phonebook.Role, persistent bool) addressData {
 	pbData := addressData{
 		networkNames:          make(map[string]bool),
-		mu:                    &deadlock.RWMutex{},
 		recentConnectionTimes: make([]time.Time, 0),
 		roles:                 phonebook.MakeRoleSet(role, persistent),
 	}
@@ -303,17 +311,15 @@ func (ps *PeerStore) deletePhonebookEntry(peerID peer.ID, networkName string) {
 		return
 	}
 	ad := data.(addressData)
-	ad.mu.Lock()
 	delete(ad.networkNames, networkName)
 	isEmpty := len(ad.networkNames) == 0
-	ad.mu.Unlock()
 	if isEmpty {
 		ps.ClearAddrs(peerID)
 		_ = ps.Put(peerID, addressDataKey, nil)
 	}
 }
 
-// AppendTime adds the current time to recentConnectionTimes in
+// appendTime adds the current time to recentConnectionTimes in
 // addressData of addr
 func (ps *PeerStore) appendTime(peerID peer.ID, t time.Time) {
 	data, _ := ps.Get(peerID, addressDataKey)
@@ -322,7 +328,7 @@ func (ps *PeerStore) appendTime(peerID peer.ID, t time.Time) {
 	_ = ps.Put(peerID, addressDataKey, ad)
 }
 
-// PopEarliestTime removes the earliest time from recentConnectionTimes in
+// popNElements removes the earliest time from recentConnectionTimes in
 // addressData for addr
 // It is expected to be later than ConnectionsRateLimitingWindow
 func (ps *PeerStore) popNElements(n int, peerID peer.ID) {