From c0ac288ee7e06d5b8c73946135412d6cde534a47 Mon Sep 17 00:00:00 2001 From: fearlessfe <505380967@qq.com> Date: Tue, 26 Dec 2023 23:44:29 +0800 Subject: [PATCH] faet: add content lookup test --- p2p/discover/portal_protocol.go | 171 +++++++++++++++++++++++++++ p2p/discover/portal_protocol_test.go | 75 ++++++++++++ 2 files changed, 246 insertions(+) diff --git a/p2p/discover/portal_protocol.go b/p2p/discover/portal_protocol.go index 53b4031337..3f848eb744 100644 --- a/p2p/discover/portal_protocol.go +++ b/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 +} diff --git a/p2p/discover/portal_protocol_test.go b/p2p/discover/portal_protocol_test.go index 6d72104321..fb5eee8a0c 100644 --- a/p2p/discover/portal_protocol_test.go +++ b/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) +}