From fb18da4dc6604c02a6e83976abf09c67a7c8c1f3 Mon Sep 17 00:00:00 2001
From: Nodari Chkuaselidze <nodar.chkuaselidze@gmail.com>
Date: Wed, 4 Jun 2025 17:29:06 +0400
Subject: [PATCH] wdb-migration: fix and test coin selection progress recovery.

---
 lib/wallet/migrations.js      | 13 ++++--
 test/wallet-migration-test.js | 86 ++++++++++++++++++++++++++++++++++-
 2 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/lib/wallet/migrations.js b/lib/wallet/migrations.js
index 3cf8950d..0c72b486 100644
--- a/lib/wallet/migrations.js
+++ b/lib/wallet/migrations.js
@@ -1405,9 +1405,16 @@ class MigrateCoinSelection extends AbstractMigration {
     await this.decodeProgress(ctx.state.inProgressData);
 
     for (const wid of wids) {
-      if (wid < this.progress.wid)
+      if (wid < this.progress.wid) {
+        this.logger.debug(
+          'Skipping wallet %d (%d/%d), already migrated.',
+          wid, this.progress.wid, wids.length);
         continue;
+      }
 
+      this.logger.info(
+        'Migrating wallet %d (%d/%d)',
+        wid, this.progress.wid, wids.length);
       await this.migrateWallet(wid, ctx);
     }
 
@@ -1473,7 +1480,7 @@ class MigrateCoinSelection extends AbstractMigration {
         this.progress.wid = wid;
         this.progress.account = account;
         this.progress.hash = hash;
-        this.progress.index = index;
+        this.progress.index = index + 1;
 
         ctx.state.inProgressData = this.encodeProgress();
         ctx.writeState(parent.root());
@@ -1488,7 +1495,7 @@ class MigrateCoinSelection extends AbstractMigration {
     this.progress.hash = consensus.ZERO_HASH;
     this.progress.index = 0;
     ctx.state.inProgressData = this.encodeProgress();
-    ctx.writeState(parent);
+    ctx.writeState(parent.root());
     await parent.write();
   }
 
diff --git a/test/wallet-migration-test.js b/test/wallet-migration-test.js
index 7cccb0a5..a6973428 100644
--- a/test/wallet-migration-test.js
+++ b/test/wallet-migration-test.js
@@ -1121,7 +1121,7 @@ describe('Wallet Migrations', function() {
     };
 
     let walletDB, ldb;
-    before(async () => {
+    beforeEach(async () => {
       WalletMigrator.migrations = {};
       await fs.mkdirp(location);
 
@@ -1133,7 +1133,7 @@ describe('Wallet Migrations', function() {
       await walletDB.close();
     });
 
-    after(async () => {
+    afterEach(async () => {
       WalletMigrator.migrations = migrationsBAK;
       await rimraf(location);
     });
@@ -1164,5 +1164,87 @@ describe('Wallet Migrations', function() {
 
       await walletDB.close();
     });
+
+    it('should resume the progress of migration if interrupted', async () => {
+      // patch the db buckets to throw after each write.
+      const patchDB = () => {
+        // throw after each bucket write.
+        const ldbBucket = walletDB.db.bucket;
+
+        walletDB.db.bucket = (prefix) => {
+          const bucket = ldbBucket.call(ldb, prefix);
+          const bucketBatch = bucket.batch;
+
+          bucket.batch = () => {
+            const batch = bucketBatch.call(bucket);
+            const originalWrite = batch.write;
+
+            batch.write = async () => {
+              await originalWrite.call(batch);
+              throw new Error('Interrupt migration');
+            };
+
+            return batch;
+          };
+
+          return bucket;
+        };
+
+        return () => {
+          walletDB.db.bucket = ldbBucket;
+        };
+      };
+
+      WalletMigrator.migrations = {
+        0: class extends Migration {
+          constructor(options) {
+            super(options);
+
+            this.batchSize = 10;
+          }
+
+          async migrate(b, ctx) {
+            const unpatch = patchDB();
+            await super.migrate(b, ctx);
+            unpatch();
+          }
+        }
+      };
+
+      await walletDB.db.open();
+
+      const migrator = new WalletMigrator({
+        walletMigrate: 0,
+        walletDB: walletDB,
+        dbVersion: 5
+      });
+
+      let err;
+
+      do {
+        try {
+          await migrator.migrate();
+          err = null;
+        } catch (e) {
+          if (e.message !== 'Interrupt migration')
+            throw e;
+
+          err = e;
+        }
+      } while (err);
+
+      await checkEntries(ldb, {
+        before: data.before,
+        after: data.after,
+        throw: true
+      });
+
+      await checkExactEntries(ldb, data.prefixes, {
+        after: data.after,
+        throw: true
+      });
+
+      await walletDB.db.close();
+    });
   });
 });