refactor: use lookupworker

p2p/discover/portal_protocol.go

@@ -13,7 +13,6 @@ import (
 	"math/big"
 	"net"
 	"sort"
-	"sync"
 	"time"
 
 	"github.com/tetratelabs/wabin/leb128"
@@ -1352,137 +1351,43 @@ func (p *PortalProtocol) collectTableNodes(rip net.IP, distances []uint, limit i
 	return nodes
 }
 
-type findContentResult struct {
+func (p *PortalProtocol) ContentLookup(contentKey []byte) ([]byte, error) {
-	flag    byte
+	lookupContext, cancel := context.WithCancel(context.Background())
-	content any
+	defer cancel()
-	src     *enode.Node
+	resChan := make(chan []byte, 1)
-	err     error
-}
 
-type ContentLookupResult struct {
+	newLookup(lookupContext, p.table, p.Self().ID(), func(n *node) ([]*node, error) {
-	Content                 []byte
+		return p.contentLookupWorker(unwrapNode(n), contentKey, resChan)
-	NodeInterestedInContent []*enode.Node
+	}).run()
-}
 
-func (p *PortalProtocol) ContentLookup(contentKey []byte) (*ContentLookupResult, error) {
+	if len(resChan) > 0 {
-	contentId := p.toContentId(contentKey)
+		return <-resChan, nil
-	closestNodes := p.Lookup(enode.ID(contentId))
-	// findContent enrs chan
-	newNodeChan := make(chan []*enode.Node)
-	// target node for findContent
-	nodeChan := p.genNodeChan(closestNodes, contentId, newNodeChan)
-
-	nodesWithoutContent := make([]*enode.Node, 0)
-	// find content result
-	stopFindContentChan := make(chan struct{}, 1)
-
-	contentResultChan := p.parallelFindContent(stopFindContentChan, nodeChan, contentKey)
-
-	defer close(newNodeChan)
-
-	for result := range contentResultChan {
-		if result.err == ContentNotFound {
-			return nil, result.err
-		}
-		switch result.flag {
-		case portalwire.ContentRawSelector:
-			content, ok := result.content.([]byte)
-			if !ok {
-				p.log.Trace("failed to assert to raw content, value is: %v", result.content)
-				continue
-			}
-			stopFindContentChan <- struct{}{}
-			return &ContentLookupResult{Content: content, NodeInterestedInContent: nodesWithoutContent}, nil
-		case portalwire.ContentEnrsSelector:
-			nodeIdStr := result.src.ID().String()
-			// Get may modify the content of srcId
-			maybeRadius := p.radiusCache.Get(nil, []byte(nodeIdStr))
-			if len(maybeRadius) > 0 {
-				radius := uint256.MustFromBig(big.NewInt(0).SetBytes(maybeRadius))
-				if inRange(result.src.ID(), radius, p.toContentId(contentKey)) {
-					nodesWithoutContent = append(nodesWithoutContent, result.src)
-				}
-			}
-			nodes, ok := result.content.([]*enode.Node)
-			if !ok {
-				p.log.Trace("failed to assert to enrs content, value is: %v", result.content)
-				continue
-			}
-			// handle new nodes
-			newNodeChan <- nodes
-		}
 	}
-
 	return nil, ContentNotFound
 }
 
-func (p *PortalProtocol) genNodeChan(initNodes []*enode.Node, contentId []byte, newNodeChan chan []*enode.Node) <-chan *enode.Node {
+func (p *PortalProtocol) contentLookupWorker(n *enode.Node, contentKey []byte, resChan chan<- []byte) ([]*node, error) {
-	nodeChan := make(chan *enode.Node, alpha)
+	wrapedNode := make([]*node, 0)
-
+	flag, content, err := p.findContent(n, contentKey)
-	closestNodes := initNodes
+	if err != nil {
-	asked := make(map[enode.ID]struct{})
+		return nil, err
-	asked[p.localNode.ID()] = struct{}{}
-	for i := 0; i < len(closestNodes) && i < alpha; i++ {
-		wrapedNode := closestNodes[i]
-		if _, ok := asked[wrapedNode.ID()]; !ok {
-			asked[wrapedNode.ID()] = struct{}{}
-			nodeChan <- wrapedNode
-		}
 	}
-	go func() {
+	switch flag {
-		var once sync.Once
+	case portalwire.ContentRawSelector:
-		for nodes := range newNodeChan {
+		content, ok := content.([]byte)
-			for _, n := range nodes {
+		if !ok {
-				if _, ok := asked[n.ID()]; !ok {
+			return wrapedNode, fmt.Errorf("failed to assert to raw content, value is: %v", content)
-					asked[n.ID()] = struct{}{}
-					p.table.addSeenNode(wrapNode(n))
-					// insert node into closestNodes with distance
-					closestNodes = insertWithDistance(closestNodes, n, enode.ID(contentId))
-					if len(closestNodes) > bucketSize {
-						closestNodes = closestNodes[:bucketSize]
-					}
-				}
-			}
-			newRequest := false
-			for _, wrapedNode := range closestNodes {
-				if _, ok := asked[wrapedNode.ID()]; !ok {
-					newRequest = true
-					asked[wrapedNode.ID()] = struct{}{}
-					nodeChan <- wrapedNode
-				}
-			}
-			if !newRequest {
-				once.Do(func() {
-					close(nodeChan)
-				})
-			}
 		}
-	}()
+		resChan <- content
-
+		return wrapedNode, err
-	return nodeChan
+	case portalwire.ContentEnrsSelector:
-}
+		nodes, ok := content.([]*enode.Node)
-
+		if !ok {
-func (p *PortalProtocol) parallelFindContent(closeChan <-chan struct{}, nodeChan <-chan *enode.Node, contentKey []byte) <-chan findContentResult {
+			return wrapedNode, fmt.Errorf("failed to assert to enrs content, value is: %v", content)
-	contentResultChan := make(chan findContentResult)
-	requestAmount := 0
-	go func() {
-		defer close(contentResultChan)
-		for {
-			select {
-			case wrapedNode, ok := <-nodeChan:
-				if wrapedNode == nil && !ok {
-					contentResultChan <- findContentResult{err: ContentNotFound}
-					return
-				}
-				flag, content, err := p.findContent(wrapedNode, contentKey)
-				requestAmount++
-				contentResultChan <- findContentResult{flag: flag, content: content, err: err, src: wrapedNode}
-			case <-closeChan:
-				return
-			}
 		}
-	}()
+		return wrapNodes(nodes), nil
-	return contentResultChan
+	}
+	return wrapedNode, nil
 }
 
 func inRange(nodeId enode.ID, nodeRadius *uint256.Int, contentId []byte) bool {
@@ -1544,20 +1449,3 @@ func getContentKeys(request *OfferRequest) [][]byte {
 		return request.Request.(*PersistOfferRequest).ContentKeys
 	}
 }
-
-func insertWithDistance(nodes []*enode.Node, newNode *enode.Node, targetId enode.ID) []*enode.Node {
-	res := make([]*enode.Node, 0, len(nodes)+1)
-	for i := 0; i < len(nodes); i++ {
-		curNode := nodes[i]
-		curDis := enode.LogDist(curNode.ID(), newNode.ID())
-		newDis := enode.LogDist(newNode.ID(), targetId)
-		if newDis < curDis {
-			res = append(res, newNode)
-			res = append(res, nodes[i:]...)
-			break
-		} else {
-			res = append(res, curNode)
-		}
-	}
-	return res
-}

p2p/discover/portal_protocol_test.go

@@ -366,11 +366,11 @@ func TestContentLookup(t *testing.T) {
 	err = node3.storage.Put(contentId, content)
 	assert.NoError(t, err)
 
-	lookupResult, err := node1.ContentLookup(contentKey)
+	res, err := node1.ContentLookup(contentKey)
 	assert.NoError(t, err)
-	assert.Equal(t, lookupResult.Content, content)
+	assert.Equal(t, res, content)
 
-	lookupResult, err = node1.ContentLookup([]byte{0x2, 0x4})
+	res, err = node1.ContentLookup([]byte{0x2, 0x4})
 	assert.Equal(t, ContentNotFound, err)
-	assert.Nil(t, lookupResult)
+	assert.Nil(t, res)
 }