network: Add capability filtered depth calculation (#1787)

nolash · zelig · commit afd1d9f8ffdf · 2019-09-20T15:40:02.000+02:00
* network: Add capability filtered depth calculation

* network: Typo in test fail error message
diff --git a/network/kademlia.go b/network/kademlia.go
@@ -229,6 +229,7 @@ type capabilityIndex struct {
 	*capability.Capability
 	conns *pot.Pot
 	addrs *pot.Pot
+	depth int
 }
 
 // NewCapabilityIndex creates a new capability index with a copy the provided capabilities array
@@ -455,6 +456,10 @@ func (k *Kademlia) setNeighbourhoodDepth() {
 		k.nDepth = nDepth
 		changed = true
 	}
+	// TODO: when hive is refactored, notifies should be made for depth change in any cap index
+	for _, idx := range k.capabilityIndex {
+		idx.depth = capabilityDepthForPot(idx, k.NeighbourhoodSize, k.base)
+	}
 	k.nDepthMu.Unlock()
 
 	if len(k.nDepthSig) > 0 && changed {
@@ -468,6 +473,7 @@ func (k *Kademlia) setNeighbourhoodDepth() {
 			}
 		}
 	}
+
 }
 
 // NeighbourhoodDepth returns the value calculated by depthForPot function
@@ -478,6 +484,16 @@ func (k *Kademlia) NeighbourhoodDepth() int {
 	return k.nDepth
 }
 
+func (k *Kademlia) NeighbourhoodDepthCapability(s string) (int, error) {
+	k.nDepthMu.RLock()
+	defer k.nDepthMu.RUnlock()
+	idx, ok := k.capabilityIndex[s]
+	if !ok {
+		return -1, fmt.Errorf("Unknown capability index %v", s)
+	}
+	return idx.depth, nil
+}
+
 // SubscribeToNeighbourhoodDepthChange returns the channel that signals
 // when neighbourhood depth value is changed. The current neighbourhood depth
 // is returned by NeighbourhoodDepth method. Returned function unsubscribes
@@ -636,6 +652,10 @@ func neighbourhoodRadiusForPot(p *pot.Pot, neighbourhoodSize int, pivotAddr []by
 	return depth
 }
 
+func capabilityDepthForPot(idx *capabilityIndex, neighbourhoodSize int, pivotAddr []byte) (depth int) {
+	return depthForPot(idx.conns, neighbourhoodSize, pivotAddr)
+}
+
 // depthForPot returns the depth for the pot
 // depth is the radius of the minimal extension of nearest neighbourhood that
 // includes all empty PO bins. I.e., depth is the deepest PO such that
diff --git a/network/kademlia_test.go b/network/kademlia_test.go
@@ -943,3 +943,57 @@ func testCapabilityIndexRemove(t *testing.T) {
 		t.Fatalf("EachAddrFiltered '42:101' expected 1 peer, got %d", c)
 	}
 }
+
+// TestCapabilityNeighbourhoodDepth tests that depth calculations filtered by capability is correct
+func TestCapabilityNeighbourhoodDepth(t *testing.T) {
+	baseAddressBytes := RandomBzzAddr().OAddr
+	kad := NewKademlia(baseAddressBytes, NewKadParams())
+	cap_both := capability.NewCapability(42, 2)
+	cap_both.Set(0)
+	cap_both.Set(1)
+	kad.RegisterCapabilityIndex("both", *cap_both)
+	cap_one := capability.NewCapability(42, 2)
+	cap_one.Set(0)
+	kad.RegisterCapabilityIndex("one", *cap_one)
+
+	baseAddress := pot.NewAddressFromBytes(baseAddressBytes)
+
+	// generate the peers
+	var peers []*Peer
+	for i := 0; i < 2; i++ {
+		addr := pot.RandomAddressAt(baseAddress, i)
+		p := newTestDiscoveryPeer(addr, kad)
+		p.BzzAddr.Capabilities.Add(cap_both)
+		peers = append(peers)
+		kad.Register(p.BzzAddr)
+		kad.On(p)
+	}
+
+	addrClosestBoth := pot.RandomAddressAt(baseAddress, 7)
+	peerClosestBoth := newTestDiscoveryPeer(addrClosestBoth, kad)
+	peerClosestBoth.BzzAddr.Capabilities.Add(cap_both)
+	kad.Register(peerClosestBoth.BzzAddr)
+	kad.On(peerClosestBoth)
+
+	addrClosestOne := pot.RandomAddressAt(baseAddress, 7)
+	peerClosestOne := newTestDiscoveryPeer(addrClosestOne, kad)
+	peerClosestOne.BzzAddr.Capabilities.Add(cap_one)
+	kad.Register(peerClosestOne.BzzAddr)
+	kad.On(peerClosestOne)
+
+	depth, err := kad.NeighbourhoodDepthCapability("both")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if depth != 1 {
+		t.Fatalf("cap 'both' expected depth 1, was %d", depth)
+	}
+
+	depth, err = kad.NeighbourhoodDepthCapability("one")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if depth != 2 {
+		t.Fatalf("cap 'one' expected depth 2, was %d", depth)
+	}
+}
