Make bitswap better

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>

exchange/bitswap/decision/engine.go

@@ -3,6 +3,7 @@ package decision
 
 import (
  "sync"
+ "time"
 
  blocks "github.com/ipfs/go-ipfs/blocks"
  bstore "github.com/ipfs/go-ipfs/blocks/blockstore"
@@ -68,7 +69,7 @@ type Engine struct {
  // peerRequestQueue is a priority queue of requests received from peers.
  // Requests are popped from the queue, packaged up, and placed in the
  // outbox.
- peerRequestQueue peerRequestQueue
+ peerRequestQueue *prq
 
  // FIXME it's a bit odd for the client and the worker to both share memory
  // (both modify the peerRequestQueue) and also to communicate over the
@@ -86,6 +87,8 @@ type Engine struct {
  lock sync.Mutex // protects the fields immediatly below
  // ledgerMap lists Ledgers by their Partner key.
  ledgerMap map[peer.ID]*ledger
+
+ ticker *time.Ticker
 }
 
 func NewEngine(ctx context.Context, bs bstore.Blockstore) *Engine {
@@ -95,6 +98,7 @@ func NewEngine(ctx context.Context, bs bstore.Blockstore) *Engine {
  peerRequestQueue: newPRQ(),
  outbox: make(chan (<-chan *Envelope), outboxChanBuffer),
  workSignal: make(chan struct{}, 1),
+ ticker: time.NewTicker(time.Millisecond * 100),
  }
  go e.taskWorker(ctx)
  return e
@@ -142,6 +146,9 @@ func (e *Engine) nextEnvelope(ctx context.Context) (*Envelope, error) {
  return nil, ctx.Err()
  case <-e.workSignal:
  nextTask = e.peerRequestQueue.Pop()
+ case <-e.ticker.C:
+ e.peerRequestQueue.thawRound()
+ nextTask = e.peerRequestQueue.Pop()
  }
  }
 
@@ -191,9 +198,6 @@ func (e *Engine) Peers() []peer.ID {
 // MessageReceived performs book-keeping. Returns error if passed invalid
 // arguments.
 func (e *Engine) MessageReceived(p peer.ID, m bsmsg.BitSwapMessage) error {
- e.lock.Lock()
- defer e.lock.Unlock()
-
  if len(m.Wantlist()) == 0 && len(m.Blocks()) == 0 {
  log.Debugf("received empty message from %s", p)
  }
@@ -206,6 +210,8 @@ func (e *Engine) MessageReceived(p peer.ID, m bsmsg.BitSwapMessage) error {
  }()
 
  l := e.findOrCreate(p)
+ l.lk.Lock()
+ defer l.lk.Unlock()
  if m.Full() {
  l.wantList = wl.New()
  }
@@ -236,10 +242,12 @@ func (e *Engine) addBlock(block blocks.Block) {
  work := false
 
  for _, l := range e.ledgerMap {
+ l.lk.Lock()
  if entry, ok := l.WantListContains(block.Key()); ok {
  e.peerRequestQueue.Push(entry, l.Partner)
  work = true
  }
+ l.lk.Unlock()
  }
 
  if work {
@@ -261,9 +269,6 @@ func (e *Engine) AddBlock(block blocks.Block) {
 // send happen atomically
 
 func (e *Engine) MessageSent(p peer.ID, m bsmsg.BitSwapMessage) error {
- e.lock.Lock()
- defer e.lock.Unlock()
-
  l := e.findOrCreate(p)
  for _, block := range m.Blocks() {
  l.SentBytes(len(block.Data()))
@@ -290,11 +295,13 @@ func (e *Engine) numBytesReceivedFrom(p peer.ID) uint64 {
 
 // ledger lazily instantiates a ledger
 func (e *Engine) findOrCreate(p peer.ID) *ledger {
+ e.lock.Lock()
  l, ok := e.ledgerMap[p]
  if !ok {
  l = newLedger(p)
  e.ledgerMap[p] = l
  }
+ e.lock.Unlock()
  return l
 }
 

exchange/bitswap/decision/ledger.go

@@ -1,6 +1,7 @@
 package decision
 
 import (
+ "sync"
  "time"
 
  key "github.com/ipfs/go-ipfs/blocks/key"
@@ -44,6 +45,8 @@ type ledger struct {
  // sentToPeer is a set of keys to ensure we dont send duplicate blocks
  // to a given peer
  sentToPeer map[key.Key]time.Time
+
+ lk sync.Mutex
 }
 
 type debtRatio struct {

exchange/bitswap/decision/peer_request_queue.go

@@ -15,14 +15,16 @@ type peerRequestQueue interface {
  Pop() *peerRequestTask
  Push(entry wantlist.Entry, to peer.ID)
  Remove(k key.Key, p peer.ID)
+
  // NB: cannot expose simply expose taskQueue.Len because trashed elements
  // may exist. These trashed elements should not contribute to the count.
 }
 
-func newPRQ() peerRequestQueue {
+func newPRQ() *prq {
  return &prq{
  taskMap: make(map[string]*peerRequestTask),
  partners: make(map[peer.ID]*activePartner),
+ frozen: make(map[peer.ID]*activePartner),
  pQueue: pq.New(partnerCompare),
  }
 }
@@ -38,6 +40,8 @@ type prq struct {
  pQueue pq.PQ
  taskMap map[string]*peerRequestTask
  partners map[peer.ID]*activePartner
+
+ frozen map[peer.ID]*activePartner
 }
 
 // Push currently adds a new peerRequestTask to the end of the list
@@ -92,7 +96,7 @@ func (tl *prq) Pop() *peerRequestTask {
  partner := tl.pQueue.Pop().(*activePartner)
 
  var out *peerRequestTask
- for partner.taskQueue.Len() > 0 {
+ for partner.taskQueue.Len() > 0 && partner.freezeVal == 0 {
  out = partner.taskQueue.Pop().(*peerRequestTask)
  delete(tl.taskMap, out.Key())
  if out.trash {
@@ -120,11 +124,47 @@ func (tl *prq) Remove(k key.Key, p peer.ID) {
  t.trash = true
 
  // having canceled a block, we now account for that in the given partner
- tl.partners[p].requests--
+ partner := tl.partners[p]
+ partner.requests--
+
+ // we now also 'freeze' that partner. If they sent us a cancel for a
+ // block we were about to send them, we should wait a short period of time
+ // to make sure we receive any other in-flight cancels before sending
+ // them a block they already potentially have
+ if partner.freezeVal == 0 {
+ tl.frozen[p] = partner
+ }
+
+ partner.freezeVal++
+ tl.pQueue.Update(partner.index)
  }
  tl.lock.Unlock()
 }
 
+func (tl *prq) fullThaw() {
+ tl.lock.Lock()
+ defer tl.lock.Unlock()
+
+ for id, partner := range tl.frozen {
+ partner.freezeVal = 0
+ delete(tl.frozen, id)
+ tl.pQueue.Update(partner.index)
+ }
+}
+
+func (tl *prq) thawRound() {
+ tl.lock.Lock()
+ defer tl.lock.Unlock()
+
+ for id, partner := range tl.frozen {
+ partner.freezeVal -= (partner.freezeVal + 1) / 2
+ if partner.freezeVal <= 0 {
+ delete(tl.frozen, id)
+ }
+ tl.pQueue.Update(partner.index)
+ }
+}
+
 type peerRequestTask struct {
  Entry wantlist.Entry
  Target peer.ID
@@ -196,6 +236,8 @@ type activePartner struct {
  // for the PQ interface
  index int
 
+ freezeVal int
+
  // priority queue of tasks belonging to this peer
  taskQueue pq.PQ
 }
@@ -208,6 +250,7 @@ func newActivePartner() *activePartner {
 }
 
 // partnerCompare implements pq.ElemComparator
+// returns true if peer 'a' has higher priority than peer 'b'
 func partnerCompare(a, b pq.Elem) bool {
  pa := a.(*activePartner)
  pb := b.(*activePartner)
@@ -220,6 +263,14 @@ func partnerCompare(a, b pq.Elem) bool {
  if pb.requests == 0 {
  return true
  }
+
+ if pa.freezeVal > pb.freezeVal {
+ return false
+ }
+ if pa.freezeVal < pb.freezeVal {
+ return true
+ }
+
  if pa.active == pb.active {
  // sorting by taskQueue.Len() aids in cleaning out trash entries faster
  // if we sorted instead by requests, one peer could potentially build up

exchange/bitswap/decision/peer_request_queue_test.go

@@ -47,6 +47,8 @@ func TestPushPop(t *testing.T) {
  prq.Remove(key.Key(consonant), partner)
  }
 
+ prq.fullThaw()
+
  var out []string
  for {
  received := prq.Pop()

exchange/bitswap/network/interface.go

@@ -25,9 +25,16 @@ type BitSwapNetwork interface {
 
  ConnectTo(context.Context, peer.ID) error
 
+ NewMessageSender(context.Context, peer.ID) (MessageSender, error)
+
  Routing
 }
 
+type MessageSender interface {
+ SendMsg(bsmsg.BitSwapMessage) error
+ Close() error
+}
+
 // Implement Receiver to receive messages from the BitSwapNetwork
 type Receiver interface {
  ReceiveMessage(

exchange/bitswap/network/ipfs_impl.go

@@ -42,6 +42,27 @@ type impl struct {
  receiver Receiver
 }
 
+type streamMessageSender struct {
+ s inet.Stream
+}
+
+func (s *streamMessageSender) Close() error {
+ return s.s.Close()
+}
+
+func (s *streamMessageSender) SendMsg(msg bsmsg.BitSwapMessage) error {
+ return msg.ToNet(s.s)
+}
+
+func (bsnet *impl) NewMessageSender(ctx context.Context, p peer.ID) (MessageSender, error) {
+ s, err := bsnet.newStreamToPeer(ctx, p)
+ if err != nil {
+ return nil, err
+ }
+
+ return &streamMessageSender{s: s}, nil
+}
+
 func (bsnet *impl) newStreamToPeer(ctx context.Context, p peer.ID) (inet.Stream, error) {
 
  // first, make sure we're connected.

exchange/bitswap/testnet/virtual.go

@@ -112,6 +112,30 @@ func (nc *networkClient) FindProvidersAsync(ctx context.Context, k key.Key, max
  return out
 }
 
+type messagePasser struct {
+ net *network
+ target peer.ID
+ local peer.ID
+ ctx context.Context
+}
+
+func (mp *messagePasser) SendMsg(m bsmsg.BitSwapMessage) error {
+ return mp.net.SendMessage(mp.ctx, mp.local, mp.target, m)
+}
+
+func (mp *messagePasser) Close() error {
+ return nil
+}
+
+func (n *networkClient) NewMessageSender(ctx context.Context, p peer.ID) (bsnet.MessageSender, error) {
+ return &messagePasser{
+ net: n.network,
+ target: p,
+ local: n.local,
+ ctx: ctx,
+ }, nil
+}
+
 // Provide provides the key to the network
 func (nc *networkClient) Provide(ctx context.Context, k key.Key) error {
  return nc.routing.Provide(ctx, k)