faet: add content lookup test

p2p/discover/portal_protocol.go

@@ -13,6 +13,7 @@ import (
 	"math/big"
 	"net"
 	"sort"
+	"sync"
 	"time"
 
 	"github.com/tetratelabs/wabin/leb128"
@@ -1351,6 +1352,159 @@ func (p *PortalProtocol) collectTableNodes(rip net.IP, distances []uint, limit i
 	return nodes
 }
 
+type contentLookupResult struct {
+	flag    byte
+	content any
+	src     *enode.Node
+	err     error
+}
+
+type ContentLookupResult struct {
+	Content                 []byte
+	NodeInterestedInContent []*enode.Node
+}
+
+func (p *PortalProtocol) ContentLookup(contentKey []byte) (res *ContentLookupResult, err error) {
+	var mutex sync.Mutex
+
+	contentId := p.toContentId(contentKey)
+	nodesByDis := p.table.findnodeByID(enode.ID(contentId), bucketSize, false)
+	closestNodes := nodesByDis.entries
+
+	asked := make(map[enode.ID]struct{})
+	seen := make(map[enode.ID]struct{})
+	asked[p.localNode.ID()] = struct{}{}
+	seen[p.localNode.ID()] = struct{}{}
+
+	for _, node := range closestNodes {
+		seen[node.ID()] = struct{}{}
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	pendingQueries := make(map[enode.ID]struct{}, 0)
+	// find content query chan
+	pendingChan := make(chan struct{}, alpha)
+
+	for i := 0; i < alpha; i++ {
+		pendingChan <- struct{}{}
+	}
+
+	requestAmount := 0
+
+	nodesWithoutContent := make([]*enode.Node, 0)
+	// find content result
+	contentResultChan := make(chan contentLookupResult)
+	stopFindContentChan := make(chan struct{})
+	// content lookup result
+	resChan := make(chan *ContentLookupResult)
+	defer close(resChan)
+	defer close(stopFindContentChan)
+
+	// get findContent
+	go func() {
+		defer close(contentResultChan)
+		for {
+			select {
+			case <-pendingChan:
+				for i := 0; i < len(closestNodes) && len(pendingQueries) < alpha; i++ {
+					wrapedNode := closestNodes[i]
+					if _, ok := asked[wrapedNode.ID()]; !ok {
+						mutex.Lock()
+						asked[wrapedNode.ID()] = struct{}{}
+						pendingQueries[wrapedNode.ID()] = struct{}{}
+						mutex.Unlock()
+						go func() {
+							flag, content, err := p.findContent(unwrapNode(wrapedNode), contentKey)
+							select {
+							case contentResultChan <- contentLookupResult{flag: flag, content: content, err: err, src: unwrapNode(wrapedNode)}:
+								return
+							case <-stopFindContentChan:
+								return
+							}
+						}()
+						requestAmount++
+						fmt.Printf("requestAmount %v", requestAmount)
+					}
+
+					if len(pendingQueries) == 0 {
+						contentResultChan <- contentLookupResult{err: ContentNotFound}
+						return
+					}
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	// handler result
+	go func() {
+		// cancel context
+		defer cancel()
+		defer close(pendingChan)
+		for result := range contentResultChan {
+			if result.err == ContentNotFound {
+				err = ContentNotFound
+				resChan <- res
+				return
+			}
+			mutex.Lock()
+			delete(pendingQueries, result.src.ID())
+			mutex.Unlock()
+			if result.err != nil {
+				p.log.Trace("find content err: %v", result.err)
+				continue
+			}
+
+			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{}{}
+				resChan <- &ContentLookupResult{Content: content, NodeInterestedInContent: nodesWithoutContent}
+				return
+
+			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
+				}
+				for _, n := range nodes {
+					if _, ok := seen[n.ID()]; !ok {
+						seen[n.ID()] = struct{}{}
+						p.table.addSeenNode(wrapNode(n))
+						// insert node into closestNodes with distance
+						closestNodes = insertWithDistance(closestNodes, wrapNode(n), enode.ID(contentId))
+						if len(closestNodes) > bucketSize {
+							closestNodes = closestNodes[:bucketSize]
+						}
+					}
+				}
+				pendingChan <- struct{}{}
+			}
+		}
+	}()
+
+	res = <-resChan
+
+	return res, err
+}
+
 func inRange(nodeId enode.ID, nodeRadius *uint256.Int, contentId []byte) bool {
 	distance := enode.LogDist(nodeId, enode.ID(contentId))
 	disBig := new(big.Int).SetInt64(int64(distance))
@@ -1410,3 +1564,20 @@ func getContentKeys(request *OfferRequest) [][]byte {
 		return request.Request.(*PersistOfferRequest).ContentKeys
 	}
 }
+
+func insertWithDistance(nodes []*node, newNode *node, targetId enode.ID) []*node {
+	res := make([]*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

@@ -1,6 +1,7 @@
 package discover
 
 import (
+	"context"
 	"crypto/rand"
 	"errors"
 	"fmt"
@@ -299,3 +300,77 @@ func TestPortalWireProtocol(t *testing.T) {
 	assert.Equal(t, testEntry2.ContentKey, contentElement.ContentKeys[1])
 	assert.Equal(t, testEntry2.Content, contentElement.Contents[1])
 }
+
+func TestCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func(ctx context.Context) {
+		defer func() {
+			t.Log("goroutine cancel")
+		}()
+
+		time.Sleep(time.Second * 5)
+	}(ctx)
+
+	cancel()
+	t.Log("after main cancel")
+
+	time.Sleep(time.Second * 3)
+}
+
+// func TestInsertWithDistance(t *testing.T) {
+// 	targetId := uint256.NewInt(0).Bytes32()
+// 	nodes := make([]*node, 0, 10)
+// 	for i := 0; i < 10; i++ {
+// 		enode := &node{
+// 			Node: enode.Node{
+// 				id: enode.ID(uint256.NewInt(uint64(i)).Bytes32()),
+// 			},
+// 		}
+// 		nodes = append(nodes, enode)
+// 	}
+// 	newNode := &node{
+// 		Node: &enode.Node{
+// 			id: uint256.NewInt(20).Bytes32(),
+// 		},
+// 	}
+// }
+
+func TestContentLookup(t *testing.T) {
+	node1, err := setupLocalPortalNode(":7777", nil)
+	assert.NoError(t, err)
+	node1.log = testlog.Logger(t, log.LvlTrace)
+	err = node1.Start()
+	assert.NoError(t, err)
+	fmt.Println(node1.localNode.Node().String())
+
+	node2, err := setupLocalPortalNode(":7778", []*enode.Node{node1.localNode.Node()})
+	assert.NoError(t, err)
+	node2.log = testlog.Logger(t, log.LvlTrace)
+	err = node2.Start()
+	assert.NoError(t, err)
+	fmt.Println(node2.localNode.Node().String())
+
+	node3, err := setupLocalPortalNode(":7779", []*enode.Node{node1.localNode.Node()})
+	assert.NoError(t, err)
+	node3.log = testlog.Logger(t, log.LvlTrace)
+	err = node3.Start()
+	assert.NoError(t, err)
+	fmt.Println(node3.localNode.Node().String())
+	time.Sleep(10 * time.Second)
+
+	contentKey := []byte{0x3, 0x4}
+	content := []byte{0x1, 0x2}
+	contentId := node1.toContentId(contentKey)
+
+	err = node3.storage.Put(contentId, content)
+	assert.NoError(t, err)
+
+	lookupResult, err := node1.ContentLookup(contentKey)
+	assert.NoError(t, err)
+	assert.Equal(t, lookupResult.Content, content)
+
+	lookupResult, err = node1.ContentLookup([]byte{0x2, 0x4})
+	assert.Equal(t, ContentNotFound, err)
+	assert.Nil(t, lookupResult)
+}