From 025a5b91389a0a4ab7c6a0efc85642553eee19e0 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Mon, 14 Nov 2016 15:45:32 -0800
Subject: [PATCH] chain: handle chain resetting more gracefully.

---
 lib/chain/chain.js     | 50 ++++++++++++++++++++++++++++++++++++------
 lib/node/fullnode.js   |  6 ++++-
 lib/node/spvnode.js    |  4 ++++
 lib/wallet/walletdb.js | 29 ++++++++++++++++++++++++
 4 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/lib/chain/chain.js b/lib/chain/chain.js
index 3218170f4..8cb213564 100644
--- a/lib/chain/chain.js
+++ b/lib/chain/chain.js
@@ -824,7 +824,7 @@ Chain.prototype.reorganizeSPV = co(function* reorganizeSPV(competitor, block) {
   // to the fork block, causing
   // us to redownload the blocks
   // on the new main chain.
-  yield this._reset(fork.hash);
+  yield this._reset(fork.hash, true);
 
   // Emit disconnection events now that
   // the chain has successfully reset.
@@ -1004,7 +1004,7 @@ Chain.prototype.saveAlternate = co(function* saveAlternate(entry, block, prev) {
 Chain.prototype.reset = co(function* reset(block) {
   var unlock = yield this.locker.lock();
   try {
-    return yield this._reset(block);
+    return yield this._reset(block, false);
   } finally {
     unlock();
   }
@@ -1017,7 +1017,7 @@ Chain.prototype.reset = co(function* reset(block) {
  * @returns {Promise}
  */
 
-Chain.prototype._reset = co(function* reset(block) {
+Chain.prototype._reset = co(function* reset(block, silent) {
   var tip = yield this.db.reset(block);
 
   // Reset state.
@@ -1028,6 +1028,9 @@ Chain.prototype._reset = co(function* reset(block) {
 
   this.emit('tip', tip);
 
+  if (!silent)
+    this.emit('reset', tip);
+
   // Reset the orphan map completely. There may
   // have been some orphans on a forked chain we
   // no longer need.
@@ -1042,6 +1045,22 @@ Chain.prototype._reset = co(function* reset(block) {
  */
 
 Chain.prototype.replay = co(function* replay(block) {
+  var unlock = yield this.locker.lock();
+  try {
+    return yield this._replay(block);
+  } finally {
+    unlock();
+  }
+});
+
+/**
+ * Reset the chain without a lock.
+ * @private
+ * @param {Hash|Number} block - hash/height
+ * @returns {Promise}
+ */
+
+Chain.prototype._replay = co(function* replay(block) {
   var entry = yield this.db.get(block);
 
   if (!entry)
@@ -1050,10 +1069,27 @@ Chain.prototype.replay = co(function* replay(block) {
   if (!(yield entry.isMainChain()))
     throw new Error('Cannot reset on alternate chain.');
 
-  if (entry.hash === this.network.genesis.hash)
-    return yield this.reset(entry.hash);
+  if (entry.isGenesis())
+    return yield this._reset(entry.hash, true);
+
+  yield this._reset(entry.prevBlock, true);
+});
+
+/**
+ * Scan the blockchain for transactions containing specified address hashes.
+ * @param {Hash} start - Block hash to start at.
+ * @param {Bloom} filter - Bloom filter containing tx and address hashes.
+ * @param {Function} iter - Iterator.
+ * @returns {Promise}
+ */
 
-  yield this.reset(entry.prevBlock);
+Chain.prototype.scan = co(function* scan(start, filter, iter) {
+  var unlock = yield this.locker.lock();
+  try {
+    return yield this.db.scan(block, filter, iter);
+  } finally {
+    unlock();
+  }
 });
 
 /**
@@ -1086,7 +1122,7 @@ Chain.prototype._resetTime = co(function* resetTime(ts) {
   if (!entry)
     return;
 
-  yield this._reset(entry.height);
+  yield this._reset(entry.height, false);
 });
 
 /**
diff --git a/lib/node/fullnode.js b/lib/node/fullnode.js
index 2336a19f8..17bcfc377 100644
--- a/lib/node/fullnode.js
+++ b/lib/node/fullnode.js
@@ -223,6 +223,10 @@ FullNode.prototype._init = function _init() {
       self.mempool.removeBlock(block).catch(onError);
   });
 
+  this.chain.on('reset', function(tip) {
+    self.walletdb.resetChain(tip).catch(onError);
+  });
+
   this.miner.on('block', function(block) {
     self.broadcast(block.toInv()).catch(onError);
   });
@@ -296,7 +300,7 @@ FullNode.prototype.watchData = function watchData(chunks) {
  */
 
 FullNode.prototype.scan = function scan(start, filter, iter) {
-  return this.chain.db.scan(start, filter, iter);
+  return this.chain.scan(start, filter, iter);
 };
 
 /**
diff --git a/lib/node/spvnode.js b/lib/node/spvnode.js
index b12cdc213..8ffaf65d5 100644
--- a/lib/node/spvnode.js
+++ b/lib/node/spvnode.js
@@ -153,6 +153,10 @@ SPVNode.prototype._init = function _init() {
   this.chain.on('disconnect', function(entry, block) {
     self.walletdb.removeBlock(entry).catch(onError);
   });
+
+  this.chain.on('reset', function(tip) {
+    self.walletdb.resetChain(tip).catch(onError);
+  });
 };
 
 /**
diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index 9905e6ca5..7f6902a13 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -2075,6 +2075,35 @@ WalletDB.prototype._unconfirm = co(function* unconfirm(tx) {
   }
 });
 
+/**
+ * Handle a chain reset.
+ * @param {ChainEntry} entry
+ * @returns {Promise}
+ */
+
+WalletDB.prototype.resetChain = co(function* resetChain(entry) {
+  var unlock = yield this.txLock.lock();
+  try {
+    return yield this._resetChain(entry);
+  } finally {
+    unlock();
+  }
+});
+
+/**
+ * Handle a chain reset without a lock.
+ * @private
+ * @param {ChainEntry} entry
+ * @returns {Promise}
+ */
+
+WalletDB.prototype._resetChain = co(function* resetChain(entry) {
+  if (entry.height > this.state.height)
+    throw new Error('WDB: Bad reset height.');
+
+  yield this.scan(entry.height);
+});
+
 /*
  * Helpers
  */