Transaction multichain fix starttracking

packages/transaction-controller/src/TransactionController.test.ts

@@ -4399,4 +4399,29 @@ describe('TransactionController', () => {
       Current tx status: ${TransactionStatus.submitted}`);
     });
   });
+  describe('startTrackinbByNetworkClientId', () => {
+    it('should start tracking in a tracking map', () => {
+      const controller = newController();
+      const trackingMap = controller.startTrackingByNetworkClientId('mainnet');
+      expect(trackingMap.get('mainnet')?.nonceTracker).toBeDefined();
+      expect(
+        trackingMap.get('mainnet')?.incomingTransactionHelper,
+      ).toBeDefined();
+      expect(
+        trackingMap.get('mainnet')?.pendingTransactionTracker,
+      ).toBeDefined();
+    });
+    it('should stop tracking in a tracking map', () => {
+      const controller = newController();
+      const trackingMap = controller.startTrackingByNetworkClientId('mainnet');
+      const incomingTransactionHelper =
+        trackingMap.get('mainnet')?.incomingTransactionHelper;
+      if (!incomingTransactionHelper) {
+        throw new Error('incomingTransactionHelper is undefined');
+      }
+      const stopSpy = jest.spyOn(incomingTransactionHelper, 'stop');
+      controller.stopTrackingByNetworkClientId('mainnet');
+      expect(stopSpy).toHaveBeenCalledTimes(1);
+    });
+  });
 });

packages/transaction-controller/src/TransactionController.ts

@@ -348,11 +348,11 @@ export class TransactionController extends BaseControllerV1<
 
   private readonly trackingMap: Map<
     NetworkClientId,
-    Set<{
+    {
       nonceTracker: NonceTracker;
       pendingTransactionTracker: PendingTransactionTracker;
       incomingTransactionHelper: IncomingTransactionHelper;
-    }>
+    }
   > = new Map();
 
   /**
@@ -400,9 +400,9 @@ export class TransactionController extends BaseControllerV1<
    * @param options.hooks.beforeCheckPendingTransaction - Additional logic to execute before checking pending transactions. Return false to prevent the broadcast of the transaction.
    * @param options.hooks.beforePublish - Additional logic to execute before publishing a transaction. Return false to prevent the broadcast of the transaction.
    * @param options.hooks.getAdditionalSignArguments - Returns additional arguments required to sign a transaction.
+   * @param options.getNetworkClientIdForDomain - Gets the network client id for the given domain.
    * @param config - Initial options used to configure this controller.
    * @param state - Initial state to set on this controller.
-   * @param options.getNetworkClientIdForDomain
    */
   constructor(
     {
@@ -915,7 +915,7 @@ export class TransactionController extends BaseControllerV1<
       txParams: newTxParams,
     });
 
-    const ethQuery = this.getEthQuery(transactionMeta.networkClientId)
+    const ethQuery = this.getEthQuery(transactionMeta.networkClientId);
     const hash = await this.publishTransaction(ethQuery, rawTx);
 
     const cancelTransactionMeta: TransactionMeta = {
@@ -1073,8 +1073,8 @@ export class TransactionController extends BaseControllerV1<
     // that the user approved them on, but it makes sense to allow cancelling
     // from any network that's also on the same chain. We will need to add a fallback
     // here to allow using networkClientIds other than the original
-    const ethQuery = this.getEthQuery(transactionMeta.networkClientId)
-    const hash = await this.publishTransaction(ethQuery, rawTx)
+    const ethQuery = this.getEthQuery(transactionMeta.networkClientId);
+    const hash = await this.publishTransaction(ethQuery, rawTx);
 
     const baseTransactionMeta: TransactionMeta = {
       ...transactionMeta,
@@ -1122,10 +1122,25 @@ export class TransactionController extends BaseControllerV1<
     this.hub.emit(`${transactionMeta.id}:speedup`, newTransactionMeta);
   }
 
+  stopTrackingByNetworkClientId(networkClientId: NetworkClientId) {
+    const trackers = this.trackingMap.get(networkClientId);
+    if (trackers) {
+      this.removePendingTransactionTrackerListeners(
+        trackers.pendingTransactionTracker,
+      );
+      trackers.incomingTransactionHelper.stop();
+
+      // doesn't seem like any cleanup is needed for nonceTracker
+      // trackers.nonceTracker
+
+      // stop not exposed for pendingTransactionTracker
+      // trackers.pendingTransactionTracker.stop();
+    }
+    this.trackingMap.delete(networkClientId);
+  }
+
   startTrackingByNetworkClientId(networkClientId: NetworkClientId) {
     const networkClient = this.getNetworkClientById(networkClientId);
-    // track using tracking map
-    this.trackingMap.set(networkClientId, new Set());
     const nonceTracker = new NonceTracker({
       provider: networkClient.provider as any,
       blockTracker: networkClient.blockTracker,
@@ -1176,11 +1191,12 @@ export class TransactionController extends BaseControllerV1<
     this.addPendingTransactionTrackerListeners(pendingTransactionTracker);
 
     // add to tracking map
-    this.trackingMap.get(networkClientId)?.add({
+    this.trackingMap.set(networkClientId, {
       nonceTracker,
       incomingTransactionHelper,
       pendingTransactionTracker,
     });
+    return this.trackingMap;
   }
 
   /**
@@ -1189,8 +1205,11 @@ export class TransactionController extends BaseControllerV1<
    * @param transaction - The transaction to estimate gas for.
    * @returns The gas and gas price.
    */
-  async estimateGas(transaction: TransactionParams, networkClientId?: NetworkClientId) {
-    const ethQuery = this.getEthQuery(networkClientId)
+  async estimateGas(
+    transaction: TransactionParams,
+    networkClientId?: NetworkClientId,
+  ) {
+    const ethQuery = this.getEthQuery(networkClientId);
     const { estimatedGas, simulationFails } = await estimateGas(
       transaction,
       ethQuery,
@@ -1205,12 +1224,13 @@ export class TransactionController extends BaseControllerV1<
    * @param transaction - The transaction params to estimate gas for.
    * @param multiplier - The multiplier to use for the gas buffer.
    */
-  async estimateGasBuffered( // NOTE(JL): Need to update SwapsController's usage of this method
+  async estimateGasBuffered(
+    // NOTE(JL): Need to update SwapsController's usage of this method
     transaction: TransactionParams,
     multiplier: number,
-    networkClientId?: NetworkClientId
+    networkClientId?: NetworkClientId,
   ) {
-    const ethQuery = this.getEthQuery(networkClientId)
+    const ethQuery = this.getEthQuery(networkClientId);
     const { blockGasLimit, estimatedGas, simulationFails } = await estimateGas(
       transaction,
       ethQuery,
@@ -1566,7 +1586,8 @@ export class TransactionController extends BaseControllerV1<
    * @param address - The hex string address for the transaction.
    * @returns object with the `nextNonce` `nonceDetails`, and the releaseLock.
    */
-  async getNonceLock(address: string): Promise<NonceLock> { // NOTE(JL): i think this should take in chainId, but not sure how to deal with networkClientId mapping
+  async getNonceLock(address: string): Promise<NonceLock> {
+    // NOTE(JL): i think this should take in chainId, but not sure how to deal with networkClientId mapping
     return this.nonceTracker.getNonceLock(address);
   }
 
@@ -1985,7 +2006,8 @@ export class TransactionController extends BaseControllerV1<
   /**
    * Create approvals for all unapproved transactions on current chain.
    */
-  private createApprovalsForUnapprovedTransactions() { // NOTE(JL): this doesn't seem to be used anywhere. Can we remove it?
+  private createApprovalsForUnapprovedTransactions() {
+    // NOTE(JL): this doesn't seem to be used anywhere. Can we remove it?
     const unapprovedTransactions = this.getCurrentChainTransactionsByStatus(
       TransactionStatus.unapproved,
     );
@@ -2220,7 +2242,7 @@ export class TransactionController extends BaseControllerV1<
         return;
       }
 
-      const ethQuery = this.getEthQuery(transactionMeta.networkClientId)
+      const ethQuery = this.getEthQuery(transactionMeta.networkClientId);
 
       if (transactionMeta.type === TransactionType.swap) {
         log('Determining pre-transaction balance');
@@ -2266,7 +2288,10 @@ export class TransactionController extends BaseControllerV1<
     }
   }
 
-  private async publishTransaction(ethQuery: EthQuery, rawTransaction: string): Promise<string> {
+  private async publishTransaction(
+    ethQuery: EthQuery,
+    rawTransaction: string,
+  ): Promise<string> {
     return await query(ethQuery, 'sendRawTransaction', [rawTransaction]);
   }
 
@@ -2547,7 +2572,7 @@ export class TransactionController extends BaseControllerV1<
    * @param transactionMeta - Nominated external transaction to be added to state.
    */
   private addExternalTransaction(transactionMeta: TransactionMeta) {
-    const { chainId } = transactionMeta
+    const { chainId } = transactionMeta;
     const { transactions } = this.state;
     const fromAddress = transactionMeta?.txParams?.from;
     const sameFromAndNetworkTransactions = transactions.filter(
@@ -2592,7 +2617,7 @@ export class TransactionController extends BaseControllerV1<
     // NOTE(JL): Should this method be exiting early if getTransaction returns no transaction object?
     const nonce = transactionMeta?.txParams?.nonce;
     const from = transactionMeta?.txParams?.from;
-    const chainId = transactionMeta?.chainId
+    const chainId = transactionMeta?.chainId;
     const sameNonceTxs = this.state.transactions.filter(
       (transaction) =>
         transaction.txParams.from === from &&
@@ -2693,6 +2718,15 @@ export class TransactionController extends BaseControllerV1<
     );
   }
 
+  private removePendingTransactionTrackerListeners(
+    pendingTransactionTracker = this.pendingTransactionTracker,
+  ) {
+    pendingTransactionTracker.hub.removeAllListeners('transaction-confirmed');
+    pendingTransactionTracker.hub.removeAllListeners('transaction-dropped');
+    pendingTransactionTracker.hub.removeAllListeners('transaction-failed');
+    pendingTransactionTracker.hub.removeAllListeners('transaction-updated');
+  }
+
   private addPendingTransactionTrackerListeners(
     pendingTransactionTracker = this.pendingTransactionTracker,
   ) {
@@ -2829,7 +2863,7 @@ export class TransactionController extends BaseControllerV1<
         return;
       }
 
-      const ethQuery = this.getEthQuery(transactionMeta.networkClientId)
+      const ethQuery = this.getEthQuery(transactionMeta.networkClientId);
       const { updatedTransactionMeta, approvalTransactionMeta } =
         await updatePostTransactionBalance(transactionMeta, {
           ethQuery,