Creating special serialized to reduce more space for token accounts

Cargo.lock

@@ -1925,8 +1925,10 @@ name = "lite-token-account-storage"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "arrayref",
  "async-trait",
  "bincode",
+ "bitflags 2.6.0",
  "dashmap",
  "futures",
  "itertools",

Cargo.toml

@@ -41,6 +41,8 @@ async-trait = "0.1.68"
 base64 = "0.21.0"
 futures = "0.3.28"
 tracing-subscriber = "0.3.16"
+bitflags = "2.6.0"
+arrayref = "0.3.7"
 
 # spl token
 spl-token = "~4.0.0"

token_account_storage/Cargo.toml

@@ -22,7 +22,8 @@ serde = {workspace = true}
 anyhow = {workspace = true}
 bincode = {workspace = true}
 lz4 = {workspace = true}
-
+bitflags = {workspace = true}
+arrayref = {workspace = true}
 
 [dev-dependencies]
 rand = "0.8.5"

token_account_storage/src/account_types.rs

@@ -11,7 +11,7 @@ pub enum Program {
     Token2022Program,
 }
 
-#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 #[repr(u8)]
 pub enum TokenProgramAccountState {
     Uninitialized,
@@ -63,19 +63,19 @@ impl From<spl_token_2022::state::AccountState> for TokenProgramAccountState {
     }
 }
 
-#[derive(Debug, Serialize, Deserialize, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct TokenAccount {
-    pub program: Program,
-    pub pubkey: Pubkey,
-    pub mint: MintIndex,
-    pub owner: Pubkey,
-    pub amount: u64,
-    pub state: TokenProgramAccountState,
-    pub delegate: Option<(Pubkey, u64)>,
-    pub is_native: Option<u64>,
-    pub close_authority: Option<Pubkey>,
-    pub additional_data: Option<Vec<u8>>,
-    pub lamports: u64,
+    pub program: Program,                 // 1 byte, offset : 0,
+    pub is_deleted: bool,                 // 1, offset : 1
+    pub pubkey: Pubkey,                   // 32, offset: 2
+    pub owner: Pubkey,                    // 32, offset : 34
+    pub mint: MintIndex,                  // 4, offset : 66
+    pub amount: u64,                      // 8, offset : 70
+    pub state: TokenProgramAccountState,  // 1
+    pub delegate: Option<(Pubkey, u64)>,  // 1
+    pub is_native: Option<u64>,           // 1
+    pub close_authority: Option<Pubkey>,  // 1
+    pub additional_data: Option<Vec<u8>>, // 0
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
@@ -101,3 +101,225 @@ pub struct MultiSig {
     pub is_initialized: bool,
     pub signers: Vec<Pubkey>,
 }
+
+const INITIALIZED_STATE_BITS: u8 = 0b00000000;
+const FROZEN_STATE_BITS: u8 = 0b00000001;
+const UNITITIALIZED_STATE_BITS: u8 = 0b00000010;
+
+// serialized adds 8 bytes to store the size of resulting array, so using custom function instead
+impl TokenAccount {
+    pub fn to_bytes(&self) -> Vec<u8> {
+        // save space in serialization, first two bits LSB is reserved for program / 2
+        // next two LSBs are reserved for state / 2
+        // next bit for has delegate / 1
+        // then next for has is_native / 1
+        // then next of has close authority / 1
+        // and last for has additional data / 1
+        let mut program_bits = 0u8;
+        if self.program == Program::Token2022Program {
+            program_bits |= 0x1; // set lsb to 1
+        }
+        match self.state {
+            TokenProgramAccountState::Uninitialized => {
+                program_bits |= UNITITIALIZED_STATE_BITS << 2
+            }
+            TokenProgramAccountState::Frozen => program_bits |= FROZEN_STATE_BITS << 2,
+            TokenProgramAccountState::Initialized => program_bits |= INITIALIZED_STATE_BITS << 2,
+        }
+        let has_delegate = self.delegate.is_some();
+        if has_delegate {
+            program_bits |= 0x1 << 4;
+        }
+
+        let has_native = self.is_native.is_some();
+        if has_native {
+            program_bits |= 0x1 << 5;
+        }
+
+        let has_close_authority = self.close_authority.is_some();
+        if has_close_authority {
+            program_bits |= 0x1 << 6;
+        }
+
+        let has_additional_data = self.additional_data.is_some();
+        if has_additional_data {
+            program_bits |= 0x1 << 7;
+        }
+
+        let mut bytes = Vec::<u8>::with_capacity(85);
+        bytes.push(program_bits);
+        bytes.push(if self.is_deleted { 1 } else { 0 });
+        bytes.extend_from_slice(&self.pubkey.to_bytes());
+        bytes.extend_from_slice(&self.owner.to_bytes());
+        bytes.extend_from_slice(&self.mint.to_le_bytes());
+        bytes.extend_from_slice(&self.amount.to_le_bytes());
+        if has_delegate {
+            let (delegate, amount) = self.delegate.unwrap();
+            bytes.extend_from_slice(&delegate.to_bytes());
+            bytes.extend_from_slice(&amount.to_le_bytes());
+        }
+
+        if has_native {
+            bytes.extend_from_slice(&self.is_native.unwrap().to_le_bytes());
+        }
+
+        if has_close_authority {
+            bytes.extend_from_slice(&self.close_authority.unwrap().to_bytes());
+        }
+
+        if has_additional_data {
+            let additional_data = self.additional_data.as_ref().unwrap();
+            let size = additional_data.len() as u32; // accounts are 10MBs max will not extend u32 space
+            bytes.extend_from_slice(&size.to_le_bytes());
+            bytes.extend_from_slice(additional_data);
+        }
+        bytes
+    }
+
+    pub fn from_bytes(v: &[u8]) -> TokenAccount {
+        let header: [u8; 78] = v[0..78].try_into().unwrap();
+        let mut others = &v[78..];
+        let (bit_flags, is_deleted, pubkey, owner, mint_index, amount) =
+            arrayref::array_refs![&header, 1, 1, 32, 32, 4, 8];
+        let bit_flags = bit_flags[0];
+        let program = match bit_flags & 0x1 {
+            0 => Program::TokenProgram,
+            1 => Program::Token2022Program,
+            _ => unreachable!(),
+        };
+        let state = match (bit_flags >> 2) & 0x3 {
+            INITIALIZED_STATE_BITS => TokenProgramAccountState::Initialized,
+            FROZEN_STATE_BITS => TokenProgramAccountState::Frozen,
+            UNITITIALIZED_STATE_BITS => TokenProgramAccountState::Uninitialized,
+            _ => unreachable!(),
+        };
+        let has_delegate = (bit_flags >> 4) & 0x1 == 0x1;
+        let has_native = (bit_flags >> 5) & 0x1 == 0x1;
+        let has_close_authority = (bit_flags >> 6) & 0x1 == 0x1;
+        let has_additional_data = (bit_flags >> 7) & 0x1 == 0x1;
+
+        let delegate = if has_delegate {
+            let delegate_bytes: &[u8; 32 + 8] = others[..32 + 8].try_into().unwrap();
+            others = &others[32 + 8..];
+            let delegate_pk = Pubkey::new_from_array(*arrayref::array_ref![delegate_bytes, 0, 32]);
+            let amount = u64::from_be_bytes(*arrayref::array_ref![delegate_bytes, 32, 8]);
+            Some((delegate_pk, amount))
+        } else {
+            None
+        };
+
+        let is_native = if has_native {
+            let native_bytes = *arrayref::array_ref![others, 0, 8];
+            others = &others[8..];
+            Some(u64::from_le_bytes(native_bytes))
+        } else {
+            None
+        };
+
+        let close_authority = if has_close_authority {
+            let native_bytes = *arrayref::array_ref![others, 0, 32];
+            others = &others[32..];
+            Some(Pubkey::new_from_array(native_bytes))
+        } else {
+            None
+        };
+
+        let additional_data = if has_additional_data {
+            let native_bytes = *arrayref::array_ref![others, 0, 4];
+            Some(others[4..4 + u32::from_le_bytes(native_bytes) as usize].to_vec())
+        } else {
+            None
+        };
+
+        TokenAccount {
+            program,
+            is_deleted: match is_deleted[0] {
+                0 => false,
+                1 => true,
+                _ => unreachable!(),
+            },
+            pubkey: Pubkey::new_from_array(*pubkey),
+            owner: Pubkey::new_from_array(*owner),
+            mint: u32::from_le_bytes(*mint_index),
+            amount: u64::from_le_bytes(*amount),
+            state,
+            delegate,
+            is_native,
+            close_authority,
+            additional_data,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use itertools::Itertools;
+    use rand::{thread_rng, Rng};
+    use solana_sdk::pubkey::Pubkey;
+
+    use super::{Program, TokenAccount, TokenProgramAccountState};
+
+    #[test]
+    pub fn test_account_serialization() {
+        let acc = TokenAccount {
+            program: super::Program::TokenProgram,
+            is_deleted: false,
+            pubkey: Pubkey::new_unique(),
+            mint: 0,
+            owner: Pubkey::new_unique(),
+            amount: 10,
+            state: TokenProgramAccountState::Initialized,
+            delegate: None,
+            is_native: None,
+            close_authority: None,
+            additional_data: None,
+        };
+        assert_eq!(acc.to_bytes().len(), 78);
+    }
+
+    #[test]
+    pub fn fuzzy_test_token_account_serialization() {
+        let mut rng = thread_rng();
+        for _ in 0..1_000_000 {
+            let acc = TokenAccount {
+                program: if rng.gen::<bool>() {
+                    Program::TokenProgram
+                } else {
+                    Program::Token2022Program
+                },
+                is_deleted: rng.gen::<bool>(),
+                pubkey: Pubkey::new_unique(),
+                mint: rng.gen(),
+                owner: Pubkey::new_unique(),
+                amount: rng.gen(),
+                state: TokenProgramAccountState::Initialized,
+                delegate: if rng.gen::<bool>() {
+                    Some((Pubkey::new_unique(), rng.gen()))
+                } else {
+                    None
+                },
+                is_native: if rng.gen::<bool>() {
+                    Some(rng.gen())
+                } else {
+                    None
+                },
+                close_authority: if rng.gen::<bool>() {
+                    Some(Pubkey::new_unique())
+                } else {
+                    None
+                },
+                additional_data: if rng.gen::<bool>() {
+                    let len = rng.gen::<usize>() % 1_000_000;
+                    let data = (0..len).map(|_| rng.gen::<u8>()).collect_vec();
+                    Some(data)
+                } else {
+                    None
+                },
+            };
+
+            let ser = acc.to_bytes();
+            let deser = TokenAccount::from_bytes(&ser);
+            assert_eq!(deser, acc);
+        }
+    }
+}

token_account_storage/src/inmemory_token_account_storage.rs

@@ -0,0 +1,114 @@
+use std::{
+    collections::{HashSet, VecDeque},
+    sync::Arc,
+};
+
+use async_trait::async_trait;
+use dashmap::DashMap;
+use itertools::Itertools;
+use lite_account_manager_common::account_store_interface::AccountLoadingError;
+use solana_sdk::pubkey::Pubkey;
+use tokio::sync::RwLock;
+
+use crate::{
+    account_types::{Program, TokenAccount, TokenAccountIndex},
+    token_account_storage_interface::TokenAccountStorageInterface,
+};
+
+#[derive(Default)]
+pub struct InmemoryTokenAccountStorage {
+    pubkey_to_index: Arc<DashMap<Pubkey, TokenAccountIndex>>,
+    token_accounts: Arc<RwLock<VecDeque<Vec<u8>>>>,
+}
+
+#[async_trait]
+impl TokenAccountStorageInterface for InmemoryTokenAccountStorage {
+    async fn contains(&self, pubkey: &Pubkey) -> Option<TokenAccountIndex> {
+        self.pubkey_to_index.get(pubkey).map(|x| *x.value())
+    }
+
+    async fn save_or_update(&self, token_account: TokenAccount) -> (TokenAccountIndex, bool) {
+        match self.pubkey_to_index.entry(token_account.pubkey) {
+            dashmap::mapref::entry::Entry::Occupied(occ) => {
+                // already present
+                let token_index = *occ.get() as usize;
+                let mut write_lk = self.token_accounts.write().await;
+                // update existing token account
+                *write_lk.get_mut(token_index).unwrap() = token_account.to_bytes();
+                (token_index as TokenAccountIndex, false)
+            }
+            dashmap::mapref::entry::Entry::Vacant(v) => {
+                // add new token account
+                let mut write_lk = self.token_accounts.write().await;
+                let token_index = write_lk.len() as TokenAccountIndex;
+                write_lk.push_back(token_account.to_bytes());
+                v.insert(token_index as TokenAccountIndex);
+                drop(write_lk);
+                (token_index, true)
+            }
+        }
+    }
+
+    async fn get_by_index(
+        &self,
+        indexes: HashSet<TokenAccountIndex>,
+    ) -> Result<Vec<TokenAccount>, AccountLoadingError> {
+        let lk = self.token_accounts.read().await;
+        Ok(indexes
+            .iter()
+            .filter_map(|index| lk.get(*index as usize))
+            .map(|x| TokenAccount::from_bytes(x))
+            .collect_vec())
+    }
+
+    async fn get_by_pubkey(&self, pubkey: &Pubkey) -> Option<TokenAccount> {
+        if let Some(acc) = self.pubkey_to_index.get(pubkey) {
+            let index = *acc.value();
+            let lk = self.token_accounts.write().await;
+            match lk.get(index as usize) {
+                Some(binary) => {
+                    let token_account: TokenAccount = TokenAccount::from_bytes(binary);
+                    Some(token_account)
+                }
+                None => None,
+            }
+        } else {
+            None
+        }
+    }
+
+    async fn delete(&self, pubkey: &Pubkey) {
+        let acc_to_remove = self.pubkey_to_index.remove(pubkey);
+        if let Some((_, index)) = acc_to_remove {
+            let mut write_lk = self.token_accounts.write().await;
+            let deleted_account = TokenAccount {
+                program: Program::TokenProgram,
+                is_deleted: true,
+                pubkey: Pubkey::default(),
+                owner: Pubkey::default(),
+                mint: 0,
+                amount: 0,
+                state: crate::account_types::TokenProgramAccountState::Uninitialized,
+                delegate: None,
+                is_native: None,
+                close_authority: None,
+                additional_data: None,
+            }
+            .to_bytes();
+            *write_lk.get_mut(index as usize).unwrap() = deleted_account;
+        }
+    }
+
+    async fn create_snapshot(&self, program: Program) -> Result<Vec<Vec<u8>>, AccountLoadingError> {
+        let program_bit = match program {
+            Program::TokenProgram => 0,
+            Program::Token2022Program => 0x1,
+        };
+        let lk = self.token_accounts.read().await;
+        Ok(lk
+            .iter()
+            .filter(|x| (*x.first().unwrap() & 0x3 == program_bit) && *x.get(1).unwrap() == 0)
+            .cloned()
+            .collect())
+    }
+}

token_account_storage/src/inmemory_token_storage.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::{BTreeMap, HashMap, VecDeque},
+    collections::{BTreeMap, HashMap, HashSet, VecDeque},
     sync::{
         atomic::{AtomicU32, AtomicU64, Ordering},
         Arc,
@@ -22,6 +22,7 @@ use tokio::sync::RwLock;
 
 use crate::{
     account_types::{MintAccount, MintIndex, MultiSig, Program, TokenAccount, TokenAccountIndex},
+    token_account_storage_interface::TokenAccountStorageInterface,
     token_program_utils::{
         get_spl_token_owner_mint_filter, get_token_program_account_type,
         token_account_to_solana_account, token_mint_to_solana_account,
@@ -293,9 +294,8 @@ pub struct InMemoryTokenStorage {
     mints_index_by_pubkey: Arc<DashMap<Pubkey, MintIndex>>,
     multisigs: Arc<DashMap<Pubkey, MultiSig>>,
     accounts_index_by_mint: Arc<DashMap<MintIndex, VecDeque<TokenAccountIndex>>>,
-    account_index_by_pubkey: Arc<DashMap<Pubkey, TokenAccountIndex>>,
     account_by_owner_pubkey: Arc<DashMap<Pubkey, Vec<TokenAccountIndex>>>,
-    token_accounts: Arc<RwLock<VecDeque<TokenAccount>>>,
+    token_accounts_storage: Arc<dyn TokenAccountStorageInterface>,
     processed_storage: ProcessedAccountStore,
     confirmed_slot: Arc<AtomicU64>,
     finalized_slot: Arc<AtomicU64>,
@@ -304,15 +304,14 @@ pub struct InMemoryTokenStorage {
 
 impl InMemoryTokenStorage {
     #[allow(clippy::new_without_default)]
-    pub fn new() -> Self {
+    pub fn new(token_accounts_storage: Arc<dyn TokenAccountStorageInterface>) -> Self {
         Self {
             mints_by_index: Arc::new(DashMap::new()),
             mints_index_by_pubkey: Arc::new(DashMap::new()),
             accounts_index_by_mint: Arc::new(DashMap::new()),
             multisigs: Arc::new(DashMap::new()),
-            account_index_by_pubkey: Arc::new(DashMap::new()),
             account_by_owner_pubkey: Arc::new(DashMap::new()),
-            token_accounts: Arc::new(RwLock::new(VecDeque::new())),
+            token_accounts_storage,
             processed_storage: ProcessedAccountStore::default(),
             confirmed_slot: Arc::new(AtomicU64::new(0)),
             finalized_slot: Arc::new(AtomicU64::new(0)),
@@ -326,46 +325,34 @@ impl InMemoryTokenStorage {
             TokenProgramAccountType::TokenAccount(token_account) => {
                 let token_account_owner = token_account.owner;
                 let mint_index = token_account.mint;
-                // find if the token account is already present
-                match self.account_index_by_pubkey.entry(token_account.pubkey) {
-                    dashmap::mapref::entry::Entry::Occupied(occ) => {
-                        // already present
-                        let token_index = *occ.get() as usize;
-                        let mut write_lk = self.token_accounts.write().await;
-                        // update existing token account
-                        *write_lk.get_mut(token_index).unwrap() = token_account;
-                    }
-                    dashmap::mapref::entry::Entry::Vacant(v) => {
-                        // add new token account
-                        let mut write_lk = self.token_accounts.write().await;
-                        let token_index = write_lk.len() as TokenAccountIndex;
-                        write_lk.push_back(token_account);
-                        v.insert(token_index as TokenAccountIndex);
-                        drop(write_lk);
-
-                        // add account index in accounts_index_by_mint map
-                        match self.accounts_index_by_mint.entry(mint_index) {
-                            dashmap::mapref::entry::Entry::Occupied(mut occ) => {
-                                occ.get_mut().push_back(token_index);
-                            }
-                            dashmap::mapref::entry::Entry::Vacant(v) => {
-                                let mut q = VecDeque::new();
-                                q.push_back(token_index);
-                                v.insert(q);
-                            }
+                let (token_index, is_added) = self
+                    .token_accounts_storage
+                    .save_or_update(token_account)
+                    .await;
+                if is_added {
+                    // add account to the indexes
+                    // add account index in accounts_index_by_mint map
+                    match self.accounts_index_by_mint.entry(mint_index) {
+                        dashmap::mapref::entry::Entry::Occupied(mut occ) => {
+                            occ.get_mut().push_back(token_index);
                         }
-
-                        // add account index in account_by_owner_pubkey map
-                        match self.account_by_owner_pubkey.entry(token_account_owner) {
-                            dashmap::mapref::entry::Entry::Occupied(mut occ) => {
-                                occ.get_mut().push(token_index);
-                            }
-                            dashmap::mapref::entry::Entry::Vacant(v) => {
-                                v.insert(vec![token_index]);
-                            }
+                        dashmap::mapref::entry::Entry::Vacant(v) => {
+                            let mut q = VecDeque::new();
+                            q.push_back(token_index);
+                            v.insert(q);
                         }
                     }
-                };
+
+                    // add account index in account_by_owner_pubkey map
+                    match self.account_by_owner_pubkey.entry(token_account_owner) {
+                        dashmap::mapref::entry::Entry::Occupied(mut occ) => {
+                            occ.get_mut().push(token_index);
+                        }
+                        dashmap::mapref::entry::Entry::Vacant(v) => {
+                            v.insert(vec![token_index]);
+                        }
+                    }
+                }
             }
             TokenProgramAccountType::Mint(mint_data) => {
                 match self.mints_index_by_pubkey.entry(mint_data.pubkey) {
@@ -390,28 +377,25 @@ impl InMemoryTokenStorage {
                     self.accounts_index_by_mint.remove(&index);
                 } else if self.mints_index_by_pubkey.remove(&account_pk).is_some() {
                     // do nothing multisig account is removed
-                } else if let Some((_, index)) = self.account_index_by_pubkey.remove(&account_pk) {
-                    let mut lk = self.token_accounts.write().await;
-                    if let Some(token_account) = lk.get_mut(index as usize) {
-                        token_account.lamports = 0;
-                        token_account.mint = 0;
-                        token_account.owner = Pubkey::default();
-                    }
+                } else {
+                    self.token_accounts_storage.delete(&account_pk).await;
                 }
             }
         }
         true
     }
 
-    pub fn is_token_program_account(&self, account_data: &AccountData) -> bool {
-        self.account_index_by_pubkey
+    pub async fn is_token_program_account(&self, account_data: &AccountData) -> bool {
+        self.mints_index_by_pubkey
             .contains_key(&account_data.pubkey)
-            || self
-                .mints_index_by_pubkey
-                .contains_key(&account_data.pubkey)
             || self.multisigs.contains_key(&account_data.pubkey)
             || account_data.account.owner == spl_token::id()
             || account_data.account.owner == spl_token_2022::id()
+            || self
+                .token_accounts_storage
+                .contains(&account_data.pubkey)
+                .await
+                .is_some()
     }
 
     pub async fn get_account_finalized(&self, account_pk: Pubkey) -> Option<AccountData> {
@@ -433,15 +417,12 @@ impl InMemoryTokenStorage {
             ));
         }
 
-        if let Some(ite) = self.account_index_by_pubkey.get(&account_pk) {
-            let token_account_index = *ite as usize;
-            let lk = self.token_accounts.read().await;
-            let token_account = lk.get(token_account_index).unwrap();
-            if token_account.lamports == 0 {
+        if let Some(token_account) = self.token_accounts_storage.get_by_pubkey(&account_pk).await {
+            if token_account.is_deleted {
                 return None;
             }
             return token_account_to_solana_account(
-                token_account,
+                &token_account,
                 self.finalized_slot.load(Ordering::Relaxed),
                 0,
                 &self.mints_by_index,
@@ -462,13 +443,18 @@ impl InMemoryTokenStorage {
             if let Some(owner) = owner {
                 match self.account_by_owner_pubkey.get(&owner) {
                     Some(token_acc_indexes) => {
-                        let lk = self.token_accounts.read().await;
-                        let indexes = token_acc_indexes.value();
+                        let indexes = token_acc_indexes
+                            .value()
+                            .iter()
+                            .cloned()
+                            .collect::<HashSet<_>>();
                         let mint =
                             mint.map(|pk| *self.mints_index_by_pubkey.get(&pk).unwrap().value());
-                        let token_accounts = indexes
+                        let token_accounts = self
+                            .token_accounts_storage
+                            .get_by_index(indexes)
+                            .await?
                             .iter()
-                            .filter_map(|index| lk.get(*index as usize))
                             .filter(|acc| {
                                 // filter by mint if necessary
                                 if let Some(mint) = mint {
@@ -495,11 +481,16 @@ impl InMemoryTokenStorage {
                 match self.mints_index_by_pubkey.get(&mint) {
                     Some(mint_index) => match self.accounts_index_by_mint.get(mint_index.value()) {
                         Some(token_acc_indexes) => {
-                            let indexes = token_acc_indexes.value();
-                            let lk = self.token_accounts.read().await;
-                            let token_accounts = indexes
+                            let indexes = token_acc_indexes
+                                .value()
+                                .iter()
+                                .cloned()
+                                .collect::<HashSet<u32>>();
+                            let token_accounts = self
+                                .token_accounts_storage
+                                .get_by_index(indexes)
+                                .await?
                                 .iter()
-                                .filter_map(|index| lk.get(*index as usize))
                                 .filter_map(|token_account| {
                                     token_account_to_solana_account(
                                         token_account,
@@ -543,13 +534,13 @@ impl InMemoryTokenStorage {
 struct TokenProgramSnapshot {
     pub mints: Vec<(MintIndex, MintAccount)>,
     pub multisigs: Vec<MultiSig>,
-    pub token_accounts: Vec<TokenAccount>,
+    pub token_accounts: Vec<Vec<u8>>,
 }
 
 #[async_trait]
 impl AccountStorageInterface for InMemoryTokenStorage {
     async fn update_account(&self, account_data: AccountData, commitment: Commitment) -> bool {
-        if !self.is_token_program_account(&account_data) {
+        if !self.is_token_program_account(&account_data).await {
             return false;
         }
 
@@ -584,7 +575,7 @@ impl AccountStorageInterface for InMemoryTokenStorage {
     }
 
     async fn initilize_or_update_account(&self, account_data: AccountData) {
-        if !self.is_token_program_account(&account_data) {
+        if !self.is_token_program_account(&account_data).await {
             return;
         }
 
@@ -759,14 +750,7 @@ impl AccountStorageInterface for InMemoryTokenStorage {
             .collect_vec();
 
         let multisigs = self.multisigs.iter().map(|ite| ite.clone()).collect_vec();
-        let token_accounts = self
-            .token_accounts
-            .read()
-            .await
-            .iter()
-            .filter(|x| x.program == program && x.lamports > 0)
-            .cloned()
-            .collect_vec();
+        let token_accounts = self.token_accounts_storage.create_snapshot(program).await?;
 
         let snapshot = TokenProgramSnapshot {
             mints,
@@ -819,39 +803,32 @@ impl AccountStorageInterface for InMemoryTokenStorage {
             }
         }
 
-        let mut lk = self.token_accounts.write().await;
-        for mut token_account in token_accounts {
+        for token_account_bytes in token_accounts {
             // use the new mint index
+            let mut token_account = TokenAccount::from_bytes(&token_account_bytes);
             let new_mint_index = *mint_mapping.get(&token_account.mint).unwrap();
             token_account.mint = new_mint_index;
-            match self.account_index_by_pubkey.get(&token_account.pubkey) {
-                Some(exists) => {
-                    // pubkey is already present, replace existing
-                    let index = *exists as usize;
-                    *lk.get_mut(index).unwrap() = token_account;
-                }
-                None => {
-                    let index = lk.len() as MintIndex;
-                    let pk = token_account.pubkey;
-                    let owner = token_account.owner;
-                    lk.push_back(token_account);
-                    self.account_index_by_pubkey.insert(pk, index);
-                    match self.account_by_owner_pubkey.get_mut(&owner) {
-                        Some(mut accs) => accs.push(index),
-                        None => {
-                            self.account_by_owner_pubkey.insert(owner, vec![index]);
-                        }
+            let owner = token_account.owner;
+            let (index, is_added) = self
+                .token_accounts_storage
+                .save_or_update(token_account)
+                .await;
+            if is_added {
+                match self.account_by_owner_pubkey.get_mut(&owner) {
+                    Some(mut accs) => accs.push(index),
+                    None => {
+                        self.account_by_owner_pubkey.insert(owner, vec![index]);
                     }
+                }
 
-                    match self.accounts_index_by_mint.get_mut(&new_mint_index) {
-                        Some(mut accs) => {
-                            accs.push_back(index);
-                        }
-                        None => {
-                            let mut q = VecDeque::new();
-                            q.push_back(index);
-                            self.accounts_index_by_mint.insert(new_mint_index, q);
-                        }
+                match self.accounts_index_by_mint.get_mut(&new_mint_index) {
+                    Some(mut accs) => {
+                        accs.push_back(index);
+                    }
+                    None => {
+                        let mut q = VecDeque::new();
+                        q.push_back(index);
+                        self.accounts_index_by_mint.insert(new_mint_index, q);
                     }
                 }
             }

token_account_storage/src/lib.rs

@@ -1,7 +1,9 @@
 use solana_sdk::pubkey::Pubkey;
 
 pub mod account_types;
-pub mod inmemory_token_storage;
+pub mod inmemory_token_account_storage; // to store only token accounts
+pub mod inmemory_token_storage; // to store all the accounts of token program
+pub mod token_account_storage_interface;
 pub mod token_program_utils;
 
 pub const TOKEN_PROGRAM_ID: Pubkey = spl_token::id();

token_account_storage/src/token_account_storage_interface.rs

@@ -0,0 +1,26 @@
+use std::collections::HashSet;
+
+use async_trait::async_trait;
+use lite_account_manager_common::account_store_interface::AccountLoadingError;
+use solana_sdk::pubkey::Pubkey;
+
+use crate::account_types::{Program, TokenAccount, TokenAccountIndex};
+
+// Interface to store token accounts
+#[async_trait]
+pub trait TokenAccountStorageInterface: Sync + Send {
+    async fn contains(&self, pubkey: &Pubkey) -> Option<TokenAccountIndex>;
+
+    async fn save_or_update(&self, token_account: TokenAccount) -> (TokenAccountIndex, bool);
+
+    async fn get_by_index(
+        &self,
+        indexes: HashSet<TokenAccountIndex>,
+    ) -> Result<Vec<TokenAccount>, AccountLoadingError>;
+
+    async fn get_by_pubkey(&self, pubkey: &Pubkey) -> Option<TokenAccount>;
+
+    async fn delete(&self, pubkey: &Pubkey);
+
+    async fn create_snapshot(&self, program: Program) -> Result<Vec<Vec<u8>>, AccountLoadingError>;
+}

token_account_storage/src/token_program_utils.rs

@@ -14,6 +14,7 @@ use solana_sdk::{
     program_option::COption,
     program_pack::Pack,
     pubkey::{Pubkey, PUBKEY_BYTES},
+    rent::Rent,
 };
 use spl_token::instruction::MAX_SIGNERS;
 
@@ -116,7 +117,7 @@ pub fn get_token_program_account_type(
                 );
                 Ok(TokenProgramAccountType::TokenAccount(TokenAccount {
                     program: crate::account_types::Program::Token2022Program,
-                    lamports: account_data.account.lamports,
+                    is_deleted: account_data.account.lamports > 0,
                     pubkey: account_data.pubkey,
                     mint: mint_index, // mint should be set inmemory account
                     owner: token_account.owner,
@@ -191,7 +192,7 @@ pub fn get_token_program_account_type(
             Ok(TokenProgramAccountType::TokenAccount(TokenAccount {
                 program: crate::account_types::Program::TokenProgram,
                 pubkey: account_data.pubkey,
-                lamports: account_data.account.lamports,
+                is_deleted: account_data.account.lamports > 0,
                 mint: mint_index, // mint should be set inmemory account
                 owner: token_account.owner,
                 amount: token_account.amount,
@@ -245,7 +246,7 @@ pub fn token_account_to_solana_account(
     write_version: u64,
     mints_by_index: &Arc<DashMap<MintIndex, MintAccount>>,
 ) -> Option<AccountData> {
-    if token_account.lamports == 0 {
+    if token_account.is_deleted {
         return None;
     }
     let (delegate, delegated_amount) = token_account.delegate.unwrap_or_default();
@@ -311,8 +312,13 @@ pub fn token_account_to_solana_account(
         }
     };
     let data_length = data.len() as u64;
+    let rent = Rent::default();
     let account = Arc::new(Account {
-        lamports: token_account.lamports,
+        lamports: if token_account.is_deleted {
+            0
+        } else {
+            rent.minimum_balance(data_length as usize)
+        },
         data: lite_account_manager_common::account_data::Data::Uncompressed(data),
         owner: match token_account.program {
             Program::TokenProgram => spl_token::id(),

token_account_storage/tests/token_program_gpa_tests.rs

@@ -1,4 +1,9 @@
-use lite_token_account_storage::inmemory_token_storage::InMemoryTokenStorage;
+use std::sync::Arc;
+
+use lite_token_account_storage::{
+    inmemory_token_account_storage::InmemoryTokenAccountStorage,
+    inmemory_token_storage::InMemoryTokenStorage,
+};
 use solana_sdk::pubkey::Pubkey;
 mod utils;
 
@@ -10,7 +15,8 @@ use lite_account_manager_common::{
 
 #[tokio::test]
 pub async fn test_gpa_token_account() {
-    let token_store = InMemoryTokenStorage::new();
+    let inmemory_token_storage = Arc::new(InmemoryTokenAccountStorage::default());
+    let token_store = InMemoryTokenStorage::new(inmemory_token_storage);
 
     let mint_1: Pubkey = Pubkey::new_unique();
     let mint_creation_params = utils::MintCreationParams::create_default(100);

token_account_storage/tests/token_program_tests.rs

@@ -6,7 +6,10 @@ use lite_account_manager_common::{
     commitment::Commitment,
     slot_info::SlotInfo,
 };
-use lite_token_account_storage::inmemory_token_storage::InMemoryTokenStorage;
+use lite_token_account_storage::{
+    inmemory_token_account_storage::InmemoryTokenAccountStorage,
+    inmemory_token_storage::InMemoryTokenStorage,
+};
 use solana_sdk::pubkey::Pubkey;
 
 mod utils;
@@ -14,7 +17,8 @@ mod utils;
 #[tokio::test]
 pub async fn test_saving_and_loading_token_account() {
     tracing_subscriber::fmt::init();
-    let token_store = InMemoryTokenStorage::new();
+    let inmemory_token_storage = Arc::new(InmemoryTokenAccountStorage::default());
+    let token_store = InMemoryTokenStorage::new(inmemory_token_storage);
 
     let mint: Pubkey = Pubkey::new_unique();
     let mint_creation_params = utils::MintCreationParams::create_default(100);