forked from ipfs/go-graphsync
-
Notifications
You must be signed in to change notification settings - Fork 0
/
linktracker.go
85 lines (75 loc) · 2.93 KB
/
linktracker.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package linktracker
import (
"github.com/ipld/go-ipld-prime"
graphsync "github.com/filecoin-project/boost-graphsync"
)
// LinkTracker records links being traversed to determine useful information
// in crafting responses for a peer. Specifically, if any in progress request
// has already sent a block for a given link, don't send it again.
// Second, keep track of whether links are missing blocks so you can determine
// at the end if a complete response has been transmitted.
type LinkTracker struct {
missingBlocks map[graphsync.RequestID]map[ipld.Link]struct{}
linksWithBlocksTraversedByRequest map[graphsync.RequestID][]ipld.Link
traversalsWithBlocksInProgress map[ipld.Link]int
}
// New makes a new link tracker
func New() *LinkTracker {
return &LinkTracker{
missingBlocks: make(map[graphsync.RequestID]map[ipld.Link]struct{}),
linksWithBlocksTraversedByRequest: make(map[graphsync.RequestID][]ipld.Link),
traversalsWithBlocksInProgress: make(map[ipld.Link]int),
}
}
// BlockRefCount returns the number of times a present block has been traversed
// by in progress requests
func (lt *LinkTracker) BlockRefCount(link ipld.Link) int {
return lt.traversalsWithBlocksInProgress[link]
}
// IsKnownMissingLink returns whether the given request recorded the given link as missing
func (lt *LinkTracker) IsKnownMissingLink(requestID graphsync.RequestID, link ipld.Link) bool {
missingBlocks, ok := lt.missingBlocks[requestID]
if !ok {
return false
}
_, ok = missingBlocks[link]
return ok
}
// RecordLinkTraversal records that we traversed a link during a request, and
// whether we had the block when we did it.
func (lt *LinkTracker) RecordLinkTraversal(requestID graphsync.RequestID, link ipld.Link, hasBlock bool) {
if hasBlock {
lt.linksWithBlocksTraversedByRequest[requestID] = append(lt.linksWithBlocksTraversedByRequest[requestID], link)
lt.traversalsWithBlocksInProgress[link]++
} else {
missingBlocks, ok := lt.missingBlocks[requestID]
if !ok {
missingBlocks = make(map[ipld.Link]struct{})
lt.missingBlocks[requestID] = missingBlocks
}
missingBlocks[link] = struct{}{}
}
}
// FinishRequest records that we have completed the given request, and returns
// true if all links traversed had blocks present.
func (lt *LinkTracker) FinishRequest(requestID graphsync.RequestID) (hasAllBlocks bool) {
_, ok := lt.missingBlocks[requestID]
hasAllBlocks = !ok
delete(lt.missingBlocks, requestID)
links, ok := lt.linksWithBlocksTraversedByRequest[requestID]
if !ok {
return
}
for _, link := range links {
lt.traversalsWithBlocksInProgress[link]--
if lt.traversalsWithBlocksInProgress[link] <= 0 {
delete(lt.traversalsWithBlocksInProgress, link)
}
}
delete(lt.linksWithBlocksTraversedByRequest, requestID)
return
}
// Empty returns true if the link tracker is empty
func (lt *LinkTracker) Empty() bool {
return len(lt.missingBlocks) == 0 && len(lt.traversalsWithBlocksInProgress) == 0
}