Address manager randomization weighted by connection failures (#1916)

* Address manager refactor stage 1

* Use a simpler weightedRand function which makes it easier parameterize to the process

* Switch back to connectionFailedCount

* Simplify selected entry deletion

* Fix function comment

Co-authored-by: Constantine Bitensky <cbitensky1@gmail.com>
Co-authored-by: Ori Newman <orinewman1@gmail.com>

infrastructure/network/addressmanager/addressmanager.go

@@ -5,26 +5,24 @@
 package addressmanager
 
 import (
+	"github.com/kaspanet/kaspad/infrastructure/db/database"
+	"github.com/kaspanet/kaspad/util/mstime"
 	"net"
 	"sync"
 	"time"
 
-	"github.com/kaspanet/kaspad/infrastructure/db/database"
-	"github.com/kaspanet/kaspad/util/mstime"
-
 	"github.com/kaspanet/kaspad/app/appmessage"
 	"github.com/pkg/errors"
 )
 
 const (
 	maxAddresses                   = 4096
-	connectionFailedCountForRemove = 3
+	connectionFailedCountForRemove = 4
 )
 
 // addressRandomizer is the interface for the randomizer needed for the AddressManager.
 type addressRandomizer interface {
-	RandomAddress(addresses []*appmessage.NetAddress) *appmessage.NetAddress
-	RandomAddresses(addresses []*appmessage.NetAddress, count int) []*appmessage.NetAddress
+	RandomAddresses(addresses []*address, count int) []*appmessage.NetAddress
 }
 
 // addressKey represents a pair of IP and port, the IP is always in V6 representation
@@ -80,7 +78,7 @@ func New(cfg *Config, database database.Database) (*AddressManager, error) {
 	return &AddressManager{
 		store:          addressStore,
 		localAddresses: localAddresses,
-		random:         NewAddressRandomize(),
+		random:         NewAddressRandomize(connectionFailedCountForRemove),
 		cfg:            cfg,
 	}, nil
 }
@@ -91,7 +89,8 @@ func (am *AddressManager) addAddressNoLock(netAddress *appmessage.NetAddress) er
 	}
 
 	key := netAddressKey(netAddress)
-	address := &address{netAddress: netAddress, connectionFailedCount: 0}
+	// We mark `connectionFailedCount` as 0 only after first success
+	address := &address{netAddress: netAddress, connectionFailedCount: 1}
 	err := am.store.add(key, address)
 	if err != nil {
 		return err
@@ -109,8 +108,7 @@ func (am *AddressManager) addAddressNoLock(netAddress *appmessage.NetAddress) er
 			}
 		}
 
-		toRemoveKey := netAddressKey(toRemove.netAddress)
-		err := am.store.remove(toRemoveKey)
+		err := am.removeAddressNoLock(toRemove.netAddress)
 		if err != nil {
 			return err
 		}
@@ -171,7 +169,6 @@ func (am *AddressManager) MarkConnectionFailure(address *appmessage.NetAddress)
 			address, entry.connectionFailedCount)
 		return am.store.remove(key)
 	}
-
 	return am.store.updateNotBanned(key, entry)
 }
 
@@ -207,19 +204,13 @@ func (am *AddressManager) BannedAddresses() []*appmessage.NetAddress {
 }
 
 // notBannedAddressesWithException returns all not banned addresses with excpetion
-func (am *AddressManager) notBannedAddressesWithException(exceptions []*appmessage.NetAddress) []*appmessage.NetAddress {
+func (am *AddressManager) notBannedAddressesWithException(exceptions []*appmessage.NetAddress) []*address {
 	am.mutex.Lock()
 	defer am.mutex.Unlock()
 
 	return am.store.getAllNotBannedNetAddressesWithout(exceptions)
 }
 
-// RandomAddress returns a random address that isn't banned and isn't in exceptions
-func (am *AddressManager) RandomAddress(exceptions []*appmessage.NetAddress) *appmessage.NetAddress {
-	validAddresses := am.notBannedAddressesWithException(exceptions)
-	return am.random.RandomAddress(validAddresses)
-}
-
 // RandomAddresses returns count addresses at random that aren't banned and aren't in exceptions
 func (am *AddressManager) RandomAddresses(count int, exceptions []*appmessage.NetAddress) []*appmessage.NetAddress {
 	validAddresses := am.notBannedAddressesWithException(exceptions)

infrastructure/network/addressmanager/addressrandomize.go

@@ -1,6 +1,7 @@
 package addressmanager
 
 import (
+	"math"
 	"math/rand"
 	"time"
 
@@ -9,38 +10,54 @@ import (
 
 // AddressRandomize implement addressRandomizer interface
 type AddressRandomize struct {
-	random *rand.Rand
+	random         *rand.Rand
+	maxFailedCount uint64
 }
 
 // NewAddressRandomize returns a new RandomizeAddress.
-func NewAddressRandomize() *AddressRandomize {
+func NewAddressRandomize(maxFailedCount uint64) *AddressRandomize {
 	return &AddressRandomize{
-		random: rand.New(rand.NewSource(time.Now().UnixNano())),
+		random:         rand.New(rand.NewSource(time.Now().UnixNano())),
+		maxFailedCount: maxFailedCount,
 	}
 }
 
-// RandomAddress returns a random address from input list
-func (amc *AddressRandomize) RandomAddress(addresses []*appmessage.NetAddress) *appmessage.NetAddress {
-	if len(addresses) > 0 {
-		randomIndex := rand.Intn(len(addresses))
-		return addresses[randomIndex]
+// weightedRand is a help function which returns a random index in the
+// range [0, len(weights)-1] with probability weighted by `weights`
+func weightedRand(weights []float32) int {
+	sum := float32(0)
+	for _, weight := range weights {
+		sum += weight
 	}
-
-	return nil
+	randPoint := rand.Float32()
+	scanPoint := float32(0)
+	for i, weight := range weights {
+		normalizedWeight := weight / sum
+		scanPoint += normalizedWeight
+		if randPoint <= scanPoint {
+			return i
+		}
+	}
+	return len(weights) - 1
 }
 
 // RandomAddresses returns count addresses at random from input list
-func (amc *AddressRandomize) RandomAddresses(addresses []*appmessage.NetAddress, count int) []*appmessage.NetAddress {
+func (amc *AddressRandomize) RandomAddresses(addresses []*address, count int) []*appmessage.NetAddress {
 	if len(addresses) < count {
 		count = len(addresses)
 	}
-
-	result := make([]*appmessage.NetAddress, 0, count)
-
-	randomIndexes := rand.Perm(len(addresses))
-	for i := 0; i < count; i++ {
-		result = append(result, addresses[randomIndexes[i]])
+	weights := make([]float32, 0, len(addresses))
+	for _, addr := range addresses {
+		weights = append(weights, float32(math.Pow(64, float64(amc.maxFailedCount-addr.connectionFailedCount))))
+	}
+	result := make([]*appmessage.NetAddress, 0, count)
+	for count > 0 {
+		i := weightedRand(weights)
+		result = append(result, addresses[i].netAddress)
+		// Zero entry i to avoid re-selection
+		weights[i] = 0
+		// Update count
+		count--
 	}
-
 	return result
 }

infrastructure/network/addressmanager/store.go

@@ -144,13 +144,13 @@ func (as *addressStore) getAllNotBannedNetAddresses() []*appmessage.NetAddress {
 	return addresses
 }
 
-func (as *addressStore) getAllNotBannedNetAddressesWithout(ignoredAddresses []*appmessage.NetAddress) []*appmessage.NetAddress {
+func (as *addressStore) getAllNotBannedNetAddressesWithout(ignoredAddresses []*appmessage.NetAddress) []*address {
 	ignoredKeys := netAddressesKeys(ignoredAddresses)
 
-	addresses := make([]*appmessage.NetAddress, 0, len(as.notBannedAddresses))
+	addresses := make([]*address, 0, len(as.notBannedAddresses))
 	for key, address := range as.notBannedAddresses {
 		if !ignoredKeys[key] {
-			addresses = append(addresses, address.netAddress)
+			addresses = append(addresses, address)
 		}
 	}
 	return addresses