kademlia: saturate based on known peers (#374)

* saturate based on connected and known peers

pkg/kademlia/export_test.go

@@ -4,5 +4,8 @@
 
 package kademlia
 
-var MaxBins = maxBins
-var TimeToRetry = &timeToRetry
+var (
+	MaxBins         = maxBins
+	TimeToRetry     = &timeToRetry
+	SaturationPeers = &saturationPeers
+)

pkg/kademlia/kademlia.go

@@ -32,9 +32,10 @@ var (
 	errOverlayMismatch         = errors.New("overlay mismatch")
 	timeToRetry                = 60 * time.Second
 	shortRetry                 = 30 * time.Second
+	saturationPeers            = 4
 )
 
-type binSaturationFunc func(bin, depth uint8, peers *pslice.PSlice) bool
+type binSaturationFunc func(bin uint8, peers, connected *pslice.PSlice) bool
 
 // Options for injecting services to Kademlia.
 type Options struct {
@@ -127,7 +128,11 @@ func (k *Kad) manage() {
 			}
 		case <-k.manageC:
 			start = time.Now()
-
+			select {
+			case <-k.quit:
+				return
+			default:
+			}
 			err := k.knownPeers.EachBinRev(func(peer swarm.Address, po uint8) (bool, bool, error) {
 				if k.connectedPeers.Exists(peer) {
 					return false, false, nil
@@ -141,7 +146,7 @@ func (k *Kad) manage() {
 				k.waitNextMu.Unlock()
 
 				currentDepth := k.NeighborhoodDepth()
-				if saturated := k.saturationFunc(po, currentDepth, k.connectedPeers); saturated {
+				if saturated := k.saturationFunc(po, k.knownPeers, k.connectedPeers); saturated {
 					return false, true, nil // bin is saturated, skip to next bin
 				}
 
@@ -180,7 +185,7 @@ func (k *Kad) manage() {
 				k.connectedPeers.Add(peer, po)
 
 				k.depthMu.Lock()
-				k.depth = k.recalcDepth()
+				k.depth = recalcDepth(k.connectedPeers)
 				k.depthMu.Unlock()
 
 				k.logger.Debugf("connected to peer: %s old depth: %d new depth: %d", peer, currentDepth, k.NeighborhoodDepth())
@@ -214,9 +219,11 @@ func (k *Kad) manage() {
 // binSaturated indicates whether a certain bin is saturated or not.
 // when a bin is not saturated it means we would like to proactively
 // initiate connections to other peers in the bin.
-func binSaturated(bin, depth uint8, peers *pslice.PSlice) bool {
+func binSaturated(bin uint8, peers, connected *pslice.PSlice) bool {
+	potentialDepth := recalcDepth(peers)
+
 	// short circuit for bins which are >= depth
-	if bin >= depth {
+	if bin >= potentialDepth {
 		return false
 	}
 
@@ -228,31 +235,31 @@ func binSaturated(bin, depth uint8, peers *pslice.PSlice) bool {
 	// gaps measurement)
 
 	size := 0
-	_ = peers.EachBin(func(_ swarm.Address, po uint8) (bool, bool, error) {
+	_ = connected.EachBin(func(_ swarm.Address, po uint8) (bool, bool, error) {
 		if po == bin {
 			size++
 		}
 		return false, false, nil
 	})
 
-	return size >= 2
+	return size >= saturationPeers
 }
 
 // recalcDepth calculates and returns the kademlia depth.
-func (k *Kad) recalcDepth() uint8 {
+func recalcDepth(peers *pslice.PSlice) uint8 {
 	// handle edge case separately
-	if k.connectedPeers.Length() <= nnLowWatermark {
+	if peers.Length() <= nnLowWatermark {
 		return 0
 	}
 	var (
-		peers                        = uint(0)
+		peersCtr                     = uint(0)
 		candidate                    = uint8(0)
-		shallowestEmpty, noEmptyBins = k.connectedPeers.ShallowestEmpty()
+		shallowestEmpty, noEmptyBins = peers.ShallowestEmpty()
 	)
 
-	_ = k.connectedPeers.EachBin(func(_ swarm.Address, po uint8) (bool, bool, error) {
-		peers++
-		if peers >= nnLowWatermark {
+	_ = peers.EachBin(func(_ swarm.Address, po uint8) (bool, bool, error) {
+		peersCtr++
+		if peersCtr >= nnLowWatermark {
 			candidate = po
 			return true, false, nil
 		}
@@ -390,7 +397,7 @@ func (k *Kad) Connected(ctx context.Context, addr swarm.Address) error {
 	k.waitNextMu.Unlock()
 
 	k.depthMu.Lock()
-	k.depth = k.recalcDepth()
+	k.depth = recalcDepth(k.connectedPeers)
 	k.depthMu.Unlock()
 
 	k.notifyPeerSig()
@@ -413,7 +420,7 @@ func (k *Kad) Disconnected(addr swarm.Address) {
 	k.waitNextMu.Unlock()
 
 	k.depthMu.Lock()
-	k.depth = k.recalcDepth()
+	k.depth = recalcDepth(k.connectedPeers)
 	k.depthMu.Unlock()
 	select {
 	case k.manageC <- struct{}{}:

pkg/kademlia/kademlia_test.go

@@ -160,11 +160,12 @@ func TestManage(t *testing.T) {
 		conns int32 // how many connect calls were made to the p2p mock
 
 		saturationVal  = false
-		saturationFunc = func(bin, depth uint8, peers *pslice.PSlice) bool {
+		saturationFunc = func(bin uint8, peers, connected *pslice.PSlice) bool {
 			return saturationVal
 		}
 		base, kad, ab, _, signer = newTestKademlia(&conns, nil, saturationFunc)
 	)
+	defer kad.Close()
 	// first, saturationFunc returns always false, this means that the bin is not saturated,
 	// hence we expect that every peer we add to kademlia will be connected to
 	for i := 0; i < 50; i++ {
@@ -200,12 +201,19 @@ func TestManage(t *testing.T) {
 // be true since depth is increasingly moving deeper, but then we add more peers
 // in shallower depth for the rest of the function to be executed
 func TestBinSaturation(t *testing.T) {
+	defer func(p int) {
+		*kademlia.SaturationPeers = p
+	}(*kademlia.SaturationPeers)
+	*kademlia.SaturationPeers = 2
+
 	var (
 		conns                    int32 // how many connect calls were made to the p2p mock
 		base, kad, ab, _, signer = newTestKademlia(&conns, nil, nil)
 		peers                    []swarm.Address
 	)
 
+	defer kad.Close()
+
 	// add two peers in a few bins to generate some depth >= 0, this will
 	// make the next iteration result in binSaturated==true, causing no new
 	// connections to be made
@@ -610,7 +618,7 @@ func TestMarshal(t *testing.T) {
 	}
 }
 
-func newTestKademlia(connCounter, failedConnCounter *int32, f func(bin, depth uint8, peers *pslice.PSlice) bool) (swarm.Address, *kademlia.Kad, addressbook.Interface, *mock.Discovery, beeCrypto.Signer) {
+func newTestKademlia(connCounter, failedConnCounter *int32, f func(bin uint8, peers, connected *pslice.PSlice) bool) (swarm.Address, *kademlia.Kad, addressbook.Interface, *mock.Discovery, beeCrypto.Signer) {
 	var (
 		base   = test.RandomAddress()                       // base address
 		ab     = addressbook.New(mockstate.NewStateStore()) // address book

pkg/p2p/libp2p/peer.go

@@ -81,7 +81,10 @@ func (r *peerRegistry) Disconnected(_ network.Network, c network.Conn) {
 func (r *peerRegistry) addStream(peerID libp2ppeer.ID, stream network.Stream, cancel context.CancelFunc) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-
+	if _, ok := r.streams[peerID]; !ok {
+		// it is possible that an addStream will be called after a disconnect
+		return
+	}
 	r.streams[peerID][stream] = cancel
 }