amm cache init refactor (#1998)

Author: moosecat2

programs/drift/src/instructions/admin.rs

@@ -1135,15 +1135,37 @@ pub fn handle_initialize_perp_market(
 
 pub fn handle_initialize_amm_cache(ctx: Context<InitializeAmmCache>) -> Result<()> {
     let amm_cache = &mut ctx.accounts.amm_cache;
-    let state = &ctx.accounts.state;
-    amm_cache
-        .cache
-        .resize_with(state.number_of_markets as usize, CacheInfo::default);
     amm_cache.bump = ctx.bumps.amm_cache;
 
     Ok(())
 }
 
+pub fn handle_resize_amm_cache(ctx: Context<ResizeAmmCache>) -> Result<()> {
+    let amm_cache = &mut ctx.accounts.amm_cache;
+    let state = &ctx.accounts.state;
+    let current_size = amm_cache.cache.len();
+    let new_size = (state.number_of_markets as usize).min(current_size + 20_usize);
+
+    msg!(
+        "resizing amm cache from {} entries to {}",
+        current_size,
+        new_size
+    );
+
+    let growth = new_size.saturating_sub(current_size);
+    validate!(
+        growth <= 20,
+        ErrorCode::DefaultError,
+        "cannot grow amm_cache by more than 20 entries in a single resize (requested +{})",
+        growth
+    )?;
+
+    amm_cache.cache.resize_with(new_size, CacheInfo::default);
+    amm_cache.validate(state)?;
+
+    Ok(())
+}
+
 #[access_control(
     perp_market_valid(&ctx.accounts.perp_market)
 )]
@@ -5510,7 +5532,7 @@ pub struct InitializeAmmCache<'info> {
     #[account(
         init,
         seeds = [AMM_POSITIONS_CACHE.as_ref()],
-        space = AmmCache::space(state.number_of_markets as usize),
+        space = AmmCache::init_space(),
         bump,
         payer = admin
     )]
@@ -5519,6 +5541,27 @@ pub struct InitializeAmmCache<'info> {
     pub system_program: Program<'info, System>,
 }
 
+#[derive(Accounts)]
+pub struct ResizeAmmCache<'info> {
+    #[account(
+        mut,
+        constraint = admin.key() == admin_hot_wallet::id() || admin.key() == state.admin
+    )]
+    pub admin: Signer<'info>,
+    pub state: Box<Account<'info, State>>,
+    #[account(
+        mut,
+        seeds = [AMM_POSITIONS_CACHE.as_ref()],
+        bump,
+        realloc = AmmCache::space(amm_cache.cache.len() + (state.number_of_markets as usize - amm_cache.cache.len()).min(20_usize)),
+        realloc::payer = admin,
+        realloc::zero = false,
+    )]
+    pub amm_cache: Box<Account<'info, AmmCache>>,
+    pub rent: Sysvar<'info, Rent>,
+    pub system_program: Program<'info, System>,
+}
+
 #[derive(Accounts)]
 pub struct DeleteInitializedPerpMarket<'info> {
     #[account(mut)]

programs/drift/src/state/amm_cache.rs

@@ -175,15 +175,19 @@ impl HasLen for AmmCacheFixed {
 }
 
 impl AmmCache {
+    pub fn init_space() -> usize {
+        8 + 8 + 4
+    }
+
     pub fn space(num_markets: usize) -> usize {
         8 + 8 + 4 + num_markets * CacheInfo::SIZE
     }
 
     pub fn validate(&self, state: &State) -> DriftResult<()> {
         validate!(
-            self.cache.len() == state.number_of_markets as usize,
+            self.cache.len() <= state.number_of_markets as usize,
             ErrorCode::DefaultError,
-            "Number of amm positions is different than number of markets"
+            "Number of amm positions is no larger than number of markets"
         )?;
         Ok(())
     }

sdk/src/adminClient.ts

@@ -507,12 +507,6 @@ export class AdminClient extends DriftClient {
 	): Promise<TransactionSignature> {
 		const currentPerpMarketIndex = this.getStateAccount().numberOfMarkets;
 
-		const ammCachePublicKey = getAmmCachePublicKey(this.program.programId);
-		const ammCacheAccount = await this.connection.getAccountInfo(
-			ammCachePublicKey
-		);
-		const mustInitializeAmmCache = ammCacheAccount?.data == null;
-
 		const initializeMarketIxs = await this.getInitializePerpMarketIx(
 			marketIndex,
 			priceOracle,
@@ -540,7 +534,6 @@ export class AdminClient extends DriftClient {
 			curveUpdateIntensity,
 			ammJitIntensity,
 			name,
-			mustInitializeAmmCache,
 			lpPoolId
 		);
 		const tx = await this.buildTransaction(initializeMarketIxs);
@@ -588,7 +581,6 @@ export class AdminClient extends DriftClient {
 		curveUpdateIntensity = 0,
 		ammJitIntensity = 0,
 		name = DEFAULT_MARKET_NAME,
-		includeInitAmmCacheIx = false,
 		lpPoolId: number = 0
 	): Promise<TransactionInstruction[]> {
 		const perpMarketPublicKey = await getPerpMarketPublicKey(
@@ -597,9 +589,6 @@ export class AdminClient extends DriftClient {
 		);
 
 		const ixs: TransactionInstruction[] = [];
-		if (includeInitAmmCacheIx) {
-			ixs.push(await this.getInitializeAmmCacheIx());
-		}
 
 		const nameBuffer = encodeName(name);
 		const initPerpIx = await this.program.instruction.initializePerpMarket(
@@ -663,9 +652,35 @@ export class AdminClient extends DriftClient {
 		return await this.program.instruction.initializeAmmCache({
 			accounts: {
 				state: await this.getStatePublicKey(),
-				admin: this.isSubscribed
-					? this.getStateAccount().admin
-					: this.wallet.publicKey,
+				admin: this.useHotWalletAdmin
+					? this.wallet.publicKey
+					: this.getStateAccount().admin,
+				ammCache: getAmmCachePublicKey(this.program.programId),
+				rent: SYSVAR_RENT_PUBKEY,
+				systemProgram: anchor.web3.SystemProgram.programId,
+			},
+		});
+	}
+
+	public async resizeAmmCache(
+		txParams?: TxParams
+	): Promise<TransactionSignature> {
+		const initializeAmmCacheIx = await this.getInitializeAmmCacheIx();
+
+		const tx = await this.buildTransaction(initializeAmmCacheIx, txParams);
+
+		const { txSig } = await this.sendTransaction(tx, [], this.opts);
+
+		return txSig;
+	}
+
+	public async getResizeAmmCacheIx(): Promise<TransactionInstruction> {
+		return await this.program.instruction.resizeAmmCache({
+			accounts: {
+				state: await this.getStatePublicKey(),
+				admin: this.useHotWalletAdmin
+					? this.wallet.publicKey
+					: this.getStateAccount().admin,
 				ammCache: getAmmCachePublicKey(this.program.programId),
 				rent: SYSVAR_RENT_PUBKEY,
 				systemProgram: anchor.web3.SystemProgram.programId,

tests/lpPool.ts

@@ -191,6 +191,8 @@ describe('LP Pool', () => {
 		solUsdLazer = getPythLazerOraclePublicKey(program.programId, 6);
 		await adminClient.initializePythLazerOracle(6);
 
+		await adminClient.initializeAmmCache();
+
 		await adminClient.initializePerpMarket(
 			0,
 			solUsd,

tests/lpPoolCUs.ts

@@ -171,6 +171,8 @@ describe('LP Pool', () => {
 
 		solUsd = await mockOracleNoProgram(bankrunContextWrapper, 200);
 
+		await adminClient.initializeAmmCache();
+
 		adminClient = new TestClient({
 			connection: bankrunContextWrapper.connection.toConnection(),
 			wallet: new anchor.Wallet(keypair),

tests/lpPoolSwap.ts

@@ -176,6 +176,8 @@ describe('LP Pool', () => {
 
 		const periodicity = new BN(0);
 
+		await adminClient.initializeAmmCache();
+
 		await adminClient.initializePerpMarket(
 			0,
 			spotMarketOracle,