refactor task queue to have queues per peer
Jeromy committed
Apr 3, 2015 at 01:07 UTC
d765c14e8fcc926b14b89e1e11b4c990628b0de4
3 files changed
+167
-15
exchange/bitswap/decision/engine.go
+8
-1
@@ -55,6 +55,9 @@ type Envelope struct {
55
Peer peer.ID
56
// Message is the payload
57
Message bsmsg.BitSwapMessage
58
+
59
+ // A callback to notify the decision queue that the task is complete
60
+ Sent func()
61
}
62
63
type Engine struct {
@@ -137,7 +140,11 @@ func (e *Engine) nextEnvelope(ctx context.Context) (*Envelope, error) {
140
141
m := bsmsg.New() // TODO: maybe add keys from our wantlist?
142
m.AddBlock(block)
140
- return &Envelope{Peer: nextTask.Target, Message: m}, nil
143
+ return &Envelope{
144
+ Peer: nextTask.Target,
145
+ Message: m,
146
+ Sent: nextTask.Done,
147
+ }, nil
148
}
149
}
150
exchange/bitswap/decision/peer_request_queue.go
+97
-10
@@ -21,8 +21,9 @@ type peerRequestQueue interface {
21
22
func newPRQ() peerRequestQueue {
23
return &prq{
24
- taskMap: make(map[string]*peerRequestTask),
25
- taskQueue: pq.New(wrapCmp(V1)),
24
+ taskMap: make(map[string]*peerRequestTask),
25
+ partners: make(map[peer.ID]*activePartner),
26
+ pQueue: pq.New(partnerCompare),
27
}
28
}
29
@@ -32,42 +33,73 @@ var _ peerRequestQueue = &prq{}
33
// to help decide how to sort tasks (on add) and how to select
34
// tasks (on getnext). For now, we are assuming a dumb/nice strategy.
35
type prq struct {
35
- lock sync.Mutex
36
- taskQueue pq.PQ
37
- taskMap map[string]*peerRequestTask
36
+ lock sync.Mutex
37
+ pQueue pq.PQ
38
+ taskMap map[string]*peerRequestTask
39
+ partners map[peer.ID]*activePartner
40
}
41
42
// Push currently adds a new peerRequestTask to the end of the list
43
func (tl *prq) Push(entry wantlist.Entry, to peer.ID) {
44
tl.lock.Lock()
45
defer tl.lock.Unlock()
46
+ partner, ok := tl.partners[to]
47
+ if !ok {
48
+ partner = &activePartner{taskQueue: pq.New(wrapCmp(V1))}
49
+ tl.pQueue.Push(partner)
50
+ tl.partners[to] = partner
51
+ }
52
+
53
if task, ok := tl.taskMap[taskKey(to, entry.Key)]; ok {
54
task.Entry.Priority = entry.Priority
46
- tl.taskQueue.Update(task.index)
55
+ partner.taskQueue.Update(task.index)
56
return
57
}
58
+
59
task := &peerRequestTask{
60
Entry: entry,
61
Target: to,
62
created: time.Now(),
63
+ Done: func() {
64
+ partner.TaskDone()
65
+ tl.lock.Lock()
66
+ tl.pQueue.Update(partner.Index())
67
+ tl.lock.Unlock()
68
+ },
69
}
54
- tl.taskQueue.Push(task)
70
+
71
+ partner.taskQueue.Push(task)
72
tl.taskMap[task.Key()] = task
73
+ partner.requests++
74
+ tl.pQueue.Update(partner.Index())
75
}
76
77
// Pop 'pops' the next task to be performed. Returns nil if no task exists.
78
func (tl *prq) Pop() *peerRequestTask {
79
tl.lock.Lock()
80
defer tl.lock.Unlock()
81
+ if tl.pQueue.Len() == 0 {
82
+ return nil
83
+ }
84
+ pElem := tl.pQueue.Pop()
85
+ if pElem == nil {
86
+ return nil
87
+ }
88
+
89
+ partner := pElem.(*activePartner)
90
+
91
var out *peerRequestTask
63
- for tl.taskQueue.Len() > 0 {
64
- out = tl.taskQueue.Pop().(*peerRequestTask)
92
+ for partner.taskQueue.Len() > 0 {
93
+ out = partner.taskQueue.Pop().(*peerRequestTask)
94
delete(tl.taskMap, out.Key())
95
if out.trash {
96
continue // discarding tasks that have been removed
97
}
98
break // and return |out|
99
}
100
+ partner.StartTask()
101
+ partner.requests--
102
+ tl.pQueue.Push(partner)
103
return out
104
}
105
@@ -80,13 +112,16 @@ func (tl *prq) Remove(k u.Key, p peer.ID) {
112
// simply mark it as trash, so it'll be dropped when popped off the
113
// queue.
114
t.trash = true
115
+ tl.partners[p].requests--
116
}
117
tl.lock.Unlock()
118
}
119
120
type peerRequestTask struct {
121
Entry wantlist.Entry
89
- Target peer.ID // required
122
+ Target peer.ID
123
+
124
+ Done func()
125
126
// trash in a book-keeping field
127
trash bool
@@ -132,3 +167,55 @@ func wrapCmp(f func(a, b *peerRequestTask) bool) func(a, b pq.Elem) bool {
167
return f(a.(*peerRequestTask), b.(*peerRequestTask))
168
}
169
}
170
+
171
+type activePartner struct {
172
+ lk sync.Mutex
173
+
174
+ // Active is the number of blocks this peer is currently being sent
175
+ active int
176
+
177
+ // requests is the number of blocks this peer is currently requesting
178
+ requests int
179
+
180
+ index int
181
+
182
+ // priority queue of
183
+ taskQueue pq.PQ
184
+}
185
+
186
+func partnerCompare(a, b pq.Elem) bool {
187
+ pa := a.(*activePartner)
188
+ pb := b.(*activePartner)
189
+
190
+ // having no blocks in their wantlist means lowest priority
191
+ if pa.requests == 0 {
192
+ return false
193
+ }
194
+ if pb.requests == 0 {
195
+ return true
196
+ }
197
+ return pa.active < pb.active
198
+}
199
+
200
+func (p *activePartner) StartTask() {
201
+ p.lk.Lock()
202
+ p.active++
203
+ p.lk.Unlock()
204
+}
205
+
206
+func (p *activePartner) TaskDone() {
207
+ p.lk.Lock()
208
+ p.active--
209
+ if p.active < 0 {
210
+ panic("more tasks finished than started!")
211
+ }
212
+ p.lk.Unlock()
213
+}
214
+
215
+func (p *activePartner) Index() int {
216
+ return p.index
217
+}
218
+
219
+func (p *activePartner) SetIndex(i int) {
220
+ p.index = i
221
+}
exchange/bitswap/decision/peer_request_queue_test.go
+62
-4
@@ -47,10 +47,68 @@ func TestPushPop(t *testing.T) {
47
prq.Remove(util.Key(consonant), partner)
48
}
49
50
- for _, expected := range vowels {
51
- received := prq.Pop().Entry.Key
52
- if received != util.Key(expected) {
53
- t.Fatal("received", string(received), "expected", string(expected))
50
+ var out []string
51
+ for {
52
+ received := prq.Pop()
53
+ if received == nil {
54
+ break
55
}
56
+
57
+ out = append(out, string(received.Entry.Key))
58
+ }
59
+
60
+ // Entries popped should already be in correct order
61
+ for i, expected := range vowels {
62
+ if out[i] != expected {
63
+ t.Fatal("received", out[i], "expected", expected)
64
+ }
65
+ }
66
+}
67
+
68
+// This test checks that peers wont starve out other peers
69
+func TestPeerRepeats(t *testing.T) {
70
+ prq := newPRQ()
71
+ a := testutil.RandPeerIDFatal(t)
72
+ b := testutil.RandPeerIDFatal(t)
73
+ c := testutil.RandPeerIDFatal(t)
74
+ d := testutil.RandPeerIDFatal(t)
75
+
76
+ // Have each push some blocks
77
+
78
+ for i := 0; i < 5; i++ {
79
+ prq.Push(wantlist.Entry{Key: util.Key(i)}, a)
80
+ prq.Push(wantlist.Entry{Key: util.Key(i)}, b)
81
+ prq.Push(wantlist.Entry{Key: util.Key(i)}, c)
82
+ prq.Push(wantlist.Entry{Key: util.Key(i)}, d)
83
+ }
84
+
85
+ // now, pop off four entries, there should be one from each
86
+ var targets []string
87
+ var tasks []*peerRequestTask
88
+ for i := 0; i < 4; i++ {
89
+ t := prq.Pop()
90
+ targets = append(targets, t.Target.Pretty())
91
+ tasks = append(tasks, t)
92
+ }
93
+
94
+ expected := []string{a.Pretty(), b.Pretty(), c.Pretty(), d.Pretty()}
95
+ sort.Strings(expected)
96
+ sort.Strings(targets)
97
+
98
+ t.Log(targets)
99
+ t.Log(expected)
100
+ for i, s := range targets {
101
+ if expected[i] != s {
102
+ t.Fatal("unexpected peer", s, expected[i])
103
+ }
104
+ }
105
+
106
+ // Now, if one of the tasks gets finished, the next task off the queue should
107
+ // be for the same peer
108
+ tasks[0].Done()
109
+
110
+ ntask := prq.Pop()
111
+ if ntask.Target != tasks[0].Target {
112
+ t.Fatal("Expected task from peer with lowest active count")
113
}
114
}