txdb: more atomicity for double-spender removal.

lib/wallet/txdb.js

@@ -633,15 +633,22 @@ TXDB.prototype._add = co(function* add(tx, info) {
   if (result)
     return true;
 
+  this.start();
+
   // Verify and get coins.
-  result = yield this.verify(tx, info);
+  // This potentially removes double-spenders.
+  try {
+    result = yield this.verify(tx, info);
+  } catch (e) {
+    this.drop();
+    throw e;
+  }
 
   if (!result)
     return false;
 
   hash = tx.hash('hex');
 
-  this.start();
   this.put(layout.t(hash), tx.toExtended());
 
   if (tx.ts === 0)
@@ -804,41 +811,38 @@ TXDB.prototype.removeConflict = co(function* removeConflict(hash, ref) {
  * @returns {Promise} - Returns Boolean.
  */
 
-TXDB.prototype.removeRecursive = co(function* removeRecursive(tx) {
+TXDB.prototype.removeRecursive = co(function* removeRecursive(tx, removed) {
   var hash = tx.hash('hex');
   var i, spent, stx, info;
 
+  if (!removed)
+    removed = {};
+
   for (i = 0; i < tx.outputs.length; i++) {
     spent = yield this.isSpent(hash, i);
 
     if (!spent)
       continue;
 
+    if (removed[spent.hash])
+      continue;
+
+    removed[spent.hash] = true;
+
     // Remove all of the spender's spenders first.
     stx = yield this.getTX(spent.hash);
 
     if (!stx)
       throw new Error('Could not find spender.');
 
-    yield this.removeRecursive(stx);
+    yield this.removeRecursive(stx, removed);
   }
 
-  this.start();
-
   // Remove the spender.
-  try {
-    info = yield this.lazyRemove(tx);
-  } catch (e) {
-    this.drop();
-    throw e;
-  }
+  info = yield this.lazyRemove(tx);
 
-  if (!info) {
-    this.drop();
+  if (!info)
     throw new Error('Cannot remove spender.');
-  }
-
-  yield this.commit();
 
   return info;
 });
@@ -1008,7 +1012,16 @@ TXDB.prototype._remove = co(function* remove(hash) {
   if (!tx)
     return;
 
-  info = yield this.removeRecursive(tx);
+  this.start();
+
+  try {
+    info = yield this.removeRecursive(tx);
+  } catch (e) {
+    this.drop();
+    throw e;
+  }
+
+  yield this.commit();
 
   if (!info)
     return;

lib/wallet/wallet.js

@@ -72,7 +72,7 @@ function Wallet(db, options) {
   this.indexCache = new LRU(10000);
   this.accountCache = new LRU(10000);
   this.pathCache = new LRU(100000);
-  this.batch = null;
+  this.current = null;
 
   this.wid = 0;
   this.id = null;

lib/wallet/walletdb.js

@@ -144,7 +144,6 @@ function WalletDB(options) {
   this.network = Network.get(options.network);
   this.fees = options.fees;
   this.logger = options.logger || Logger.global;
-  this.batches = {};
   this.wallets = {};
 
   this.tip = this.network.genesis.hash;
@@ -292,9 +291,9 @@ WalletDB.prototype.getDepth = co(function* getDepth() {
  */
 
 WalletDB.prototype.start = function start(wallet) {
-  assert(!wallet.batch, 'Batch already started.');
-  wallet.batch = this.db.batch();
-  return wallet.batch;
+  assert(!wallet.current, 'Batch already started.');
+  wallet.current = this.db.batch();
+  return wallet.current;
 };
 
 /**
@@ -305,7 +304,7 @@ WalletDB.prototype.start = function start(wallet) {
 
 WalletDB.prototype.drop = function drop(wallet) {
   var batch = this.batch(wallet);
-  wallet.batch = null;
+  wallet.current = null;
   batch.clear();
 };
 
@@ -317,8 +316,8 @@ WalletDB.prototype.drop = function drop(wallet) {
  */
 
 WalletDB.prototype.batch = function batch(wallet) {
-  assert(wallet.batch, 'Batch does not exist.');
-  return wallet.batch;
+  assert(wallet.current, 'Batch does not exist.');
+  return wallet.current;
 };
 
 /**
@@ -329,8 +328,8 @@ WalletDB.prototype.batch = function batch(wallet) {
  */
 
 WalletDB.prototype.commit = function commit(wallet) {
-  var batch = wallet.batch;
-  wallet.batch = null;
+  var batch = wallet.current;
+  wallet.current = null;
   return batch.write();
 };