Merge branch 'storage-leakfix' of https://github.com/astro/webtorrent into astro-storage-leakfix

3 files changed, 25 insertions, 9 deletions

diff --git a/bin/cmd.js b/bin/cmd.js
index 81990e4..fadf0dd 100755
--- a/bin/cmd.js
+++ b/bin/cmd.js
@@ -11,6 +11,7 @@ var networkAddress = require('network-address')
 var parseTorrent = require('parse-torrent')
 var path = require('path')
 var prettyBytes = require('pretty-bytes')
+var Storage = require('../lib/storage')
 var WebTorrent = require('../')
 var zeroFill = require('zero-fill')
 
@@ -523,20 +524,33 @@ function drawTorrent (torrent) {
 
     if (argv.verbose) {
       var pieces = torrent.storage.pieces
+      var storageMem = 0
       for (var i = 0; i < pieces.length; i++) {
         var piece = pieces[i]
-        if (piece.verified || (piece.blocksWritten === 0 && !piece.blocks[0])) {
-          continue
-        }
+        if (piece.buffer) storageMem += piece.buffer.length
+        if (piece.verified || (piece.blocksWritten === 0 && !piece.blocks[0])) continue
         var bar = ''
         for (var j = 0; j < piece.blocks.length; j++) {
-          bar += piece.blocks[j] ? (piece.blocks[j] === 1 ? '{blue:█}' : '{green:█}') : '{red:█}'
+          switch (piece.blocks[j]) {
+            case Storage.BLOCK_BLANK:
+              bar += '{red:█}'
+              break
+            case Storage.BLOCK_RESERVED:
+              bar += '{blue:█}'
+              break
+            case Storage.BLOCK_WRITTEN:
+              bar += '{green:█}'
+              break
+          }
         }
         clivas.line('{4+cyan:' + i + '} ' + bar)
         linesRemaining -= 1
       }
+      clivas.line(
+        '{red:storage mem:} {bold:' + prettyBytes(storageMem) + ' KB}  '
+      )
       clivas.line('{80:}')
-      linesRemaining -= 1
+      linesRemaining -= 2
     }
 
     torrent.swarm.wires.every(function (wire) {
diff --git a/lib/fs-storage.js b/lib/fs-storage.js
index 51e5ea5..14b423e 100644
--- a/lib/fs-storage.js
+++ b/lib/fs-storage.js
@@ -156,6 +156,7 @@ FSStorage.prototype._onPieceDone = function (piece) {
   var writeToNextFile = function (err) {
     if (err) return self.emit('error', err)
     if (i >= end) {
+      delete piece.buffer
       return cb()
     }
 
diff --git a/lib/storage.js b/lib/storage.js
index 805c498..d13e116 100644
--- a/lib/storage.js
+++ b/lib/storage.js
@@ -1,4 +1,4 @@
-module.exports = Storage
+module.exports = exports = Storage
 
 var BitField = require('bitfield')
 var BlockStream = require('block-stream2')
@@ -16,9 +16,10 @@ var sha1 = require('simple-sha1')
 
 var BLOCK_LENGTH = 16 * 1024
 
-var BLOCK_BLANK = 0
-var BLOCK_RESERVED = 1
-var BLOCK_WRITTEN = 2
+var BLOCK_BLANK = exports.BLOCK_BLANK = 0
+var BLOCK_RESERVED = exports.BLOCK_RESERVED = 1
+var BLOCK_WRITTEN = exports.BLOCK_WRITTEN = 2
+
 function noop () {}
 
 inherits(Piece, EventEmitter)