Repository: libp2p/go-libp2p
Pull request #3187 (2025-02-13 15:42:55)

autonatv2: allow multiple concurrent requests per peer
This fixes the flaky test:
TestServerDialData(https://github.com/libp2p/go-libp2p/issues/3179) The exact sequence of events that flakiness is:

The client rejects the dial data request by resetting the stream. The client makes a new request to the server. The server rejects this request as it hasn't read the previous stream's reset yet.

It is better to allow multiple concurrent requests per peer. 5 concurrent requests per peer is okay.

sukunrt: No science to this. I guess 2 is okay too. 

Patch:
diff --git a/p2p/protocol/autonatv2/options.go b/p2p/protocol/autonatv2/options.go
index 15268f62..76cd3735 100644
--- a/p2p/protocol/autonatv2/options.go
+++ b/p2p/protocol/autonatv2/options.go
@@ -8,6 +8,7 @@ type autoNATSettings struct {
 	serverRPM                            int
 	serverPerPeerRPM                     int
 	serverDialDataRPM                    int
+	maxConcurrentRequestsPerPeer         int
 	dataRequestPolicy                    dataRequestPolicyFunc
 	now                                  func() time.Time
 	amplificatonAttackPreventionDialWait time.Duration
@@ -20,6 +21,7 @@ func defaultSettings() *autoNATSettings {
 		serverRPM:                            60, // 1 every second
 		serverPerPeerRPM:                     12, // 1 every 5 seconds
 		serverDialDataRPM:                    12, // 1 every 5 seconds
+		maxConcurrentRequestsPerPeer:         2,
 		dataRequestPolicy:                    amplificationAttackPrevention,
 		amplificatonAttackPreventionDialWait: 3 * time.Second,
 		now:                                  time.Now,
@@ -28,11 +30,12 @@ func defaultSettings() *autoNATSettings {
 
 type AutoNATOption func(s *autoNATSettings) error
 
-func WithServerRateLimit(rpm, perPeerRPM, dialDataRPM int) AutoNATOption {
+func WithServerRateLimit(rpm, perPeerRPM, dialDataRPM int, maxConcurrentRequestsPerPeer int) AutoNATOption {
 	return func(s *autoNATSettings) error {
 		s.serverRPM = rpm
 		s.serverPerPeerRPM = perPeerRPM
 		s.serverDialDataRPM = dialDataRPM
+		s.maxConcurrentRequestsPerPeer = maxConcurrentRequestsPerPeer
 		return nil
 	}
 }
diff --git a/p2p/protocol/autonatv2/server.go b/p2p/protocol/autonatv2/server.go
index d28cfc93..05e0bdd9 100644
--- a/p2p/protocol/autonatv2/server.go
+++ b/p2p/protocol/autonatv2/server.go
@@ -67,10 +67,11 @@ func newServer(host, dialer host.Host, s *autoNATSettings) *server {
 		amplificatonAttackPreventionDialWait: s.amplificatonAttackPreventionDialWait,
 		allowPrivateAddrs:                    s.allowPrivateAddrs,
 		limiter: &rateLimiter{
-			RPM:         s.serverRPM,
-			PerPeerRPM:  s.serverPerPeerRPM,
-			DialDataRPM: s.serverDialDataRPM,
-			now:         s.now,
+			RPM:                          s.serverRPM,
+			PerPeerRPM:                   s.serverPerPeerRPM,
+			DialDataRPM:                  s.serverDialDataRPM,
+			MaxConcurrentRequestsPerPeer: s.maxConcurrentRequestsPerPeer,
+			now:                          s.now,
 		},
 		now:           s.now,
 		metricsTracer: s.metricsTracer,
@@ -391,16 +392,17 @@ type rateLimiter struct {
 	RPM int
 	// DialDataRPM is the rate limit for requests that require dial data
 	DialDataRPM int
+	// MaxConcurrentRequestsPerPeer is the maximum number of concurrent requests per peer
+	MaxConcurrentRequestsPerPeer int
 
 	mu           sync.Mutex
 	closed       bool
 	reqs         []entry
 	peerReqs     map[peer.ID][]time.Time
 	dialDataReqs []time.Time
-	// ongoingReqs tracks in progress requests. This is used to disallow multiple concurrent requests by the
-	// same peer
-	// TODO: Should we allow a few concurrent requests per peer?
-	ongoingReqs map[peer.ID]struct{}
+	// inProgressReqs tracks in progress requests. This is used to limit multiple
+	// concurrent requests by the same peer.
+	inProgressReqs map[peer.ID]int
 
 	now func() time.Time // for tests
 }
@@ -410,28 +412,31 @@ type entry struct {
 	Time   time.Time
 }
 
+func (r *rateLimiter) init() {
+	if r.peerReqs == nil {
+		r.peerReqs = make(map[peer.ID][]time.Time)
+		r.inProgressReqs = make(map[peer.ID]int)
+	}
+}
+
 func (r *rateLimiter) Accept(p peer.ID) bool {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if r.closed {
 		return false
 	}
-	if r.peerReqs == nil {
-		r.peerReqs = make(map[peer.ID][]time.Time)
-		r.ongoingReqs = make(map[peer.ID]struct{})
-	}
-
+	r.init()
 	nw := r.now()
 	r.cleanup(nw)
 
-	if _, ok := r.ongoingReqs[p]; ok {
+	if r.inProgressReqs[p] >= r.MaxConcurrentRequestsPerPeer {
 		return false
 	}
 	if len(r.reqs) >= r.RPM || len(r.peerReqs[p]) >= r.PerPeerRPM {
 		return false
 	}
 
-	r.ongoingReqs[p] = struct{}{}
+	r.inProgressReqs[p]++
 	r.reqs = append(r.reqs, entry{PeerID: p, Time: nw})
 	r.peerReqs[p] = append(r.peerReqs[p], nw)
 	return true
@@ -443,10 +448,7 @@ func (r *rateLimiter) AcceptDialDataRequest(p peer.ID) bool {
 	if r.closed {
 		return false
 	}
-	if r.peerReqs == nil {
-		r.peerReqs = make(map[peer.ID][]time.Time)
-		r.ongoingReqs = make(map[peer.ID]struct{})
-	}
+	r.init()
 	nw := r.now()
 	r.cleanup(nw)
 	if len(r.dialDataReqs) >= r.DialDataRPM {
@@ -495,7 +497,13 @@ func (r *rateLimiter) cleanup(now time.Time) {
 func (r *rateLimiter) CompleteRequest(p peer.ID) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
-	delete(r.ongoingReqs, p)
+	r.inProgressReqs[p]--
+	if r.inProgressReqs[p] <= 0 {
+		delete(r.inProgressReqs, p)
+		if r.inProgressReqs[p] < 0 {
+			log.Errorf("BUG: negative in progress requests for peer %s", p)
+		}
+	}
 }
 
 func (r *rateLimiter) Close() {
@@ -503,7 +511,7 @@ func (r *rateLimiter) Close() {
 	defer r.mu.Unlock()
 	r.closed = true
 	r.peerReqs = nil
-	r.ongoingReqs = nil
+	r.inProgressReqs = nil
 	r.dialDataReqs = nil
 }
 


Test patch:
diff --git a/p2p/protocol/autonatv2/server_test.go b/p2p/protocol/autonatv2/server_test.go
index 0b40f275..6664f1c3 100644
--- a/p2p/protocol/autonatv2/server_test.go
+++ b/p2p/protocol/autonatv2/server_test.go
@@ -149,7 +149,7 @@ func TestServerDataRequest(t *testing.T) {
 			}
 			return false
 		}),
-		WithServerRateLimit(10, 10, 10),
+		WithServerRateLimit(10, 10, 10, 2),
 		withAmplificationAttackPreventionDialWait(0),
 	)
 	defer an.Close()
@@ -187,6 +187,69 @@ func TestServerDataRequest(t *testing.T) {
 	_, err = c.GetReachability(context.Background(), []Request{{Addr: quicAddr, SendDialData: true}, {Addr: tcpAddr}})
 	require.Error(t, err)
 }
+
+func TestServerMaxConcurrentRequestsPerPeer(t *testing.T) {
+	const concurrentRequests = 5
+
+	// server will skip all tcp addresses
+	dialer := bhost.NewBlankHost(swarmt.GenSwarm(t, swarmt.OptDisableTCP))
+
+	doneChan := make(chan struct{})
+	an := newAutoNAT(t, dialer, allowPrivateAddrs, withDataRequestPolicy(
+		// stall all allowed requests
+		func(s network.Stream, dialAddr ma.Multiaddr) bool {
+			<-doneChan
+			return true
+		}),
+		WithServerRateLimit(10, 10, 10, concurrentRequests),
+		withAmplificationAttackPreventionDialWait(0),
+	)
+	defer an.Close()
+	defer an.host.Close()
+
+	c := newAutoNAT(t, nil, allowPrivateAddrs)
+	defer c.Close()
+	defer c.host.Close()
+
+	idAndWait(t, c, an)
+
+	errChan := make(chan error)
+	const N = 10
+	// num concurrentRequests will stall and N will fail
+	for i := 0; i < concurrentRequests+N; i++ {
+		go func() {
+			_, err := c.GetReachability(context.Background(), []Request{{Addr: c.host.Addrs()[0], SendDialData: false}})
+			errChan <- err
+		}()
+	}
+
+	// check N failures
+	for i := 0; i < N; i++ {
+		select {
+		case err := <-errChan:
+			require.Error(t, err)
+		case <-time.After(10 * time.Second):
+			t.Fatalf("expected %d errors: got: %d", N, i)
+		}
+	}
+
+	// check concurrentRequests failures, as we won't send dial data
+	close(doneChan)
+	for i := 0; i < concurrentRequests; i++ {
+		select {
+		case err := <-errChan:
+			require.Error(t, err)
+		case <-time.After(5 * time.Second):
+			t.Fatalf("expected %d errors: got: %d", concurrentRequests, i)
+		}
+	}
+	select {
+	case err := <-errChan:
+		t.Fatalf("expected no more errors: got: %v", err)
+	default:
+	}
+}
+
 func TestServerDataRequestJitter(t *testing.T) {
 	// server will skip all tcp addresses
 	dialer := bhost.NewBlankHost(swarmt.GenSwarm(t, swarmt.OptDisableTCP))
@@ -198,7 +261,7 @@ func TestServerDataRequestJitter(t *testing.T) {
 			}
 			return false
 		}),
-		WithServerRateLimit(10, 10, 10),
+		WithServerRateLimit(10, 10, 10, 2),
 		withAmplificationAttackPreventionDialWait(5*time.Second),
 	)
 	defer an.Close()
@@ -238,7 +301,7 @@ func TestServerDataRequestJitter(t *testing.T) {
 }
 
 func TestServerDial(t *testing.T) {
-	an := newAutoNAT(t, nil, WithServerRateLimit(10, 10, 10), allowPrivateAddrs)
+	an := newAutoNAT(t, nil, WithServerRateLimit(10, 10, 10, 2), allowPrivateAddrs)
 	defer an.Close()
 	defer an.host.Close()
 
@@ -295,7 +358,7 @@ func TestServerDial(t *testing.T) {
 
 func TestRateLimiter(t *testing.T) {
 	cl := test.NewMockClock()
-	r := rateLimiter{RPM: 3, PerPeerRPM: 2, DialDataRPM: 1, now: cl.Now}
+	r := rateLimiter{RPM: 3, PerPeerRPM: 2, DialDataRPM: 1, now: cl.Now, MaxConcurrentRequestsPerPeer: 1}
 
 	require.True(t, r.Accept("peer1"))
 
@@ -333,12 +396,37 @@ func TestRateLimiter(t *testing.T) {
 
 	cl.AdvanceBy(10 * time.Second)
 	require.True(t, r.Accept("peer3"))
+
+}
+
+func TestRateLimiterConcurrentRequests(t *testing.T) {
+	const N = 5
+	const Peers = 5
+	for concurrentRequests := 1; concurrentRequests <= N; concurrentRequests++ {
+		cl := test.NewMockClock()
+		r := rateLimiter{RPM: 10 * Peers * N, PerPeerRPM: 10 * Peers * N, DialDataRPM: 10 * Peers * N, now: cl.Now, MaxConcurrentRequestsPerPeer: concurrentRequests}
+		for p := 0; p < Peers; p++ {
+			for i := 0; i < concurrentRequests; i++ {
+				require.True(t, r.Accept(peer.ID(fmt.Sprintf("peer-%d", p))))
+			}
+			require.False(t, r.Accept(peer.ID(fmt.Sprintf("peer-%d", p))))
+			// Now complete the requests
+			for i := 0; i < concurrentRequests; i++ {
+				r.CompleteRequest(peer.ID(fmt.Sprintf("peer-%d", p)))
+			}
+			// Now we should be able to accept new requests
+			for i := 0; i < concurrentRequests; i++ {
+				require.True(t, r.Accept(peer.ID(fmt.Sprintf("peer-%d", p))))
+			}
+			require.False(t, r.Accept(peer.ID(fmt.Sprintf("peer-%d", p))))
+		}
+	}
 }
 
 func TestRateLimiterStress(t *testing.T) {
 	cl := test.NewMockClock()
 	for i := 0; i < 10; i++ {
-		r := rateLimiter{RPM: 20 + i, PerPeerRPM: 10 + i, DialDataRPM: i, now: cl.Now}
+		r := rateLimiter{RPM: 20 + i, PerPeerRPM: 10 + i, DialDataRPM: i, MaxConcurrentRequestsPerPeer: 1, now: cl.Now}
 
 		peers := make([]peer.ID, 10+i)
 		for i := 0; i < len(peers); i++ {
@@ -386,7 +474,7 @@ func TestRateLimiterStress(t *testing.T) {
 		require.Equal(t, len(r.peerReqs), 1)
 		require.Equal(t, len(r.peerReqs[peers[0]]), 1)
 		require.Equal(t, len(r.dialDataReqs), 0)
-		require.Equal(t, len(r.ongoingReqs), 1)
+		require.Equal(t, len(r.inProgressReqs), 1)
 	}
 }
 
@@ -433,7 +521,7 @@ func TestReadDialData(t *testing.T) {
 }
 
 func FuzzServerDialRequest(f *testing.F) {
-	a := newAutoNAT(f, nil, allowPrivateAddrs, WithServerRateLimit(math.MaxInt32, math.MaxInt32, math.MaxInt32))
+	a := newAutoNAT(f, nil, allowPrivateAddrs, WithServerRateLimit(math.MaxInt32, math.MaxInt32, math.MaxInt32, 2))
 	c := newAutoNAT(f, nil)
 	idAndWait(f, c, a)
 	// reduce the streamTimeout before running this. TODO: fix this