@cryptotaxi247 / kubo / commits / e43d1317b

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 }