feat(walletManager): add shouldRecordUsage option and tests

src/index.ts

@@ -235,7 +235,8 @@ app.post("/", async (req, res, next) => {
       } else {
         Logger.debug(req.transactionId, "Finding unlock that does not revert the bundle...");
 
-        const unlock = await findUnlock(flashbotsBundleProvider, backrunTxs, targetBlock, req);
+        // FindUnlock with recordWalletUsage set to false to prevent recording wallet usage in WalletManager.
+        const unlock = await findUnlock(flashbotsBundleProvider, backrunTxs, targetBlock, req, false);
         if (!unlock) {
           Logger.debug(req.transactionId, "No valid unlock found!");
           await handleUnsupportedRequest(req, res, "No valid unlock found"); // Pass through if no unlock is found.

src/lib/bundleUtils.ts

@@ -46,17 +46,19 @@ export const createUnlockLatestValueTx = async (
 };
 
 // Prepare unlockLatestValue transaction for a given Oval instance and simulate the bundle with the unlock transaction prepended if simulate is true.
+// If shouldRecordWalletUsage is true, the wallet usage will be recorded in the WalletManager.
 export const prepareUnlockTransaction = async (
   flashbotsBundleProvider: FlashbotsBundleProvider,
   backrunTxs: string[],
   targetBlock: number,
   ovalAddress: string,
   req: express.Request,
   simulate = true,
+  shouldRecordWalletUsage = true,
 ) => {
   const provider = getProvider();
   const [baseFee, network] = await Promise.all([getBaseFee(provider, req), provider.getNetwork()]);
-  const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId);
+  const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock, req.transactionId, shouldRecordWalletUsage);
   const isSharedWallet = isOvalSharedUnlockerKey(unlockerWallet.address);
 
   // Encode the unlockLatestValue function call depending on whether the unlocker is a shared wallet or not.
@@ -75,11 +77,11 @@ export const prepareUnlockTransaction = async (
     target,
   );
 
-  if (!simulate) return { ovalAddress, unlockTxHash, signedUnlockTx };
+  if (!simulate) return { ovalAddress, unlockTxHash, signedUnlockTx, unlockerWallet };
 
   const simulationResponse = await flashbotsBundleProvider.simulate([signedUnlockTx, ...backrunTxs], targetBlock);
 
-  return { ovalAddress, unlockTxHash, signedUnlockTx, simulationResponse };
+  return { ovalAddress, unlockTxHash, signedUnlockTx, simulationResponse, unlockerWallet };
 };
 
 export const getUnlockBundlesFromOvalAddresses = async (
@@ -99,7 +101,7 @@ export const getUnlockBundlesFromOvalAddresses = async (
       targetBlock,
       ovalAddress,
       req,
-      false,
+      false, // Do not simulate
     );
 
     // Construct the inner bundle with call to Oval to unlock the latest value.
@@ -117,16 +119,19 @@ export const getUnlockBundlesFromOvalAddresses = async (
 };
 
 // Simulate calls to unlockLatestValue on each Oval instance bundled with original backrun transactions. Tries to find
-// the first unlock that doesn't revert the bundle.
+// the first unlock that doesn't revert the bundle. If updateWalletUsage is true, the wallet usage will be recorded in the WalletManager.
 export const findUnlock = async (
   flashbotsBundleProvider: FlashbotsBundleProvider,
   backrunTxs: string[],
   targetBlock: number,
   req: express.Request,
+  updateWalletUsage = true
 ) => {
   const unlocks = await Promise.all(
     getOvalAddresses().map(async (ovalAddress) =>
-      prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req),
+      // Do not record wallet usage here as we will record it in the loop below once we find a valid unlock.
+      // See shouldRecordWalletUsage parameter in WalletManager.getWallet.
+      prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req, true, false),
     ),
   );
 
@@ -137,6 +142,9 @@ export const findUnlock = async (
       !("error" in unlock.simulationResponse) &&
       !unlock.simulationResponse.firstRevert
     ) {
+      if (updateWalletUsage) {
+        WalletManager.getInstance().updateWalletUsage(unlock.ovalAddress, unlock.unlockerWallet, targetBlock);
+      }
       return {
         // Spread in order to preserve inferred SimulationResponseSuccess type.
         ...unlock,

src/lib/walletManager.ts

@@ -49,50 +49,57 @@ export class WalletManager {
     }
 
     // Get a wallet for a given address
-    public getWallet(address: string, targetBlock: number, transactionId: string): Wallet {
+    // If shouldRecordUsage is true, the wallet usage will be recorded in the WalletManager.
+    public getWallet(address: string, targetBlock: number, transactionId: string, shouldRecordUsage = true): Wallet {
         if (!this.provider) {
             throw new Error("Provider is not initialized");
         }
         const checkSummedAddress = getAddress(address);
         const wallet = this.wallets[checkSummedAddress];
         if (!wallet) {
-            return this.getSharedWallet(address, targetBlock, transactionId);
+            return this.getSharedWallet(address, targetBlock, transactionId, shouldRecordUsage);
         }
         return wallet.connect(this.provider);
     }
 
     // Get a shared wallet for a given Oval instance and target block
-    private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string): Wallet {
+    private getSharedWallet(ovalInstance: string, targetBlock: number, transactionId: string, shouldRecordUsage = true): Wallet {
         if (!this.provider) {
             throw new Error("Provider is not initialized");
         }
+
         if (!this.ovalDiscovery.isOval(ovalInstance)) {
             throw new Error(`Oval instance ${ovalInstance} is not found`);
         }
 
-        let selectedWallet: Wallet | undefined;
-
-        // Check if a wallet has already been assigned to this Oval instance
-        for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) {
-            for (const [_, record] of instanceUsage.entries()) {
-                if (record.ovalInstances && record.ovalInstances.has(ovalInstance)) {
-                    selectedWallet = this.sharedWallets.get(walletPubKey)!.connect(this.provider!);
-                }
-            }
-        }
+        let selectedWallet = this.findAssignedWallet(ovalInstance);
 
         // If no wallet has been assigned, find the least used wallet
         if (!selectedWallet) {
             selectedWallet = this.findLeastUsedWallet(transactionId);
         }
 
-        // Update the usage of the selected wallet
-        if (selectedWallet) {
+        if (!selectedWallet) {
+            throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`);
+        }
+
+        // Update the usage of the selected wallet if recording is enabled
+        if (shouldRecordUsage) {
             this.updateWalletUsage(ovalInstance, selectedWallet, targetBlock);
-            return selectedWallet.connect(this.provider);
         }
 
-        throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`);
+        return selectedWallet.connect(this.provider);
+    }
+
+    private findAssignedWallet(ovalInstance: string): Wallet | undefined {
+        for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) {
+            for (const record of instanceUsage.values()) {
+                if (record.ovalInstances?.has(ovalInstance)) {
+                    return this.sharedWallets.get(walletPubKey)?.connect(this.provider!);
+                }
+            }
+        }
+        return undefined;
     }
 
     public isOvalSharedUnlocker(unlockerPublicKey: string): boolean {
@@ -190,12 +197,11 @@ export class WalletManager {
         if (minInstances !== Infinity && minInstances !== 0) {
             Logger.error(transactionId, `Public key ${selectedWallet?.address} is reused in multiple Oval instances because no free wallets are available.`);
         }
-
         return selectedWallet;
     }
 
     // Update the usage statistics for a wallet
-    private async updateWalletUsage(ovalInstance: string, wallet: Wallet, targetBlock: number): Promise<void> {
+    public async updateWalletUsage(ovalInstance: string, wallet: Wallet, targetBlock: number): Promise<void> {
         const walletPubKey = await wallet.getAddress();
         const instanceUsage = this.sharedWalletUsage.get(walletPubKey) || new Map();
         const existingRecord = instanceUsage.get(targetBlock);

test/walletManager.ts

@@ -141,6 +141,42 @@ describe('WalletManager Tests', () => {
         expect(errorSpy.args[0][1]).to.include(`Public key ${wallet2.address} is reused in multiple Oval instances because no free wallets are available.`);
     });
 
+    it('should not record usage when shouldRecordUsage is false', async () => {
+        const unlockerRandom = getRandomAddressAndKey();
+        const sharedConfigs: OvalConfigsShared = [
+            { unlockerKey: unlockerRandom.privateKey },
+        ];
+        const walletManager = WalletManager.getInstance();
+        await walletManager.initialize(mockProvider, {}, sharedConfigs);
+
+        const ovalInstance = getRandomAddressAndKey().address;
+        const targetBlock = 123;
+
+        const wallet = walletManager.getWallet(ovalInstance, targetBlock, "transactionId", false);
+
+        const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress());
+        expect(sharedWalletUsage).to.be.undefined;
+    });
+
+    it('should record usage when shouldRecordUsage is true', async () => {
+        const unlockerRandom = getRandomAddressAndKey();
+        const sharedConfigs: OvalConfigsShared = [
+            { unlockerKey: unlockerRandom.privateKey },
+        ];
+        const walletManager = WalletManager.getInstance();
+        await walletManager.initialize(mockProvider, {}, sharedConfigs);
+
+        const ovalInstance = getRandomAddressAndKey().address;
+        const targetBlock = 123;
+
+        const wallet = await walletManager.getWallet(ovalInstance, targetBlock, "transactionId");
+
+        const sharedWalletUsage = walletManager['sharedWalletUsage'].get(await wallet.getAddress());
+        expect(sharedWalletUsage).to.not.be.undefined;
+        expect(sharedWalletUsage?.get(targetBlock)?.count).to.equal(1);
+        expect(sharedWalletUsage?.get(targetBlock)?.ovalInstances.has(ovalInstance)).to.be.true;
+    });
+
     it('should cleanup old records correctly', async () => {
         const unlockerRandom = getRandomAddressAndKey();
         const sharedConfigs: OvalConfigsShared = [