fix: cache node

This commit is contained in:
thinkAfCod 2024-10-13 12:30:19 +08:00 committed by Chen Kai
parent 95c2637065
commit d28ff8ff6d
2 changed files with 41 additions and 27 deletions

View file

@ -169,7 +169,7 @@ func DefaultPortalProtocolConfig() *PortalProtocolConfig {
type PortalProtocol struct {
table *Table
cachedIdsLock sync.Mutex
cachedIds map[string]enode.ID
cachedIds map[string]*enode.Node
protocolId string
protocolName string
@ -213,7 +213,7 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId portalwire.Proto
closeCtx, cancelCloseCtx := context.WithCancel(context.Background())
protocol := &PortalProtocol{
cachedIds: make(map[string]enode.ID),
cachedIds: make(map[string]*enode.Node),
protocolId: string(protocolId),
protocolName: protocolId.Name(),
ListenAddr: config.ListenAddr,
@ -330,10 +330,10 @@ func (p *PortalProtocol) setupUDPListening() error {
p.cachedIdsLock.Lock()
defer p.cachedIdsLock.Unlock()
if id, ok := p.cachedIds[addr.String()]; ok {
if n, ok := p.cachedIds[addr.String()]; ok {
//_, err := p.DiscV5.TalkRequestToID(id, addr, string(portalwire.UTPNetwork), buf)
req := &v5wire.TalkRequest{Protocol: string(portalwire.Utp), Message: buf}
p.DiscV5.sendFromAnotherThread(id, netip.AddrPortFrom(netutil.IPToAddr(addr.IP), uint16(addr.Port)), req)
p.DiscV5.sendFromAnotherThreadWithNode(n, netip.AddrPortFrom(netutil.IPToAddr(addr.IP), uint16(addr.Port)), req)
return len(buf), err
} else {
@ -388,21 +388,22 @@ func (p *PortalProtocol) setupDiscV5AndTable() error {
return nil
}
func (p *PortalProtocol) putCacheNodeId(node *enode.Node) {
func (p *PortalProtocol) cacheNode(node *enode.Node) {
p.cachedIdsLock.Lock()
defer p.cachedIdsLock.Unlock()
addr := &net.UDPAddr{IP: node.IP(), Port: node.UDP()}
if _, ok := p.cachedIds[addr.String()]; !ok {
p.cachedIds[addr.String()] = node.ID()
p.cachedIds[addr.String()] = node
}
}
func (p *PortalProtocol) putCacheId(id enode.ID, addr *net.UDPAddr) {
p.cachedIdsLock.Lock()
defer p.cachedIdsLock.Unlock()
if _, ok := p.cachedIds[addr.String()]; !ok {
p.cachedIds[addr.String()] = id
}
func (p *PortalProtocol) cacheNodeById(id enode.ID, addr *net.UDPAddr) {
go func() {
if _, ok := p.cachedIds[addr.String()]; !ok {
n := p.ResolveNodeId(id)
p.cacheNode(n)
}
}()
}
func (p *PortalProtocol) ping(node *enode.Node) (uint64, error) {
@ -566,7 +567,7 @@ func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request *
}
p.Log.Info("will process Offer", "id", target.ID(), "ip", target.IP().To4().String(), "port", target.UDP())
p.putCacheNodeId(target)
p.cacheNode(target)
accept := &portalwire.Accept{}
err = accept.UnmarshalSSZ(resp[1:])
@ -704,7 +705,7 @@ func (p *PortalProtocol) processContent(target *enode.Node, resp []byte) (byte,
}
p.Log.Info("will process content", "id", target.ID(), "ip", target.IP().To4().String(), "port", target.UDP())
p.putCacheNodeId(target)
p.cacheNode(target)
switch resp[1] {
case portalwire.ContentRawSelector:
@ -910,21 +911,14 @@ func (p *PortalProtocol) processPong(target *enode.Node, resp []byte) (*portalwi
}
func (p *PortalProtocol) handleUtpTalkRequest(id enode.ID, addr *net.UDPAddr, msg []byte) []byte {
if n := p.DiscV5.getNode(id); n != nil {
p.table.addInboundNode(n)
}
p.putCacheId(id, addr)
p.cacheNodeById(id, addr)
p.Log.Trace("receive utp data", "addr", addr, "msg-length", len(msg))
p.packetRouter.ReceiveMessage(msg, addr)
return []byte("")
}
func (p *PortalProtocol) handleTalkRequest(id enode.ID, addr *net.UDPAddr, msg []byte) []byte {
if n := p.DiscV5.getNode(id); n != nil {
p.table.addInboundNode(n)
}
p.putCacheId(id, addr)
p.cacheNodeById(id, addr)
msgCode := msg[0]
@ -1107,7 +1101,7 @@ func (p *PortalProtocol) handleFindContent(id enode.ID, addr *net.UDPAddr, reque
return nil, err
}
p.putCacheId(id, addr)
p.cacheNodeById(id, addr)
if errors.Is(err, ContentNotFound) {
closestNodes := p.findNodesCloseToContent(contentId, portalFindnodesResultLimit)
@ -1303,7 +1297,7 @@ func (p *PortalProtocol) handleOffer(id enode.ID, addr *net.UDPAddr, request *po
}
}
p.putCacheId(id, addr)
p.cacheNodeById(id, addr)
idBuffer := make([]byte, 2)
if contentKeyBitlist.Count() != 0 {

View file

@ -102,6 +102,7 @@ type UDPv5 struct {
type sendRequest struct {
destID enode.ID
destNode *enode.Node
destAddr netip.AddrPort
msg v5wire.Packet
}
@ -596,7 +597,19 @@ func (t *UDPv5) dispatch() {
t.sendNextCall(c.id)
case r := <-t.sendCh:
t.send(r.destID, r.destAddr, r.msg, nil)
c := &callV5{id: r.destID, addr: r.destAddr}
c.node = r.destNode
c.packet = r.msg
c.reqid = make([]byte, 8)
c.ch = make(chan v5wire.Packet, 1)
c.err = make(chan error, 1)
// Assign request ID.
crand.Read(c.reqid)
c.packet.SetRequestID(c.reqid)
newNonce, _ := t.send(c.id, c.addr, c.packet, nil)
c.nonce = newNonce
t.activeCallByAuth[newNonce] = c
t.startResponseTimeout(c)
case p := <-t.packetInCh:
t.handlePacket(p.Data, p.Addr)
@ -681,7 +694,14 @@ func (t *UDPv5) sendResponse(toID enode.ID, toAddr netip.AddrPort, packet v5wire
func (t *UDPv5) sendFromAnotherThread(toID enode.ID, toAddr netip.AddrPort, packet v5wire.Packet) {
select {
case t.sendCh <- sendRequest{toID, toAddr, packet}:
case t.sendCh <- sendRequest{toID, nil, toAddr, packet}:
case <-t.closeCtx.Done():
}
}
func (t *UDPv5) sendFromAnotherThreadWithNode(node *enode.Node, toAddr netip.AddrPort, packet v5wire.Packet) {
select {
case t.sendCh <- sendRequest{node.ID(), node, toAddr, packet}:
case <-t.closeCtx.Done():
}
}