rework how refcounted wantlists work
License: MIT Signed-off-by: Jeromy <jeromyj@gmail.com>
Jeromy committed
Apr 27, 2017 at 17:38 UTC
e43d1317bb7cf4e6ad0ad68a2c8f0b1e8422f680
7 files changed
+211
-52
core/commands/bitswap.go
+5
-1
@@ -64,7 +64,11 @@ var unwantCmd = &cmds.Command{
64
ks = append(ks, c)
65
}
66
67
- bs.CancelWants(ks)
67
+ // TODO: This should maybe find *all* sessions for this request and cancel them?
68
+ // (why): in reality, i think this command should be removed. Its
69
+ // messing with the internal state of bitswap. You should cancel wants
70
+ // by killing the command that caused the want.
71
+ bs.CancelWants(ks, 0)
72
},
73
}
74
exchange/bitswap/bitswap.go
+20
-8
@@ -169,6 +169,9 @@ type Bitswap struct {
169
// Sessions
170
sessions []*Session
171
sessLk sync.Mutex
172
+
173
+ sessID uint64
174
+ sessIDLk sync.Mutex
175
}
176
177
type blockRequest struct {
@@ -219,7 +222,9 @@ func (bs *Bitswap) GetBlocks(ctx context.Context, keys []*cid.Cid) (<-chan block
222
log.Event(ctx, "Bitswap.GetBlockRequest.Start", k)
223
}
224
222
- bs.wm.WantBlocks(ctx, keys, nil)
225
+ mses := bs.getNextSessionID()
226
+
227
+ bs.wm.WantBlocks(ctx, keys, nil, mses)
228
229
// NB: Optimization. Assumes that providers of key[0] are likely to
230
// be able to provide for all keys. This currently holds true in most
@@ -241,7 +246,7 @@ func (bs *Bitswap) GetBlocks(ctx context.Context, keys []*cid.Cid) (<-chan block
246
defer close(out)
247
defer func() {
248
// can't just defer this call on its own, arguments are resolved *when* the defer is created
244
- bs.CancelWants(remaining.Keys())
249
+ bs.CancelWants(remaining.Keys(), mses)
250
}()
251
for {
252
select {
@@ -250,6 +255,7 @@ func (bs *Bitswap) GetBlocks(ctx context.Context, keys []*cid.Cid) (<-chan block
255
return
256
}
257
258
+ bs.CancelWants([]*cid.Cid{blk.Cid()}, mses)
259
remaining.Remove(blk.Cid())
260
select {
261
case out <- blk:
@@ -270,9 +276,16 @@ func (bs *Bitswap) GetBlocks(ctx context.Context, keys []*cid.Cid) (<-chan block
276
}
277
}
278
279
+func (bs *Bitswap) getNextSessionID() uint64 {
280
+ bs.sessIDLk.Lock()
281
+ defer bs.sessIDLk.Unlock()
282
+ bs.sessID++
283
+ return bs.sessID
284
+}
285
+
286
// CancelWant removes a given key from the wantlist
274
-func (bs *Bitswap) CancelWants(cids []*cid.Cid) {
275
- bs.wm.CancelWants(context.Background(), cids, nil)
287
+func (bs *Bitswap) CancelWants(cids []*cid.Cid, ses uint64) {
288
+ bs.wm.CancelWants(context.Background(), cids, nil, ses)
289
}
290
291
// HasBlock announces the existance of a block to this bitswap service. The
@@ -314,7 +327,7 @@ func (bs *Bitswap) SessionsForBlock(c *cid.Cid) []*Session {
327
328
var out []*Session
329
for _, s := range bs.sessions {
317
- if s.InterestedIn(c) {
330
+ if s.interestedIn(c) {
331
out = append(out, s)
332
}
333
}
@@ -346,8 +359,6 @@ func (bs *Bitswap) ReceiveMessage(ctx context.Context, p peer.ID, incoming bsmsg
359
keys = append(keys, block.Cid())
360
}
361
349
- bs.wm.CancelWants(context.Background(), keys, nil)
350
-
362
wg := sync.WaitGroup{}
363
for _, block := range iblocks {
364
wg.Add(1)
@@ -360,7 +371,8 @@ func (bs *Bitswap) ReceiveMessage(ctx context.Context, p peer.ID, incoming bsmsg
371
log.Event(ctx, "Bitswap.GetBlockRequest.End", k)
372
373
for _, ses := range bs.SessionsForBlock(k) {
363
- ses.ReceiveBlock(p, b)
374
+ ses.receiveBlockFrom(p, b)
375
+ bs.CancelWants([]*cid.Cid{k}, ses.id)
376
}
377
log.Debugf("got block %s from %s", b, p)
378
if err := bs.HasBlock(b); err != nil {
exchange/bitswap/bitswap_test.go
+5
@@ -332,6 +332,11 @@ func TestBasicBitswap(t *testing.T) {
332
t.Fatal(err)
333
}
334
335
+ time.Sleep(time.Millisecond * 20)
336
+ if len(instances[1].Exchange.GetWantlist()) != 0 {
337
+ t.Fatal("shouldnt have anything in wantlist")
338
+ }
339
+
340
st0, err := instances[0].Exchange.Stat()
341
if err != nil {
342
t.Fatal(err)
exchange/bitswap/session.go
+17
-5
@@ -16,6 +16,9 @@ import (
16
17
const activeWantsLimit = 16
18
19
+// Session holds state for an individual bitswap transfer operation.
20
+// This allows bitswap to make smarter decisions about who to send wantlist
21
+// info to, and who to request blocks from
22
type Session struct {
23
ctx context.Context
24
tofetch []*cid.Cid
@@ -40,8 +43,12 @@ type Session struct {
43
notif notifications.PubSub
44
45
uuid logging.Loggable
46
+
47
+ id uint64
48
}
49
50
+// NewSession creates a new bitswap session whose lifetime is bounded by the
51
+// given context
52
func (bs *Bitswap) NewSession(ctx context.Context) *Session {
53
s := &Session{
54
activePeers: make(map[peer.ID]struct{}),
@@ -54,6 +61,7 @@ func (bs *Bitswap) NewSession(ctx context.Context) *Session {
61
notif: notifications.New(),
62
uuid: loggables.Uuid("GetBlockRequest"),
63
baseTickDelay: time.Millisecond * 500,
64
+ id: bs.getNextSessionID(),
65
}
66
67
cache, _ := lru.New(2048)
@@ -73,11 +81,11 @@ type blkRecv struct {
81
blk blocks.Block
82
}
83
76
-func (s *Session) ReceiveBlock(from peer.ID, blk blocks.Block) {
84
+func (s *Session) receiveBlockFrom(from peer.ID, blk blocks.Block) {
85
s.incoming <- blkRecv{from: from, blk: blk}
86
}
87
80
-func (s *Session) InterestedIn(c *cid.Cid) bool {
88
+func (s *Session) interestedIn(c *cid.Cid) bool {
89
return s.interest.Contains(c.KeyString())
90
}
91
@@ -134,14 +142,14 @@ func (s *Session) run(ctx context.Context) {
142
143
case <-s.tick.C:
144
var live []*cid.Cid
137
- for c, _ := range s.liveWants {
145
+ for c := range s.liveWants {
146
cs, _ := cid.Cast([]byte(c))
147
live = append(live, cs)
148
s.liveWants[c] = time.Now()
149
}
150
151
// Broadcast these keys to everyone we're connected to
144
- s.bs.wm.WantBlocks(ctx, live, nil)
152
+ s.bs.wm.WantBlocks(ctx, live, nil, s.id)
153
154
if len(live) > 0 {
155
go func() {
@@ -181,7 +189,7 @@ func (s *Session) wantBlocks(ctx context.Context, ks []*cid.Cid) {
189
for _, c := range ks {
190
s.liveWants[c.KeyString()] = time.Now()
191
}
184
- s.bs.wm.WantBlocks(ctx, ks, s.activePeersArr)
192
+ s.bs.wm.WantBlocks(ctx, ks, s.activePeersArr, s.id)
193
}
194
195
func (s *Session) cancel(keys []*cid.Cid) {
@@ -211,11 +219,15 @@ func (s *Session) fetch(ctx context.Context, keys []*cid.Cid) {
219
}
220
}
221
222
+// GetBlocks fetches a set of blocks within the context of this session and
223
+// returns a channel that found blocks will be returned on. No order is
224
+// guaranteed on the returned blocks.
225
func (s *Session) GetBlocks(ctx context.Context, keys []*cid.Cid) (<-chan blocks.Block, error) {
226
ctx = logging.ContextWithLoggable(ctx, s.uuid)
227
return getBlocksImpl(ctx, keys, s.notif, s.fetch, s.cancelWants)
228
}
229
230
+// GetBlock fetches a single block
231
func (s *Session) GetBlock(parent context.Context, k *cid.Cid) (blocks.Block, error) {
232
return getBlock(parent, k, s.GetBlocks)
233
}
exchange/bitswap/wantlist/wantlist.go
+67
-25
@@ -10,8 +10,8 @@ import (
10
)
11
12
type ThreadSafe struct {
13
- lk sync.RWMutex
14
- Wantlist Wantlist
13
+ lk sync.RWMutex
14
+ set map[string]*Entry
15
}
16
17
// not threadsafe
@@ -23,7 +23,16 @@ type Entry struct {
23
Cid *cid.Cid
24
Priority int
25
26
- RefCnt int
26
+ SesTrk map[uint64]struct{}
27
+}
28
+
29
+// NewRefEntry creates a new reference tracked wantlist entry
30
+func NewRefEntry(c *cid.Cid, p int) *Entry {
31
+ return &Entry{
32
+ Cid: c,
33
+ Priority: p,
34
+ SesTrk: make(map[uint64]struct{}),
35
+ }
36
}
37
38
type entrySlice []*Entry
@@ -34,7 +43,7 @@ func (es entrySlice) Less(i, j int) bool { return es[i].Priority > es[j].Priorit
43
44
func NewThreadSafe() *ThreadSafe {
45
return &ThreadSafe{
37
- Wantlist: *New(),
46
+ set: make(map[string]*Entry),
47
}
48
}
49
@@ -44,46 +53,86 @@ func New() *Wantlist {
53
}
54
}
55
47
-func (w *ThreadSafe) Add(k *cid.Cid, priority int) bool {
56
+func (w *ThreadSafe) Add(c *cid.Cid, priority int, ses uint64) bool {
57
w.lk.Lock()
58
defer w.lk.Unlock()
50
- return w.Wantlist.Add(k, priority)
59
+ k := c.KeyString()
60
+ if e, ok := w.set[k]; ok {
61
+ e.SesTrk[ses] = struct{}{}
62
+ return false
63
+ }
64
+
65
+ w.set[k] = &Entry{
66
+ Cid: c,
67
+ Priority: priority,
68
+ SesTrk: map[uint64]struct{}{ses: struct{}{}},
69
+ }
70
+
71
+ return true
72
}
73
53
-func (w *ThreadSafe) AddEntry(e *Entry) bool {
74
+func (w *ThreadSafe) AddEntry(e *Entry, ses uint64) bool {
75
w.lk.Lock()
76
defer w.lk.Unlock()
56
- return w.Wantlist.AddEntry(e)
77
+ k := e.Cid.KeyString()
78
+ if ex, ok := w.set[k]; ok {
79
+ ex.SesTrk[ses] = struct{}{}
80
+ return false
81
+ }
82
+ w.set[k] = e
83
+ e.SesTrk[ses] = struct{}{}
84
+ return true
85
}
86
59
-func (w *ThreadSafe) Remove(k *cid.Cid) bool {
87
+func (w *ThreadSafe) Remove(c *cid.Cid, ses uint64) bool {
88
w.lk.Lock()
89
defer w.lk.Unlock()
62
- return w.Wantlist.Remove(k)
90
+ k := c.KeyString()
91
+ e, ok := w.set[k]
92
+ if !ok {
93
+ return false
94
+ }
95
+
96
+ delete(e.SesTrk, ses)
97
+ if len(e.SesTrk) == 0 {
98
+ delete(w.set, k)
99
+ return true
100
+ }
101
+ return false
102
}
103
104
func (w *ThreadSafe) Contains(k *cid.Cid) (*Entry, bool) {
105
w.lk.RLock()
106
defer w.lk.RUnlock()
68
- return w.Wantlist.Contains(k)
107
+ e, ok := w.set[k.KeyString()]
108
+ return e, ok
109
}
110
111
func (w *ThreadSafe) Entries() []*Entry {
112
w.lk.RLock()
113
defer w.lk.RUnlock()
74
- return w.Wantlist.Entries()
114
+ var es entrySlice
115
+ for _, e := range w.set {
116
+ es = append(es, e)
117
+ }
118
+ return es
119
}
120
121
func (w *ThreadSafe) SortedEntries() []*Entry {
122
w.lk.RLock()
123
defer w.lk.RUnlock()
80
- return w.Wantlist.SortedEntries()
124
+ var es entrySlice
125
+ for _, e := range w.set {
126
+ es = append(es, e)
127
+ }
128
+ sort.Sort(es)
129
+ return es
130
}
131
132
func (w *ThreadSafe) Len() int {
133
w.lk.RLock()
134
defer w.lk.RUnlock()
86
- return w.Wantlist.Len()
135
+ return len(w.set)
136
}
137
138
func (w *Wantlist) Len() int {
@@ -92,15 +141,13 @@ func (w *Wantlist) Len() int {
141
142
func (w *Wantlist) Add(c *cid.Cid, priority int) bool {
143
k := c.KeyString()
95
- if e, ok := w.set[k]; ok {
96
- e.RefCnt++
144
+ if _, ok := w.set[k]; ok {
145
return false
146
}
147
148
w.set[k] = &Entry{
149
Cid: c,
150
Priority: priority,
103
- RefCnt: 1,
151
}
152
153
return true
@@ -108,8 +155,7 @@ func (w *Wantlist) Add(c *cid.Cid, priority int) bool {
155
156
func (w *Wantlist) AddEntry(e *Entry) bool {
157
k := e.Cid.KeyString()
111
- if ex, ok := w.set[k]; ok {
112
- ex.RefCnt++
158
+ if _, ok := w.set[k]; ok {
159
return false
160
}
161
w.set[k] = e
@@ -118,16 +164,12 @@ func (w *Wantlist) AddEntry(e *Entry) bool {
164
165
func (w *Wantlist) Remove(c *cid.Cid) bool {
166
k := c.KeyString()
121
- e, ok := w.set[k]
167
+ _, ok := w.set[k]
168
if !ok {
169
return false
170
}
171
126
- e.RefCnt--
127
- if e.RefCnt <= 0 {
128
- delete(w.set, k)
129
- return true
130
- }
172
+ delete(w.set, k)
173
return false
174
}
175
exchange/bitswap/wantlist/wantlist_test.go
new
+87
@@ -0,0 +1,87 @@
1
+package wantlist
2
+
3
+import (
4
+ "testing"
5
+
6
+ cid "gx/ipfs/QmYhQaCYEcaPPjxJX7YcPcVKkQfRy6sJ7B3XmGFk82XYdQ/go-cid"
7
+)
8
+
9
+var testcids []*cid.Cid
10
+
11
+func init() {
12
+ strs := []string{
13
+ "QmQL8LqkEgYXaDHdNYCG2mmpow7Sp8Z8Kt3QS688vyBeC7",
14
+ "QmcBDsdjgSXU7BP4A4V8LJCXENE5xVwnhrhRGVTJr9YCVj",
15
+ "QmQakgd2wDxc3uUF4orGdEm28zUT9Mmimp5pyPG2SFS9Gj",
16
+ }
17
+ for _, s := range strs {
18
+ c, err := cid.Decode(s)
19
+ if err != nil {
20
+ panic(err)
21
+ }
22
+ testcids = append(testcids, c)
23
+ }
24
+
25
+}
26
+
27
+type wli interface {
28
+ Contains(*cid.Cid) (*Entry, bool)
29
+}
30
+
31
+func assertHasCid(t *testing.T, w wli, c *cid.Cid) {
32
+ e, ok := w.Contains(c)
33
+ if !ok {
34
+ t.Fatal("expected to have ", c)
35
+ }
36
+ if !e.Cid.Equals(c) {
37
+ t.Fatal("returned entry had wrong cid value")
38
+ }
39
+}
40
+
41
+func assertNotHasCid(t *testing.T, w wli, c *cid.Cid) {
42
+ _, ok := w.Contains(c)
43
+ if ok {
44
+ t.Fatal("expected not to have ", c)
45
+ }
46
+}
47
+
48
+func TestBasicWantlist(t *testing.T) {
49
+ wl := New()
50
+
51
+ wl.Add(testcids[0], 5)
52
+ assertHasCid(t, wl, testcids[0])
53
+ wl.Add(testcids[1], 4)
54
+ assertHasCid(t, wl, testcids[0])
55
+ assertHasCid(t, wl, testcids[1])
56
+
57
+ if wl.Len() != 2 {
58
+ t.Fatal("should have had two items")
59
+ }
60
+
61
+ wl.Add(testcids[1], 4)
62
+ assertHasCid(t, wl, testcids[0])
63
+ assertHasCid(t, wl, testcids[1])
64
+
65
+ if wl.Len() != 2 {
66
+ t.Fatal("should have had two items")
67
+ }
68
+
69
+ wl.Remove(testcids[0])
70
+ assertHasCid(t, wl, testcids[1])
71
+ if _, has := wl.Contains(testcids[0]); has {
72
+ t.Fatal("shouldnt have this cid")
73
+ }
74
+}
75
+
76
+func TestSesRefWantlist(t *testing.T) {
77
+ wl := NewThreadSafe()
78
+
79
+ wl.Add(testcids[0], 5, 1)
80
+ assertHasCid(t, wl, testcids[0])
81
+ wl.Remove(testcids[0], 2)
82
+ assertHasCid(t, wl, testcids[0])
83
+ wl.Add(testcids[0], 5, 1)
84
+ assertHasCid(t, wl, testcids[0])
85
+ wl.Remove(testcids[0], 1)
86
+ assertNotHasCid(t, wl, testcids[0])
87
+}
exchange/bitswap/wantmanager.go
+10
-13
@@ -71,34 +71,31 @@ type msgQueue struct {
71
done chan struct{}
72
}
73
74
-func (pm *WantManager) WantBlocks(ctx context.Context, ks []*cid.Cid, peers []peer.ID) {
74
+func (pm *WantManager) WantBlocks(ctx context.Context, ks []*cid.Cid, peers []peer.ID, ses uint64) {
75
log.Infof("want blocks: %s", ks)
76
- pm.addEntries(ctx, ks, peers, false)
76
+ pm.addEntries(ctx, ks, peers, false, ses)
77
}
78
79
-func (pm *WantManager) CancelWants(ctx context.Context, ks []*cid.Cid, peers []peer.ID) {
80
- pm.addEntries(context.Background(), ks, peers, true)
79
+func (pm *WantManager) CancelWants(ctx context.Context, ks []*cid.Cid, peers []peer.ID, ses uint64) {
80
+ pm.addEntries(context.Background(), ks, peers, true, ses)
81
}
82
83
type wantSet struct {
84
entries []*bsmsg.Entry
85
targets []peer.ID
86
+ from uint64
87
}
88
88
-func (pm *WantManager) addEntries(ctx context.Context, ks []*cid.Cid, targets []peer.ID, cancel bool) {
89
+func (pm *WantManager) addEntries(ctx context.Context, ks []*cid.Cid, targets []peer.ID, cancel bool, ses uint64) {
90
var entries []*bsmsg.Entry
91
for i, k := range ks {
92
entries = append(entries, &bsmsg.Entry{
93
Cancel: cancel,
93
- Entry: &wantlist.Entry{
94
- Cid: k,
95
- Priority: kMaxPriority - i,
96
- RefCnt: 1,
97
- },
94
+ Entry: wantlist.NewRefEntry(k, kMaxPriority-i),
95
})
96
}
97
select {
101
- case pm.incoming <- &wantSet{entries: entries, targets: targets}:
98
+ case pm.incoming <- &wantSet{entries: entries, targets: targets, from: ses}:
99
case <-pm.ctx.Done():
100
case <-ctx.Done():
101
}
@@ -290,11 +287,11 @@ func (pm *WantManager) Run() {
287
// add changes to our wantlist
288
for _, e := range ws.entries {
289
if e.Cancel {
293
- if pm.wl.Remove(e.Cid) {
290
+ if pm.wl.Remove(e.Cid, ws.from) {
291
pm.wantlistGauge.Dec()
292
}
293
} else {
297
- if pm.wl.AddEntry(e.Entry) {
294
+ if pm.wl.AddEntry(e.Entry, ws.from) {
295
pm.wantlistGauge.Inc()
296
}
297
}