Add _streamSelectionsfix/deselect

3 files changed, 43 insertions, 11 deletions

diff --git a/lib/file-stream.js b/lib/file-stream.js
index 7a579b6..6ed743a 100644
--- a/lib/file-stream.js
+++ b/lib/file-stream.js
@@ -34,9 +34,9 @@ class FileStream extends stream.Readable {
     this._notifying = false
     this._criticalLength = Math.min((1024 * 1024 / pieceLength) | 0, 2)
 
-    this._torrent.select(this._startPiece, this._endPiece, true, () => {
+    this._torrent._select(this._startPiece, this._endPiece, true, () => {
       this._notify()
-    })
+    }, true)
 
     // Ensure that cleanup happens even if destroy() is never called (readable-stream v3 currently doesn't call it automaticallly)
     eos(this, (err) => {
@@ -97,7 +97,7 @@ class FileStream extends stream.Readable {
 
   _destroy (err, cb) {
     if (!this._torrent.destroyed) {
-      this._torrent.deselect(this._startPiece, this._endPiece)
+      this._torrent._deselect(this._startPiece, this._endPiece, true)
     }
     cb(err)
   }
diff --git a/lib/selection-list.js b/lib/selection-list.js
index fba1468..cd5e419 100644
--- a/lib/selection-list.js
+++ b/lib/selection-list.js
@@ -49,7 +49,11 @@ class SelectionList {
       this._onTreeDelete(interval)
     })
 
+    // Note: notifications can overlap
     this._notifications = new IBplusTree()
+
+    // Overlapping selections to handle file.createReadStream()
+    this._streamSelections = []
   }
 
   get (i) {
@@ -61,15 +65,34 @@ class SelectionList {
     this._list[i] = elem
   }
 
-  add (from, to, priority, notify) {
-    this._tree.add(from, to, priority)
+  add (from, to, priority, notify, isStreamSelection) {
+    if (isStreamSelection) {
+      const s = { from, to, priority }
+      this._streamSelections.push(s)
+      // Sync with list
+      this._onTreeInsert(s)
+    } else {
+      this._tree.add(from, to, priority)
+    }
 
     // Add (optional) notify
     if (notify) this._notifications.insert(new NotificationInterval(from, to, notify))
   }
 
-  remove (from, to) {
-    this._tree.remove(from, to)
+  remove (from, to, isStreamSelection) {
+    if (isStreamSelection) {
+      for (let i = 0; i < this._streamSelections.length; ++i) {
+        const s = this._streamSelections[i]
+        if (s.from === from && s.to === to) {
+          this._streamSelections.splice(i, 1)
+          // Sync with list
+          this._onTreeDelete(s)
+          break
+        }
+      }
+    } else {
+      this._tree.remove(from, to)
+    }
 
     // Remove notifications from the range
     const notificationIntervals = this._notifications.containedRangeSearch(from, to)
@@ -119,6 +142,7 @@ class SelectionList {
     this._list = null
     this._tree = null
     this._notifications = null
+    this._streamSelections = []
   }
 }
 
diff --git a/lib/torrent.js b/lib/torrent.js
index a28c871..3a25b32 100644
--- a/lib/torrent.js
+++ b/lib/torrent.js
@@ -925,7 +925,7 @@ class Torrent extends EventEmitter {
     this._drain()
   }
 
-  select (start, end, priority, notify) {
+  _select (start, end, priority, notify, isStreamSelection) {
     if (this.destroyed) throw new Error('torrent is destroyed')
 
     if (start < 0 || end < start || this.pieces.length <= end) {
@@ -935,21 +935,29 @@ class Torrent extends EventEmitter {
 
     this._debug('select %s-%s (priority %s)', start, end, priority)
 
-    this._selections.add(start, end, priority, notify)
+    this._selections.add(start, end, priority, notify, isStreamSelection)
 
     this._updateSelections()
   }
 
-  deselect (start, end) {
+  select (start, end, priority, notify) {
+    this._select(start, end, priority, notify, false)
+  }
+
+  _deselect (start, end, isStreamSelection) {
     if (this.destroyed) throw new Error('torrent is destroyed')
 
     this._debug('deselect %s-%s', start, end)
 
-    this._selections.remove(start, end)
+    this._selections.remove(start, end, isStreamSelection)
 
     this._updateSelections()
   }
 
+  deselect (start, end, priority, notify) {
+    this._deselect(start, end, priority, notify, false)
+  }
+
   critical (start, end) {
     if (this.destroyed) throw new Error('torrent is destroyed')