fix infinite loop if not enough new addrs

p2p/pex/addrbook.go

@@ -417,6 +417,11 @@ ADDRS_LOOP:
 		pickFromOldBucket := int((float64(selectionIndex)/float64(numAddresses))*100) >= biasTowardsNewAddrs
 		pickFromOldBucket = (pickFromOldBucket && a.nOld > 0 && len(oldBucketToAddrsMap) < a.nOld) || a.nNew == 0
 
+		// if there's not enough new addrs to pick from, just return early.
+		if !pickFromOldBucket && len(newBucketToAddrsMap) >= a.nNew {
+			return selection
+		}
+
 		bucket := make(map[string]*knownAddress)
 
 		// loop until we pick a random non-empty bucket

p2p/pex/addrbook_test.go

@@ -36,11 +36,10 @@ func deleteTempFile(fname string) {
 	}
 }
 
-func TestAddrBookGetSelectionWithMarkedGood(t *testing.T) {
+func TestAddrBookGetSelectionWithOneMarkedGood(t *testing.T) {
 	fname := createTempFileName("addrbook_test")
 	defer deleteTempFile(fname)
 
-	// 0 addresses
 	book := NewAddrBook(fname, true)
 	book.SetLogger(log.TestingLogger())
 	assert.Zero(t, book.Size())
@@ -50,11 +49,42 @@ func TestAddrBookGetSelectionWithMarkedGood(t *testing.T) {
 	for _, addr := range randAddrs {
 		book.AddAddress(addr.addr, addr.src)
 	}
+
+	// mark one addr as good
 	for _, addr := range randAddrs[:1] {
 		book.MarkGood(addr.addr)
 	}
 
-	addrs := book.GetSelectionWithBias(10)
+	// ensure we can get selection.
+	// this used to result in an infinite loop (!)
+	addrs := book.GetSelectionWithBias(biasToSelectNewPeers)
+	assert.NotNil(t, addrs, "expected an address")
+}
+
+func TestAddrBookGetSelectionWithOneNotMarkedGood(t *testing.T) {
+	fname := createTempFileName("addrbook_test")
+	defer deleteTempFile(fname)
+
+	// 0 addresses
+	book := NewAddrBook(fname, true)
+	book.SetLogger(log.TestingLogger())
+	assert.Zero(t, book.Size())
+
+	N := 10
+	randAddrs := randNetAddressPairs(t, N)
+	// mark all addrs but one as good
+	for _, addr := range randAddrs[:N-1] {
+		book.AddAddress(addr.addr, addr.src)
+		book.MarkGood(addr.addr)
+	}
+	// add the last addr, don't mark as good
+	for _, addr := range randAddrs[N-1:] {
+		book.AddAddress(addr.addr, addr.src)
+	}
+
+	// ensure we can get selection.
+	// this used to result in an infinite loop (!)
+	addrs := book.GetSelectionWithBias(biasToSelectNewPeers)
 	assert.NotNil(t, addrs, "expected an address")
 }