diff --git a/device/allowedips_test.go b/device/allowedips_test.go index 005df4899..8f95faf98 100644 --- a/device/allowedips_test.go +++ b/device/allowedips_test.go @@ -258,3 +258,49 @@ func TestTrieIPv6(t *testing.T) { assertEQ(h, 0x24046800, 0x40040800, 0x10101010, 0x10101010) assertEQ(a, 0x24046800, 0x40040800, 0xdeadbeef, 0xdeadbeef) } + +// Repeatedly insert and remove IPv4 /32s from AllowedIPs. +func BenchmarkAllowedIPsInsertRemove(b *testing.B) { + // First, make 64k peers, all with unique IPs, in a seeded pseudo-random order. + const num = 256 * 256 + var peers [num]*Peer + var ips [num]net.IP + for i := range peers { + peers[i] = new(Peer) + ips[i] = net.IPv4(100, 64, byte(i>>8), byte(i)).To4() + } + rand.Seed(1) + rand.Shuffle(num, func(i, j int) { ips[i], ips[j] = ips[j], ips[i] }) + + // Then repeatedly add one and remove one that was insert 32k inserts back. + b.ResetTimer() + var a AllowedIPs + for i := 0; i < b.N; i++ { + a.Insert(ips[i%num], 32, peers[i%num]) + a.RemoveByPeer(peers[(i+num/2)%num]) + } + + // Finally, some stats & validity checks. + nodes, numPeer := 0, 0 + peersSeen := map[*Peer]bool{} + foreachEntry(a.IPv4, func(n *trieEntry) { + nodes++ + if n.peer != nil { + numPeer++ + peersSeen[n.peer] = true + } + }) + b.Logf("for N=%v: nodes=%v, peers=%v, unique_peers=%v", b.N, nodes, numPeer, len(peersSeen)) + if numPeer != len(peersSeen) { + b.Errorf("walked %d nodes with peers != %d unique peers seen", numPeer, len(peersSeen)) + } +} + +func foreachEntry(n *trieEntry, f func(*trieEntry)) { + if n == nil { + return + } + f(n) + foreachEntry(n.child[0], f) + foreachEntry(n.child[1], f) +}