Repository: libp2p/go-libp2p
Pull request #2844 (2024-06-20 16:13:44)

swarm: improve dial worker performance for common case
This makes the common case O(1) compared to O(N) before. Updating an element is a separate op and we check the element from the back of the slice as most inserts are sorted. 


Patch:
diff --git a/p2p/net/swarm/dial_worker.go b/p2p/net/swarm/dial_worker.go
index 0cac6e4f..360a99e2 100644
--- a/p2p/net/swarm/dial_worker.go
+++ b/p2p/net/swarm/dial_worker.go
@@ -254,7 +254,7 @@ loop:
 						if simConnect, _, _ := network.GetSimultaneousConnect(ad.ctx); !simConnect {
 							ad.ctx = network.WithSimultaneousConnect(ad.ctx, isClient, reason)
 							// update the element in dq to use the simultaneous connect delay.
-							dq.Add(network.AddrDelay{
+							dq.UpdateOrAdd(network.AddrDelay{
 								Addr:  ad.addr,
 								Delay: addrDelay[string(ad.addr.Bytes())],
 							})
@@ -436,12 +436,31 @@ type dialQueue struct {
 
 // newDialQueue returns a new dialQueue
 func newDialQueue() *dialQueue {
-	return &dialQueue{q: make([]network.AddrDelay, 0, 16)}
+	return &dialQueue{
+		q: make([]network.AddrDelay, 0, 16),
+	}
 }
 
-// Add adds adelay to the queue. If another element exists in the queue with
-// the same address, it replaces that element.
+// Add adds a new element to the dialQueue. To update an element use UpdateOrAdd.
 func (dq *dialQueue) Add(adelay network.AddrDelay) {
+	for i := dq.Len() - 1; i >= 0; i-- {
+		if dq.q[i].Delay <= adelay.Delay {
+			// insert at pos i+1
+			dq.q = append(dq.q, network.AddrDelay{}) // extend the slice
+			copy(dq.q[i+2:], dq.q[i+1:])
+			dq.q[i+1] = adelay
+			return
+		}
+	}
+	// insert at position 0
+	dq.q = append(dq.q, network.AddrDelay{}) // extend the slice
+	copy(dq.q[1:], dq.q[0:])
+	dq.q[0] = adelay
+}
+
+// UpdateOrAdd updates the elements with address adelay.Addr to the new delay
+// Useful when hole punching
+func (dq *dialQueue) UpdateOrAdd(adelay network.AddrDelay) {
 	for i := 0; i < dq.Len(); i++ {
 		if dq.q[i].Addr.Equal(adelay.Addr) {
 			if dq.q[i].Delay == adelay.Delay {
@@ -451,19 +470,9 @@ func (dq *dialQueue) Add(adelay network.AddrDelay) {
 			// remove the element
 			copy(dq.q[i:], dq.q[i+1:])
 			dq.q = dq.q[:len(dq.q)-1]
-			break
-		}
-	}
-
-	for i := 0; i < dq.Len(); i++ {
-		if dq.q[i].Delay > adelay.Delay {
-			dq.q = append(dq.q, network.AddrDelay{}) // extend the slice
-			copy(dq.q[i+1:], dq.q[i:])
-			dq.q[i] = adelay
-			return
 		}
 	}
-	dq.q = append(dq.q, adelay)
+	dq.Add(adelay)
 }
 
 // NextBatch returns all the elements in the queue with the highest priority


Test patch:
diff --git a/p2p/net/swarm/dial_worker_test.go b/p2p/net/swarm/dial_worker_test.go
index 9e9465fc..ed4f00ff 100644
--- a/p2p/net/swarm/dial_worker_test.go
+++ b/p2p/net/swarm/dial_worker_test.go
@@ -401,9 +401,10 @@ func TestDialQueueNextBatch(t *testing.T) {
 		addrs = append(addrs, ma.StringCast(fmt.Sprintf("/ip4/1.2.3.4/tcp/%d", i)))
 	}
 	testcase := []struct {
-		name   string
-		input  []network.AddrDelay
-		output [][]ma.Multiaddr
+		name       string
+		input      []network.AddrDelay
+		output     [][]ma.Multiaddr
+		hasUpdates bool
 	}{
 		{
 			name: "next batch",
@@ -456,6 +457,7 @@ func TestDialQueueNextBatch(t *testing.T) {
 				{addrs[4]},
 				{},
 			},
+			hasUpdates: true,
 		},
 		{
 			name:  "null input",
@@ -470,7 +472,11 @@ func TestDialQueueNextBatch(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			q := newDialQueue()
 			for i := 0; i < len(tc.input); i++ {
-				q.Add(tc.input[i])
+				if tc.hasUpdates {
+					q.UpdateOrAdd(tc.input[i])
+				} else {
+					q.Add(tc.input[i])
+				}
 			}
 			for _, batch := range tc.output {
 				b := q.NextBatch()
@@ -1140,3 +1146,51 @@ func TestDialWorkerLoopTCPConnUpgradeWait(t *testing.T) {
 		t.Errorf("expected a fail response")
 	}
 }
+
+func BenchmarkDialRanker(b *testing.B) {
+	const N = 10000
+
+	benchDialQueue := func(adelays []network.AddrDelay) {
+		dq := newDialQueue()
+		for _, a := range adelays {
+			dq.Add(a)
+		}
+		for {
+			batch := dq.NextBatch()
+			if len(batch) == 0 {
+				return
+			}
+		}
+	}
+	addrs := make([]ma.Multiaddr, N)
+	for i := 0; i < N; i++ {
+		addrs[i] = ma.StringCast(fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", i))
+	}
+
+	b.Run("equal delay", func(b *testing.B) {
+		b.ReportAllocs()
+		addrDelays := make([]network.AddrDelay, N)
+		for i := 0; i < N; i++ {
+			addrDelays[i] = network.AddrDelay{
+				Addr:  addrs[i],
+				Delay: 0,
+			}
+		}
+		for i := 0; i < b.N; i++ {
+			benchDialQueue(addrDelays)
+		}
+	})
+	b.Run("sorted delay", func(b *testing.B) {
+		b.ReportAllocs()
+		addrDelays := make([]network.AddrDelay, N)
+		for i := 0; i < N; i++ {
+			addrDelays[i] = network.AddrDelay{
+				Addr:  addrs[i],
+				Delay: time.Millisecond * time.Duration(i),
+			}
+		}
+		for i := 0; i < b.N; i++ {
+			benchDialQueue(addrDelays)
+		}
+	})
+}