Skip to content

Prevent OOM when close account ix included for swap#2148

Open
tmns wants to merge 1 commit intomasterfrom
methias/prevent-swap-oom
Open

Prevent OOM when close account ix included for swap#2148
tmns wants to merge 1 commit intomasterfrom
methias/prevent-swap-oom

Conversation

@tmns
Copy link
Collaborator

@tmns tmns commented Mar 17, 2026

We recently started allowing for swap txns to include close account ixs. Turns out including the ix causes an out of memory error:

[
    "Program log: Instruction: BeginSwap",
    ...
    "Program log: Instruction: Transfer",
    "Program ... success",
    "Program log: Error: memory allocation failed, out of memory",
    "Program ... failed: SBF program panicked"
]

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_checked was 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

    • Enabled a full spot-swap test that explicitly creates the user’s associated token account (ATA), executes the swap, and verifies the ATA is closed.
  • Chores

    • Improved instruction parsing and validation to reduce memory use and enforce stricter post-swap checks.
  • Documentation

    • Added an Unreleased changelog entry noting a fix to prevent out-of-memory during certain swap+close flows.

@tmns tmns self-assigned this Mar 17, 2026
@coderabbitai
Copy link

coderabbitai bot commented Mar 17, 2026

Walkthrough

Refactors 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

Cohort / File(s) Summary
Changelog
CHANGELOG.md
Adds a new "Unreleased" Fix entry: "program: prevent OOM when close account ix included for swap" (text-only change).
Instruction Sysvar Parsing Refactor
programs/drift/src/instructions/user.rs
Adds InstructionSysvarView<'a> plus zero-copy readers (read_u16_le, read_pubkey, read_slice, load_instruction_sysvar_view_at, etc.). Replaces full Instruction deserialization with slice-based access in swap flows, updates is_token_close_account_for_swap_ix signature, tightens account/meta validations and error paths.
Spot Swap Test (ATA handling)
tests/spotSwap.ts
Adds imports from @solana/spl-token, computes taker USDC ATA, inserts createAssociatedTokenAccountIdempotentInstruction before swap instructions, and enables the test path that opens an ATA, performs the swap, then closes the ATA.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

Poem

🐰 I nibbled at bytes with nimble paws,

Zero-copy tunnels through sysvar laws.
An ATA opened, a swap set free,
Close hops tidy as a clean marquee.
Hooray — a rabbit's compact glee.

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title directly and clearly captures the main objective of the PR: preventing an out-of-memory issue that occurs when close account instructions are included in swap transactions, matching the core change of replacing instruction deserialization with zero-copy parsing.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

📝 Coding Plan
  • Generate coding plan for human review comments

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tests/spotSwap.ts (1)

755-821: Avoid shadowing takerUSDC inside the test.

The local const takerUSDC hides the suite-level takerUSDC, 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

📥 Commits

Reviewing files that changed from the base of the PR and between 8be7a58 and 49aa508.

📒 Files selected for processing (3)
  • CHANGELOG.md
  • programs/drift/src/instructions/user.rs
  • tests/spotSwap.ts

@tmns tmns force-pushed the methias/prevent-swap-oom branch from 49aa508 to 9b6ff44 Compare March 17, 2026 07:45
@tmns tmns force-pushed the methias/prevent-swap-oom branch from 9b6ff44 to 6d2a378 Compare March 17, 2026 08:02
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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

📥 Commits

Reviewing files that changed from the base of the PR and between 9b6ff44 and 6d2a378.

📒 Files selected for processing (3)
  • CHANGELOG.md
  • programs/drift/src/instructions/user.rs
  • tests/spotSwap.ts
🚧 Files skipped from review as they are similar to previous changes (1)
  • CHANGELOG.md

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant