]> Sergey Matveev's repositories - btrtrc.git/blobdiff - pex.go
Drop support for go 1.20
[btrtrc.git] / pex.go
diff --git a/pex.go b/pex.go
index a6295f361104e474abb6ad7b128ac1212db3fc3b..a0a5f49f5babfb54a2ca11b532d7a017babf4bb8 100644 (file)
--- a/pex.go
+++ b/pex.go
@@ -2,8 +2,10 @@ package torrent
 
 import (
        "net"
+       "net/netip"
+       "sync"
+       "time"
 
-       "github.com/anacrolix/dht/v2/krpc"
        pp "github.com/anacrolix/torrent/peer_protocol"
 )
 
@@ -24,14 +26,16 @@ const (
 // represents a single connection (t=pexAdd) or disconnection (t=pexDrop) event
 type pexEvent struct {
        t    pexEventType
-       addr net.Addr
+       addr netip.AddrPort
        f    pp.PexPeerFlags
+       next *pexEvent // event feed list
 }
 
 // facilitates efficient de-duplication while generating PEX messages
 type pexMsgFactory struct {
-       added   map[addrKey]pexEvent
-       dropped map[addrKey]pexEvent
+       msg     pp.PexMsg
+       added   map[netip.AddrPort]struct{}
+       dropped map[netip.AddrPort]struct{}
 }
 
 func (me *pexMsgFactory) DeltaLen() int {
@@ -40,42 +44,90 @@ func (me *pexMsgFactory) DeltaLen() int {
                int64(len(me.dropped))))
 }
 
-type addrKey string
+type addrKey = netip.AddrPort
 
 // Returns the key to use to identify a given addr in the factory.
-func (me *pexMsgFactory) addrKey(addr net.Addr) addrKey {
-       return addrKey(addr.String())
+func (me *pexMsgFactory) addrKey(addr netip.AddrPort) addrKey {
+       return addr
 }
 
 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
 // won't hit the limit consuming this event).
 func (me *pexMsgFactory) add(e pexEvent) {
        key := me.addrKey(e.addr)
-       if _, ok := me.dropped[key]; ok {
-               delete(me.dropped, key)
+       if _, ok := me.added[key]; ok {
                return
        }
        if me.added == nil {
-               me.added = make(map[addrKey]pexEvent, pexMaxDelta)
+               me.added = make(map[addrKey]struct{}, pexMaxDelta)
+       }
+       addr := krpcNodeAddrFromAddrPort(e.addr)
+       m := &me.msg
+       switch {
+       case addr.IP.To4() != nil:
+               if _, ok := me.dropped[key]; ok {
+                       if i := m.Dropped.Index(addr); i >= 0 {
+                               m.Dropped = append(m.Dropped[:i], m.Dropped[i+1:]...)
+                       }
+                       delete(me.dropped, key)
+                       return
+               }
+               m.Added = append(m.Added, addr)
+               m.AddedFlags = append(m.AddedFlags, e.f)
+       case len(addr.IP) == net.IPv6len:
+               if _, ok := me.dropped[key]; ok {
+                       if i := m.Dropped6.Index(addr); i >= 0 {
+                               m.Dropped6 = append(m.Dropped6[:i], m.Dropped6[i+1:]...)
+                       }
+                       delete(me.dropped, key)
+                       return
+               }
+               m.Added6 = append(m.Added6, addr)
+               m.Added6Flags = append(m.Added6Flags, e.f)
+       default:
+               panic(addr)
        }
-       me.added[key] = e
+       me.added[key] = struct{}{}
 }
 
 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
 // won't hit the limit consuming this event).
 func (me *pexMsgFactory) drop(e pexEvent) {
+       addr := krpcNodeAddrFromAddrPort(e.addr)
        key := me.addrKey(e.addr)
-       if _, ok := me.added[key]; ok {
-               delete(me.added, key)
+       if me.dropped == nil {
+               me.dropped = make(map[addrKey]struct{}, pexMaxDelta)
+       }
+       if _, ok := me.dropped[key]; ok {
                return
        }
-       if me.dropped == nil {
-               me.dropped = make(map[addrKey]pexEvent, pexMaxDelta)
+       m := &me.msg
+       switch {
+       case addr.IP.To4() != nil:
+               if _, ok := me.added[key]; ok {
+                       if i := m.Added.Index(addr); i >= 0 {
+                               m.Added = append(m.Added[:i], m.Added[i+1:]...)
+                               m.AddedFlags = append(m.AddedFlags[:i], m.AddedFlags[i+1:]...)
+                       }
+                       delete(me.added, key)
+                       return
+               }
+               m.Dropped = append(m.Dropped, addr)
+       case len(addr.IP) == net.IPv6len:
+               if _, ok := me.added[key]; ok {
+                       if i := m.Added6.Index(addr); i >= 0 {
+                               m.Added6 = append(m.Added6[:i], m.Added6[i+1:]...)
+                               m.Added6Flags = append(m.Added6Flags[:i], m.Added6Flags[i+1:]...)
+                       }
+                       delete(me.added, key)
+                       return
+               }
+               m.Dropped6 = append(m.Dropped6, addr)
        }
-       me.dropped[key] = e
+       me.dropped[key] = struct{}{}
 }
 
-func (me *pexMsgFactory) addEvent(event pexEvent) {
+func (me *pexMsgFactory) append(event pexEvent) {
        switch event.t {
        case pexAdd:
                me.add(event)
@@ -86,80 +138,58 @@ func (me *pexMsgFactory) addEvent(event pexEvent) {
        }
 }
 
-func (me *pexMsgFactory) PexMsg() (ret pp.PexMsg) {
-       for key, added := range me.added {
-               addr, ok := nodeAddr(added.addr)
-               if !ok {
-                       continue
-               }
-               switch len(addr.IP) {
-               case net.IPv4len:
-                       ret.Added = append(ret.Added, addr)
-                       ret.AddedFlags = append(ret.AddedFlags, added.f)
-               case net.IPv6len:
-                       ret.Added6 = append(ret.Added6, addr)
-                       ret.Added6Flags = append(ret.Added6Flags, added.f)
-               default:
-                       panic(key)
-               }
-       }
-       for key, dropped := range me.dropped {
-               addr, ok := nodeAddr(dropped.addr)
-               if !ok {
-                       continue
-               }
-               switch len(addr.IP) {
-               case net.IPv4len:
-                       ret.Dropped = append(ret.Dropped, addr)
-               case net.IPv6len:
-                       ret.Dropped6 = append(ret.Dropped6, addr)
-               default:
-                       panic(key)
-               }
-       }
-       return
-}
-
-// Convert an arbitrary torrent peer Addr into one that can be represented by the compact addr
-// format.
-func nodeAddr(addr net.Addr) (_ krpc.NodeAddr, ok bool) {
-       ipport, ok := tryIpPortFromNetAddr(addr)
-       if !ok {
-               return
-       }
-       return krpc.NodeAddr{IP: shortestIP(ipport.IP), Port: ipport.Port}, true
-}
-
-// mainly for the krpc marshallers
-func shortestIP(ip net.IP) net.IP {
-       if ip4 := ip.To4(); ip4 != nil {
-               return ip4
-       }
-       return ip
+func (me *pexMsgFactory) PexMsg() *pp.PexMsg {
+       return &me.msg
 }
 
 // Per-torrent PEX state
 type pexState struct {
-       ev   []pexEvent // event feed, append-only
+       sync.RWMutex
+       tail *pexEvent  // event feed list
        hold []pexEvent // delayed drops
-       nc   int        // net number of alive conns
+       // Torrent-wide cooldown deadline on inbound. This exists to prevent PEX from drowning out other
+       // peer address sources, until that is fixed.
+       rest time.Time
+       nc   int           // net number of alive conns
+       msg0 pexMsgFactory // initial message
 }
 
+// Reset wipes the state clean, releasing resources. Called from Torrent.Close().
 func (s *pexState) Reset() {
-       s.ev = nil
+       s.Lock()
+       defer s.Unlock()
+       s.tail = nil
        s.hold = nil
        s.nc = 0
+       s.rest = time.Time{}
+       s.msg0 = pexMsgFactory{}
+}
+
+func (s *pexState) append(e *pexEvent) {
+       if s.tail != nil {
+               s.tail.next = e
+       }
+       s.tail = e
+       s.msg0.append(*e)
 }
 
 func (s *pexState) Add(c *PeerConn) {
+       e, err := c.pexEvent(pexAdd)
+       if err != nil {
+               return
+       }
+       s.Lock()
+       defer s.Unlock()
        s.nc++
        if s.nc >= pexTargAdded {
-               s.ev = append(s.ev, s.hold...)
+               for _, e := range s.hold {
+                       ne := e
+                       s.append(&ne)
+               }
                s.hold = s.hold[:0]
        }
-       e := c.pexEvent(pexAdd)
-       s.ev = append(s.ev, e)
        c.pex.Listed = true
+       s.append(&e)
 }
 
 func (s *pexState) Drop(c *PeerConn) {
@@ -167,26 +197,50 @@ func (s *pexState) Drop(c *PeerConn) {
                // skip connections which were not previously Added
                return
        }
-       e := c.pexEvent(pexDrop)
+       e, err := c.pexEvent(pexDrop)
+       if err != nil {
+               return
+       }
+       s.Lock()
+       defer s.Unlock()
        s.nc--
        if s.nc < pexTargAdded && len(s.hold) < pexMaxHold {
                s.hold = append(s.hold, e)
        } else {
-               s.ev = append(s.ev, e)
+               s.append(&e)
        }
 }
 
-// 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) {
-       var factory pexMsgFactory
-       n := start
-       for _, e := range s.ev[start:] {
-               if start > 0 && factory.DeltaLen() >= pexMaxDelta {
+// Generate a PEX message based on the event feed.
+// Also returns a pointer to pass to the subsequent calls
+// to produce incremental deltas.
+func (s *pexState) Genmsg(start *pexEvent) (pp.PexMsg, *pexEvent) {
+       s.RLock()
+       defer s.RUnlock()
+       if start == nil {
+               return *s.msg0.PexMsg(), s.tail
+       }
+       var msg pexMsgFactory
+       last := start
+       for e := start.next; e != nil; e = e.next {
+               if msg.DeltaLen() >= pexMaxDelta {
                        break
                }
-               factory.addEvent(e)
-               n++
+               msg.append(*e)
+               last = e
        }
-       return factory.PexMsg(), n
+       return *msg.PexMsg(), last
+}
+
+// The same as Genmsg but just counts up the distinct events that haven't been sent.
+func (s *pexState) numPending(start *pexEvent) (num int) {
+       s.RLock()
+       defer s.RUnlock()
+       if start == nil {
+               return s.msg0.PexMsg().Len()
+       }
+       for e := start.next; e != nil; e = e.next {
+               num++
+       }
+       return
 }