feat: add unwantBlocks method

Author: dignifiedquire

API.md

@@ -27,6 +27,14 @@ Fetch a single block.
 
 Fetch multiple blocks.
 
+### `unwantBlocks(keys)`
+
+- `keys: []Multihash`
+
+Cancel previously requested keys, forcefully. That means they are removed from the
+wantlist independent of how many other resources requested these keys. Callbacks
+attached to `getBlock` are errored with `Error('manual unwant: key')`.
+
 ### `cancelWants(keys)`
 
 - `keys: []Multihash`

src/index.js

@@ -137,9 +137,9 @@ module.exports = class Bitwap {
       cb(err, block)
     }
 
-    this.getBlocks([key], (err, res) => {
-      if (err) {
-        return done(err)
+    this.getBlocks([key], (errs, res) => {
+      if (errs) {
+        return done(errs[0])
       }
 
       done(null, res[0])
@@ -153,36 +153,80 @@ module.exports = class Bitwap {
 
   getBlocks (keys, cb) {
     const blocks = []
-    const finish = (block) => {
-      blocks.push(block)
-      if (blocks.length === keys.length) {
+    const errs = []
+    const unwantListeners = {}
+    const blockListeners = {}
+    const unwantEvent = (key) => `unwant:${key.toString('hex')}`
+    const blockEvent = (key) => `block:${key.toString('hex')}`
+
+    const cleanupListeners = () => {
+      keys.forEach((key) => {
+        this.notifications.removeListener(unwantEvent(key), unwantListeners[key])
+        this.notifications.removeListener(blockEvent(key), blockListeners[key])
+      })
+    }
+
+    const addListeners = () => {
+      keys.forEach((key) => {
+        unwantListeners[key] = () => {
+          finish(new Error(`manual unwant: ${key.toString('hex')}`))
+        }
+
+        blockListeners[key] = (block) => {
+          finish(null, block)
+        }
+
+        this.notifications.once(unwantEvent(key), unwantListeners[key])
+        this.notifications.once(blockEvent(key), blockListeners[key])
+      })
+    }
+
+    const finish = (err, block) => {
+      if (err) {
+        errs.push(err)
+      }
+      if (block) {
+        blocks.push(block)
+      }
+
+      if (blocks.length + errs.length === keys.length) {
+        cleanupListeners()
+        if (errs.length > 0) {
+          return cb(errs, blocks)
+        }
         cb(null, blocks)
       }
     }
 
+    addListeners()
+
     keys.forEach((key) => {
       // Sanity check, we don't want to announce looking for blocks
       // when we might have them ourselves
       this.datastore.get(key, (err, res) => {
         if (!err && res) {
           this.wm.cancelWants([key])
-          finish(res)
+          finish(null, res)
           return
         }
 
         if (err) {
           log('error in datastore.get: ', err.message)
         }
-
-        this.notifications.once(`block:${key.toString('hex')}`, (block) => {
-          finish(block)
-        })
       })
     })
 
     this.wm.wantBlocks(keys)
   }
 
+  // removes the given keys from the want list independent of any ref counts
+  unwantBlocks (keys) {
+    this.wm.unwantBlocks(keys)
+    keys.forEach((key) => {
+      this.notifications.emit(`unwant:${key.toString('hex')}`)
+    })
+  }
+
   // removes the given keys from the want list
   cancelWants (keys) {
     this.wm.cancelWants(keys)

src/wantlist/index.js

@@ -35,6 +35,12 @@ class Wantlist {
     this.set.delete(key)
   }
 
+  removeForce (key) {
+    if (this.set.has(key)) {
+      this.set.delete(key)
+    }
+  }
+
   entries () {
     return this.set.entries()
   }

src/wantmanager/index.js

@@ -23,7 +23,7 @@ module.exports = class Wantmanager {
     return new MsgQueue(peerId, this.network)
   }
 
-  _addEntries (keys, cancel) {
+  _addEntries (keys, cancel, force) {
     let i = -1
     _(keys)
       .map((key) => {
@@ -33,7 +33,11 @@ module.exports = class Wantmanager {
       .tap((e) => {
         // add changes to our wantlist
         if (e.cancel) {
-          this.wl.remove(e.key)
+          if (force) {
+            this.wl.removeForce(e.key)
+          } else {
+            this.wl.remove(e.key)
+          }
         } else {
           this.wl.add(e.key, e.priority)
         }
@@ -90,6 +94,12 @@ module.exports = class Wantmanager {
     this._addEntries(keys, false)
   }
 
+  // remove blocks of all the given keys without respecting refcounts
+  unwantBlocks (keys) {
+    log('unwant blocks:', keys)
+    this._addEntries(keys, true, true)
+  }
+
   // cancel wanting all of the given keys
   cancelWants (keys) {
     log('cancel wants: ', keys)