Merge pull request #177 from ipfs/refactor/use-global-pubsub-notifier

Refactor: use global pubsub notifier

Author: Stebalien

bitswap.go

@@ -16,6 +16,7 @@ import (
 	bsmsg "github.com/ipfs/go-bitswap/message"
 	bsmq "github.com/ipfs/go-bitswap/messagequeue"
 	bsnet "github.com/ipfs/go-bitswap/network"
+	notifications "github.com/ipfs/go-bitswap/notifications"
 	bspm "github.com/ipfs/go-bitswap/peermanager"
 	bspqm "github.com/ipfs/go-bitswap/providerquerymanager"
 	bssession "github.com/ipfs/go-bitswap/session"
@@ -116,16 +117,18 @@ func New(parent context.Context, network bsnet.BitSwapNetwork,
 	pqm := bspqm.New(ctx, network)
 
 	sessionFactory := func(ctx context.Context, id uint64, pm bssession.PeerManager, srs bssession.RequestSplitter,
+		notif notifications.PubSub,
 		provSearchDelay time.Duration,
 		rebroadcastDelay delay.D) bssm.Session {
-		return bssession.New(ctx, id, wm, pm, srs, provSearchDelay, rebroadcastDelay)
+		return bssession.New(ctx, id, wm, pm, srs, notif, provSearchDelay, rebroadcastDelay)
 	}
 	sessionPeerManagerFactory := func(ctx context.Context, id uint64) bssession.PeerManager {
 		return bsspm.New(ctx, id, network.ConnectionManager(), pqm)
 	}
 	sessionRequestSplitterFactory := func(ctx context.Context) bssession.RequestSplitter {
 		return bssrs.New(ctx)
 	}
+	notif := notifications.New()
 
 	bs := &Bitswap{
 		blockstore:       bstore,
@@ -136,7 +139,8 @@ func New(parent context.Context, network bsnet.BitSwapNetwork,
 		provideKeys:      make(chan cid.Cid, provideKeysBufferSize),
 		wm:               wm,
 		pqm:              pqm,
-		sm:               bssm.New(ctx, sessionFactory, sessionPeerManagerFactory, sessionRequestSplitterFactory),
+		sm:               bssm.New(ctx, sessionFactory, sessionPeerManagerFactory, sessionRequestSplitterFactory, notif),
+		notif:            notif,
 		counters:         new(counters),
 		dupMetric:        dupHist,
 		allMetric:        allHist,
@@ -163,6 +167,7 @@ func New(parent context.Context, network bsnet.BitSwapNetwork,
 	go func() {
 		<-px.Closing() // process closes first
 		cancelFunc()
+		notif.Shutdown()
 	}()
 	procctx.CloseAfterContext(px, ctx) // parent cancelled first
 
@@ -187,6 +192,9 @@ type Bitswap struct {
 	// NB: ensure threadsafety
 	blockstore blockstore.Blockstore
 
+	// manages channels of outgoing blocks for sessions
+	notif notifications.PubSub
+
 	// newBlocks is a channel for newly added blocks to be provided to the
 	// network.  blocks pushed down this channel get buffered and fed to the
 	// provideKeys channel later on to avoid too much network activity
@@ -307,18 +315,38 @@ func (bs *Bitswap) receiveBlocksFrom(from peer.ID, blks []blocks.Block) error {
 	// to the same node. We should address this soon, but i'm not going to do
 	// it now as it requires more thought and isnt causing immediate problems.
 
-	// Send all blocks (including duplicates) to any sessions that want them.
+	allKs := make([]cid.Cid, 0, len(blks))
+	for _, b := range blks {
+		allKs = append(allKs, b.Cid())
+	}
+
+	wantedKs := allKs
+	if len(blks) != len(wanted) {
+		wantedKs = make([]cid.Cid, 0, len(wanted))
+		for _, b := range wanted {
+			wantedKs = append(wantedKs, b.Cid())
+		}
+	}
+
+	// Send all block keys (including duplicates) to any sessions that want them.
 	// (The duplicates are needed by sessions for accounting purposes)
-	bs.sm.ReceiveBlocksFrom(from, blks)
+	bs.sm.ReceiveFrom(from, allKs)
 
-	// Send wanted blocks to decision engine
-	bs.engine.AddBlocks(wanted)
+	// Send wanted block keys to decision engine
+	bs.engine.AddBlocks(wantedKs)
+
+	// Publish the block to any Bitswap clients that had requested blocks.
+	// (the sessions use this pubsub mechanism to inform clients of received
+	// blocks)
+	for _, b := range wanted {
+		bs.notif.Publish(b)
+	}
 
 	// If the reprovider is enabled, send wanted blocks to reprovider
 	if bs.provideEnabled {
-		for _, b := range wanted {
+		for _, k := range wantedKs {
 			select {
-			case bs.newBlocks <- b.Cid():
+			case bs.newBlocks <- k:
 				// send block off to be reprovided
 			case <-bs.process.Closing():
 				return bs.process.Close()

decision/engine.go

@@ -10,7 +10,6 @@ import (
 	"github.com/google/uuid"
 	bsmsg "github.com/ipfs/go-bitswap/message"
 	wl "github.com/ipfs/go-bitswap/wantlist"
-	blocks "github.com/ipfs/go-block-format"
 	cid "github.com/ipfs/go-cid"
 	bstore "github.com/ipfs/go-ipfs-blockstore"
 	logging "github.com/ipfs/go-log"
@@ -312,13 +311,13 @@ func (e *Engine) MessageReceived(p peer.ID, m bsmsg.BitSwapMessage) {
 	}
 }
 
-func (e *Engine) addBlocks(blocks []blocks.Block) {
+func (e *Engine) addBlocks(ks []cid.Cid) {
 	work := false
 
 	for _, l := range e.ledgerMap {
 		l.lk.Lock()
-		for _, block := range blocks {
-			if entry, ok := l.WantListContains(block.Cid()); ok {
+		for _, k := range ks {
+			if entry, ok := l.WantListContains(k); ok {
 				e.peerRequestQueue.PushBlock(l.Partner, peertask.Task{
 					Identifier: entry.Cid,
 					Priority:   entry.Priority,
@@ -337,11 +336,11 @@ func (e *Engine) addBlocks(blocks []blocks.Block) {
 // AddBlocks is called when new blocks are received and added to a block store,
 // meaning there may be peers who want those blocks, so we should send the blocks
 // to them.
-func (e *Engine) AddBlocks(blocks []blocks.Block) {
+func (e *Engine) AddBlocks(ks []cid.Cid) {
 	e.lock.Lock()
 	defer e.lock.Unlock()
 
-	e.addBlocks(blocks)
+	e.addBlocks(ks)
 }
 
 // TODO add contents of m.WantList() to my local wantlist? NB: could introduce

getter/getter.go

@@ -61,40 +61,56 @@ func SyncGetBlock(p context.Context, k cid.Cid, gb GetBlocksFunc) (blocks.Block,
 type WantFunc func(context.Context, []cid.Cid)
 
 // AsyncGetBlocks take a set of block cids, a pubsub channel for incoming
-// blocks, a want function, and a close function,
-// and returns a channel of incoming blocks.
-func AsyncGetBlocks(ctx context.Context, keys []cid.Cid, notif notifications.PubSub, want WantFunc, cwants func([]cid.Cid)) (<-chan blocks.Block, error) {
+// blocks, a want function, and a close function, and returns a channel of
+// incoming blocks.
+func AsyncGetBlocks(ctx context.Context, sessctx context.Context, keys []cid.Cid, notif notifications.PubSub,
+	want WantFunc, cwants func([]cid.Cid)) (<-chan blocks.Block, error) {
+
+	// If there are no keys supplied, just return a closed channel
 	if len(keys) == 0 {
 		out := make(chan blocks.Block)
 		close(out)
 		return out, nil
 	}
 
+	// Use a PubSub notifier to listen for incoming blocks for each key
 	remaining := cid.NewSet()
 	promise := notif.Subscribe(ctx, keys...)
 	for _, k := range keys {
 		log.Event(ctx, "Bitswap.GetBlockRequest.Start", k)
 		remaining.Add(k)
 	}
 
+	// Send the want request for the keys to the network
 	want(ctx, keys)
 
 	out := make(chan blocks.Block)
-	go handleIncoming(ctx, remaining, promise, out, cwants)
+	go handleIncoming(ctx, sessctx, remaining, promise, out, cwants)
 	return out, nil
 }
 
-func handleIncoming(ctx context.Context, remaining *cid.Set, in <-chan blocks.Block, out chan blocks.Block, cfun func([]cid.Cid)) {
+// Listens for incoming blocks, passing them to the out channel.
+// If the context is cancelled or the incoming channel closes, calls cfun with
+// any keys corresponding to blocks that were never received.
+func handleIncoming(ctx context.Context, sessctx context.Context, remaining *cid.Set,
+	in <-chan blocks.Block, out chan blocks.Block, cfun func([]cid.Cid)) {
+
 	ctx, cancel := context.WithCancel(ctx)
+
+	// Clean up before exiting this function, and call the cancel function on
+	// any remaining keys
 	defer func() {
 		cancel()
 		close(out)
 		// can't just defer this call on its own, arguments are resolved *when* the defer is created
 		cfun(remaining.Keys())
 	}()
+
 	for {
 		select {
 		case blk, ok := <-in:
+			// If the channel is closed, we're done (note that PubSub closes
+			// the channel once all the keys have been received)
 			if !ok {
 				return
 			}
@@ -104,9 +120,13 @@ func handleIncoming(ctx context.Context, remaining *cid.Set, in <-chan blocks.Bl
 			case out <- blk:
 			case <-ctx.Done():
 				return
+			case <-sessctx.Done():
+				return
 			}
 		case <-ctx.Done():
 			return
+		case <-sessctx.Done():
+			return
 		}
 	}
 }

notifications/notifications.go

@@ -60,8 +60,8 @@ func (ps *impl) Shutdown() {
 }
 
 // Subscribe returns a channel of blocks for the given |keys|. |blockChannel|
-// is closed if the |ctx| times out or is cancelled, or after sending len(keys)
-// blocks.
+// is closed if the |ctx| times out or is cancelled, or after receiving the blocks
+// corresponding to |keys|.
 func (ps *impl) Subscribe(ctx context.Context, keys ...cid.Cid) <-chan blocks.Block {
 
 	blocksCh := make(chan blocks.Block, len(keys))
@@ -82,6 +82,8 @@ func (ps *impl) Subscribe(ctx context.Context, keys ...cid.Cid) <-chan blocks.Bl
 	default:
 	}
 
+	// AddSubOnceEach listens for each key in the list, and closes the channel
+	// once all keys have been received
 	ps.wrapped.AddSubOnceEach(valuesCh, toStrings(keys)...)
 	go func() {
 		defer func() {