feat(walletManager): add shouldRecordUsage option and tests

Signed-off-by: Pablo Maldonado <pablo@umaproject.org>

src/index.ts

@@ -151,6 +151,7 @@ app.post("/", async (req, res, next) => {
           await handleUnsupportedRequest(req, res, "No valid unlock found"); // Pass through if no unlock is found.
           return;
         }
+
         // Dynamically adjust refund percent so that builder nets at least configured minimum. We don't need to consider
         // refund gas costs as the builder is deducting them from refund and should not include the bundle if refund gas
         // costs exceed refund value.

src/lib/bundleUtils.ts

@@ -53,10 +53,11 @@ export const prepareUnlockTransaction = async (
   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 +76,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 +100,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.
@@ -123,10 +124,11 @@ export const findUnlock = async (
   backrunTxs: string[],
   targetBlock: number,
   req: express.Request,
+  updateWalletUsage = false
 ) => {
   const unlocks = await Promise.all(
     getOvalAddresses().map(async (ovalAddress) =>
-      prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req),
+      prepareUnlockTransaction(flashbotsBundleProvider, backrunTxs, targetBlock, ovalAddress, req, true, false),
     ),
   );
 
@@ -137,6 +139,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,56 @@ export class WalletManager {
     }
 
     // Get a wallet for a given address
-    public getWallet(address: string, targetBlock: number, transactionId: string): Wallet {
+    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 +196,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 = [