diff --git a/go.mod b/go.mod index fb2afb9..c6cbefc 100644 --- a/go.mod +++ b/go.mod @@ -27,7 +27,7 @@ require ( github.com/stretchr/testify v1.10.0 github.com/tetratelabs/wabin v0.0.0-20230304001439-f6f874872834 github.com/urfave/cli/v2 v2.27.6 - github.com/zen-eth/utp-go v0.0.0-20250517113239-5d962dd66394 + github.com/zen-eth/utp-go v0.0.0-20250601092006-0508310c08c3 go.uber.org/automaxprocs v1.6.0 golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 golang.org/x/sync v0.14.0 @@ -117,4 +117,4 @@ require ( rsc.io/tmplfunc v0.0.3 // indirect ) -replace github.com/ethereum/go-ethereum => github.com/optimism-java/shisui v1.14.6-0.20250516133529-e5d979e5825f +replace github.com/ethereum/go-ethereum => github.com/optimism-java/shisui v1.14.6-0.20250526144834-7569b8f3f6b5 diff --git a/go.sum b/go.sum index 0fbf337..ae65aff 100644 --- a/go.sum +++ b/go.sum @@ -210,8 +210,8 @@ github.com/onsi/ginkgo v1.14.0/go.mod h1:iSB4RoI2tjJc9BBv4NKIKWKya62Rps+oPG/Lv9k github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= github.com/onsi/gomega v1.10.1 h1:o0+MgICZLuZ7xjH7Vx6zS/zcu93/BEp1VwkIW1mEXCE= github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= -github.com/optimism-java/shisui v1.14.6-0.20250516133529-e5d979e5825f h1:/qGqNuNi9Q+/WZAikhZEHvwhm8TIwVxgamxk9G8mkfg= -github.com/optimism-java/shisui v1.14.6-0.20250516133529-e5d979e5825f/go.mod h1:Cak0mDbSoH9geLA691ioescz1vXoXchz7rxF1DcNizE= +github.com/optimism-java/shisui v1.14.6-0.20250526144834-7569b8f3f6b5 h1:pu6xWJifKseNAUB37VC4mn0xceUwQu2ogECgXCHtwEg= +github.com/optimism-java/shisui v1.14.6-0.20250526144834-7569b8f3f6b5/go.mod h1:Cak0mDbSoH9geLA691ioescz1vXoXchz7rxF1DcNizE= github.com/optimism-java/zrnt v0.32.4-0.20250528142456-bc543d07ddb2 h1:F40qRmm51C04XA8zZWSd1dZDu57OOPKTyiKpUFuLw68= github.com/optimism-java/zrnt v0.32.4-0.20250528142456-bc543d07ddb2/go.mod h1:A0fezkp9Tt3GBLATSPIbuY4ywYESyAuc/FFmPKg8Lqs= github.com/orisano/pixelmatch v0.0.0-20220722002657-fb0b55479cde/go.mod h1:nZgzbfBr3hhjoZnS66nKrHmduYNpc34ny7RK4z5/HM0= @@ -300,8 +300,8 @@ github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1/go.mod h1:Ohn+xnUBi github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= -github.com/zen-eth/utp-go v0.0.0-20250517113239-5d962dd66394 h1:iBectudmNrjGXexs3R/GEYbJHZ5RNdpUOoBc0GYFRGU= -github.com/zen-eth/utp-go v0.0.0-20250517113239-5d962dd66394/go.mod h1:j1tRbjVA455+fpF7B9yh+qkld0OQCXwMzb9dyuRWFGs= +github.com/zen-eth/utp-go v0.0.0-20250601092006-0508310c08c3 h1:VT3Obx3YsaqB9sKHFDflbIsmFp2+zPvDSH+LvxIYA88= +github.com/zen-eth/utp-go v0.0.0-20250601092006-0508310c08c3/go.mod h1:j1tRbjVA455+fpF7B9yh+qkld0OQCXwMzb9dyuRWFGs= go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs= go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= diff --git a/portal/node.go b/portal/node.go index ed59953..e217207 100644 --- a/portal/node.go +++ b/portal/node.go @@ -294,10 +294,11 @@ func (n *Node) initUDP() error { // initDiscV5 initializes the discV5 protocol and local node func (n *Node) initDiscV5() error { discCfg := discover.Config{ - PrivateKey: n.config.PrivateKey, - NetRestrict: n.config.PortalProtocolConfig.NetRestrict, - Bootnodes: n.config.PortalProtocolConfig.BootstrapNodes, - Log: log.New("protocol", "discV5"), + PrivateKey: n.config.PrivateKey, + NetRestrict: n.config.PortalProtocolConfig.NetRestrict, + Bootnodes: n.config.PortalProtocolConfig.BootstrapNodes, + V5RespTimeout: n.config.PortalProtocolConfig.Discv5RespTimeout, + Log: log.New("protocol", "discV5"), } nodeDB, err := enode.OpenDB(n.config.PortalProtocolConfig.NodeDBPath) diff --git a/portalwire/api.go b/portalwire/api.go index 2353bf5..af38dd4 100644 --- a/portalwire/api.go +++ b/portalwire/api.go @@ -1,6 +1,7 @@ package portalwire import ( + "context" "errors" "github.com/ethereum/go-ethereum/common/hexutil" @@ -410,7 +411,7 @@ func (p *PortalProtocolAPI) FindContent(enr string, contentKey string) (interfac return nil, err } - flag, findContent, err := p.portalProtocol.findContent(n, contentKeyBytes) + flag, findContent, err := p.portalProtocol.findContent(context.Background(), n, contentKeyBytes) if err != nil { return nil, err } @@ -474,7 +475,7 @@ func (p *PortalProtocolAPI) Offer(enr string, contentItems [][2]string) (string, Kind: TransientOfferRequestKind, Request: transientOfferRequest, } - accept, err := p.portalProtocol.offer(n, offerReq, &NoPermit{}) + accept, err := p.portalProtocol.offer(n, offerReq, PermitNotLimit) if err != nil { return "", err } @@ -519,7 +520,7 @@ func (p *PortalProtocolAPI) TraceOffer(enr string, key string, value string) (in Request: transientOfferRequestWithResult, } - _, err = p.portalProtocol.offer(n, offerReq, &NoPermit{}) + _, err = p.portalProtocol.offer(n, offerReq, PermitNotLimit) if err != nil { return nil, err } @@ -554,7 +555,6 @@ func (p *PortalProtocolAPI) RecursiveFindContent(contentKeyHex string) (*Content UtpTransfer: false, }, nil } - p.portalProtocol.Log.Warn("find content err", "contextKey", hexutil.Encode(contentKey), "err", err) content, utpTransfer, err := p.portalProtocol.ContentLookup(contentKey, contentId) @@ -683,7 +683,7 @@ func (p *PortalProtocolAPI) PutContent(contentKeyHex, contentHex string) (*PutCo continue } - _, offerErr := p.portalProtocol.offer(nodeToOffer, offerReq, &NoPermit{}) + _, offerErr := p.portalProtocol.offer(nodeToOffer, offerReq, PermitNotLimit) if offerErr != nil { p.portalProtocol.Log.Warn("Failed to offer content to lookup node", "node", nodeToOffer.ID(), "err", offerErr) continue diff --git a/portalwire/lookup.go b/portalwire/lookup.go index 8674e29..dde8ec3 100644 --- a/portalwire/lookup.go +++ b/portalwire/lookup.go @@ -18,9 +18,14 @@ package portalwire import ( "context" + "encoding/hex" "errors" + "sync" + "sync/atomic" "time" + "github.com/ethereum/go-ethereum/common/hexutil" + "github.com/ethereum/go-ethereum/p2p/enode" ) @@ -206,3 +211,347 @@ func (it *lookupIterator) Next() bool { func (it *lookupIterator) Close() { it.cancel() } + +// nodeQueue 按距离排序的节点队列 +type nodeQueue struct { + target enode.ID + nodes []*enode.Node +} + +func newNodeQueue(target enode.ID) *nodeQueue { + return &nodeQueue{target: target, nodes: make([]*enode.Node, 0)} +} + +func (nq *nodeQueue) push(n *enode.Node) { + // 按距离插入排序 + dist := enode.LogDist(nq.target, n.ID()) + insertPos := 0 + for i, existing := range nq.nodes { + if enode.LogDist(nq.target, existing.ID()) > dist { + insertPos = i + break + } + insertPos = i + 1 + } + + // 插入节点 + nq.nodes = append(nq.nodes, nil) + copy(nq.nodes[insertPos+1:], nq.nodes[insertPos:]) + nq.nodes[insertPos] = n + + // 限制队列大小 + if len(nq.nodes) > bucketSize { + nq.nodes = nq.nodes[:bucketSize] + } +} + +func (nq *nodeQueue) pop() *enode.Node { + if len(nq.nodes) == 0 { + return nil + } + node := nq.nodes[0] + nq.nodes = nq.nodes[1:] + return node +} + +func (nq *nodeQueue) len() int { + return len(nq.nodes) +} + +// ContentLookupResult contains found result +type ContentLookupResult struct { + Content []byte + UtpTransfer bool + FoundAt *enode.Node +} + +// contentLookupState hold lookup state +type contentLookupState struct { + target enode.ID + contentKey []byte + protocol *PortalProtocol + ctx context.Context + cancel context.CancelFunc + + contacted map[enode.ID]bool + pending map[enode.ID]bool + candidates *nodeQueue + + // concurrency control + resultChan chan *ContentLookupResult + newNodeChan chan *enode.Node // new node found notify + queryDoneChan chan *enode.Node // query done notify + mu sync.Mutex + found atomic.Bool + activeQueries int + + trace *Trace +} + +func newContentLookupState(ctx context.Context, cancel context.CancelFunc, p *PortalProtocol, contentId []byte, contentKey []byte, trace *Trace) *contentLookupState { + return &contentLookupState{ + target: enode.ID(contentId), + contentKey: contentKey, + protocol: p, + ctx: ctx, + cancel: cancel, + contacted: make(map[enode.ID]bool), + pending: make(map[enode.ID]bool), + candidates: newNodeQueue(enode.ID(contentId)), + resultChan: make(chan *ContentLookupResult, 1), + newNodeChan: make(chan *enode.Node, 100), + queryDoneChan: make(chan *enode.Node, 100), + trace: trace, + } +} + +// run completely event drive for content look up +func (state *contentLookupState) run() { + defer func() { + state.protocol.Log.Info("content lookup task done", "found", state.found.Load(), "ctx.err", state.ctx.Err(), "contentKey", hex.EncodeToString(state.contentKey)) + state.cancel() + }() + state.initCandidates() + const maxConcurrent = 1 + activeChan := make(chan struct{}, maxConcurrent) + + // 启动初始查询 + state.startAvailableQueries(activeChan) + + // 主事件循环 + for { + select { + case <-state.ctx.Done(): + return + + case newNode := <-state.newNodeChan: + state.addNewCandidate(newNode) + + case completedNode := <-state.queryDoneChan: + state.onQueryCompleted(completedNode) + <-activeChan + if state.isLookupComplete() { + // It will not affect the normal process + // If no content is found, it read a nil value + // state.resultChan will not be closed anytime + select { + case state.resultChan <- nil: + default: + } + return + } + // try start new queries + state.startAvailableQueries(activeChan) + } + } +} + +// startAvailableQueries start a available query +func (state *contentLookupState) startAvailableQueries(activeChan chan struct{}) { + for { + select { + case <-state.ctx.Done(): + return + default: + } + if state.found.Load() { + state.protocol.Log.Info("content has been found, will not start a new goroutine", "contentKey", hexutil.Encode(state.contentKey)) + return + } + state.mu.Lock() + node := state.candidates.pop() + if node == nil { + state.mu.Unlock() + return + } + + if state.contacted[node.ID()] || state.pending[node.ID()] { + state.mu.Unlock() + continue + } + + if state.activeQueries >= 1 { + state.protocol.Log.Info("exceed active query limit", "contentKey", hexutil.Encode(state.contentKey), "state.activeQueries", state.activeQueries) + state.mu.Unlock() + return + } + state.activeQueries++ + state.pending[node.ID()] = true + state.mu.Unlock() + + go func(n *enode.Node) { + select { + case activeChan <- struct{}{}: + state.queryNode(n) + case <-state.ctx.Done(): + return + } + }(node) + } +} + +func (state *contentLookupState) queryNode(node *enode.Node) { + // notify query done + defer func() { + select { + case state.queryDoneChan <- node: + case <-state.ctx.Done(): + } + }() + select { + case <-state.ctx.Done(): + return + default: + } + if state.found.Load() { + return + } + state.protocol.Log.Info("findContent method start", "node", node.ID(), "contentKey", hexutil.Encode(state.contentKey)) + flag, content, err := state.protocol.findContent(state.ctx, node, state.contentKey) + if err != nil { + state.protocol.Log.Error("content lookup query failed", + "node", node.ID(), "contentKey", hexutil.Encode(state.contentKey), "err", err) + return + } + state.protocol.Log.Info("findContent method end", "node", node.ID(), "contentKey", hexutil.Encode(state.contentKey)) + + hexId := "0x" + node.ID().String() + if state.trace != nil { + state.mu.Lock() + + dis := state.protocol.Distance(node.ID(), state.target) + + state.trace.Metadata[hexId] = &NodeMetadata{ + Enr: node.String(), + Distance: hexutil.Encode(dis[:]), + } + state.mu.Unlock() + } + + switch flag { + case ContentRawSelector, ContentConnIdSelector: + contentBytes, ok := content.([]byte) + if !ok { + state.protocol.Log.Error("invalid content type", + "node", node.ID(), "content", content) + return + } + + result := &ContentLookupResult{ + Content: contentBytes, + UtpTransfer: flag == ContentConnIdSelector, + FoundAt: node, + } + + if state.found.CompareAndSwap(false, true) { + state.protocol.Log.Info("found content", "node", node.ID(), "contentKey", hexutil.Encode(state.contentKey)) + select { + case state.resultChan <- result: + default: + } + } + + case ContentEnrsSelector: + nodes, ok := content.([]*enode.Node) + if !ok { + state.protocol.Log.Error("invalid enrs type", + "node", node.ID(), "content", content) + return + } + if state.trace != nil { + state.handleEnrsWithTrace(hexId, nodes) + return + } + for _, newNode := range nodes { + if newNode.ID() == state.protocol.Self().ID() { + continue + } + select { + case state.newNodeChan <- newNode: + case <-state.ctx.Done(): + return + } + } + } +} + +func (state *contentLookupState) handleEnrsWithTrace(fromHexId string, nodes []*enode.Node) { + respByNode := RespByNode{ + RespondedWith: make([]string, 0, len(nodes)), + } + state.mu.Lock() + defer state.mu.Unlock() + + for _, newNode := range nodes { + idInner := "0x" + newNode.ID().String() + respByNode.RespondedWith = append(respByNode.RespondedWith, idInner) + + if _, ok := state.trace.Metadata[idInner]; !ok { + dis := state.protocol.Distance(newNode.ID(), state.target) + state.trace.Metadata[idInner] = &NodeMetadata{ + Enr: newNode.String(), + Distance: hexutil.Encode(dis[:]), + } + } + if newNode.ID() == state.protocol.Self().ID() { + continue + } + select { + case state.newNodeChan <- newNode: + case <-state.ctx.Done(): + return + } + } + state.trace.Responses[fromHexId] = respByNode +} + +func (state *contentLookupState) addNewCandidate(node *enode.Node) { + if state.found.Load() { + return + } + if !state.contacted[node.ID()] && !state.pending[node.ID()] { + state.candidates.push(node) + } +} + +// onQueryCompleted 处理查询完成 +func (state *contentLookupState) onQueryCompleted(node *enode.Node) { + state.mu.Lock() + defer state.mu.Unlock() + + delete(state.pending, node.ID()) + state.contacted[node.ID()] = true + state.activeQueries-- +} + +// isLookupComplete 检查查找是否完成 +func (state *contentLookupState) isLookupComplete() bool { + state.mu.Lock() + defer state.mu.Unlock() + state.protocol.Log.Trace("check lookup complete", "found", state.found.Load(), "candidates.len", state.candidates.len(), "activeQueries", state.activeQueries) + return state.found.Load() || (state.candidates.len() == 0 && state.activeQueries == 0) +} + +// hasLocalResult 检查本地是否有结果 +func (state *contentLookupState) hasLocalResult() bool { + if content, err := state.protocol.storage.Get(state.contentKey, state.target[:]); err == nil { + result := &ContentLookupResult{ + Content: content, + UtpTransfer: false, + FoundAt: state.protocol.Self(), + } + state.resultChan <- result + return true + } + return false +} + +// initCandidates 初始化候选节点 +func (state *contentLookupState) initCandidates() { + closestNodes := state.protocol.findNodesCloseToContent(state.target[:], bucketSize) + for _, node := range closestNodes { + if node.ID() != state.protocol.Self().ID() { + state.candidates.push(node) + } + } +} diff --git a/portalwire/portal_protocol.go b/portalwire/portal_protocol.go index 56f581b..96fc026 100644 --- a/portalwire/portal_protocol.go +++ b/portalwire/portal_protocol.go @@ -9,14 +9,16 @@ import ( "encoding/binary" "errors" "fmt" + "io" "math/big" "net" "slices" "sort" - "sync" "sync/atomic" "time" + "github.com/panjf2000/ants/v2" + "github.com/VictoriaMetrics/fastcache" "github.com/ethereum/go-ethereum/common" "github.com/ethereum/go-ethereum/common/hexutil" @@ -60,7 +62,9 @@ const ( defaultUTPReadTimeout = 60 * time.Second - DefaultUtpConnSize = 50 + DefaultUtpConnSize = 64 + + DefaultDiscV5RespTimeout = 3 * time.Second // These are the concurrent offers per Portal wire protocol that is running. // Using the `offerQueue` allows for limiting the amount of offers send and @@ -80,6 +84,8 @@ const ( lookupRequestLimit = 3 // max requests against a single node during lookup + lookupContentTimeout = 30 * time.Second + maxPacketSize = 1280 ) @@ -100,21 +106,18 @@ const ( Failed ) -var expirationVersionMinutes = 5 * time.Minute // cache versionsCache expiration time in minutes - -var versionsCacheSize = nBuckets * (bucketSize + maxReplacements) // VersionsCacheSize ideally should have the buckets plus the replacement Buckets size - -type protocolVersions []uint8 - -func (pv protocolVersions) ENRKey() string { return "pv" } - -var Versions protocolVersions = protocolVersions{0, 1} //protocol network versions defined here +const Tag ClientTag = "shisui" -type ClientTag string +const ( + contentOverhead = 1 + 1 // msg id + SSZ Union selector + enrOverhead = 4 // per added ENR, 4 bytes offset overheadvar expirationVersionMinutes = 5 * time.Minute // cache versionsCache expiration time in minutes + expirationVersionMinutes = 5 * time.Minute // cache versionsCache expiration time in minutes + DefaultLookupPoolSize = 200 +) -func (c ClientTag) ENRKey() string { return "c" } +var maxPayloadSize = maxPacketSize - talkRespOverhead - contentOverhead -const Tag ClientTag = "shisui" +var versionsCacheSize = nBuckets * (bucketSize + maxReplacements) // VersionsCacheSize ideally should have the buckets plus the replacement Buckets size var MaxDistance = hexutil.MustDecode("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff") @@ -145,6 +148,16 @@ var ( errLowPort = errors.New("low port") ) +type protocolVersions []uint8 + +func (pv protocolVersions) ENRKey() string { return "pv" } + +var Versions protocolVersions = protocolVersions{0, 1} //protocol network versions defined here + +type ClientTag string + +func (c ClientTag) ENRKey() string { return "c" } + type ContentElement struct { Node enode.ID ContentKeys [][]byte @@ -177,8 +190,6 @@ type OfferRequest struct { type OfferRequestWithNode struct { Request *OfferRequest Node *enode.Node - - permit Permit } type ContentInfoResp struct { @@ -186,13 +197,6 @@ type ContentInfoResp struct { UtpTransfer bool } -type traceContentInfoResp struct { - Node *enode.Node - Flag byte - Content any - UtpTransfer bool -} - // OfferTrace Define the type that can hold any variant type OfferTrace struct { Type OfferTraceType @@ -217,6 +221,7 @@ type PortalProtocolConfig struct { clock mclock.Clock TrustedBlockRoot []byte MaxUtpConnSize int + Discv5RespTimeout time.Duration } func DefaultPortalProtocolConfig() *PortalProtocolConfig { @@ -234,6 +239,7 @@ func DefaultPortalProtocolConfig() *PortalProtocolConfig { clock: mclock.System{}, TrustedBlockRoot: make([]byte, 0), MaxUtpConnSize: DefaultUtpConnSize, + Discv5RespTimeout: DefaultDiscV5RespTimeout, } } @@ -277,6 +283,8 @@ type PortalProtocol struct { transferringKeyCache *fastcache.Cache versionsCache cache.Cache[*enode.Node, uint8] + + lookupContentPool *ants.Pool } func defaultContentIdFunc(contentKey []byte) []byte { @@ -299,7 +307,11 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId ProtocolId, priv } closeCtx, cancelCloseCtx := context.WithCancel(context.Background()) - + pool, err := ants.NewPool(DefaultLookupPoolSize, ants.WithPreAlloc(true)) + if err != nil { + cancelCloseCtx() + return nil, err + } protocol := &PortalProtocol{ protocolId: string(protocolId), protocolName: protocolId.Name(), @@ -326,6 +338,7 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId ProtocolId, priv Utp: utp, currentVersions: currentVersions, transferringKeyCache: fastcache.New(config.ContentKeyCacheSize), + lookupContentPool: pool, } for _, setOpt := range setOpts { @@ -582,12 +595,12 @@ func (p *PortalProtocol) findNodes(node *enode.Node, distances []uint) ([]*enode return p.processNodes(node, talkResp, distances) } -func (p *PortalProtocol) findContent(node *enode.Node, contentKey []byte) (byte, interface{}, error) { +func (p *PortalProtocol) findContent(ctx context.Context, node *enode.Node, contentKey []byte) (byte, interface{}, error) { findContent := &FindContent{ ContentKey: contentKey, } - p.Log.Trace(">> FIND_CONTENT/"+p.protocolName, "id", node.ID(), "findContent", findContent) + p.Log.Info(">> FIND_CONTENT/"+p.protocolName, "id", node.ID(), "findContent", findContent) if metrics.Enabled() { p.portalMetrics.messagesSentFindContent.Mark(1) } @@ -601,13 +614,19 @@ func (p *PortalProtocol) findContent(node *enode.Node, contentKey []byte) (byte, copy(talkRequestBytes[1:], findContentBytes) talkResp, err := p.DiscV5.TalkRequest(node, p.protocolId, talkRequestBytes) if err != nil { - return 0xff, nil, err + return 0xff, nil, fmt.Errorf("find content in discv5 has err: %s", err.Error()) } return p.processContent(node, talkResp) } func (p *PortalProtocol) offer(node *enode.Node, offerRequest *OfferRequest, permit Permit) ([]byte, error) { + permitReleaseInProcessOffer := false + defer func() { + if !permitReleaseInProcessOffer { + p.Utp.releasePermit(permit) + } + }() contentKeys := getContentKeys(offerRequest) offer := &Offer{ @@ -631,14 +650,15 @@ func (p *PortalProtocol) offer(node *enode.Node, offerRequest *OfferRequest, per return nil, err } + permitReleaseInProcessOffer = true return p.processOffer(node, talkResp, offerRequest, permit) } func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request *OfferRequest, permit Permit) ([]byte, error) { - notStartedUtp := true + permitReleaseInGo := false defer func() { - if notStartedUtp { - permit.Release() + if !permitReleaseInGo { + p.Utp.releasePermit(permit) } }() var err error @@ -701,10 +721,10 @@ func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request * p.Log.Debug("trace offer declined", "keys", accept.GetContentKeys()) return accept.GetContentKeys(), nil } - notStartedUtp = false + permitReleaseInGo = true connId := binary.BigEndian.Uint16(accept.GetConnectionId()) - go func(ctx context.Context) { - defer permit.Release() + go func(ctx context.Context, pr Permit) { + defer p.Utp.releasePermit(pr) var conn *utp.UtpStream for { select { @@ -756,7 +776,7 @@ func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request * Type: Failed, } } - p.Log.Error("failed to dial utp connection", "err", err) + p.Log.Error("failed to dial utp connection", "destId", target.ID().String(), "ip", target.IP().String(), "port", target.UDP(), "err", err) return } @@ -791,7 +811,7 @@ func (p *PortalProtocol) processOffer(target *enode.Node, resp []byte, request * return } } - }(p.closeCtx) + }(p.closeCtx, permit) return accept.GetContentKeys(), nil } @@ -859,13 +879,13 @@ func (p *PortalProtocol) processContent(target *enode.Node, resp []byte) (byte, defer readCancel() var data []byte n, err := conn.ReadToEOF(readCtx, &data) - if err != nil { + if err != nil && !errors.Is(err, io.EOF) { if metrics.Enabled() { p.portalMetrics.utpInFailRead.Inc(1) } return 0xff, nil, err } - p.Log.Trace("<< CONTENT/"+p.protocolName, "id", target.ID(), "size", n, "data", data) + p.Log.Trace("<< CONTENT/"+p.protocolName, "id", target.ID(), "size", n) data, err = p.decodeUtpContent(target, data) if err != nil { if metrics.Enabled() { @@ -1411,10 +1431,47 @@ func (p *PortalProtocol) handleFindNodes(fromAddr *net.UDPAddr, request *FindNod return talkRespBytes, nil } +func (p *PortalProtocol) closestNodeToContent(n *enode.Node, addr *net.UDPAddr, contentId []byte, limit int) ([]byte, error) { + closestNodes := p.findNodesCloseToContent(contentId, limit) + for i, closeNode := range closestNodes { + if closeNode.ID() == n.ID() { + closestNodes = append(closestNodes[:i], closestNodes[i+1:]...) + break + } + } + + enrs := p.truncateNodes(closestNodes, maxPayloadSize, enrOverhead) + // TODO fix when no content and no enrs found + if len(enrs) == 0 { + enrs = nil + } + + enrsMsg := &Enrs{ + Enrs: enrs, + } + + p.Log.Trace(">> CONTENT_ENRS/"+p.protocolName, "protocol", p.protocolName, "source", addr, "enrs.size", len(enrsMsg.Enrs)) + if metrics.Enabled() { + p.portalMetrics.messagesSentContent.Mark(1) + } + var enrsMsgBytes []byte + enrsMsgBytes, err := enrsMsg.MarshalSSZ() + if err != nil { + return nil, err + } + + contentMsgBytes := make([]byte, 0, len(enrsMsgBytes)+1) + contentMsgBytes = append(contentMsgBytes, ContentEnrsSelector) + contentMsgBytes = append(contentMsgBytes, enrsMsgBytes...) + + talkRespBytes := make([]byte, 0, len(contentMsgBytes)+1) + talkRespBytes = append(talkRespBytes, CONTENT) + talkRespBytes = append(talkRespBytes, contentMsgBytes...) + + return talkRespBytes, nil +} + func (p *PortalProtocol) handleFindContent(n *enode.Node, addr *net.UDPAddr, request *FindContent) ([]byte, error) { - contentOverhead := 1 + 1 // msg id + SSZ Union selector - maxPayloadSize := maxPacketSize - talkRespOverhead - contentOverhead - enrOverhead := 4 // per added ENR, 4 bytes offset overhead var err error contentKey := request.ContentKey contentId := p.toContentId(contentKey) @@ -1429,49 +1486,13 @@ func (p *PortalProtocol) handleFindContent(n *enode.Node, addr *net.UDPAddr, req } if errors.Is(err, ErrContentNotFound) { - closestNodes := p.findNodesCloseToContent(contentId, portalFindnodesResultLimit) - for i, closeNode := range closestNodes { - if closeNode.ID() == n.ID() { - closestNodes = append(closestNodes[:i], closestNodes[i+1:]...) - break - } - } - - enrs := p.truncateNodes(closestNodes, maxPayloadSize, enrOverhead) - // TODO fix when no content and no enrs found - if len(enrs) == 0 { - enrs = nil - } - - enrsMsg := &Enrs{ - Enrs: enrs, - } - - p.Log.Trace(">> CONTENT_ENRS/"+p.protocolName, "protocol", p.protocolName, "source", addr, "enrs.size", len(enrsMsg.Enrs)) - if metrics.Enabled() { - p.portalMetrics.messagesSentContent.Mark(1) - } - var enrsMsgBytes []byte - enrsMsgBytes, err = enrsMsg.MarshalSSZ() - if err != nil { - return nil, err - } - - contentMsgBytes := make([]byte, 0, len(enrsMsgBytes)+1) - contentMsgBytes = append(contentMsgBytes, ContentEnrsSelector) - contentMsgBytes = append(contentMsgBytes, enrsMsgBytes...) - - talkRespBytes := make([]byte, 0, len(contentMsgBytes)+1) - talkRespBytes = append(talkRespBytes, CONTENT) - talkRespBytes = append(talkRespBytes, contentMsgBytes...) - - return talkRespBytes, nil + return p.closestNodeToContent(n, addr, contentId, portalFindnodesResultLimit) } else if len(content) <= maxPayloadSize { rawContentMsg := &Content{ Content: content, } - p.Log.Trace(">> CONTENT_RAW/"+p.protocolName, "protocol", p.protocolName, "source", addr, "content", rawContentMsg) + p.Log.Trace(">> CONTENT_RAW/"+p.protocolName, "protocol", p.protocolName, "source", addr, "content.len", len(rawContentMsg.Content)) if metrics.Enabled() { p.portalMetrics.messagesSentContent.Mark(1) } @@ -1494,7 +1515,12 @@ func (p *PortalProtocol) handleFindContent(n *enode.Node, addr *net.UDPAddr, req } else { connectionId := p.Utp.CidWithAddr(n, addr, false) + permit := p.Utp.getOutboundLimit() + if permit != PermitOutbound { + return p.closestNodeToContent(n, addr, contentId, portalFindnodesResultLimit) + } go func(bctx context.Context, connId *utp.ConnectionId) { + defer p.Utp.releasePermit(permit) var conn *utp.UtpStream var connectCtx context.Context var cancel context.CancelFunc @@ -1511,7 +1537,7 @@ func (p *PortalProtocol) handleFindContent(n *enode.Node, addr *net.UDPAddr, req if metrics.Enabled() { p.portalMetrics.utpOutFailConn.Inc(1) } - p.Log.Error("failed to accept utp connection for handle find content", "connId", connectionId.Send, "err", err) + p.Log.Error("failed to accept utp connection for handle find content", "destId", n.ID().String(), "addr", addr.String(), "connId", connectionId.Send, "err", err) return } @@ -1525,26 +1551,24 @@ func (p *PortalProtocol) handleFindContent(n *enode.Node, addr *net.UDPAddr, req p.Log.Error("encode utp content failed", "err", err) return } - var n int - n, err = conn.Write(writeCtx, content) + + _, err = conn.Write(writeCtx, content) conn.Close() if err != nil { if metrics.Enabled() { p.portalMetrics.utpOutFailWrite.Inc(1) } - p.Log.Error("failed to write content to utp connection", "err", err) + p.Log.Error("failed to write content to utp connection", "err", err, "destId", n.ID().String(), "destAddr", addr.String(), "connId", connectionId.Send) return } if metrics.Enabled() { p.portalMetrics.utpOutSuccess.Inc(1) } - p.Log.Trace("wrote content size to utp connection", "n", n) return } } }(p.closeCtx, connectionId) - idBuffer := make([]byte, 2) binary.BigEndian.PutUint16(idBuffer, connectionId.Send) connIdMsg := &ConnectionId{ @@ -1586,8 +1610,8 @@ func (p *PortalProtocol) handleOffer(node *enode.Node, addr *net.UDPAddr, reques idBuffer := make([]byte, 2) idValue := uint16(0) if len(contentKeys) > 0 { - permit, getPermit := p.Utp.GetInboundPermit() - if !getPermit { + permit := p.Utp.getInboundLimit() + if permit != PermitInbound { p.Log.Debug("utp rate limited") if acceptV1, isV1 := accept.(*AcceptV1); isV1 { keysLen := len(acceptV1.ContentKeys) @@ -1599,9 +1623,11 @@ func (p *PortalProtocol) handleOffer(node *enode.Node, addr *net.UDPAddr, reques } } else { connectionId := p.Utp.CidWithAddr(node, addr, false) - go func(bctx context.Context, connId *utp.ConnectionId, releasePermit Permit) { - defer releasePermit.Release() - defer p.deleteTransferringContentKeys(contentKeys) + go func(bctx context.Context, connId *utp.ConnectionId, permitValue Permit) { + defer func() { + p.Utp.releasePermit(permitValue) + p.deleteTransferringContentKeys(contentKeys) + }() var conn *utp.UtpStream for { select { @@ -1617,7 +1643,7 @@ func (p *PortalProtocol) handleOffer(node *enode.Node, addr *net.UDPAddr, reques if metrics.Enabled() { p.portalMetrics.utpInFailConn.Inc(1) } - p.Log.Error("failed to accept utp connection for handle offer", "connId", connectionId.Send, "err", err) + p.Log.Error("failed to accept utp connection for handle offer", "destId", node.ID().String(), "addr", addr.String(), "connId", connectionId.Send, "err", err) return } @@ -1628,9 +1654,14 @@ func (p *PortalProtocol) handleOffer(node *enode.Node, addr *net.UDPAddr, reques defer readCancel() n, err = conn.ReadToEOF(readCtx, &data) conn.Close() - // release permit fast - releasePermit.Release() - p.Log.Trace("<< OFFER_CONTENT/"+p.protocolName, "id", node.ID(), "size", n, "data", data) + if err != nil && !errors.Is(err, io.EOF) { + if metrics.Enabled() { + p.portalMetrics.utpInFailRead.Inc(1) + } + p.Log.Error("failed to read from the accepted utp connection for handling offer", "addr", addr.String(), "connId", connectionId.Send, "err", err) + return + } + p.Log.Trace("<< OFFER_CONTENT/"+p.protocolName, "id", node.ID(), "size", n) if metrics.Enabled() { p.portalMetrics.messagesReceivedContent.Mark(1) } @@ -1644,6 +1675,7 @@ func (p *PortalProtocol) handleOffer(node *enode.Node, addr *net.UDPAddr, reques if metrics.Enabled() { p.portalMetrics.utpInSuccess.Inc(1) } + return } } }(p.closeCtx, connectionId, permit) @@ -1806,9 +1838,14 @@ func (p *PortalProtocol) offerWorker() { return case offerRequestWithNode := <-p.offerQueue: p.Log.Trace("offerWorker", "offerRequestWithNode", offerRequestWithNode) - _, err := p.offer(offerRequestWithNode.Node, offerRequestWithNode.Request, offerRequestWithNode.permit) + permit := p.Utp.getOutboundLimit() + if permit == PermitReject { + p.Log.Debug("utp rate limited, rejecting offer", "offerRequestWithNode", offerRequestWithNode) + continue + } + _, err := p.offer(offerRequestWithNode.Node, offerRequestWithNode.Request, permit) if err != nil { - p.Log.Error("failed to offer", "err", err) + p.Log.Error("failed to offer", "destNodeId", offerRequestWithNode.Node.ID().String(), "ip", offerRequestWithNode.Node.IP().String(), "port", offerRequestWithNode.Node.UDP(), "err", err) } } } @@ -1933,50 +1970,38 @@ func (p *PortalProtocol) collectTableNodes(rip net.IP, distances []uint, limit i return nodes } +var concurrencyLookup = atomic.Int32{} + func (p *PortalProtocol) ContentLookup(contentKey, contentId []byte) ([]byte, bool, error) { - lookupContext, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), lookupContentTimeout) defer cancel() + start := time.Now() + state := newContentLookupState(ctx, cancel, p, contentId, contentKey, nil) + defer concurrencyLookup.Add(-1) + if err := p.lookupContentPool.Submit(state.run); err != nil { + return nil, false, ErrContentNotFound + } + c := concurrencyLookup.Add(1) + p.Log.Info("submitted a state to contentLookup pool", "concurrency", c) - resChan := make(chan *traceContentInfoResp, alpha) - hasResult := int32(0) - - result := ContentInfoResp{} - - var wg sync.WaitGroup - wg.Add(1) - - go func() { - defer wg.Done() - for res := range resChan { - if res.Flag != ContentEnrsSelector { - result.Content = res.Content.([]byte) - result.UtpTransfer = res.UtpTransfer - } + // 等待结果 + select { + case result := <-state.resultChan: + p.Log.Info("get content result", "duration", time.Since(start).Milliseconds(), "contentKey", hexutil.Encode(contentKey)) + if result != nil { + return result.Content, result.UtpTransfer, nil } - }() - - newLookup(lookupContext, p.table, enode.ID(contentId), func(n *enode.Node) ([]*enode.Node, error) { - return p.contentLookupWorker(n, contentKey, resChan, cancel, &hasResult) - }).run() - close(resChan) - - wg.Wait() - if hasResult == 1 { - return result.Content, result.UtpTransfer, nil + return nil, false, ErrContentNotFound + case <-ctx.Done(): + return nil, false, fmt.Errorf("content lookup timeout: %w", ctx.Err(), "contentKey", hexutil.Encode(contentKey)) } - - return nil, false, ErrContentNotFound } func (p *PortalProtocol) TraceContentLookup(contentKey, contentId []byte) (*TraceContentResult, error) { - lookupContext, cancel := context.WithCancel(context.Background()) - // resp channel - resChan := make(chan *traceContentInfoResp, alpha) - - hasResult := int32(0) + ctx, cancel := context.WithTimeout(context.Background(), lookupContentTimeout) + defer cancel() traceContentRes := &TraceContentResult{} - selfHexId := "0x" + p.Self().ID().String() trace := &Trace{ @@ -1987,130 +2012,55 @@ func (p *PortalProtocol) TraceContentLookup(contentKey, contentId []byte) (*Trac Metadata: make(map[string]*NodeMetadata), Cancelled: make([]string, 0), } + p.setupLocalTraceInfo(trace, selfHexId, contentId) - nodes := p.table.findnodeByID(enode.ID(contentId), bucketSize, false) + state := newContentLookupState(ctx, cancel, p, contentId, contentKey, trace) + + if err := p.lookupContentPool.Submit(state.run); err != nil { + return nil, ErrContentNotFound + } + + // 等待结果 + select { + case result := <-state.resultChan: + if result != nil { + hexId := "0x" + result.FoundAt.ID().String() + trace.ReceivedFrom = hexId + traceContentRes.Content = hexutil.Encode(result.Content) + traceContentRes.UtpTransfer = result.UtpTransfer + + // 确保找到内容的节点有一个空的响应记录（表示它有内容，不返回其他节点） + if _, exists := trace.Responses[hexId]; !exists { + trace.Responses[hexId] = RespByNode{ + RespondedWith: nil, + } + } + } + case <-ctx.Done(): + return nil, fmt.Errorf("content lookup timeout: %w", ctx.Err()) + } + traceContentRes.Trace = *trace + return traceContentRes, nil +} +func (p *PortalProtocol) setupLocalTraceInfo(trace *Trace, selfHexId string, contentId []byte) { + nodes := p.table.findnodeByID(enode.ID(contentId), bucketSize, false) localResponse := make([]string, 0, len(nodes.entries)) for _, node := range nodes.entries { id := "0x" + node.ID().String() localResponse = append(localResponse, id) } + trace.Responses[selfHexId] = RespByNode{ DurationMs: 0, RespondedWith: localResponse, } dis := p.Distance(p.Self().ID(), enode.ID(contentId)) - trace.Metadata[selfHexId] = &NodeMetadata{ Enr: p.Self().String(), Distance: hexutil.Encode(dis[:]), } - - var wg sync.WaitGroup - wg.Add(1) - - go func() { - defer wg.Done() - for res := range resChan { - node := res.Node - hexId := "0x" + node.ID().String() - dis := p.Distance(node.ID(), enode.ID(contentId)) - p.Log.Debug("reveice res", "id", hexId, "flag", res.Flag) - trace.Metadata[hexId] = &NodeMetadata{ - Enr: node.String(), - Distance: hexutil.Encode(dis[:]), - } - // no content return - if traceContentRes.Content == "" { - if res.Flag == ContentRawSelector || res.Flag == ContentConnIdSelector { - trace.ReceivedFrom = hexId - content := res.Content.([]byte) - traceContentRes.Content = hexutil.Encode(content) - traceContentRes.UtpTransfer = res.UtpTransfer - trace.Responses[hexId] = RespByNode{} - } else { - nodes := res.Content.([]*enode.Node) - respByNode := RespByNode{ - RespondedWith: make([]string, 0, len(nodes)), - } - for _, node := range nodes { - idInner := "0x" + node.ID().String() - respByNode.RespondedWith = append(respByNode.RespondedWith, idInner) - if _, ok := trace.Metadata[idInner]; !ok { - dis := p.Distance(node.ID(), enode.ID(contentId)) - trace.Metadata[idInner] = &NodeMetadata{ - Enr: node.String(), - Distance: hexutil.Encode(dis[:]), - } - } - trace.Responses[hexId] = respByNode - } - } - } else { - trace.Cancelled = append(trace.Cancelled, hexId) - } - } - }() - - lookup := newLookup(lookupContext, p.table, enode.ID(contentId), func(n *enode.Node) ([]*enode.Node, error) { - return p.contentLookupWorker(n, contentKey, resChan, cancel, &hasResult) - }) - lookup.run() - close(resChan) - - wg.Wait() - if hasResult == 0 { - cancel() - } - traceContentRes.Trace = *trace - - return traceContentRes, nil -} - -func (p *PortalProtocol) contentLookupWorker(n *enode.Node, contentKey []byte, resChan chan<- *traceContentInfoResp, cancel context.CancelFunc, done *int32) ([]*enode.Node, error) { - wrapedNode := make([]*enode.Node, 0) - flag, content, err := p.findContent(n, contentKey) - if err != nil { - return nil, err - } - p.Log.Debug("traceContentLookupWorker reveice response", "ip", n.IP().String(), "flag", flag) - - switch flag { - case ContentRawSelector, ContentConnIdSelector: - content, ok := content.([]byte) - if !ok { - return wrapedNode, fmt.Errorf("failed to assert to raw content, value is: %v", content) - } - res := &traceContentInfoResp{ - Node: n, - Flag: flag, - Content: content, - UtpTransfer: false, - } - if flag == ContentConnIdSelector { - res.UtpTransfer = true - } - if atomic.CompareAndSwapInt32(done, 0, 1) { - p.Log.Debug("contentLookupWorker find content", "ip", n.IP().String(), "port", n.UDP()) - resChan <- res - cancel() - } - return wrapedNode, err - case ContentEnrsSelector: - nodes, ok := content.([]*enode.Node) - if !ok { - return wrapedNode, fmt.Errorf("failed to assert to enrs content, value is: %v", content) - } - resChan <- &traceContentInfoResp{ - Node: n, - Flag: flag, - Content: content, - UtpTransfer: false, - } - return nodes, nil - } - return wrapedNode, nil } func (p *PortalProtocol) ToContentId(contentKey []byte) []byte { @@ -2198,11 +2148,6 @@ func (p *PortalProtocol) GossipAndReturnPeers(srcNodeId *enode.ID, contentKeys [ } for _, n := range finalGossipNodes { - permit, ok := p.Utp.GetOutboundPermit() - if !ok { - p.Log.Debug("reached utp conn limit, will drop this content", "network", p.protocolName, "nodeId", n.ID(), "addr", n.IPAddr().String()) - continue - } transientOfferRequest := &TransientOfferRequest{ Contents: contentList, } @@ -2215,7 +2160,6 @@ func (p *PortalProtocol) GossipAndReturnPeers(srcNodeId *enode.ID, contentKeys [ offerRequestWithNode := &OfferRequestWithNode{ Node: n, Request: offerRequest, - permit: permit, } select { case p.offerQueue <- offerRequestWithNode: diff --git a/portalwire/portal_protocol_test.go b/portalwire/portal_protocol_test.go index b95a13d..1ee38a3 100644 --- a/portalwire/portal_protocol_test.go +++ b/portalwire/portal_protocol_test.go @@ -10,6 +10,7 @@ import ( "net" "os" "sync" + "sync/atomic" "testing" "time" @@ -219,7 +220,7 @@ func TestPortalWireProtocolUdp(t *testing.T) { defer connId2Conn.Close() var buf []byte n, err := connId2Conn.ReadToEOF(context.Background(), &buf) - assert.NoError(t, err) + assert.Equal(t, io.EOF, err) assert.Equal(t, len(cliSendMsgWithCid2)+len(largeTestContent), n) assert.True(t, bytes.Equal([]byte(cliSendMsgWithCid2), buf[:len(cliSendMsgWithCid2)])) assert.True(t, bytes.Equal(largeTestContent, buf[len(cliSendMsgWithCid2):])) @@ -311,12 +312,12 @@ func TestPortalWireProtocol(t *testing.T) { err = node1.storage.Put(nil, node1.toContentId([]byte("test_key")), []byte("test_value")) assert.NoError(t, err) - flag, content, err := node2.findContent(node1.localNode.Node(), []byte("test_key")) + flag, content, err := node2.findContent(context.Background(), node1.localNode.Node(), []byte("test_key")) assert.NoError(t, err) assert.Equal(t, ContentRawSelector, flag) assert.Equal(t, []byte("test_value"), content) - flag, content, err = node2.findContent(node3.localNode.Node(), []byte("test_key")) + flag, content, err = node2.findContent(context.Background(), node3.localNode.Node(), []byte("test_key")) assert.NoError(t, err) assert.Equal(t, ContentEnrsSelector, flag) assert.Equal(t, 1, len(content.([]*enode.Node))) @@ -331,7 +332,7 @@ func TestPortalWireProtocol(t *testing.T) { err = node1.storage.Put(nil, node1.toContentId([]byte("large_test_key")), largeTestContent) assert.NoError(t, err) - flag, content, err = node2.findContent(node1.localNode.Node(), []byte("large_test_key")) + flag, content, err = node2.findContent(context.Background(), node1.localNode.Node(), []byte("large_test_key")) assert.NoError(t, err) assert.Equal(t, largeTestContent, content) assert.Equal(t, ContentConnIdSelector, flag) @@ -355,7 +356,7 @@ func TestPortalWireProtocol(t *testing.T) { Request: testTransientOfferRequest, } - contentKeys, err := node1.offer(node3.localNode.Node(), offerRequest, &NoPermit{}) + contentKeys, err := node1.offer(node3.localNode.Node(), offerRequest, PermitNotLimit) assert.Equal(t, uint64(2), bitfield.Bitlist(contentKeys).Count()) assert.NoError(t, err) @@ -402,7 +403,7 @@ func TestPortalWireProtocol(t *testing.T) { Request: testTransientOfferRequestWithResult, } - _, err = node1.offer(node3.localNode.Node(), traceOfferRequest, &NoPermit{}) + _, err = node1.offer(node3.localNode.Node(), traceOfferRequest, PermitNotLimit) assert.NoError(t, err) offerTrace := <-testTransientOfferRequestWithResult.Result @@ -421,7 +422,7 @@ func TestPortalWireProtocol(t *testing.T) { err = node3.storage.Put(nil, node3.toContentId(testTraceEntry.ContentKey), testTraceEntry.Content) assert.NoError(t, err) - _, err = node1.offer(node3.localNode.Node(), traceOfferRequest1, &NoPermit{}) + _, err = node1.offer(node3.localNode.Node(), traceOfferRequest1, PermitNotLimit) assert.NoError(t, err) offerTrace1 := <-testTransientOfferRequestWithResult1.Result @@ -480,7 +481,6 @@ func TestContentLookup(t *testing.T) { node2.Log = testlog.Logger(t, log.LvlInfo) err = node2.Start() assert.NoError(t, err) - fmt.Println(node2.localNode.Node().String()) node3, err := setupLocalPortalNode(":17779", []*enode.Node{node1.localNode.Node(), node2.localNode.Node()}, DefaultUtpConnSize) assert.NoError(t, err) @@ -508,7 +508,7 @@ func TestContentLookup(t *testing.T) { _, err = node2.ping(node3.localNode.Node()) assert.NoError(t, err) - res, _, err := node1.ContentLookup(contentKey, contentId) + res, _, err := node2.ContentLookup(contentKey, contentId) assert.NoError(t, err) assert.Equal(t, res, content) @@ -547,20 +547,20 @@ func TestTraceContentLookup(t *testing.T) { content := []byte{0x1, 0x2} contentId := node1.toContentId(contentKey) - err = node1.storage.Put(nil, contentId, content) - assert.NoError(t, err) - node1Id := hexutil.Encode(node1.Self().ID().Bytes()) node2Id := hexutil.Encode(node2.Self().ID().Bytes()) node3Id := hexutil.Encode(node3.Self().ID().Bytes()) + err = node1.storage.Put(nil, contentId, content) + assert.NoError(t, err) + res, err := node3.TraceContentLookup(contentKey, contentId) assert.NoError(t, err) - assert.Equal(t, res.Content, hexutil.Encode(content)) - assert.Equal(t, res.UtpTransfer, false) - assert.Equal(t, res.Trace.Origin, node3Id) - assert.Equal(t, res.Trace.TargetId, hexutil.Encode(contentId)) - assert.Equal(t, res.Trace.ReceivedFrom, node1Id) + assert.Equal(t, hexutil.Encode(content), res.Content) + assert.Equal(t, false, res.UtpTransfer) + assert.Equal(t, node3Id, res.Trace.Origin) + assert.Equal(t, hexutil.Encode(contentId), res.Trace.TargetId) + assert.Equal(t, node1Id, res.Trace.ReceivedFrom) // check nodeMeta node1Meta := res.Trace.Metadata[node1Id] @@ -730,7 +730,7 @@ func TestOfferV1(t *testing.T) { Request: testTransientOfferRequest, } // all accept - contentKeys, err := node1.offer(node2.localNode.Node(), offerRequest, &NoPermit{}) + contentKeys, err := node1.offer(node2.localNode.Node(), offerRequest, PermitNotLimit) assert.NoError(t, err) assert.Equal(t, 2, len(contentKeys)) for _, val := range contentKeys { @@ -740,7 +740,7 @@ func TestOfferV1(t *testing.T) { // one reject node1.storage.Put(testEntry1.ContentKey, node2.toContentId(testEntry1.ContentKey), testEntry1.Content) node1.transferringKeyCache.Set(testEntry2.ContentKey, EmptyBytes) - acceptCodes, err := node2.offer(node1.localNode.Node(), offerRequest, &NoPermit{}) + acceptCodes, err := node2.offer(node1.localNode.Node(), offerRequest, PermitNotLimit) assert.NoError(t, err) assert.Equal(t, 2, len(acceptCodes)) assert.Equal(t, uint8(AlreadyStored), acceptCodes[0]) @@ -980,6 +980,61 @@ func TestGetOrStoreHighestVersionOverflow(t *testing.T) { assert.Equal(t, 1, node1.versionsCache.Len()) } +func TestHandleFindContent_Ratelimit(t *testing.T) { + original := maxPayloadSize + maxPayloadSize = 10 + defer func() { + maxPayloadSize = original + }() + + node1, err := setupLocalPortalNode(":3321", nil, 0, 0, 1) + assert.NoError(t, err) + node1.Log = testlog.Logger(t, log.LevelCrit) + err = node1.Start() + assert.NoError(t, err) + defer stopNode(node1) + + node2, err := setupLocalPortalNode(":3322", []*enode.Node{node1.localNode.Node()}, 64, 0, 1) + assert.NoError(t, err) + node2.Log = testlog.Logger(t, log.LevelCrit) + err = node2.Start() + assert.NoError(t, err) + defer stopNode(node2) + + for i := 0; i < 10; i++ { + key := []byte(fmt.Sprintf("test_entry_key_%d", i)) + value := []byte(fmt.Sprintf("test_entry_value_%d", i)) + node2.storage.Put(key, node2.toContentId(key), value) + } + + addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", node1.localNode.Node().IPAddr().String(), node1.localNode.Node().UDP())) + assert.NoError(t, err) + testKey := &FindContent{ + ContentKey: []byte("test_entry_key_1"), + } + var permitCount atomic.Int32 + permitCountPtr := &permitCount + + var wg sync.WaitGroup + var waitGoroutines sync.WaitGroup + wg.Add(1000) + waitGoroutines.Add(1000) + for i := 0; i < 1000; i++ { + go func() { + defer wg.Done() + waitGoroutines.Wait() + content, err2 := node2.handleFindContent(node1.localNode.Node(), addr, testKey) + assert.NoError(t, err2) + if len(content) != 2 { + permitCountPtr.Add(1) + } + }() + waitGoroutines.Done() + } + wg.Wait() + assert.Equal(t, int32(64), permitCountPtr.Load()) +} + func TestAcceptCode_Ratelmit(t *testing.T) { node1, err := setupLocalPortalNode(":3321", nil, 0, 0, 1) assert.NoError(t, err) @@ -1017,7 +1072,7 @@ func TestAcceptCode_Ratelmit(t *testing.T) { Request: testTransientOfferRequest, } // all accept - contentKeys, err := node1.offer(node2.localNode.Node(), offerRequest, &NoPermit{}) + contentKeys, err := node1.offer(node2.localNode.Node(), offerRequest, PermitReject) assert.NoError(t, err) assert.Len(t, contentKeys, 2, "excepted: 2, but got: %d", len(contentKeys)) for _, val := range contentKeys { diff --git a/portalwire/utp_transport.go b/portalwire/utp_transport.go index ca20623..f780c04 100644 --- a/portalwire/utp_transport.go +++ b/portalwire/utp_transport.go @@ -24,14 +24,17 @@ var ( ) type UtpTransportService struct { - startOnce sync.Once - ctx context.Context - log log.Logger - discV5 *discover.UDPv5 - socket *utp.UtpSocket - socketConfig *utp.ConnectionConfig - ListenAddr string - utpController *utpController + startOnce sync.Once + ctx context.Context + log log.Logger + discV5 *discover.UDPv5 + socket *utp.UtpSocket + socketConfig *utp.ConnectionConfig + ListenAddr string + utpControllerRef *utpController + getInboundLimit func() Permit + getOutboundLimit func() Permit + releasePermit func(permit Permit) } type UtpPeer struct { @@ -76,27 +79,14 @@ type packetItem struct { data []byte } -// ReleasePermit is a interface type that releases a UTP connection permit. -type Permit interface { - Release() -} - -type NoPermit struct{} - -func (n *NoPermit) Release() {} +type Permit int -// ReleasePermit is a function type that releases a UTP connection permit. -// It is returned when a permit is successfully acquired and must be called to release the permit. -type ReleasePermit struct { - released atomic.Bool - action func() -} - -func (p *ReleasePermit) Release() { - if p.released.CompareAndSwap(false, true) { - p.action() - } -} +const ( + PermitNotLimit Permit = iota + PermitReject + PermitInbound + PermitOutbound +) type utpController struct { inboundLimit *semaphore.Weighted @@ -112,36 +102,38 @@ func newUtpController(maxLimit int) *utpController { // GetInboundPermit tries to acquire a permit for inbound UTP connections. // It is returned when a permit is successfully acquired and must be called to release the permit. -func (u *utpController) GetInboundPermit() (Permit, bool) { - if ok := u.inboundLimit.TryAcquire(1); !ok { - return &NoPermit{}, false +func (u *utpController) GetInboundPermit() Permit { + ok := u.inboundLimit.TryAcquire(1) + if !ok { + return PermitReject } - return &ReleasePermit{ - action: func() { - u.inboundLimit.Release(1) - }, - }, true + return PermitInbound } // GetInboundPermit tries to acquire a permit for outbound UTP connections. // It is returned when a permit is successfully acquired and must be called to release the permit. -func (u *utpController) GetOutboundPermit() (Permit, bool) { - if ok := u.outboundLimit.TryAcquire(1); !ok { - return &NoPermit{}, false +func (u *utpController) GetOutboundPermit() Permit { + ok := u.outboundLimit.TryAcquire(1) + if !ok { + return PermitReject + } + return PermitOutbound +} + +func (u *utpController) Release(permit Permit) { + switch permit { + case PermitInbound: + u.inboundLimit.Release(1) + case PermitOutbound: + u.outboundLimit.Release(1) } - return &ReleasePermit{ - action: func() { - u.outboundLimit.Release(1) - }, - }, true } type discv5Conn struct { - logger log.Logger - receive chan *packetItem - conn *discover.UDPv5 - closed *atomic.Bool - UtpController *utpController + logger log.Logger + receive chan *packetItem + conn *discover.UDPv5 + closed *atomic.Bool } func newDiscv5Conn(conn *discover.UDPv5, logger log.Logger) *discv5Conn { @@ -192,13 +184,17 @@ func (c *discv5Conn) Close() error { } func NewZenEthUtp(ctx context.Context, config *PortalProtocolConfig, discV5 *discover.UDPv5, conn discover.UDPConn) *UtpTransportService { + utpControllerRef := newUtpController(config.MaxUtpConnSize) uts := &UtpTransportService{ - ctx: ctx, - log: log.New("protocol", "utp", "local", conn.LocalAddr().String()), - discV5: discV5, - socketConfig: utp.NewConnectionConfig(), - ListenAddr: config.ListenAddr, - utpController: newUtpController(config.MaxUtpConnSize), + ctx: ctx, + log: log.New("protocol", "utp", "local", conn.LocalAddr().String()), + discV5: discV5, + socketConfig: utp.NewConnectionConfig(), + ListenAddr: config.ListenAddr, + utpControllerRef: utpControllerRef, + getInboundLimit: utpControllerRef.GetInboundPermit, + getOutboundLimit: utpControllerRef.GetOutboundPermit, + releasePermit: utpControllerRef.Release, } return uts } @@ -212,14 +208,6 @@ func (z *UtpTransportService) Start() error { return nil } -func (z *UtpTransportService) GetOutboundPermit() (Permit, bool) { - return z.utpController.GetOutboundPermit() -} - -func (z *UtpTransportService) GetInboundPermit() (Permit, bool) { - return z.utpController.GetInboundPermit() -} - func (z *UtpTransportService) DialWithCid(ctx context.Context, dest *enode.Node, connId uint16) (*utp.UtpStream, error) { cid := z.SendId(dest, connId) stream, err := z.socket.ConnectWithCid(ctx, cid, z.socketConfig) diff --git a/portalwire/utp_transport_test.go b/portalwire/utp_transport_test.go index 72f37ba..c2d59f8 100644 --- a/portalwire/utp_transport_test.go +++ b/portalwire/utp_transport_test.go @@ -7,46 +7,41 @@ import ( "testing" ) -func testUtpControllerPermitAcquisition(t *testing.T, getPermit func() (Permit, bool)) { - firstPermit, ok := getPermit() - require.True(t, ok) +func testUtpControllerPermitAcquisition(t *testing.T, getPermit func() Permit, releasePermit func(Permit)) { + firstPermit := getPermit() require.NotNil(t, firstPermit) - noPermit, ok := getPermit() - require.False(t, ok) - require.NotNil(t, noPermit) - noPermit.Release() - stillLimited, ok := getPermit() - require.False(t, ok) - require.NotNil(t, stillLimited) + noPermit := getPermit() + require.Equal(t, PermitReject, noPermit) + releasePermit(noPermit) + stillLimited := getPermit() + require.Equal(t, PermitReject, stillLimited) // after release permit, should be able to get permit again - firstPermit.Release() - secondPermit, ok := getPermit() - require.True(t, ok) + releasePermit(firstPermit) + secondPermit := getPermit() require.NotNil(t, secondPermit) // should not be able to get permit - noPermit, ok = getPermit() - require.False(t, ok) - require.NotNil(t, noPermit) + noPermit = getPermit() + require.Equal(t, PermitReject, noPermit) } func TestUtpController_GetPermit(t *testing.T) { utpCtrl := newUtpController(1) - testUtpControllerPermitAcquisition(t, utpCtrl.GetInboundPermit) - testUtpControllerPermitAcquisition(t, utpCtrl.GetOutboundPermit) + testUtpControllerPermitAcquisition(t, utpCtrl.GetInboundPermit, utpCtrl.Release) + testUtpControllerPermitAcquisition(t, utpCtrl.GetOutboundPermit, utpCtrl.Release) } -func testUtpControllerConcurrencyGetPermitAcquisition(t *testing.T, getPermit func() (Permit, bool)) { +func testUtpControllerConcurrencyGetPermitAcquisition(t *testing.T, getPermit func() Permit) { var permitCount atomic.Int32 var wg sync.WaitGroup wg.Add(10) for i := 0; i < 10; i++ { go func() { defer wg.Done() - _, ok := getPermit() - if ok { + permit := getPermit() + if permit == PermitInbound || permit == PermitOutbound { permitCount.Add(1) } }()