Skip to content

Commit

Permalink
feat: improve wallet manager selection criteria
Browse files Browse the repository at this point in the history 
Signed-off-by: Pablo Maldonado <[email protected]>
  • Loading branch information
@md0x
md0x committed Jul 19, 2024
1 parent 28ba4fc commit 8773050
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 43 deletions.
2 changes: 1 addition & 1 deletion src/lib/bundleUtils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ export const prepareUnlockTransaction = async (
simulate = true,
) => {
const provider = getProvider();
const unlockerWallet = WalletManager.getInstance(provider).getWallet(ovalAddress, targetBlock);
const unlockerWallet = WalletManager.getInstance().getWallet(ovalAddress, targetBlock);
const [baseFee, network] = await Promise.all([getBaseFee(provider, req), provider.getNetwork()]);
const data = ovalInterface.encodeFunctionData("unlockLatestValue");
const { unlockTxHash, signedUnlockTx } = await createUnlockLatestValueTx(
Expand Down
8 changes: 6 additions & 2 deletions src/lib/helpers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {
import { Request } from "express";
import { FlashbotsBundleProvider } from "flashbots-ethers-v6-provider-bundle";
import { JSONRPCRequest } from "json-rpc-2.0";
import { OvalDiscovery } from "./";
import { OvalDiscovery, WalletManager } from "./";
import { flashbotsSupportedNetworks, supportedNetworks } from "./constants";
import { env } from "./env";
import { Logger } from "./logging";
Expand Down Expand Up @@ -281,7 +281,11 @@ export function getOvalConfigsShared(input: string): OvalConfigsShared {

export const getOvalAddresses = (): string[] => {
const factoryInstances = OvalDiscovery.getInstance().getOvalFactoryInstances();
return [...factoryInstances, ...Object.keys(env.ovalConfigs)].map(getAddress);
return [...Object.keys(env.ovalConfigs), ...factoryInstances].map(getAddress);
};

export const isOvalSharedUnlockerKey = (unlockerKey: string): boolean => {
return WalletManager.getInstance().isOvalSharedUnlocker(unlockerKey);
};

export const getOvalRefundConfig = (ovalAddress: string): RefundConfig => {
Expand Down
78 changes: 50 additions & 28 deletions src/lib/walletManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,9 @@ type WalletConfig = {
};

type WalletUsed = {
walletPubKey: string;
count: number;
walletPubKey: string;
ovalInstances: Set<string>;
};

// WalletManager class to handle wallet operations.
Expand All @@ -22,32 +23,35 @@ export class WalletManager {
private wallets: Record<string, Wallet> = {};
private sharedWallets: Map<string, Wallet> = new Map();
private sharedWalletUsage: Map<string, Map<number, WalletUsed>> = new Map();
private provider: JsonRpcProvider;
private provider: JsonRpcProvider | undefined;

private constructor(provider: JsonRpcProvider) {
this.provider = provider;
private constructor() {
this.ovalDiscovery = OvalDiscovery.getInstance();
this.setupCleanupInterval();
}

// Singleton pattern to get an instance of WalletManager
public static getInstance(provider: JsonRpcProvider): WalletManager {
public static getInstance(): WalletManager {
if (!WalletManager.instance) {
WalletManager.instance = new WalletManager(provider);
WalletManager.instance = new WalletManager();
}
return WalletManager.instance;
}

// Initialize wallets with configurations
public async initialize(ovalConfigs: OvalConfigs, sharedConfigs?: OvalConfigsShared): Promise<void> {
public async initialize(provider: JsonRpcProvider, ovalConfigs: OvalConfigs, sharedConfigs?: OvalConfigsShared): Promise<void> {
this.provider = provider;
await this.initializeWallets(ovalConfigs);
if (sharedConfigs) {
await this.initializeSharedWallets(sharedConfigs);
}
this.setupCleanupInterval();
}

// Get a wallet for a given address
public getWallet(address: string, targetBlock: number): Wallet {
if (!this.provider) {
throw new Error("Provider is not initialized");
}
const checkSummedAddress = getAddress(address);
const wallet = this.wallets[checkSummedAddress];
if (!wallet) {
Expand All @@ -59,15 +63,20 @@ export class WalletManager {

// Get a shared wallet for a given Oval instance and target block
private getSharedWallet(ovalInstance: string, targetBlock: number): 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`);
}

// Check if a wallet has already been assigned to this Oval instance
const instanceUsage = this.sharedWalletUsage.get(ovalInstance);
if (instanceUsage) {
const [_, usage] = instanceUsage.entries().next().value;
return this.sharedWallets.get(usage.walletPubKey)!.connect(this.provider);
for (const [walletPubKey, instanceUsage] of this.sharedWalletUsage.entries()) {
for (const [block, record] of instanceUsage.entries()) {
if (record.ovalInstances && record.ovalInstances.has(ovalInstance)) {
return this.sharedWallets.get(record.walletPubKey)!.connect(this.provider!);
}
}
}

// If no wallet has been assigned, find the least used wallet
Expand All @@ -80,10 +89,15 @@ export class WalletManager {
throw new Error(`No available shared wallets for Oval instance ${ovalInstance} at block ${targetBlock}`);
}

public isOvalSharedUnlocker(unlockerPublicKey: string): boolean {
return this.sharedWallets.has(unlockerPublicKey);
}

// Private helper methods
private setupCleanupInterval(): void {
if (isMochaTest()) return;
if (isMochaTest() || !this.provider) return;
setInterval(async () => {
if (!this.provider) return;
const currentBlock = await this.provider.getBlockNumber();
this.cleanupOldRecords(currentBlock);
}, env.sharedWalletUsageCleanupInterval * 1000);
Expand All @@ -101,7 +115,6 @@ export class WalletManager {
if (wallet) {
const walletPubKey = await wallet.getAddress();
this.sharedWallets.set(walletPubKey, wallet);
this.sharedWalletUsage.set(walletPubKey, new Map());
}
}
}
Expand All @@ -123,26 +136,34 @@ export class WalletManager {

private findLeastUsedWallet(): Wallet | undefined {
let selectedWallet: Wallet | undefined;
const usageCount = new Map<string, number>();

// Initialize usage counts for each wallet
this.sharedWallets.forEach((_, walletPubKey) => {
usageCount.set(walletPubKey, 0);
const totalUsage = new Map<string, {
totalCount: number;
ovalInstances: Set<string>;
}>();

// Initialize total usage counts for each wallet
this.sharedWallets.forEach((wallet) => {
totalUsage.set(wallet.address, { totalCount: 0, ovalInstances: new Set() });
});

// Sum usage counts for each wallet
this.sharedWalletUsage.forEach((instanceUsage) => {
instanceUsage.forEach((record) => {
const count = usageCount.get(record.walletPubKey) || 0;
usageCount.set(record.walletPubKey, count + record.count);
const totalInstanceUsage = totalUsage.get(record.walletPubKey)!;
totalInstanceUsage.totalCount += record.count;
record.ovalInstances.forEach(instance => totalInstanceUsage.ovalInstances.add(instance));
totalUsage.set(record.walletPubKey, totalInstanceUsage);
});
});

// Find the wallet with the least usage
let minInstances = Infinity;
let minUsage = Infinity;
usageCount.forEach((count, walletPubKey) => {
if (count < minUsage) {
minUsage = count;
totalUsage.forEach((usage, walletPubKey) => {
const instanceCount = usage.ovalInstances.size;
if (instanceCount < minInstances || (instanceCount === minInstances && usage.totalCount < minUsage)) {
minInstances = instanceCount;
minUsage = usage.totalCount;
selectedWallet = this.sharedWallets.get(walletPubKey);
}
});
Expand All @@ -152,16 +173,17 @@ export class WalletManager {

private async updateWalletUsage(ovalInstance: string, wallet: Wallet, targetBlock: number): Promise<void> {
const walletPubKey = await wallet.getAddress();
const instanceUsage = this.sharedWalletUsage.get(ovalInstance) || new Map();
const instanceUsage = this.sharedWalletUsage.get(walletPubKey) || new Map();
const existingRecord = instanceUsage.get(targetBlock);

if (existingRecord) {
existingRecord.count += 1;
(existingRecord as WalletUsed).count += 1;
(existingRecord as WalletUsed).ovalInstances.add(ovalInstance);
} else {
instanceUsage.set(targetBlock, { walletPubKey, count: 1 });
instanceUsage.set(targetBlock, { walletPubKey, count: 1, ovalInstances: new Set([ovalInstance]) });
}

this.sharedWalletUsage.set(ovalInstance, instanceUsage);
this.sharedWalletUsage.set(walletPubKey, instanceUsage);
}

private cleanupOldRecords(currentBlock: number): void {
Expand Down
28 changes: 16 additions & 12 deletions test/walletManager.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,15 +26,17 @@ describe('WalletManager Tests', () => {
findOval: sinon.stub().resolves()
};
sinon.stub(ovalDiscovery.OvalDiscovery, 'getInstance').returns(ovalDiscoveryInstance as any);
// Cleanup old records
WalletManager.getInstance()['cleanupOldRecords'](Infinity);
});

afterEach(() => {
sinon.restore();
});

it('should return a singleton instance', () => {
const instance1 = WalletManager.getInstance(mockProvider);
const instance2 = WalletManager.getInstance(mockProvider);
const instance1 = WalletManager.getInstance();
const instance2 = WalletManager.getInstance();
expect(instance1).to.equal(instance2);
});

Expand All @@ -49,8 +51,8 @@ describe('WalletManager Tests', () => {
[oval2]: { gckmsKeyId: 'gckmsKeyId456', refundAddress: refundRandom, refundPercent: 20 },
};
sinon.stub(gckms, 'retrieveGckmsKey').resolves(gckmsRandom.privateKey);
const walletManager = WalletManager.getInstance(mockProvider);
await walletManager.initialize(ovalConfigs);
const walletManager = WalletManager.getInstance();
await walletManager.initialize(mockProvider, ovalConfigs);

const walletRandom = walletManager.getWallet(oval1, 123);
expect(walletRandom?.privateKey).to.equal(unlockerRandom.privateKey);
Expand All @@ -67,8 +69,8 @@ describe('WalletManager Tests', () => {
{ gckmsKeyId: 'gckmsKeyId456' },
];
sinon.stub(gckms, 'retrieveGckmsKey').resolves(gckmsRandom.privateKey);
const walletManager = WalletManager.getInstance(mockProvider);
await walletManager.initialize({}, sharedConfigs);
const walletManager = WalletManager.getInstance();
await walletManager.initialize(mockProvider, {}, sharedConfigs);

// Check if shared wallets are initialized
const sharedWallets = Array.from(walletManager['sharedWallets'].values());
Expand All @@ -80,12 +82,14 @@ describe('WalletManager Tests', () => {
const sharedConfigs: OvalConfigsShared = [
{ unlockerKey: unlockerRandom.privateKey },
];
const walletManager = WalletManager.getInstance(mockProvider);
await walletManager.initialize({}, sharedConfigs);
const walletManager = WalletManager.getInstance();

await walletManager.initialize(mockProvider, {}, sharedConfigs);

const ovalInstance = 'ovalInstance1';
const targetBlock = 123;


const wallet1 = await walletManager['getSharedWallet'](ovalInstance, targetBlock);
const wallet2 = await walletManager['getSharedWallet'](ovalInstance, targetBlock + 1);

Expand All @@ -99,8 +103,8 @@ describe('WalletManager Tests', () => {
{ unlockerKey: unlockerRandom1.privateKey },
{ unlockerKey: unlockerRandom2.privateKey },
];
const walletManager = WalletManager.getInstance(mockProvider);
await walletManager.initialize({}, sharedConfigs);
const walletManager = WalletManager.getInstance();
await walletManager.initialize(mockProvider, {}, sharedConfigs);

const ovalInstance1 = 'ovalInstance1';
const ovalInstance2 = 'ovalInstance2';
Expand All @@ -118,8 +122,8 @@ describe('WalletManager Tests', () => {
const sharedConfigs: OvalConfigsShared = [
{ unlockerKey: unlockerRandom.privateKey },
];
const walletManager = WalletManager.getInstance(mockProvider);
await walletManager.initialize({}, sharedConfigs);
const walletManager = WalletManager.getInstance();
await walletManager.initialize(mockProvider, {}, sharedConfigs);

const ovalInstance = 'ovalInstance1';
const targetBlock = 123;
Expand Down

0 comments on commit 8773050

Please sign in to comment.