p2p: refactor GetSelectionWithBias for addressbook (#3475)

Why submit this pr:

    we have suffered from infinite loop in addrbook bug which takes us a long time to find out why process become a zombie peer. It have been fixed in #3232. But the ADDRS_LOOP is still there, risk of infinite loop is still exist.
    The algorithm that to random pick a bucket is not stable, which means the peer may unluckily always choose the wrong bucket for a long time, the time and cpu cost is meaningless.

A simple improvement:
shuffle bucketsNew and bucketsOld, and pick necessary number of address from them. A stable
algorithm.

p2p/pex/addrbook.go

@@ -9,6 +9,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math"
+	"math/rand"
 	"net"
 	"sync"
 	"time"
@@ -405,89 +406,11 @@ func (a *addrBook) GetSelectionWithBias(biasTowardsNewAddrs int) []*p2p.NetAddre
 		bookSize*getSelectionPercent/100)
 	numAddresses = cmn.MinInt(maxGetSelection, numAddresses)
 
-	selection := make([]*p2p.NetAddress, numAddresses)
-
-	oldBucketToAddrsMap := make(map[int]map[string]struct{})
-	var oldIndex int
-	newBucketToAddrsMap := make(map[int]map[string]struct{})
-	var newIndex int
-
-	// initialize counters used to count old and new added addresses.
-	// len(oldBucketToAddrsMap) cannot be used as multiple addresses can endup in the same bucket.
-	var oldAddressesAdded int
-	var newAddressesAdded int
-
 	// number of new addresses that, if possible, should be in the beginning of the selection
-	numRequiredNewAdd := percentageOfNum(biasTowardsNewAddrs, numAddresses)
-
-	selectionIndex := 0
-ADDRS_LOOP:
-	for selectionIndex < numAddresses {
-		// biasedTowardsOldAddrs indicates if the selection can switch to old addresses
-		biasedTowardsOldAddrs := selectionIndex >= numRequiredNewAdd
-		// An old addresses is selected if:
-		// - the bias is for old and old addressees are still available or,
-		// - there are no new addresses or all new addresses have been selected.
-		// numAddresses <= a.nOld + a.nNew therefore it is guaranteed that there are enough
-		// addresses to fill the selection
-		pickFromOldBucket :=
-			(biasedTowardsOldAddrs && oldAddressesAdded < a.nOld) ||
-				a.nNew == 0 || newAddressesAdded >= a.nNew
-
-		bucket := make(map[string]*knownAddress)
-
-		// loop until we pick a random non-empty bucket
-		for len(bucket) == 0 {
-			if pickFromOldBucket {
-				oldIndex = a.rand.Intn(len(a.bucketsOld))
-				bucket = a.bucketsOld[oldIndex]
-			} else {
-				newIndex = a.rand.Intn(len(a.bucketsNew))
-				bucket = a.bucketsNew[newIndex]
-			}
-		}
-
-		// pick a random index
-		randIndex := a.rand.Intn(len(bucket))
-
-		// loop over the map to return that index
-		var selectedAddr *p2p.NetAddress
-		for _, ka := range bucket {
-			if randIndex == 0 {
-				selectedAddr = ka.Addr
-				break
-			}
-			randIndex--
-		}
-
-		// if we have selected the address before, restart the loop
-		// otherwise, record it and continue
-		if pickFromOldBucket {
-			if addrsMap, ok := oldBucketToAddrsMap[oldIndex]; ok {
-				if _, ok = addrsMap[selectedAddr.String()]; ok {
-					continue ADDRS_LOOP
-				}
-			} else {
-				oldBucketToAddrsMap[oldIndex] = make(map[string]struct{})
-			}
-			oldBucketToAddrsMap[oldIndex][selectedAddr.String()] = struct{}{}
-			oldAddressesAdded++
-		} else {
-			if addrsMap, ok := newBucketToAddrsMap[newIndex]; ok {
-				if _, ok = addrsMap[selectedAddr.String()]; ok {
-					continue ADDRS_LOOP
-				}
-			} else {
-				newBucketToAddrsMap[newIndex] = make(map[string]struct{})
-			}
-			newBucketToAddrsMap[newIndex][selectedAddr.String()] = struct{}{}
-			newAddressesAdded++
-		}
-
-		selection[selectionIndex] = selectedAddr
-		selectionIndex++
-	}
-
+	// if there are no enough old addrs, will choose new addr instead.
+	numRequiredNewAdd := cmn.MaxInt(percentageOfNum(biasTowardsNewAddrs, numAddresses), numAddresses-a.nOld)
+	selection := a.randomPickAddresses(bucketTypeNew, numRequiredNewAdd)
+	selection = append(selection, a.randomPickAddresses(bucketTypeOld, numAddresses-len(selection))...)
 	return selection
 }
 
@@ -726,6 +649,44 @@ func (a *addrBook) addAddress(addr, src *p2p.NetAddress) error {
 	return nil
 }
 
+func (a *addrBook) randomPickAddresses(bucketType byte, num int) []*p2p.NetAddress {
+	var buckets []map[string]*knownAddress
+	switch bucketType {
+	case bucketTypeNew:
+		buckets = a.bucketsNew
+	case bucketTypeOld:
+		buckets = a.bucketsOld
+	default:
+		panic("unexpected bucketType")
+	}
+	total := 0
+	for _, bucket := range buckets {
+		total = total + len(bucket)
+	}
+	addresses := make([]*knownAddress, 0, total)
+	for _, bucket := range buckets {
+		for _, ka := range bucket {
+			addresses = append(addresses, ka)
+		}
+	}
+	selection := make([]*p2p.NetAddress, 0, num)
+	chosenSet := make(map[string]bool, num)
+	rand.Shuffle(total, func(i, j int) {
+		addresses[i], addresses[j] = addresses[j], addresses[i]
+	})
+	for _, addr := range addresses {
+		if chosenSet[addr.Addr.String()] {
+			continue
+		}
+		chosenSet[addr.Addr.String()] = true
+		selection = append(selection, addr.Addr)
+		if len(selection) >= num {
+			return selection
+		}
+	}
+	return selection
+}
+
 // Make space in the new buckets by expiring the really bad entries.
 // If no bad entries are available we remove the oldest.
 func (a *addrBook) expireNew(bucketIdx int) {

p2p/pex/addrbook_test.go

@@ -435,12 +435,12 @@ func TestPrivatePeers(t *testing.T) {
 
 func testAddrBookAddressSelection(t *testing.T, bookSize int) {
 	// generate all combinations of old (m) and new addresses
-	for nOld := 0; nOld <= bookSize; nOld++ {
-		nNew := bookSize - nOld
-		dbgStr := fmt.Sprintf("book of size %d (new %d, old %d)", bookSize, nNew, nOld)
+	for nBookOld := 0; nBookOld <= bookSize; nBookOld++ {
+		nBookNew := bookSize - nBookOld
+		dbgStr := fmt.Sprintf("book of size %d (new %d, old %d)", bookSize, nBookNew, nBookOld)
 
 		// create book and get selection
-		book, fname := createAddrBookWithMOldAndNNewAddrs(t, nOld, nNew)
+		book, fname := createAddrBookWithMOldAndNNewAddrs(t, nBookOld, nBookNew)
 		defer deleteTempFile(fname)
 		addrs := book.GetSelectionWithBias(biasToSelectNewPeers)
 		assert.NotNil(t, addrs, "%s - expected a non-nil selection", dbgStr)
@@ -460,27 +460,25 @@ func testAddrBookAddressSelection(t *testing.T, bookSize int) {
 		// Given:
 		// n - num new addrs, m - num old addrs
 		// k - num new addrs expected in the beginning (based on bias %)
-		// i=min(n, k), aka expFirstNew
+		// i=min(n, max(k,r-m)), aka expNew
 		// j=min(m, r-i), aka expOld
 		//
 		// We expect this layout:
-		// indices:      0...i-1   i...i+j-1    i+j...r
-		// addresses:    N0..Ni-1  O0..Oj-1     Ni...
+		// indices:      0...i-1   i...i+j-1
+		// addresses:    N0..Ni-1  O0..Oj-1
 		//
 		// There is at least one partition and at most three.
 		var (
-			k           = percentageOfNum(biasToSelectNewPeers, nAddrs)
-			expFirstNew = cmn.MinInt(nNew, k)
-			expOld      = cmn.MinInt(nOld, nAddrs-expFirstNew)
-			expNew      = nAddrs - expOld
-			expLastNew  = expNew - expFirstNew
+			k      = percentageOfNum(biasToSelectNewPeers, nAddrs)
+			expNew = cmn.MinInt(nNew, cmn.MaxInt(k, nAddrs-nBookOld))
+			expOld = cmn.MinInt(nOld, nAddrs-expNew)
 		)
 
 		// Verify that the number of old and new addresses are as expected
-		if nNew < expNew || nNew > expNew {
+		if nNew != expNew {
 			t.Fatalf("%s - expected new addrs %d, got %d", dbgStr, expNew, nNew)
 		}
-		if nOld < expOld || nOld > expOld {
+		if nOld != expOld {
 			t.Fatalf("%s - expected old addrs %d, got %d", dbgStr, expOld, nOld)
 		}
 
@@ -499,15 +497,12 @@ func testAddrBookAddressSelection(t *testing.T, bookSize int) {
 		case expOld == 0: // all new addresses
 			expSeqLens = []int{nAddrs}
 			expSeqTypes = []int{1}
-		case expFirstNew == 0: // all old addresses
+		case expNew == 0: // all old addresses
 			expSeqLens = []int{nAddrs}
 			expSeqTypes = []int{2}
-		case nAddrs-expFirstNew-expOld == 0: // new addresses, old addresses
-			expSeqLens = []int{expFirstNew, expOld}
+		case nAddrs-expNew-expOld == 0: // new addresses, old addresses
+			expSeqLens = []int{expNew, expOld}
 			expSeqTypes = []int{1, 2}
-		default: // new addresses, old addresses, new addresses
-			expSeqLens = []int{expFirstNew, expOld, expLastNew}
-			expSeqTypes = []int{1, 2, 1}
 		}
 
 		assert.Equal(t, expSeqLens, seqLens,