optimise generation of the initial PEX

diff --git a/pex.go b/pex.go
index a6295f361104e474abb6ad7b128ac1212db3fc3b..7d2ba208d5ae21aa8289e2e97e46a0f0af2b2933 100644 (file)
--- a/pex.go
+++ b/pex.go
@@ -2,6 +2,7 @@ package torrent
 
 import (
        "net"
+       "sync"
 
        "github.com/anacrolix/dht/v2/krpc"
        pp "github.com/anacrolix/torrent/peer_protocol"
@@ -140,15 +141,23 @@ func shortestIP(ip net.IP) net.IP {
 
 // Per-torrent PEX state
 type pexState struct {
-       ev   []pexEvent // event feed, append-only
-       hold []pexEvent // delayed drops
-       nc   int        // net number of alive conns
+       ev        []pexEvent    // event feed, append-only
+       hold      []pexEvent    // delayed drops
+       nc        int           // net number of alive conns
+       initCache pexMsgFactory // last generated initial message
+       initSeq   int           // number of events which went into initCache
+       initLock  sync.RWMutex  // serialise access to initCache and initSeq
 }
 
+// Reset wipes the state clean, releasing resources. Called from Torrent.Close().
 func (s *pexState) Reset() {
        s.ev = nil
        s.hold = nil
        s.nc = 0
+       s.initLock.Lock()
+       s.initCache = pexMsgFactory{}
+       s.initSeq = 0
+       s.initLock.Unlock()
 }
 
 func (s *pexState) Add(c *PeerConn) {
@@ -179,6 +188,10 @@ func (s *pexState) Drop(c *PeerConn) {
 // Generate a PEX message based on the event feed. Also returns an index to pass to the subsequent
 // calls, producing incremental deltas.
 func (s *pexState) Genmsg(start int) (pp.PexMsg, int) {
+       if start == 0 {
+               return s.genmsg0()
+       }
+
        var factory pexMsgFactory
        n := start
        for _, e := range s.ev[start:] {
@@ -190,3 +203,17 @@ func (s *pexState) Genmsg(start int) (pp.PexMsg, int) {
        }
        return factory.PexMsg(), n
 }
+
+func (s *pexState) genmsg0() (pp.PexMsg, int) {
+       s.initLock.Lock()
+       for _, e := range s.ev[s.initSeq:] {
+               s.initCache.addEvent(e)
+               s.initSeq++
+       }
+       s.initLock.Unlock()
+       s.initLock.RLock()
+       n := s.initSeq
+       msg := s.initCache.PexMsg()
+       s.initLock.RUnlock()
+       return msg, n
+}

diff --git a/pex_test.go b/pex_test.go
index 006644c2c8bf7eddc30e4405a95f5e5f5a747c81..96b1ed71ca82f42dacc95c8b506bdb4fa3c9650a 100644 (file)
--- a/pex_test.go
+++ b/pex_test.go
@@ -307,11 +307,11 @@ func TestPexGenmsg(t *testing.T) {
 func addrgen(n int) chan net.Addr {
        c := make(chan net.Addr)
        go func() {
+               defer close(c)
                for i := 4747; i < 65535 && n > 0; i++ {
                        c <- &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: i}
                        n--
                }
-               close(c)
        }()
        return c
 }
@@ -335,6 +335,24 @@ func TestPexInitialNoCutoff(t *testing.T) {
        require.EqualValues(t, 0, len(m.Dropped6))
 }
 
+func benchmarkPexInitialN(b *testing.B, npeers int) {
+       for i := 0; i < b.N; i++ {
+               var s pexState
+               c := addrgen(npeers)
+               for addr := range c {
+                       s.Add(&PeerConn{peer: peer{RemoteAddr: addr}})
+                       s.Genmsg(0)
+               }
+       }
+}
+
+// obtain at least 5 points, e.g. to plot a graph
+func BenchmarkPexInitial4(b *testing.B)   { benchmarkPexInitialN(b, 4) }
+func BenchmarkPexInitial50(b *testing.B)  { benchmarkPexInitialN(b, 50) }
+func BenchmarkPexInitial100(b *testing.B) { benchmarkPexInitialN(b, 100) }
+func BenchmarkPexInitial200(b *testing.B) { benchmarkPexInitialN(b, 200) }
+func BenchmarkPexInitial400(b *testing.B) { benchmarkPexInitialN(b, 400) }
+
 func TestPexAdd(t *testing.T) {
        t.Run("ipv4", func(t *testing.T) {
                addrs := addrs4