enhancement: prevent unsub/sub race condition

[thlorenz]

Right now we do unsubscribe accounts that we only query but don't clone at the end of the task that needed that information.
One example are delegation records.

However there is a potential race condition such that if a delegation record is cloned + subscribed to in the meantime we would still unsubscribe from it in the end.

More details and a plan to fix this issue are in the comments.

Since the cases for which this matters are very rare this is a lower priority problem.

[thlorenz]

Race Condition Analysis: subscribe/unsubscribe Race in magicblock-chainlink

1. Executive Summary

Yes, the race condition is real and can lead to lost subscriptions.

The RemoteAccountProvider performs subscribe and unsubscribe as two non-atomic
steps: (1) a pubsub network operation and (2) an LRU cache update. These steps
can interleave with concurrent operations on the same pubkey, producing states
where a pubkey is subscribed in pubsub but absent from the LRU cache (or vice
versa). The periodic reconciler can then "correct" this transient inconsistency
by unsubscribing the account from pubsub—destroying a subscription that was
actively being established.

Four distinct race scenarios are identified below. Scenario D (reconciliation
racing with active registration) is the most severe and the hardest to guard
against with the current architecture.

---

2. Architecture Overview

Subscription State: Two Independent Sources of Truth

Component	Role	Synchronization
Pubsub clients (WebSocket/gRPC)	Actual network subscriptions to chain data	Async operations via SubMuxClient → AccountSubscriptionTask::process()
LRU cache (AccountsLruCache)	Tracks which accounts should be subscribed; source of truth for is_watching()	parking_lot::Mutex, locked per-operation (no cross-operation atomicity)

Subscribe Flow (RemoteAccountProvider::subscribe)

subscribe(pubkey)
  ├─ is_watching(pubkey)?  [LRU.contains — lock/unlock]
  │   ├─ YES → promote in LRU, return Ok
  │   └─ NO  → register_subscription(pubkey)
  │              ├─ Step 1: pubsub_client.subscribe(pubkey)   [ASYNC, ~2s timeout]
  │              │           ↑ GAP: pubkey in pubsub, NOT in LRU
  │              └─ Step 2: lrucache.add(pubkey)              [SYNC, fast]
  │                         └─ may evict LRU victim → unsubscribe victim
  └─ return Ok

File: remote_account_provider/mod.rs, lines 890–953

Unsubscribe Flow (RemoteAccountProvider::unsubscribe)

unsubscribe(pubkey)
  ├─ can_evict(pubkey)?     [check never-evict set]
  │   └─ NO → return Ok (no-op)
  ├─ LRU.contains(pubkey)?  [lock/unlock]
  │   └─ NO → return Ok (no-op, warns)
  ├─ Step 1: pubsub_client.unsubscribe(pubkey)  [ASYNC]
  │           + send removal notification
  │           ↑ GAP: pubkey removed from pubsub, still in LRU
  └─ Step 2: LRU.remove(pubkey)                 [SYNC, fast]

File: remote_account_provider/mod.rs, lines 972–997

Reconciler (reconcile_subscriptions)

Runs periodically (every ACTIVE_SUBSCRIPTIONS_UPDATE_INTERVAL_MS):

reconcile_subscriptions()
  ├─ Snapshot LRU pubkeys
  ├─ Snapshot pubsub subscriptions (union + intersection across clients)
  ├─ extra_in_lru  = LRU - pubsub_intersection  → resubscribe
  └─ extra_in_pubsub = pubsub_union - LRU       → UNSUBSCRIBE + notify removal

File: remote_account_provider/subscription_reconciler.rs, lines 72–160

Entry Points (Concurrent Callers)

Entry Point	Trigger	Can Subscribe?	Can Unsubscribe?
ensure_transaction_accounts	Every incoming transaction	✅ (via fetch_and_clone_accounts)	✅ (via cancel_subs)
process_subscription_update	Each pubsub update	❌	✅ (via unsubscribe_from_delegated_account)
undelegation_requested	Committor service	✅ (via subscribe_to_account)	❌
reconcile_subscriptions	Periodic timer	✅ (resubscribe)	✅ (stale cleanup)

---

3. Race Condition Scenarios

Scenario A: Concurrent Pipeline Cleanup Cancels Another Pipeline's Subscription

Actors:

• Thread A: fetch_and_clone_accounts_with_dedup for pubkeys {P, Q}
• Thread B: fetch_and_clone_accounts_with_dedup for pubkeys {P, R}

Interleaving:

Step	Thread A	Thread B	LRU State	Pubsub State
1	subscribe(P) → pubsub ✅, LRU ✅		{P}	{P}
2	subscribe(Q) → pubsub ✅, LRU ✅		{P, Q}	{P, Q}
3		subscribe(P) → is_watching(P) = true → promote, return	{P, Q}	{P, Q}
4		subscribe(R) → pubsub ✅, LRU ✅	{P, Q, R}	{P, Q, R}
5		Fetch fails for R → cancel_subs(CancelStrategy::New{new_subs={P,R}, existing_subs=...})
6		P was NOT in existing_subs for Thread B (checked before subscribe)
7		unsubscribe(P) → pubsub removed, LRU removed	{Q, R}	{Q, R}
8	Pipeline A completed, P subscribed by A but now gone		P LOST	P LOST

Key code locations:

• cancel_subs: fetch_cloner/subscription.rs:104–158
• CancelStrategy::New difference: subscription.rs:121 — new_subs.difference(&existing_subs)
• build_existing_subs captures the existing subs snapshot before subscribing: pipeline.rs:33–60

Guard analysis: is_pending check at subscription.rs:146 only checks
fetching_accounts (the dedup map for RPC fetches). If Thread A has already
completed its fetch and cleared pending_requests, is_pending(P) returns
false and the guard does not protect P.

Severity: MEDIUM — Requires two concurrent pipelines sharing a pubkey where
one encounters an error mid-pipeline. The subscription is lost but would be
re-established on the next transaction for P.

---

Scenario B: unsubscribe_from_delegated_account Racing with subscribe

Actors:

• Task A: process_subscription_update receives an update showing P is delegated → calls unsubscribe_from_delegated_account(P)
• Task B: undelegation_requested(P) → calls subscribe_to_account(P) (committor service schedules undelegation)

Interleaving:

Step	Task A (unsub)	Task B (sub)	LRU	Pubsub
1		subscribe(P): is_watching(P) → true (P was previously subscribed) → promote, return	{..P..}	{..P..}
2	unsubscribe(P): can_evict(P) → true		{..P..}	{..P..}
3	contains(P) → true		{..P..}	{..P..}
4	pubsub_client.unsubscribe(P) ✅		{..P..}	{..}
5	LRU.remove(P)		{..}	{..}
6		P is gone — undelegation tracking fails silently	P LOST	P LOST

This is the canonical case described in the original question: subscribe
completes, then unsubscribe races and wins.

The opposite order is also problematic:

Step	Task A (unsub)	Task B (sub)	LRU	Pubsub
1	contains(P) → true, start async pubsub_client.unsubscribe(P)		{..P..}	{..P..}
2		is_watching(P) → true → promote, return Ok	{..P..}	{..P..}
3	pubsub unsubscribe completes ✅		{..P..}	{..}
4	LRU.remove(P)		{..}	{..}
5		subscribe returned Ok, but P is no longer subscribed	P LOST	P LOST

In this interleaving, subscribe returns Ok(()) at step 2 because
is_watching sees P in the LRU. But the concurrent unsubscribe then removes P
from both pubsub and LRU. The caller of subscribe believes P is being
watched, but it isn't.

Key code locations:

• unsubscribe_from_delegated_account: fetch_cloner/mod.rs:400–410
• subscribe_to_account: fetch_cloner/mod.rs:1415–1427
• undelegation_requested: chainlink/mod.rs:486–504
• process_subscription_update: fetch_cloner/mod.rs:257–388, specifically line 294 and 355

Severity: HIGH — This directly breaks undelegation tracking. When a commit
with undelegation is in flight, the validator needs to watch the account so it
can clone the new state once undelegation completes. A lost subscription means
the account gets stuck in "undelegating" state in the bank. This scenario is
realistic because process_subscription_update and undelegation_requested are
invoked from separate async tasks that can easily overlap.

---

Scenario C: LRU Temporal Gap During register_subscription

Actors:

• Task A: register_subscription(P) — in the gap between pubsub subscribe and LRU add
• Task B: unsubscribe(P) — concurrent

Interleaving:

Step	Task A (register)	Task B (unsub)	LRU	Pubsub
1	pubsub_client.subscribe(P) ✅		{}	{P}
2	— gap: P in pubsub but NOT in LRU —	contains(P) → false	{}	{P}
3		Returns early: "not subscribed in LRU" (logs warning)	{}	{P}
4	LRU.add(P)		{P}	{P}
5	Result: P is subscribed (both pubsub and LRU)

In this specific interleaving, the race actually results in P remaining
subscribed — the opposite problem from what was asked. P survives because the
unsubscribe was a no-op (LRU guard rejected it).

However, the concern is what happens next: if the caller of unsubscribe
expected P to be unsubscribed but it wasn't, they may not retry.

Inverted concern: If unsubscribe was called after LRU.add(P) at step 4
but before the caller of register_subscription acts on P's existence, then
unsubscribe succeeds and P is removed — but the caller of subscribe believes it
succeeded.

Severity: LOW — The LRU containment check protects against the most
dangerous interleaving in this specific scenario. The inverted case (P survives
when it shouldn't) is less harmful because the reconciler will eventually clean
it up.

---

Scenario D: Reconciliation Racing with Active Registration

Actors:

• Task A: register_subscription(P) — in progress
• Task R: periodic reconcile_subscriptions — runs concurrently

Interleaving:

Step	Task A (register)	Task R (reconcile)	LRU	Pubsub
1	pubsub_client.subscribe(P) ✅		{}	{P}
2	— gap —	Snapshot: LRU = {}, pubsub_union = {..P..}	{}	{P}
3		extra_in_pubsub = {P} (P in pubsub but not LRU)	{}	{P}
4	LRU.add(P)		{P}	{P}
5	register_subscription returns Ok		{P}	{P}
6		unsubscribe_and_notify_removal(P) → pubsub removes P	{P}	{}
7		Sends removal notification → bank removes P's data	P in LRU	P NOT in pubsub

Result: P is in the LRU cache but NOT subscribed in pubsub. The caller
believes P is being watched (is_watching(P) returns true), but no updates
will be received for P.

On the next reconciliation cycle, the reconciler will see P in LRU but not in
pubsub_intersection and will resubscribe it (extra_in_lru path, line 130).
So the state is self-healing — but there is a window (one reconciliation
interval, default ACTIVE_SUBSCRIPTIONS_UPDATE_INTERVAL_MS) during which P
receives no updates.

Worse variant: If the removal notification at step 7 triggers bank eviction
of P's data, the account state is lost from the bank. Even after resubscription,
the account would need to be re-fetched and re-cloned.

Key code locations:

• register_subscription gap: remote_account_provider/mod.rs:895–900
• Reconciler unsubscribe path: subscription_reconciler.rs:141–154
• Reconciler invocation: remote_account_provider/mod.rs:244–251

Severity: HIGH — This is the most dangerous scenario because:

1. It requires no special conditions — just the periodic reconciler running
   during a normal subscription operation.
2. The gap between pubsub subscribe (line 895) and LRU add (line 900) is the
   duration of an async network call (up to 2 seconds timeout), which is a large
   window.
3. The removal notification causes active bank data loss.
4. Although self-healing via the next reconciliation cycle, the interim data loss
   can cause transaction failures or stale reads.

---

4. Root Cause Analysis

Fundamental Problem: No Atomicity Between LRU and Pubsub

The subscribe and unsubscribe operations span two independent state stores (LRU
cache and pubsub client) without any mechanism to make the compound operation
atomic:

// register_subscription (mod.rs:890-914)
async fn register_subscription(&self, pubkey: &Pubkey) {
    // ──── BEGIN NON-ATOMIC WINDOW ────
    self.pubsub_client.subscribe(*pubkey, None).await?;  // Step 1: network
    //                                                      ↑
    //  P is in pubsub but NOT in LRU for the entire duration of the await
    //                                                      ↓
    self.lrucache_subscribed_accounts.add(*pubkey);       // Step 2: local
    // ──── END NON-ATOMIC WINDOW ────
}

The LRU cache uses parking_lot::Mutex which provides per-operation atomicity
but no cross-operation transactional guarantees. Each call to add(),
contains(), and remove() independently acquires and releases the lock
(lru_cache.rs:56–108).

No Per-Pubkey Serialization

There is no per-pubkey lock or serialization mechanism. Multiple async tasks can
concurrently call subscribe(P) and unsubscribe(P) on the same pubkey, and
these operations interleave freely at every .await point.

is_pending Guard Is Insufficient

The is_pending check in cancel_subs (subscription.rs:146) only checks the
fetching_accounts map (mod.rs:934–937), which tracks in-flight RPC fetch
operations. It does NOT protect against:

• Subscribe operations that have completed their fetch but not yet cleaned up
• Subscribe operations triggered by undelegation_requested (which calls
  subscribe_to_account directly, bypassing the fetch/dedup pipeline)
• Reconciler-driven resubscriptions

---

5. Severity Summary

Scenario	Description	Severity	Self-Healing?	Data Loss?
A	Concurrent pipelines sharing a pubkey	Medium	Yes (next tx re-subscribes)	No
B	Undelegation subscribe vs delegation unsubscribe	High	No (requires manual retry)	Possible (stuck undelegation)
C	LRU temporal gap (inverse: P survives)	Low	Yes (reconciler)	No
D	Reconciler vs active registration	High	Partial (next reconcile cycle)	Yes (bank removal notification)

---

6. Conclusion

The subscribe/unsubscribe race condition is confirmed and exploitable under
normal concurrent operation. The core issue is the lack of atomicity between the
pubsub network operation and the LRU cache update in both register_subscription
and unsubscribe.

Scenario B is the most practically dangerous: when a committor service
triggers undelegation_requested (subscribe) and a concurrent subscription
update triggers unsubscribe_from_delegated_account (unsubscribe) for the same
pubkey, the subscription can be silently lost. This causes the account to get
stuck in "undelegating" state, which is a correctness bug visible to users.

Scenario D is the most architecturally concerning: the periodic reconciler—
designed to fix inconsistencies—can actually create them by unsubscribing
accounts that are in the middle of being registered. The 2-second async gap in
register_subscription makes this window substantial.

The is_pending guard in cancel_subs provides partial protection for
Scenario A but is ineffective against Scenarios B, C, and D because it only
checks the fetch dedup map, not the subscribe operation itself.

A proper fix would require one or more of:

1. Per-pubkey subscribe/unsubscribe serialization (e.g., a
   DashMap<Pubkey, Mutex<()>> acting as a key-level lock)
2. Add-to-LRU-first strategy: add the pubkey to the LRU before the async
   pubsub subscribe, so concurrent unsubscribe and reconciler operations see it
   immediately. Rollback the LRU entry if pubsub subscribe fails.
3. In-flight set: a separate HashSet<Pubkey> of "currently subscribing"
   pubkeys that the unsubscribe path and reconciler check before acting, similar
   to is_pending but covering subscribe operations.
4. Reconciler exclusion window: skip unsubscribing pubkeys from pubsub that
   were added in the last N seconds, giving register_subscription time to
   complete its LRU update.

[thlorenz]

Fix Plan: Subscribe/Unsubscribe Race Condition in magicblock-chainlink

NOTE: we should indeed solve this via Option B

Chosen Approach: Option B — Per-Pubkey Lock via scc::HashMap

Justification

Option B (per-pubkey lock) is the best fit because:

1. Fine-grained concurrency: Only serializes operations on the same pubkey. Different pubkeys can subscribe/unsubscribe concurrently without contention.
2. scc::HashMap already in use: The crate already depends on scc (used in fetch_cloner/mod.rs). No new dependency needed.
3. Async-safe: Uses tokio::sync::Mutex (already used as AsyncMutex in pubsub_connection_pool.rs), which can be held across .await points without blocking the executor.
4. Simple mental model: "Before touching pubkey P's subscription state, acquire P's lock." All race scenarios are eliminated by construction.
5. Low overhead: Lock entries are lightweight (Arc<tokio::sync::Mutex<()>>). Cleanup is straightforward.

Why not other options:

• Option A (actor/channel): Serializes ALL sub/unsub operations globally, unnecessary bottleneck.
• Option C (refcount/generation): Complex, hard to reason about correctness, easy to introduce subtle bugs.
• Option D (global tokio::sync::Mutex): Too coarse — blocks all pubkeys when any one is in-flight.

---

Detailed Implementation Plan

Step 1: Add per-pubkey lock map to RemoteAccountProvider

File: magicblock-chainlink/src/remote_account_provider/mod.rs

Add a new field to RemoteAccountProvider:

use tokio::sync::Mutex as AsyncMutex;

pub struct RemoteAccountProvider<T: ChainRpcClient, U: ChainPubsubClient> {
    // ... existing fields ...

    /// Per-pubkey lock to serialize subscribe/unsubscribe operations
    /// on the same pubkey, preventing race conditions between
    /// LRU cache and pubsub client state changes.
    subscription_locks: scc::HashMap<Pubkey, Arc<AsyncMutex<()>>>,
}

Initialize it in RemoteAccountProvider::new() (around line 262):

subscription_locks: scc::HashMap::new(),

Step 2: Add helper to acquire per-pubkey lock

File: magicblock-chainlink/src/remote_account_provider/mod.rs

Add a private helper method:

/// Acquires the per-pubkey subscription lock, creating it if needed.
/// Returns the guard that must be held for the duration of the
/// subscribe or unsubscribe operation.
async fn acquire_subscription_lock(
    &self,
    pubkey: &Pubkey,
) -> tokio::sync::OwnedMutexGuard<()> {
    let mutex = self
        .subscription_locks
        .entry(*pubkey)
        .or_insert_with(|| Arc::new(AsyncMutex::new(())))
        .get()
        .clone();
    mutex.lock_owned().await
}

Step 3: Modify subscribe() to use per-pubkey lock

File: magicblock-chainlink/src/remote_account_provider/mod.rs
Lines: ~941-953

Change from:

pub async fn subscribe(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    if self.is_watching(pubkey) {
        self.lrucache_subscribed_accounts.add(*pubkey);
        return Ok(());
    }
    self.register_subscription(pubkey).await?;
    Ok(())
}

To:

pub async fn subscribe(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    let _guard = self.acquire_subscription_lock(pubkey).await;
    if self.is_watching(pubkey) {
        // Promote in LRU cache even if already subscribed
        self.lrucache_subscribed_accounts.add(*pubkey);
        return Ok(());
    }
    self.register_subscription(pubkey).await?;
    Ok(())
}

What this fixes: With the lock held, no concurrent unsubscribe(P) can proceed between the is_watching check and the completion of register_subscription. This prevents Scenario 3 where T3's subscribe sees P in LRU (and short-circuits) while T2's unsubscribe is in-flight and about to remove P from both pubsub and LRU.

Step 4: Modify unsubscribe() to use per-pubkey lock

File: magicblock-chainlink/src/remote_account_provider/mod.rs
Lines: ~972-997

Change from:

pub async fn unsubscribe(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    if !self.lrucache_subscribed_accounts.can_evict(pubkey) {
        warn!(...);
        return Ok(());
    }
    if !self.lrucache_subscribed_accounts.contains(pubkey) {
        warn!(...);
        return Ok(());
    }

    let success = subscription_reconciler::unsubscribe_and_notify_removal(
        *pubkey,
        &self.pubsub_client,
        &self.removed_account_tx,
    )
    .await;

    if success {
        self.lrucache_subscribed_accounts.remove(pubkey);
    }
    Ok(())
}

To:

pub async fn unsubscribe(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    if !self.lrucache_subscribed_accounts.can_evict(pubkey) {
        warn!(...);
        return Ok(());
    }

    let _guard = self.acquire_subscription_lock(pubkey).await;

    // Re-check under lock — another task may have already removed it
    if !self.lrucache_subscribed_accounts.contains(pubkey) {
        warn!(...);
        return Ok(());
    }

    let success = subscription_reconciler::unsubscribe_and_notify_removal(
        *pubkey,
        &self.pubsub_client,
        &self.removed_account_tx,
    )
    .await;

    if success {
        self.lrucache_subscribed_accounts.remove(pubkey);
    }
    Ok(())
}

Key detail: The can_evict check remains outside the lock (it's a pure check against an immutable HashSet — no race). The contains check moves inside the lock so it's atomic with the subsequent unsubscribe + remove.

What this fixes: Together with Step 3, this eliminates Scenario 3. If T2 holds the lock for unsubscribe(P), T3's subscribe(P) will block until T2 finishes. When T3 resumes, is_watching(P) returns false (T2 removed it), so T3 calls register_subscription — re-subscribing in both pubsub and LRU. No lost subscription.

Step 5: Modify register_subscription() to remove from lock map on eviction

File: magicblock-chainlink/src/remote_account_provider/mod.rs
Lines: ~890-914

No structural change needed. The eviction path calls pubsub_client.unsubscribe(evicted) directly (not through self.unsubscribe()), so no deadlock risk. However, we should also acquire the evicted pubkey's lock to prevent a concurrent subscribe(evicted) from racing with the eviction:

Change from:

async fn register_subscription(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    self.pubsub_client.subscribe(*pubkey, None).await?;

    if let Some(evicted) = self.lrucache_subscribed_accounts.add(*pubkey) {
        trace!(evicted = %evicted, "Evicting account");
        if let Err(err) = self.pubsub_client.unsubscribe(evicted).await {
            warn!(...);
        }
        self.send_removal_update(evicted).await?;
    }
    Ok(())
}

To:

async fn register_subscription(
    &self,
    pubkey: &Pubkey,
) -> RemoteAccountProviderResult<()> {
    // NOTE: caller (subscribe) already holds the lock for `pubkey`
    self.pubsub_client.subscribe(*pubkey, None).await?;

    if let Some(evicted) = self.lrucache_subscribed_accounts.add(*pubkey) {
        trace!(evicted = %evicted, "Evicting account");

        // Acquire lock for the evicted pubkey to prevent concurrent
        // subscribe(evicted) from racing with this eviction
        let _evicted_guard =
            self.acquire_subscription_lock(&evicted).await;

        if let Err(err) = self.pubsub_client.unsubscribe(evicted).await {
            warn!(...);
        }
        self.send_removal_update(evicted).await?;
    }
    Ok(())
}

Deadlock analysis: The caller holds the lock for pubkey. We acquire the lock for evicted. Since evicted != pubkey (enforced by the debug_assert_ne! in AccountsLruCache::add), there is no self-deadlock. Two-key deadlocks are possible only if two tasks each hold one key and wait for the other — but subscribe only ever holds one key's lock at a time before acquiring a second, and the second is always the evicted key (determined by LRU order), making circular waits practically impossible. If paranoia is warranted, we can use try_lock with a fallback, but this is unnecessary given the constraints.

Step 6: Fix cancel_subs to handle Scenario 5

File: magicblock-chainlink/src/chainlink/fetch_cloner/subscription.rs

The current is_pending check only guards against active RPC fetches but doesn't account for another pipeline call that has already completed its fetch and is now relying on the subscription. The per-pubkey lock from Steps 3-4 prevents race conditions within individual subscribe/unsubscribe calls, but Scenario 5 is a higher-level logical bug: cancel_subs unsubscribes a pubkey that another concurrent pipeline needs.

The root cause is that cancel_subs with CancelStrategy::Hybrid { all: {P} } unconditionally unsubscribes P without knowing that another caller also needs P's subscription.

Fix: Add a reference count to AccountsLruCache tracking how many active pipeline operations are using each subscription. cancel_subs should only unsubscribe if the refcount is 1 (only the cancelling pipeline is using it).

Step 6a: Add refcount tracking to AccountsLruCache

File: magicblock-chainlink/src/remote_account_provider/lru_cache.rs

Change the LRU value type from () to u32 (reference count):

pub struct AccountsLruCache {
    /// Tracks which accounts are currently subscribed to,
    /// value is the number of active users of this subscription
    subscribed_accounts: Mutex<LruCache<Pubkey, u32>>,
    accounts_to_never_evict: HashSet<Pubkey>,
}

Modify add():

pub fn add(&self, pubkey: Pubkey) -> Option<Pubkey> {
    if self.accounts_to_never_evict.contains(&pubkey) {
        return None;
    }

    let mut subs = self.subscribed_accounts.lock();
    if let Some(count) = subs.get_mut(&pubkey) {
        // Already present — promote and increment refcount
        *count += 1;
        trace!(pubkey = %pubkey, refcount = *count, "Account promoted");
        return None;
    }

    // New entry with refcount = 1
    let evicted = subs
        .push(pubkey, 1)
        .map(|(evicted_pubkey, _)| evicted_pubkey);

    if let Some(evicted_pubkey) = evicted {
        inc_evicted_accounts_count();
        debug_assert_ne!(evicted_pubkey, pubkey);
        trace!(evicted_pubkey = %evicted_pubkey, "Evict candidate");
    }

    evicted
}

Add a release() method that decrements the refcount:

/// Decrements the refcount for a pubkey. Returns true if the refcount
/// reached 0 (caller should unsubscribe). Returns false if other users
/// remain or the pubkey is not in the cache.
pub fn release(&self, pubkey: &Pubkey) -> bool {
    let mut subs = self.subscribed_accounts.lock();
    if let Some(count) = subs.get_mut(pubkey) {
        *count = count.saturating_sub(1);
        if *count == 0 {
            subs.pop(pubkey);
            trace!(pubkey = %pubkey, "Released last reference, removed");
            return true;
        }
        trace!(pubkey = %pubkey, refcount = *count, "Released reference");
    }
    false
}

Keep remove() as a force-remove (for reconciliation and eviction paths).

Step 6b: Update subscribe() to always increment refcount

The current subscribe() calls lrucache.add() which now increments the refcount. This means each call to subscribe(P) increments P's refcount. This is correct — each pipeline that subscribes to P contributes a reference.

Step 6c: Update unsubscribe() to use release() instead of remove()

File: magicblock-chainlink/src/remote_account_provider/mod.rs

In unsubscribe(), replace:

if success {
    self.lrucache_subscribed_accounts.remove(pubkey);
}

With:

if success {
    // Only actually removed if this was the last reference
    if !self.lrucache_subscribed_accounts.release(pubkey) {
        // Other callers still need this subscription — resubscribe
        if let Err(err) =
            self.pubsub_client.subscribe(*pubkey, None).await
        {
            warn!(
                pubkey = %pubkey, error = ?err,
                "Failed to resubscribe after premature unsubscribe"
            );
        }
    }
}

Wait — this is getting complex. Let me reconsider.

Simpler approach for Scenario 5: Instead of refcounting in the LRU (which changes a lot of semantics), keep the LRU as-is (no refcount) and fix cancel_subs to only cancel subscriptions that are genuinely no longer needed.

Revised Step 6: Fix cancel_subs by checking is_watching under lock

The real issue in Scenario 5 is that cancel_subs with All { P } force-unsubscribes P even though another pipeline re-subscribed P (which promoted it in LRU). The per-pubkey lock already ensures that unsubscribe(P) and subscribe(P) are serialized. The remaining issue is ordering: if the re-subscribe happened before cancel_subs runs, then P is correctly in LRU and pubsub — but cancel_subs removes it anyway.

The fix: In cancel_subs, before calling unsubscribe, check if the account is still is_pending (for the calling pipeline's batch). But the real fix is deeper — cancel_subs should not be called with pubkeys that other concurrent pipelines are actively using.

Practical fix: The is_pending guard in cancel_subs (line 146) already prevents unsubscribing accounts with pending fetch requests. For accounts that have completed fetching but are actively used by another pipeline, we need the per-pubkey lock to ensure atomicity. With the per-pubkey lock in place:

1. If T2's pipeline calls subscribe(P) before T1's cancel_subs runs unsubscribe(P):

   • T2's subscribe(P) acquires lock for P, sees P in LRU (promotes), releases lock.
   • T1's cancel_subs → unsubscribe(P) acquires lock for P, checks LRU → P is there → unsubscribes from pubsub → removes from LRU.
   • Still a bug: T2 thinks P is subscribed but it isn't.

2. The fundamental issue is that cancel_subs should not unsubscribe P if it was promoted (re-used) by another pipeline after the cancel decision was made.

Root-cause fix for Scenario 5: Add a generation counter per pubkey in LRU. When subscribe() is called, record the current generation. When cancel_subs runs, only unsubscribe if the generation hasn't changed (meaning no one else re-subscribed).

This is getting complex. Let me propose a simpler, more practical approach:

Final Step 6: Snapshot-based cancel guard

File: magicblock-chainlink/src/remote_account_provider/mod.rs

Add a method subscribe_generation that tracks when a pubkey was last subscribed:

Actually, the simplest fix: Don't cancel subscriptions for pubkeys in CancelStrategy::All/Hybrid.all that have been re-promoted in the LRU since the cancel decision was made.

In practice, this means cancel_subs should re-check is_watching right before calling unsubscribe. But is_watching returning true doesn't help — it just means the pubkey is in LRU, which it always is until we remove it.

The correct fix: Change the LRU cache value from () to a use counter that tracks how many concurrent pipeline operations are referencing this subscription. This is cleaner than a generation counter.

Revised Step 6: Reference-counted subscriptions (simplified)

After careful analysis, the cleanest fix for Scenario 5 that does NOT over-complicate the LRU:

Add a separate refcount map (not in LRU) that tracks active pipeline usage:

File: magicblock-chainlink/src/remote_account_provider/mod.rs

pub struct RemoteAccountProvider<T: ChainRpcClient, U: ChainPubsubClient> {
    // ... existing fields ...

    /// Per-pubkey lock for subscribe/unsubscribe atomicity
    subscription_locks: scc::HashMap<Pubkey, Arc<AsyncMutex<()>>>,

    /// Tracks the number of active pipeline operations using each
    /// subscription. Unsubscribe is skipped if refcount > 0.
    active_subscription_users: Arc<scc::HashMap<Pubkey, u32>>,
}

Add methods:

/// Increment the active user count for a pubkey.
/// Called when a pipeline starts relying on a subscription.
pub fn retain_subscription(&self, pubkey: &Pubkey) {
    self.active_subscription_users
        .entry(*pubkey)
        .and_modify(|count| *count += 1)
        .or_insert(1);
}

/// Decrement the active user count for a pubkey.
/// Returns true if the count reached 0 (safe to unsubscribe).
pub fn release_subscription(&self, pubkey: &Pubkey) -> bool {
    let mut should_remove = false;
    self.active_subscription_users.entry(*pubkey).and_modify(|count| {
        *count = count.saturating_sub(1);
        if *count == 0 {
            should_remove = true;
        }
    });
    if should_remove {
        let _ = self.active_subscription_users.remove(pubkey);
    }
    should_remove || !self.active_subscription_users.contains(pubkey)
}

/// Check if any pipeline is actively using this subscription.
pub fn has_active_users(&self, pubkey: &Pubkey) -> bool {
    self.active_subscription_users
        .get(pubkey)
        .map(|entry| *entry.get() > 0)
        .unwrap_or(false)
}

File: magicblock-chainlink/src/chainlink/fetch_cloner/subscription.rs

Update cancel_subs to check active users:

for pubkey in subs_to_cancel {
    let provider_clone = provider.clone();
    joinset.spawn(async move {
        // Check if there are pending requests OR active pipeline
        // users for this account before unsubscribing
        if provider_clone.is_pending(&pubkey)
            || provider_clone.has_active_users(&pubkey)
        {
            debug!(
                pubkey = %pubkey,
                "Skipping unsubscribe — has pending requests \
                 or active users"
            );
            return;
        }

        if let Err(err) = provider_clone.unsubscribe(&pubkey).await {
            warn!(pubkey = %pubkey, error = ?err, "Failed to unsubscribe");
        }
    });
}

File: magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs (callers)

At the point where a pipeline starts using a subscription result, call retain_subscription. When the pipeline completes (or the subscription result is consumed), call release_subscription. The exact call sites need to be identified in the fetch_cloner pipeline — likely around fetch_and_clone_accounts_with_dedup.

---

Summary of All Changes

File 1: magicblock-chainlink/src/remote_account_provider/mod.rs

Change	Lines (approx)	Description
Add import	1-8	use tokio::sync::Mutex as AsyncMutex;
Add fields	99-145	subscription_locks: scc::HashMap<Pubkey, Arc<AsyncMutex<()>>> and active_subscription_users: Arc<scc::HashMap<Pubkey, u32>>
Init fields	~290	Initialize both new fields in new()
Add helper	~935	acquire_subscription_lock() method
Add methods	~935	retain_subscription(), release_subscription(), has_active_users()
Modify subscribe()	941-953	Acquire per-pubkey lock before check-and-subscribe
Modify unsubscribe()	972-997	Acquire per-pubkey lock; move contains check inside lock
Modify register_subscription()	890-914	Acquire evicted pubkey's lock before unsubscribing evicted account

File 2: magicblock-chainlink/src/chainlink/fetch_cloner/subscription.rs

Change	Lines (approx)	Description
Modify cancel_subs	146-149	Add has_active_users check alongside existing is_pending

File 3: magicblock-chainlink/src/chainlink/fetch_cloner/mod.rs

Change	Lines (approx)	Description
Call retain_subscription	TBD	When a pipeline starts relying on subscription results
Call release_subscription	TBD	When a pipeline is done with subscription results

File 4: magicblock-chainlink/src/remote_account_provider/lru_cache.rs

No changes needed. The LRU cache remains a simple presence cache. Race protection is handled by the per-pubkey lock at the RemoteAccountProvider level.

File 5: magicblock-chainlink/src/remote_account_provider/subscription_reconciler.rs

No structural changes needed. The reconciler runs periodically and operates on snapshots of LRU + pubsub state. The per-pubkey locks ensure that any concurrent subscribe/unsubscribe operations are atomic, so the reconciler's snapshots are consistent within each pubkey. The reconciler acts as a safety net — it will catch any remaining inconsistencies.

---

Test Scenarios to Verify the Fix

Test 1: Scenario 3 — Subscribe sees LRU present while unsubscribe is in-flight

#[tokio::test]
async fn test_concurrent_unsubscribe_resubscribe_no_lost_subscription() {
    // 1. Subscribe(P) — P is in LRU + pubsub
    // 2. Spawn unsubscribe(P) — starts async pubsub.unsubscribe
    // 3. Spawn subscribe(P) — should block on per-pubkey lock
    // 4. unsubscribe completes — P removed from LRU + pubsub, lock released
    // 5. subscribe resumes — sees P NOT in LRU — calls register_subscription
    // 6. Verify: P is in LRU AND pubsub
}

Test 2: Scenario 5 — cancel_subs skips actively-used subscriptions

#[tokio::test]
async fn test_cancel_subs_preserves_subscription_with_active_users() {
    // 1. Subscribe(P), retain_subscription(P) for pipeline T2
    // 2. cancel_subs with All({P})
    // 3. Verify: P is NOT unsubscribed (has_active_users returns true)
    // 4. release_subscription(P)
    // 5. cancel_subs with All({P}) again
    // 6. Verify: P IS unsubscribed now
}

Test 3: Concurrent subscribes to same pubkey don't deadlock

#[tokio::test]
async fn test_concurrent_subscribes_same_pubkey_no_deadlock() {
    // 1. Spawn 10 tasks all calling subscribe(P) concurrently
    // 2. Verify: all complete within timeout (no deadlock)
    // 3. Verify: P is in LRU AND pubsub exactly once
}

Test 4: Subscribe + LRU eviction doesn't deadlock

#[tokio::test]
async fn test_subscribe_eviction_no_deadlock() {
    // LRU capacity = 2
    // 1. subscribe(A), subscribe(B) — cache full
    // 2. subscribe(C) — evicts A
    // 3. Concurrent subscribe(A) — should succeed after eviction lock released
    // 4. Verify: A, B or A, C depending on order — no deadlock
}

Test 5: Unsubscribe idempotency under lock

#[tokio::test]
async fn test_double_unsubscribe_is_noop() {
    // 1. subscribe(P)
    // 2. Spawn two unsubscribe(P) tasks concurrently
    // 3. First one removes P from LRU + pubsub
    // 4. Second one acquires lock, sees P not in LRU, returns early
    // 5. Verify: no errors, P not in LRU or pubsub
}

Test 6: Reconciler doesn't conflict with locked operations

#[tokio::test]
async fn test_reconciler_during_active_subscribe() {
    // 1. Start a slow subscribe(P) (mock pubsub with delay)
    // 2. Run reconcile_subscriptions while subscribe is in-flight
    // 3. Verify: no panic, no inconsistent state after both complete
}

---

Migration/Backward Compatibility Notes

1. No public API changes: subscribe() and unsubscribe() signatures are unchanged. The per-pubkey lock is internal.

2. New public methods: retain_subscription(), release_subscription(), has_active_users() are added to RemoteAccountProvider. These need to be called by pipeline code in fetch_cloner. All callers are internal to the crate.

3. Performance impact: Minimal. Per-pubkey locks only add contention when two tasks operate on the same pubkey concurrently, which is the exact scenario we want to serialize. The scc::HashMap is lock-free for reads and uses fine-grained locking for writes.

4. Memory: subscription_locks grows with unique pubkeys. Since it mirrors the LRU cache size (bounded by capacity), memory growth is bounded. Stale lock entries can be cleaned up lazily during unsubscribe (remove entry after releasing lock if no waiters). Alternatively, they can be cleaned up during reconciliation.

5. Lock entry cleanup: After unsubscribe completes and the lock is released, the entry in subscription_locks can be removed. Since scc::HashMap handles concurrent access safely, this is straightforward:

   // After _guard is dropped (end of unsubscribe)
   // Only remove if no one else is waiting
   if let Some(entry) = self.subscription_locks.get(pubkey) {
       if Arc::strong_count(entry.get()) == 1 {
           let _ = self.subscription_locks.remove(pubkey);
       }
   }

   This cleanup is optional — the memory overhead of keeping stale entries is negligible given LRU bounds.

6. Reconciler safety: The reconciler doesn't acquire per-pubkey locks (it operates directly on pubsub_client and LRU snapshots). This is intentional — the reconciler is a periodic safety net, not a concurrent participant. If a reconciler operation races with a locked subscribe/unsubscribe, the next reconciliation pass will fix any inconsistency. This is acceptable because reconciliation is already an eventually-consistent mechanism.

---

Implementation Order

1. Add subscription_locks field + acquire_subscription_lock helper — build should pass (unused field)
2. Wrap subscribe() with lock — build + existing tests pass
3. Wrap unsubscribe() with lock — build + existing tests pass
4. Wrap eviction in register_subscription() with evicted key's lock — build + existing tests pass
5. Add active_subscription_users field + methods — build passes
6. Update cancel_subs to check has_active_users — build passes
7. Add retain_subscription/release_subscription calls in pipeline — requires identifying exact call sites in fetch_cloner
8. Add new tests for all scenarios above
9. Run full test suite: cargo nextest run --package magicblock-chainlink
10. Run lint: make lint

[Rhovian]

Hey, can i take this one?