fix: slow address generation

db/migrations/20241203185227-add-address-wallet_id-index.js

@@ -0,0 +1,18 @@
+'use strict';
+
+/** @type {import('sequelize-cli').Migration} */
+module.exports = {
+  async up(queryInterface) {
+    await queryInterface.addIndex(
+      'address',
+      ['wallet_id', 'index'],
+      {
+        name: 'idx_wallet_address_index',
+      }
+    );
+  },
+
+  async down(queryInterface) {
+    await queryInterface.removeIndex('address', 'idx_wallet_address_index');
+  }
+};

packages/daemon/__tests__/db/index.test.ts

@@ -12,13 +12,13 @@ import {
   addUtxos,
   fetchAddressBalance,
   fetchAddressTxHistorySum,
-  generateAddresses,
   getAddressWalletInfo,
   getBestBlockHeight,
   getDbConnection,
   getExpiredTimelocksUtxos,
   getLastSyncedEvent,
   getLockedUtxoFromInputs,
+  getMaxIndicesForWallets,
   getMinersList,
   getTokenInformation,
   getTokenSymbols,
@@ -73,6 +73,7 @@ import { DbTxOutput, StringMap, TokenInfo, WalletStatus } from '../../src/types'
 import { Authorities, TokenBalanceMap } from '@wallet-service/common';
 // @ts-ignore
 import { constants } from '@hathor/wallet-lib';
+import { generateAddresses } from '../../src/utils';
 
 // Use a single mysql connection for all tests
 let mysql: Connection;
@@ -782,84 +783,69 @@ describe('address and wallet related tests', () => {
     await expect(checkWalletBalanceTable(mysql, 3, wallet1, tokenId, 25, 5, now, 5, 0b11, 0b01)).resolves.toBe(true);
   });
 
-  test('generateAddresses', async () => {
+  test('should generate addresses correctly', async () => {
     expect.hasAssertions();
+
     const maxGap = 5;
     const address0 = ADDRESSES[0];
+    const address1 = ADDRESSES[1];
+    const address4 = ADDRESSES[4];
 
     // check first with no addresses on database, so it should return only maxGap addresses
-    let addressesInfo = await generateAddresses(mysql, XPUBKEY, maxGap);
+    let addresses = await generateAddresses('mainnet', XPUBKEY, 0, maxGap);
 
-    expect(addressesInfo.addresses).toHaveLength(maxGap);
-    expect(addressesInfo.existingAddresses).toStrictEqual({});
-    expect(Object.keys(addressesInfo.newAddresses)).toHaveLength(maxGap);
-    expect(addressesInfo.addresses[0]).toBe(address0);
+    expect(Object.keys(addresses).length).toBe(maxGap);
+    expect(addresses[address0]).toBe(0);
 
-    // add first address with no transactions. As it's not used, we should still only generate maxGap addresses
+    // add address0 to the database with no transactions
     await addToAddressTable(mysql, [{
       address: address0,
       index: 0,
-      walletId: null,
+      walletId: 'wallet1',
       transactions: 0,
     }]);
-    addressesInfo = await generateAddresses(mysql, XPUBKEY, maxGap);
-    expect(addressesInfo.addresses).toHaveLength(maxGap);
-    expect(addressesInfo.existingAddresses).toStrictEqual({ [address0]: 0 });
-    expect(addressesInfo.lastUsedAddressIndex).toStrictEqual(-1);
 
-    let totalLength = Object.keys(addressesInfo.addresses).length;
-    let existingLength = Object.keys(addressesInfo.existingAddresses).length;
-    expect(Object.keys(addressesInfo.newAddresses)).toHaveLength(totalLength - existingLength);
-    expect(addressesInfo.addresses[0]).toBe(address0);
+    addresses = await generateAddresses('mainnet', XPUBKEY, 0, maxGap);
+    expect(Object.keys(addresses).length).toBe(maxGap);
+    expect(addresses[address0]).toBe(0);
 
-    // mark address as used and check again
+    // now mark address0 as used
     let usedIndex = 0;
     await mysql.query('UPDATE `address` SET `transactions` = ? WHERE `address` = ?', [1, address0]);
-    addressesInfo = await generateAddresses(mysql, XPUBKEY, maxGap);
-    expect(addressesInfo.addresses).toHaveLength(maxGap + usedIndex + 1);
-    expect(addressesInfo.existingAddresses).toStrictEqual({ [address0]: 0 });
-    expect(addressesInfo.lastUsedAddressIndex).toStrictEqual(0);
-
-    totalLength = Object.keys(addressesInfo.addresses).length;
-    existingLength = Object.keys(addressesInfo.existingAddresses).length;
-    expect(Object.keys(addressesInfo.newAddresses)).toHaveLength(totalLength - existingLength);
+    addresses = await generateAddresses('mainnet', XPUBKEY, 0, maxGap + usedIndex + 1);
+    expect(Object.keys(addresses).length).toBe(maxGap + usedIndex + 1);
+    expect(addresses[address0]).toBe(0);
 
-    // add address with index 1 as used
+    // add address1 to the database with transactions
     usedIndex = 1;
-    const address1 = ADDRESSES[1];
     await addToAddressTable(mysql, [{
       address: address1,
-      index: usedIndex,
-      walletId: null,
+      index: 1,
+      walletId: 'wallet1',
       transactions: 1,
     }]);
-    addressesInfo = await generateAddresses(mysql, XPUBKEY, maxGap);
-    expect(addressesInfo.addresses).toHaveLength(maxGap + usedIndex + 1);
-    expect(addressesInfo.existingAddresses).toStrictEqual({ [address0]: 0, [address1]: 1 });
-    expect(addressesInfo.lastUsedAddressIndex).toStrictEqual(1);
-    totalLength = Object.keys(addressesInfo.addresses).length;
-    existingLength = Object.keys(addressesInfo.existingAddresses).length;
-    expect(Object.keys(addressesInfo.newAddresses)).toHaveLength(totalLength - existingLength);
-
-    // add address with index 4 as used
+
+    addresses = await generateAddresses('mainnet', XPUBKEY, 0, maxGap + usedIndex + 1);
+    expect(Object.keys(addresses).length).toBe(maxGap + usedIndex + 1);
+    expect(addresses[address0]).toBe(0);
+    expect(addresses[address1]).toBe(1);
+
+    // add address4 to the database with transactions
     usedIndex = 4;
-    const address4 = ADDRESSES[4];
     await addToAddressTable(mysql, [{
       address: address4,
-      index: usedIndex,
-      walletId: null,
+      index: 4,
+      walletId: 'wallet1',
       transactions: 1,
     }]);
-    addressesInfo = await generateAddresses(mysql, XPUBKEY, maxGap);
-    expect(addressesInfo.addresses).toHaveLength(maxGap + usedIndex + 1);
-    expect(addressesInfo.existingAddresses).toStrictEqual({ [address0]: 0, [address1]: 1, [address4]: 4 });
-    expect(addressesInfo.lastUsedAddressIndex).toStrictEqual(4);
-    totalLength = Object.keys(addressesInfo.addresses).length;
-    existingLength = Object.keys(addressesInfo.existingAddresses).length;
-    expect(Object.keys(addressesInfo.newAddresses)).toHaveLength(totalLength - existingLength);
+
+    addresses = await generateAddresses('mainnet', XPUBKEY, 0, maxGap + usedIndex + 1);
+    expect(Object.keys(addresses).length).toBe(maxGap + usedIndex + 1);
+    expect(addresses[address0]).toBe(0);
+    expect(addresses[address4]).toBe(4);
 
     // make sure no address was skipped from being generated
-    for (const [index, address] of addressesInfo.addresses.entries()) {
+    for (const [address, index] of Object.entries(addresses)) {
       expect(ADDRESSES[index]).toBe(address);
     }
   }, 15000);
@@ -1273,3 +1259,90 @@ describe('voidTransaction', () => {
     await expect(voidTransaction(mysql, 'mysterious-transaction', {})).rejects.toThrow('Tried to void a transaction that is not in the database.');
   });
 });
+
+describe('address generation and index methods', () => {
+  test('generateAddresses should generate correct addresses', async () => {
+    expect.hasAssertions();
+
+    const startIndex = 0;
+    const count = 3;
+    const addresses = await generateAddresses('mainnet', XPUBKEY, startIndex, count);
+
+    // Check if we got the expected number of addresses
+    expect(Object.keys(addresses).length).toBe(count);
+
+    // Check if the addresses are mapped to correct indices
+    Object.entries(addresses).forEach(([address, index]) => {
+      expect(typeof address).toBe('string');
+      expect(index).toBeGreaterThanOrEqual(startIndex);
+      expect(index).toBeLessThan(startIndex + count);
+    });
+  });
+
+  test('getMaxIndicesForWallets should return correct indices for multiple wallets', async () => {
+    expect.hasAssertions();
+
+    const wallet1 = 'wallet1';
+    const wallet2 = 'wallet2';
+    const addresses1 = ['addr1', 'addr2', 'addr3'];
+    const addresses2 = ['addr4', 'addr5'];
+    const indices1 = [5, 10, 15];
+    const indices2 = [7, 12];
+
+    // Add addresses for wallet1
+    const entries1 = addresses1.map((address, i) => ({
+      address,
+      index: indices1[i],
+      walletId: wallet1,
+      transactions: 0,
+    }));
+    await addToAddressTable(mysql, entries1);
+
+    // Add addresses for wallet2
+    const entries2 = addresses2.map((address, i) => ({
+      address,
+      index: indices2[i],
+      walletId: wallet2,
+      transactions: 0,
+    }));
+    await addToAddressTable(mysql, entries2);
+
+    // Test getting indices for both wallets
+    const walletData = [
+      { walletId: wallet1, addresses: addresses1 },
+      { walletId: wallet2, addresses: addresses2 },
+    ];
+    const indices = await getMaxIndicesForWallets(mysql, walletData);
+
+    // Check wallet1 indices
+    const wallet1Indices = indices.get(wallet1);
+    expect(wallet1Indices).toBeDefined();
+    expect(wallet1Indices?.maxAmongAddresses).toBe(15);
+    expect(wallet1Indices?.maxWalletIndex).toBe(15);
+
+    // Check wallet2 indices
+    const wallet2Indices = indices.get(wallet2);
+    expect(wallet2Indices).toBeDefined();
+    expect(wallet2Indices?.maxAmongAddresses).toBe(12);
+    expect(wallet2Indices?.maxWalletIndex).toBe(12);
+
+    // Test with empty wallet data
+    const emptyIndices = await getMaxIndicesForWallets(mysql, []);
+    expect(emptyIndices.size).toBe(0);
+
+    // Test with non-existent wallet
+    const nonExistentIndices = await getMaxIndicesForWallets(mysql, [
+      { walletId: 'nonexistent', addresses: ['addr1'] }
+    ]);
+    expect(nonExistentIndices.size).toBe(0);
+
+    // Test with subset of addresses
+    const subsetIndices = await getMaxIndicesForWallets(mysql, [
+      { walletId: wallet1, addresses: addresses1.slice(0, 2) }
+    ]);
+    const subsetWallet1 = subsetIndices.get(wallet1);
+    expect(subsetWallet1).toBeDefined();
+    expect(subsetWallet1?.maxAmongAddresses).toBe(10);
+    expect(subsetWallet1?.maxWalletIndex).toBe(15);
+  });
+});

packages/daemon/__tests__/services/services.test.ts

@@ -19,8 +19,8 @@ import {
   getUtxosLockedAtHeight,
   addOrUpdateTx,
   getAddressWalletInfo,
-  generateAddresses,
   storeTokenInformation,
+  getMaxIndicesForWallets,
 } from '../../src/db';
 import {
   fetchInitialState,
@@ -39,6 +39,7 @@ import {
   getFullnodeHttpUrl,
   invokeOnTxPushNotificationRequestedLambda,
   getWalletBalancesForTx,
+  generateAddresses,
 } from '../../src/utils';
 import getConfig from '../../src/config';
 
@@ -83,6 +84,9 @@ jest.mock('../../src/db', () => ({
   generateAddresses: jest.fn(),
   addNewAddresses: jest.fn(),
   updateWalletTablesWithTx: jest.fn(),
+  getMaxIndicesForWallets: jest.fn(() => new Map([
+    ['wallet1', { maxAmongAddresses: 10, maxWalletIndex: 15 }]
+  ])),
 }));
 
 jest.mock('../../src/utils', () => ({
@@ -102,6 +106,7 @@ jest.mock('../../src/utils', () => ({
   invokeOnTxPushNotificationRequestedLambda: jest.fn(),
   sendMessageSQS: jest.fn(),
   getWalletBalancesForTx: jest.fn(),
+  generateAddresses: jest.fn(),
 }));
 
 beforeEach(() => {
@@ -454,11 +459,18 @@ describe('handleVertexAccepted', () => {
     jest.clearAllMocks();
 
     (getDbConnection as jest.Mock).mockResolvedValue(mockDb);
-    (getAddressWalletInfo as jest.Mock).mockResolvedValue({});
+    (getAddressWalletInfo as jest.Mock).mockResolvedValue({
+      address1: { walletId: 'wallet1', xpubkey: 'xpubkey1', maxGap: 10 }
+    });
+
     (generateAddresses as jest.Mock).mockResolvedValue({
-      newAddresses: ['mockAddress1', 'mockAddress2'],
-      lastUsedAddressIndex: 1
+      'new-address-1': 16,
+      'new-address-2': 17,
     });
+
+    (getMaxIndicesForWallets as jest.Mock).mockResolvedValue(new Map([
+      ['wallet1', { maxAmongAddresses: 10, maxWalletIndex: 15 }]
+    ]));
   });
 
   it('should handle vertex accepted successfully', async () => {