NPeersForCpl and collapse empty buckets

table.go

@@ -85,6 +85,58 @@ func (rt *RoutingTable) Close() error {
 	return nil
 }
 
+// NPeersForCPL returns the number of peers we have for a given Cpl
+func (rt *RoutingTable) NPeersForCpl(cpl uint) int {
+	rt.tabLock.RLock()
+	defer rt.tabLock.RUnlock()
+
+	// it's in the last bucket
+	if int(cpl) >= len(rt.buckets)-1 {
+		count := 0
+		b := rt.buckets[len(rt.buckets)-1]
+		for _, p := range b.peerIds() {
+			if CommonPrefixLen(rt.local, ConvertPeerID(p)) == int(cpl) {
+				count++
+			}
+		}
+		return count
+	} else {
+		return rt.buckets[cpl].len()
+	}
+}
+
+// GetPeersForCpl returns the peers in the Routing Table with this cpl.
+func (rt *RoutingTable) GetPeersForCpl(cpl uint) []peer.ID {
+	rt.tabLock.RLock()
+	defer rt.tabLock.RUnlock()
+
+	var peers []peer.ID
+
+	// it's in the last bucket
+	if int(cpl) >= len(rt.buckets)-1 {
+		b := rt.buckets[len(rt.buckets)-1]
+		for _, p := range b.peerIds() {
+			if CommonPrefixLen(rt.local, ConvertPeerID(p)) == int(cpl) {
+				peers = append(peers, p)
+			}
+		}
+	} else {
+		for _, p := range rt.buckets[cpl].peerIds() {
+			peers = append(peers, p)
+		}
+	}
+
+	return peers
+}
+
+// IsBucketFull returns true if the Logical bucket for a given Cpl is full
+func (rt *RoutingTable) IsBucketFull(cpl uint) bool {
+	rt.tabLock.RLock()
+	defer rt.tabLock.RUnlock()
+
+	return rt.NPeersForCpl(cpl) == rt.bucketsize
+}
+
 // TryAddPeer tries to add a peer to the Routing table. If the peer ALREADY exists in the Routing Table, this call is a no-op.
 // If the peer is a queryPeer i.e. we queried it or it queried us, we set the LastSuccessfulOutboundQuery to the current time.
 // If the peer is just a peer that we connect to/it connected to us without any DHT query, we consider it as having
@@ -117,6 +169,11 @@ func (rt *RoutingTable) addPeer(p peer.ID, queryPeer bool) (bool, error) {
 
 	// peer already exists in the Routing Table.
 	if peer := bucket.getPeer(p); peer != nil {
+		// if we're querying the peer first time after adding it, let's give it a
+		// usefulness bump. This will ONLY happen once.
+		if peer.LastUsefulAt.IsZero() && queryPeer {
+			peer.LastUsefulAt = lastUsefulAt
+		}
 		return false, nil
 	}
 
@@ -230,7 +287,21 @@ func (rt *RoutingTable) removePeer(p peer.ID) {
 	if bucket.remove(p) {
 		// peer removed callback
 		rt.PeerRemoved(p)
-		return
+
+		// remove this bucket if it was the last bucket and it's now empty
+		// provided it isn't the ONLY bucket we have.
+		if len(rt.buckets) > 1 && bucketID == len(rt.buckets)-1 && len(bucket.peers()) == 0 {
+			rt.buckets[bucketID] = nil
+			rt.buckets = rt.buckets[:bucketID]
+			return
+		}
+
+		// if the second last bucket just became empty, remove and replace it with the last bucket.
+		if bucketID == len(rt.buckets)-2 && len(bucket.peers()) == 0 {
+			rt.buckets[bucketID] = rt.buckets[bucketID+1]
+			rt.buckets[bucketID+1] = nil
+			rt.buckets = rt.buckets[:bucketID+1]
+		}
 	}
 }
 

table_test.go

@@ -80,6 +80,162 @@ func TestBucket(t *testing.T) {
 	}
 }
 
+func TestNPeersForCpl(t *testing.T) {
+	t.Parallel()
+	local := test.RandPeerIDFatal(t)
+	m := pstore.NewMetrics()
+	rt, err := NewRoutingTable(2, ConvertPeerID(local), time.Hour, m, NoOpThreshold)
+	require.NoError(t, err)
+
+	require.Equal(t, 0, rt.NPeersForCpl(0))
+	require.Equal(t, 0, rt.NPeersForCpl(1))
+
+	// one peer with cpl 1
+	p, _ := rt.GenRandPeerID(1)
+	rt.TryAddPeer(p, true)
+	require.Equal(t, 0, rt.NPeersForCpl(0))
+	require.Equal(t, 1, rt.NPeersForCpl(1))
+	require.Equal(t, 0, rt.NPeersForCpl(2))
+
+	// one peer with cpl 0
+	p, _ = rt.GenRandPeerID(0)
+	rt.TryAddPeer(p, true)
+	require.Equal(t, 1, rt.NPeersForCpl(0))
+	require.Equal(t, 1, rt.NPeersForCpl(1))
+	require.Equal(t, 0, rt.NPeersForCpl(2))
+
+	// split the bucket with a peer with cpl 1
+	p, _ = rt.GenRandPeerID(1)
+	rt.TryAddPeer(p, true)
+	require.Equal(t, 1, rt.NPeersForCpl(0))
+	require.Equal(t, 2, rt.NPeersForCpl(1))
+	require.Equal(t, 0, rt.NPeersForCpl(2))
+
+	p, _ = rt.GenRandPeerID(0)
+	rt.TryAddPeer(p, true)
+	require.Equal(t, 2, rt.NPeersForCpl(0))
+}
+
+func TestGetPeersForCpl(t *testing.T) {
+	t.Parallel()
+	local := test.RandPeerIDFatal(t)
+	m := pstore.NewMetrics()
+	rt, err := NewRoutingTable(2, ConvertPeerID(local), time.Hour, m, NoOpThreshold)
+	require.NoError(t, err)
+
+	require.Empty(t, rt.GetPeersForCpl(0))
+	require.Empty(t, rt.GetPeersForCpl(1))
+
+	// one peer with cpl 1
+	p1, _ := rt.GenRandPeerID(1)
+	rt.TryAddPeer(p1, true)
+	require.Empty(t, rt.GetPeersForCpl(0))
+	require.Len(t, rt.GetPeersForCpl(1), 1)
+	require.Contains(t, rt.GetPeersForCpl(1), p1)
+	require.Empty(t, rt.GetPeersForCpl(2))
+
+	// one peer with cpl 0
+	p2, _ := rt.GenRandPeerID(0)
+	rt.TryAddPeer(p2, true)
+	require.Len(t, rt.GetPeersForCpl(0), 1)
+	require.Contains(t, rt.GetPeersForCpl(0), p2)
+	require.Len(t, rt.GetPeersForCpl(1), 1)
+	require.Contains(t, rt.GetPeersForCpl(1), p1)
+	require.Empty(t, rt.GetPeersForCpl(2))
+
+	// split the bucket with a peer with cpl 1
+	p3, _ := rt.GenRandPeerID(1)
+	rt.TryAddPeer(p3, true)
+	require.Len(t, rt.GetPeersForCpl(0), 1)
+	require.Contains(t, rt.GetPeersForCpl(0), p2)
+
+	require.Len(t, rt.GetPeersForCpl(1), 2)
+	require.Contains(t, rt.GetPeersForCpl(1), p1)
+	require.Contains(t, rt.GetPeersForCpl(1), p3)
+	require.Empty(t, rt.GetPeersForCpl(2))
+
+	p0, _ := rt.GenRandPeerID(0)
+	rt.TryAddPeer(p0, true)
+	require.Len(t, rt.GetPeersForCpl(0), 2)
+}
+
+func TestEmptyBucketCollapse(t *testing.T) {
+	t.Parallel()
+	local := test.RandPeerIDFatal(t)
+
+	m := pstore.NewMetrics()
+	rt, err := NewRoutingTable(1, ConvertPeerID(local), time.Hour, m, NoOpThreshold)
+	require.NoError(t, err)
+
+	// generate peers with cpl 0,1,2 & 3
+	p1, _ := rt.GenRandPeerID(0)
+	p2, _ := rt.GenRandPeerID(1)
+	p3, _ := rt.GenRandPeerID(2)
+	p4, _ := rt.GenRandPeerID(3)
+
+	// remove peer on an empty bucket should not panic.
+	rt.RemovePeer(p1)
+
+	// add peer with cpl 0 and remove it..bucket should still exist as it's the ONLY bucket we have
+	b, err := rt.TryAddPeer(p1, true)
+	require.True(t, b)
+	require.NoError(t, err)
+	rt.RemovePeer(p1)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 1)
+	rt.tabLock.Unlock()
+	require.Empty(t, rt.ListPeers())
+
+	// add peer with cpl 0 and cpl 1 and verify we have two buckets.
+	b, err = rt.TryAddPeer(p1, true)
+	require.True(t, b)
+	b, err = rt.TryAddPeer(p2, true)
+	require.True(t, b)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 2)
+	rt.tabLock.Unlock()
+
+	// removing a peer from the last bucket collapses it.
+	rt.RemovePeer(p2)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 1)
+	rt.tabLock.Unlock()
+	require.Len(t, rt.ListPeers(), 1)
+	require.Contains(t, rt.ListPeers(), p1)
+
+	// add p2 again
+	b, err = rt.TryAddPeer(p2, true)
+	require.True(t, b)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 2)
+	rt.tabLock.Unlock()
+
+	// now remove a peer from the second-last i.e. first bucket and ensure it collapses
+	rt.RemovePeer(p1)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 1)
+	rt.tabLock.Unlock()
+	require.Len(t, rt.ListPeers(), 1)
+	require.Contains(t, rt.ListPeers(), p2)
+
+	// let's have a total of 4 buckets now
+	rt.TryAddPeer(p1, true)
+	rt.TryAddPeer(p2, true)
+	rt.TryAddPeer(p3, true)
+	rt.TryAddPeer(p4, true)
+
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 4)
+	rt.tabLock.Unlock()
+
+	// removing a peer from the middle bucket does not collapse any bucket
+	rt.RemovePeer(p2)
+	rt.tabLock.Lock()
+	require.Len(t, rt.buckets, 4)
+	rt.tabLock.Unlock()
+	require.NotContains(t, rt.ListPeers(), p2)
+}
+
 func TestRemovePeer(t *testing.T) {
 	t.Parallel()
 	local := test.RandPeerIDFatal(t)