Prevent OOM when close account ix included for swap#2148
Conversation
WalkthroughRefactors Solana instruction sysvar parsing to a zero-copy InstructionSysvarView in user.rs, updates swap handling and account validations accordingly, adds explicit ATA creation in the spot swap test, and adds a changelog entry. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Program
participant InstructionSysvar
participant TokenProgram
Client->>Program: begin_swap transaction (includes swap + possible close)
Program->>InstructionSysvar: load_instruction_sysvar_view_at(sysvar)
InstructionSysvar-->>Program: InstructionSysvarView (slice-based)
Program->>InstructionSysvar: inspect accounts_len / account_meta_bytes_at
alt token-close detected for swap
Program->>TokenProgram: validate close-account ix via view
TokenProgram-->>Program: ok / error
end
Program-->>Client: proceed or abort swap/end_swap
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 📝 Coding Plan
Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/spotSwap.ts (1)
755-821: Avoid shadowingtakerUSDCinside the test.The local
const takerUSDChides the suite-leveltakerUSDC, which makes future edits easy to misread. Renaming the local ATA variable improves clarity.Suggested patch
- const takerUSDC = getAssociatedTokenAddressSync( + const takerUsdcAta = getAssociatedTokenAddressSync( usdcMint.publicKey, takerDriftClient.wallet.publicKey ); const { beginSwapIx, endSwapIx } = await takerDriftClient.getSwapIx({ amountIn, inMarketIndex: 0, outMarketIndex: 1, - inTokenAccount: takerUSDC, + inTokenAccount: takerUsdcAta, outTokenAccount: takerWSOL, }); const openIx = createAssociatedTokenAccountIdempotentInstruction( takerDriftClient.wallet.publicKey, - takerUSDC, + takerUsdcAta, takerDriftClient.wallet.publicKey, usdcMint.publicKey, TOKEN_PROGRAM_ID ); const transferIn = createTransferInstruction( - takerUSDC, + takerUsdcAta, makerUSDC.publicKey, takerDriftClient.wallet.publicKey, amountIn.toNumber() ); const closeIx = createCloseAccountInstruction( - takerUSDC, + takerUsdcAta, takerDriftClient.wallet.publicKey, takerDriftClient.wallet.publicKey, undefined, TOKEN_PROGRAM_ID ); const accountInfo = await bankrunContextWrapper.connection.getAccountInfo( - takerUSDC + takerUsdcAta );🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/spotSwap.ts` around lines 755 - 821, The local const takerUSDC declared in this test shadows the suite-level takerUSDC; rename the local ATA variable (e.g., takerUSDCAta or takerUSDCAccount) and update all its usages in this snippet — the getSwapIx call (inTokenAccount), createAssociatedTokenAccountIdempotentInstruction, createTransferInstruction (transferIn), createCloseAccountInstruction, the Transaction.add sequence, and the final getAccountInfo/assert check — so the suite-level wallet variable remains unshadowed and the test references the ATA by the new name consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@CHANGELOG.md`:
- Line 19: Remove the duplicated changelog bullet that repeats "program: skip
isolated positions when checking for cross margin bankruptcy
[`#2123`](https://github.com/drift-labs/protocol-v2/pull/2123)" under the
Unreleased > Fixes section in CHANGELOG.md by deleting the duplicate line so
only one instance of that bullet remains.
---
Nitpick comments:
In `@tests/spotSwap.ts`:
- Around line 755-821: The local const takerUSDC declared in this test shadows
the suite-level takerUSDC; rename the local ATA variable (e.g., takerUSDCAta or
takerUSDCAccount) and update all its usages in this snippet — the getSwapIx call
(inTokenAccount), createAssociatedTokenAccountIdempotentInstruction,
createTransferInstruction (transferIn), createCloseAccountInstruction, the
Transaction.add sequence, and the final getAccountInfo/assert check — so the
suite-level wallet variable remains unshadowed and the test references the ATA
by the new name consistently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: f77d2883-d176-405d-9e5a-3c87356085e3
📒 Files selected for processing (3)
CHANGELOG.mdprograms/drift/src/instructions/user.rstests/spotSwap.ts
49aa508 to
9b6ff44
Compare
9b6ff44 to
6d2a378
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
programs/drift/src/instructions/user.rs (1)
4205-4214: Avoid per-iteration whitelist allocations in the swap scan loop.At Line 4205,
WHITELISTED_SWAP_PROGRAMS.to_vec()still allocates each iteration. In this hot path, use membership checks without cloning to reduce heap churn further.♻️ Proposed allocation-free check
- let mut whitelisted_programs = WHITELISTED_SWAP_PROGRAMS.to_vec(); - if !delegate_is_signer { - whitelisted_programs.push(AssociatedToken::id()); - whitelisted_programs.push(Token::id()); - whitelisted_programs.push(Token2022::id()); - whitelisted_programs.push(marinade_mainnet::ID); - } - validate!( - whitelisted_programs.contains(&ix.program_id), + let is_whitelisted = WHITELISTED_SWAP_PROGRAMS.contains(&ix.program_id) + || (!delegate_is_signer + && (ix.program_id == AssociatedToken::id() + || ix.program_id == Token::id() + || ix.program_id == Token2022::id() + || ix.program_id == marinade_mainnet::ID)); + validate!( + is_whitelisted, ErrorCode::InvalidSwap, "only allowed to pass in ixs to ATA, openbook, Jupiter v3/v4/v6, dflow, or titan programs" )?;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@programs/drift/src/instructions/user.rs` around lines 4205 - 4214, WHITELISTED_SWAP_PROGRAMS.to_vec() allocates each loop; replace the per-iteration vector allocation with an allocation-free membership check: implement a small conditional check (or helper like is_whitelisted_swap_program(program_id, delegate_is_signer)) that first tests WHITELISTED_SWAP_PROGRAMS.contains(&ix.program_id) and, when !delegate_is_signer, also checks equality against AssociatedToken::id(), Token::id(), Token2022::id(), and marinade_mainnet::ID, then use that result in the existing validate!(..., ErrorCode::InvalidSwap) call so no temporary Vec is created each iteration.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@programs/drift/src/instructions/user.rs`:
- Around line 4205-4214: WHITELISTED_SWAP_PROGRAMS.to_vec() allocates each loop;
replace the per-iteration vector allocation with an allocation-free membership
check: implement a small conditional check (or helper like
is_whitelisted_swap_program(program_id, delegate_is_signer)) that first tests
WHITELISTED_SWAP_PROGRAMS.contains(&ix.program_id) and, when
!delegate_is_signer, also checks equality against AssociatedToken::id(),
Token::id(), Token2022::id(), and marinade_mainnet::ID, then use that result in
the existing validate!(..., ErrorCode::InvalidSwap) call so no temporary Vec is
created each iteration.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ad990ed4-83f9-44eb-8cfa-47ec1a247ac4
📒 Files selected for processing (3)
CHANGELOG.mdprograms/drift/src/instructions/user.rstests/spotSwap.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- CHANGELOG.md
We recently started allowing for swap txns to include close account ixs. Turns out including the ix causes an out of memory error:
First, per-iteration
WHITELISTED_SWAP_PROGRAMS.to_vec()allocations were replaced with static membership checks to try to reduce memory alloc but it had no effect.So instead
load_instruction_at_checkedwas replaced with a custom parser that avoids the former's tendency toward memory growth from repeated deserializations. Not sure if it's enough to work in practice but it's got the test passing.Summary by CodeRabbit
Tests
Chores
Documentation