feat(wallet): implement UTxO lock/unlock functionality, fix utxo sync

packages/wallet/package.json

@@ -30,9 +30,10 @@
     "@cardano-ogmios/schema": "4.1.0",
     "@cardano-sdk/cip2": "0.1.3",
     "@cardano-sdk/core": "0.1.3",
-    "lodash-es": "^4.17.21",
-    "isomorphic-bip39": "^3.0.5",
     "buffer": "^6.0.3",
+    "emittery": "^0.10.0",
+    "isomorphic-bip39": "^3.0.5",
+    "lodash-es": "^4.17.21",
     "ts-custom-error": "^3.2.0",
     "ts-log": "^2.2.3"
   }

packages/wallet/src/InMemoryUtxoRepository.ts

@@ -5,30 +5,51 @@ import { CardanoProvider, Ogmios, CardanoSerializationLib, CSL } from '@cardano-
 import { dummyLogger, Logger } from 'ts-log';
 import { ImplicitCoin, InputSelector, SelectionConstraints, SelectionResult } from '@cardano-sdk/cip2';
 import { KeyManager } from './KeyManagement';
+import { OnTransactionArgs, TransactionTracker, UtxoRepositoryEvents } from '.';
+import { cslToOgmios } from '@cardano-sdk/core/src/Ogmios';
+import Emittery from 'emittery';
 
-export class InMemoryUtxoRepository implements UtxoRepository {
+export interface InMemoryUtxoRepositoryProps {
+  csl: CardanoSerializationLib;
+  provider: CardanoProvider;
+  keyManager: KeyManager;
+  inputSelector: InputSelector;
+  txTracker: TransactionTracker;
+  logger?: Logger;
+}
+
+// Review: is comparing txIn enough to identify unique utxo?
+const utxoEquals = ([txIn1]: [Schema.TxIn, Schema.TxOut], [txIn2]: [Schema.TxIn, Schema.TxOut]): boolean =>
+  txIn1.txId === txIn2.txId && txIn1.index === txIn2.index;
+export class InMemoryUtxoRepository extends Emittery<UtxoRepositoryEvents> implements UtxoRepository {
   #csl: CardanoSerializationLib;
   #delegationAndRewards: Schema.DelegationsAndRewards;
   #inputSelector: InputSelector;
   #keyManager: KeyManager;
   #logger: Logger;
   #provider: CardanoProvider;
   #utxoSet: Set<[TxIn, TxOut]>;
+  #lockedUtxoSet: Set<[TxIn, TxOut]> = new Set();
 
-  constructor(
-    csl: CardanoSerializationLib,
-    provider: CardanoProvider,
-    keyManager: KeyManager,
-    inputSelector: InputSelector,
-    logger?: Logger
-  ) {
+  constructor({
+    csl,
+    logger = dummyLogger,
+    provider,
+    inputSelector,
+    keyManager,
+    txTracker
+  }: InMemoryUtxoRepositoryProps) {
+    super();
     this.#csl = csl;
-    this.#logger = logger ?? dummyLogger;
+    this.#logger = logger;
     this.#provider = provider;
     this.#utxoSet = new Set();
     this.#delegationAndRewards = { rewards: undefined, delegate: undefined };
     this.#inputSelector = inputSelector;
     this.#keyManager = keyManager;
+    txTracker.on('transaction', (args) => {
+      this.#onTransaction(args);
+    });
   }
 
   public async sync(): Promise<void> {
@@ -39,11 +60,17 @@ export class InMemoryUtxoRepository implements UtxoRepository {
     );
     this.#logger.trace(result);
     for (const utxo of result.utxo) {
-      if (!this.#utxoSet.has(utxo)) {
+      if (!this.allUtxos.some((oldUtxo) => utxoEquals(utxo, oldUtxo))) {
         this.#utxoSet.add(utxo);
         this.#logger.debug('New UTxO', utxo);
       }
     }
+    for (const utxo of this.#utxoSet) {
+      if (!result.utxo.some((newUtxo) => utxoEquals(utxo, newUtxo))) {
+        this.#utxoSet.delete(utxo);
+        this.#logger.debug('UTxO is gone', utxo);
+      }
+    }
     if (this.#delegationAndRewards.delegate !== result.delegationAndRewards.delegate) {
       this.#delegationAndRewards.delegate = result.delegationAndRewards.delegate;
       this.#logger.debug('Delegation stored', result.delegationAndRewards.delegate);
@@ -75,11 +102,35 @@ export class InMemoryUtxoRepository implements UtxoRepository {
     return [...this.#utxoSet.values()];
   }
 
+  public get availableUtxos(): Schema.Utxo {
+    return this.allUtxos.filter((utxo) => !this.#lockedUtxoSet.has(utxo));
+  }
+
   public get rewards(): Schema.Lovelace | null {
     return this.#delegationAndRewards.rewards ?? null;
   }
 
   public get delegation(): Schema.PoolId | null {
     return this.#delegationAndRewards.delegate ?? null;
   }
+
+  async #onTransaction({ transaction, confirmed }: OnTransactionArgs) {
+    const utxoLockedByTx: Schema.Utxo = [];
+    const inputs = transaction.body().inputs();
+    for (let inputIdx = 0; inputIdx < inputs.len(); inputIdx++) {
+      const { txId, index } = cslToOgmios.txIn(inputs.get(inputIdx));
+      const utxo = this.allUtxos.find(([txIn]) => txIn.txId === txId && txIn.index === index)!;
+      this.#lockedUtxoSet.add(utxo);
+      utxoLockedByTx.push(utxo);
+    }
+    // We don't care what the error is here. Unlock UTxO and resync on both success and failure.
+    await confirmed.catch((): void => undefined);
+    for (const utxo of utxoLockedByTx) this.#lockedUtxoSet.delete(utxo);
+    try {
+      await this.sync();
+    } catch {
+      this.emit('outOfSync', void 0);
+      this.#logger.debug('Failed to resync InMemoryUtxoRepository after transaction. Emitting "outOfSync"');
+    }
+  }
 }

packages/wallet/src/types.ts

@@ -1,9 +1,12 @@
 import Schema from '@cardano-ogmios/schema';
 import { ImplicitCoin, SelectionConstraints, SelectionResult } from '@cardano-sdk/cip2';
 import { CSL } from '@cardano-sdk/core';
+import Emittery from 'emittery';
 
-export interface UtxoRepository {
+export type UtxoRepositoryEvents = { outOfSync: void };
+export interface UtxoRepository extends Emittery<UtxoRepositoryEvents> {
   allUtxos: Schema.Utxo;
+  availableUtxos: Schema.Utxo;
   rewards: Schema.Lovelace | null;
   delegation: Schema.PoolId | null;
   sync: () => Promise<void>;
@@ -13,3 +16,19 @@ export interface UtxoRepository {
     implicitCoin?: ImplicitCoin
   ) => Promise<SelectionResult>;
 }
+
+export interface OnTransactionArgs {
+  transaction: CSL.Transaction;
+  /**
+   * Resolves when transaction is confirmed.
+   * Rejects if transaction fails to submit or validate.
+   */
+  confirmed: Promise<void>;
+}
+export type TransactionTrackerEvents = { transaction: OnTransactionArgs };
+export interface TransactionTracker extends Emittery<TransactionTrackerEvents> {
+  /**
+   * Track a new transaction
+   */
+  trackTransaction(transaction: CSL.Transaction): void;
+}

packages/wallet/test/InMemoryUtxoRepository.test.ts

@@ -1,19 +1,22 @@
 import { roundRobinRandomImprove, InputSelector } from '@cardano-sdk/cip2';
-import { loadCardanoSerializationLib, CardanoSerializationLib, CSL, CardanoProvider, Ogmios } from '@cardano-sdk/core';
-import { SelectionConstraints } from '@cardano-sdk/util-dev';
-import { providerStub, delegate, rewards } from './ProviderStub';
+import { loadCardanoSerializationLib, CardanoSerializationLib, CSL, Ogmios } from '@cardano-sdk/core';
+import { flushPromises, SelectionConstraints } from '@cardano-sdk/util-dev';
+import { providerStub, delegate, rewards, ProviderStub, utxo, delegationAndRewards } from './ProviderStub';
 import { InMemoryUtxoRepository, KeyManagement, UtxoRepository } from '../src';
+import { MockTransactionTracker } from './mockTransactionTracker';
+import { ogmiosToCsl } from '@cardano-sdk/core/src/Ogmios';
 
 const addresses = [
   'addr_test1qq585l3hyxgj3nas2v3xymd23vvartfhceme6gv98aaeg9muzcjqw982pcftgx53fu5527z2cj2tkx2h8ux2vxsg475q2g7k3g'
 ];
 
 describe('InMemoryUtxoRepository', () => {
   let utxoRepository: UtxoRepository;
-  let provider: CardanoProvider;
+  let provider: ProviderStub;
   let inputSelector: InputSelector;
   let csl: CardanoSerializationLib;
   let outputs: Set<CSL.TransactionOutput>;
+  let txTracker: MockTransactionTracker;
 
   beforeEach(async () => {
     provider = providerStub();
@@ -35,7 +38,8 @@ describe('InMemoryUtxoRepository', () => {
         value: { coins: 2_000_000 }
       })
     ]);
-    utxoRepository = new InMemoryUtxoRepository(csl, provider, keyManager, inputSelector);
+    txTracker = new MockTransactionTracker();
+    utxoRepository = new InMemoryUtxoRepository({ csl, provider, keyManager, inputSelector, txTracker });
   });
 
   test('constructed state', async () => {
@@ -49,6 +53,16 @@ describe('InMemoryUtxoRepository', () => {
     await expect(utxoRepository.allUtxos.length).toBe(3);
     await expect(utxoRepository.rewards).toBe(rewards);
     await expect(utxoRepository.delegation).toBe(delegate);
+    const identicalUtxo = [{ ...utxo[1][0] }, { ...utxo[1][1] }] as const; // clone UTxO
+    provider.utxoDelegationAndRewards.mockResolvedValueOnce({
+      utxo: [utxo[0], identicalUtxo],
+      delegationAndRewards
+    });
+    await utxoRepository.sync();
+    await expect(utxoRepository.allUtxos.length).toBe(2);
+    // Verify we're not replacing the object with an identical one in the UTxO set
+    await expect(utxoRepository.allUtxos).not.toContain(identicalUtxo);
+    await expect(utxoRepository.allUtxos).toContain(utxo[1]);
   });
 
   describe('selectInputs', () => {
@@ -62,4 +76,47 @@ describe('InMemoryUtxoRepository', () => {
       await expect(result.selection.change.size).toBe(2);
     });
   });
+
+  describe('availableUtxos', () => {
+    // TODO: this test is a little difficult to read, try to refactor.
+    // eslint-disable-next-line unicorn/consistent-function-scoping
+    const testUtxoLock = (confirmSuccessful: boolean, failToResync: boolean) => async () => {
+      await utxoRepository.sync();
+      let confirm: Function;
+      const confirmed = new Promise<void>((resolve, reject) => (confirm = confirmSuccessful ? resolve : reject));
+      const transactionUtxo = utxo[0];
+      expect(utxoRepository.availableUtxos).toHaveLength(utxoRepository.allUtxos.length);
+      expect(utxoRepository.availableUtxos).toContain(transactionUtxo);
+      await txTracker.emit('transaction', {
+        transaction: {
+          body: () => ({
+            inputs: () => ({
+              len: () => 1,
+              get: () => ogmiosToCsl(csl).txIn(transactionUtxo[0])
+            })
+          })
+        } as unknown as CSL.Transaction,
+        confirmed
+      });
+      expect(utxoRepository.availableUtxos).toHaveLength(utxoRepository.allUtxos.length - 1);
+      expect(utxoRepository.availableUtxos).not.toContain(transactionUtxo);
+      if (failToResync) {
+        provider.utxoDelegationAndRewards.mockRejectedValueOnce(new Error('doesnt matter'));
+        confirm!();
+        await flushPromises();
+        expect(utxoRepository.availableUtxos).toContain(transactionUtxo);
+      } else {
+        provider.utxoDelegationAndRewards.mockResolvedValueOnce({ utxo: utxo.slice(1), delegationAndRewards });
+        confirm!();
+        await flushPromises();
+        expect(utxoRepository.availableUtxos).not.toContain(transactionUtxo);
+      }
+      expect(utxoRepository.availableUtxos).toHaveLength(utxoRepository.allUtxos.length);
+    };
+
+    it('confirmed, resync success', testUtxoLock(true, false));
+    it('confirmed, resync failure', testUtxoLock(true, true));
+    it('rejected, resync success', testUtxoLock(false, false));
+    it('rejected, resync failure', testUtxoLock(false, true));
+  });
 });

packages/wallet/test/ProviderStub.ts

@@ -1,5 +1,4 @@
 /* eslint-disable max-len */
-import { CardanoProvider } from '@cardano-sdk/core';
 import * as Schema from '@cardano-ogmios/schema';
 
 export const stakeKeyHash = 'stake_test1up7pvfq8zn4quy45r2g572290p9vf99mr9tn7r9xrgy2l2qdsf58d';
@@ -51,14 +50,14 @@ export const utxo: Schema.Utxo = [
 
 export const delegate = 'pool185g59xpqzt7gf0ljr8v8f3akl95qnmardf2f8auwr3ffx7atjj5';
 export const rewards = 33_333;
+export const delegationAndRewards = { delegate, rewards };
 
 /**
  * Provider stub for testing
  *
- * @returns {CardanoProvider} CardanoProvider
+ * returns CardanoProvider-compatible object
  */
-
-export const providerStub = (): CardanoProvider => ({
+export const providerStub = () => ({
   ledgerTip: async () => ({
     blockNo: 1_111_111,
     hash: '10d64cc11e9b20e15b6c46aa7b1fed11246f437e62225655a30ea47bf8cc22d0',
@@ -92,14 +91,7 @@ export const providerStub = (): CardanoProvider => ({
       retiring: 5
     }
   }),
-  utxoDelegationAndRewards: async () => {
-    const delegationAndRewards = {
-      delegate,
-      rewards
-    };
-
-    return { utxo, delegationAndRewards };
-  },
+  utxoDelegationAndRewards: jest.fn().mockResolvedValue({ utxo, delegationAndRewards }),
   queryTransactionsByAddresses: async () =>
     Promise.resolve([
       {
@@ -145,3 +137,5 @@ export const providerStub = (): CardanoProvider => ({
     coinsPerUtxoWord: 34_482
   })
 });
+
+export type ProviderStub = ReturnType<typeof providerStub>;

packages/wallet/test/SingleAddressWallet.test.ts

@@ -10,6 +10,7 @@ import {
   SingleAddressWalletDependencies,
   UtxoRepository
 } from '../src';
+import { txTracker } from './mockTransactionTracker';
 
 describe('Wallet', () => {
   const name = 'Test Wallet';
@@ -30,7 +31,7 @@ describe('Wallet', () => {
     });
     provider = providerStub();
     inputSelector = roundRobinRandomImprove(csl);
-    utxoRepository = new InMemoryUtxoRepository(csl, provider, keyManager, inputSelector);
+    utxoRepository = new InMemoryUtxoRepository({ csl, provider, keyManager, inputSelector, txTracker });
     walletDependencies = { csl, keyManager, provider, utxoRepository };
   });
 

packages/wallet/test/Transaction/createTransactionInternals.test.ts

@@ -12,6 +12,7 @@ import { KeyManager } from '../../src/KeyManagement';
 import { testKeyManager } from '../testKeyManager';
 import { UtxoRepository } from '../../src/types';
 import { InMemoryUtxoRepository } from '../../src/InMemoryUtxoRepository';
+import { txTracker } from '../mockTransactionTracker';
 
 const address =
   'addr_test1qq585l3hyxgj3nas2v3xymd23vvartfhceme6gv98aaeg9muzcjqw982pcftgx53fu5527z2cj2tkx2h8ux2vxsg475q2g7k3g';
@@ -53,7 +54,7 @@ describe('Transaction.createTransactionInternals', () => {
         value: { coins: 2_000_000 }
       })
     ]);
-    utxoRepository = new InMemoryUtxoRepository(csl, provider, keyManager, inputSelector);
+    utxoRepository = new InMemoryUtxoRepository({ csl, provider, keyManager, inputSelector, txTracker });
   });
 
   test('simple transaction', async () => {

packages/wallet/test/mockTransactionTracker.ts

@@ -0,0 +1,8 @@
+import Emittery from 'emittery';
+import { TransactionTrackerEvents } from '../src';
+
+export class MockTransactionTracker extends Emittery<TransactionTrackerEvents> {
+  trackTransaction = jest.fn();
+}
+
+export const txTracker = new MockTransactionTracker();

yarn.lock

@@ -3119,6 +3119,11 @@ email-validator@^2.0.4:
   resolved "https://registry.yarnpkg.com/email-validator/-/email-validator-2.0.4.tgz#b8dfaa5d0dae28f1b03c95881d904d4e40bfe7ed"
   integrity sha512-gYCwo7kh5S3IDyZPLZf6hSS0MnZT8QmJFqYvbqlDZSbwdZlY6QZWxJ4i/6UhITOJ4XzyI647Bm2MXKCLqnJ4nQ==
 
+emittery@^0.10.0:
+  version "0.10.0"
+  resolved "https://registry.yarnpkg.com/emittery/-/emittery-0.10.0.tgz#bb373c660a9d421bb44706ec4967ed50c02a8026"
+  integrity sha512-AGvFfs+d0JKCJQ4o01ASQLGPmSCxgfU9RFXvzPvZdjKK8oscynksuJhWrSTSw7j7Ep/sZct5b5ZhYCi8S/t0HQ==
+
 emittery@^0.8.1:
   version "0.8.1"
   resolved "https://registry.yarnpkg.com/emittery/-/emittery-0.8.1.tgz#bb23cc86d03b30aa75a7f734819dee2e1ba70860"