blockworks-foundation
/
solana-lite-account-manager
mirror of https://github.com/blockworks-foundation/solana-lite-account-manager.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

creating partial pubkey to save space used for keys

This commit is contained in:
godmodegalactus 2024-07-09 14:22:48 +02:00
parent 211d32e196
commit d71878da64
No known key found for this signature in database
GPG Key ID: 22DA4A30887FDA3C
5 changed files with 119 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
common/src/lib.rs
View File

@ -5,5 +5,6 @@ pub mod account_store_interface;
pub mod accounts_source_interface;
pub mod commitment;
pub mod mutable_filter_store;
pub mod pubkey_container_utils;
pub mod simple_filter_store;
pub mod slot_info;

22
common/src/pubkey_container_utils.rs Normal file
View File

@ -0,0 +1,22 @@
use solana_sdk::pubkey::Pubkey;
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
pub struct PartialPubkey<const SIZE: usize> {
partial_pubkey_bytes: [u8; SIZE],
}
impl<const SIZE: usize> From<Pubkey> for PartialPubkey<SIZE> {
fn from(value: Pubkey) -> Self {
Self {
partial_pubkey_bytes: value.to_bytes()[0..SIZE].try_into().unwrap(),
}
}
}
impl<const SIZE: usize> From<&Pubkey> for PartialPubkey<SIZE> {
fn from(value: &Pubkey) -> Self {
Self {
partial_pubkey_bytes: value.to_bytes()[0..SIZE].try_into().unwrap(),
}
}
}

20
token_account_storage/src/account_types.rs
View File

@ -249,6 +249,21 @@ impl TokenAccount {
additional_data,
}
}
pub fn get_pubkey_from_binary(binary: &[u8]) -> Pubkey {
let pubkey_array: [u8; 32] = binary[2..34].try_into().unwrap();
Pubkey::new_from_array(pubkey_array)
}
pub fn get_owner_from_binary(binary: &[u8]) -> Pubkey {
let pubkey_array: [u8; 32] = binary[34..66].try_into().unwrap();
Pubkey::new_from_array(pubkey_array)
}
pub fn get_mint_index_from_binary(binary: &[u8]) -> u32 {
let bytes: [u8; 4] = binary[66..66 + 4].try_into().unwrap();
u32::from_le_bytes(bytes)
}
}
#[cfg(test)]
@ -317,8 +332,11 @@ mod test {
None
},
};
let ser = acc.to_bytes();
assert_eq!(TokenAccount::get_pubkey_from_binary(&ser), acc.pubkey);
assert_eq!(TokenAccount::get_owner_from_binary(&ser), acc.owner);
assert_eq!(TokenAccount::get_mint_index_from_binary(&ser), acc.mint);
let deser = TokenAccount::from_bytes(&ser);
assert_eq!(deser, acc);
}

82
token_account_storage/src/inmemory_token_account_storage.rs
View File

@ -6,7 +6,9 @@ use std::{
use async_trait::async_trait;
use dashmap::DashMap;
use itertools::Itertools;
use lite_account_manager_common::account_store_interface::AccountLoadingError;
use lite_account_manager_common::{
account_store_interface::AccountLoadingError, pubkey_container_utils::PartialPubkey,
};
use solana_sdk::pubkey::Pubkey;
use tokio::sync::RwLock;
@ -15,26 +17,61 @@ use crate::{
token_account_storage_interface::TokenAccountStorageInterface,
};
const PARTIAL_PUBKEY_SIZE: usize = 8;
type InmemoryPubkey = PartialPubkey<PARTIAL_PUBKEY_SIZE>;
#[derive(Default)]
pub struct InmemoryTokenAccountStorage {
pubkey_to_index: Arc<DashMap<Pubkey, TokenAccountIndex>>,
// use only first 8 bytes of pubkey to create an index
pubkey_to_index: Arc<DashMap<InmemoryPubkey, Vec<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())
match self.pubkey_to_index.get(&pubkey.into()) {
Some(x) => {
let value = x.value();
if value.len() > 1 {
let lk = self.token_accounts.read().await;
for index in value {
if TokenAccount::get_pubkey_from_binary(lk.get(*index as usize).unwrap())
== *pubkey
{
return Some(*index);
}
}
None
} else {
Some(value[0])
}
}
None => None,
}
}
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) => {
match self.pubkey_to_index.entry(token_account.pubkey.into()) {
dashmap::mapref::entry::Entry::Occupied(mut 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();
{
let token_indexes = occ.get();
// update existing token account
for token_index in token_indexes {
let binary = write_lk.get_mut(*token_index as usize).unwrap();
if TokenAccount::get_pubkey_from_binary(binary) == token_account.pubkey {
*write_lk.get_mut(*token_index as usize).unwrap() =
token_account.to_bytes();
return (*token_index as TokenAccountIndex, false);
}
}
}
let token_indexes = occ.get_mut();
let token_index = write_lk.len() as TokenAccountIndex;
write_lk.push_back(token_account.to_bytes());
token_indexes.push(token_index);
(token_index as TokenAccountIndex, false)
}
dashmap::mapref::entry::Entry::Vacant(v) => {
@ -42,7 +79,7 @@ impl TokenAccountStorageInterface for InmemoryTokenAccountStorage {
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);
v.insert(vec![token_index as TokenAccountIndex]);
drop(write_lk);
(token_index, true)
}
@ -63,24 +100,25 @@ impl TokenAccountStorageInterface for InmemoryTokenAccountStorage {
}
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) => {
if let Some(acc) = self.pubkey_to_index.get(&pubkey.into()) {
let indexes = acc.value();
let lk = self.token_accounts.read().await;
for index in indexes {
let binary = lk.get(*index as usize).unwrap();
if TokenAccount::get_pubkey_from_binary(binary) == *pubkey {
let token_account: TokenAccount = TokenAccount::from_bytes(binary);
Some(token_account)
return Some(token_account);
}
None => 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 acc_to_remove = self.pubkey_to_index.remove(&pubkey.into());
if let Some((_, indexes)) = acc_to_remove {
let mut write_lk = self.token_accounts.write().await;
let deleted_account = TokenAccount {
program: Program::TokenProgram,
@ -96,7 +134,13 @@ impl TokenAccountStorageInterface for InmemoryTokenAccountStorage {
additional_data: None,
}
.to_bytes();
*write_lk.get_mut(index as usize).unwrap() = deleted_account;
for index in indexes {
let binary = write_lk.get_mut(index as usize).unwrap();
if TokenAccount::get_pubkey_from_binary(binary) == *pubkey {
*binary = deleted_account;
return;
}
}
}
}

21
token_account_storage/src/inmemory_token_storage.rs
View File

@ -14,6 +14,7 @@ use lite_account_manager_common::{
account_filter::AccountFilterType,
account_store_interface::{AccountLoadingError, AccountStorageInterface},
commitment::Commitment,
pubkey_container_utils::PartialPubkey,
slot_info::SlotInfo,
};
use serde::{Deserialize, Serialize};
@ -31,6 +32,8 @@ use crate::{
},
};
const PARTIAL_PUBKEY_SIZE: usize = 6;
type InmemoryPubkey = PartialPubkey<PARTIAL_PUBKEY_SIZE>;
#[derive(Clone)]
struct ProcessedAccount {
pub processed_account: TokenProgramAccountType,
@ -294,7 +297,7 @@ 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_by_owner_pubkey: Arc<DashMap<Pubkey, Vec<TokenAccountIndex>>>,
account_by_owner_pubkey: Arc<DashMap<InmemoryPubkey, Vec<TokenAccountIndex>>>,
token_accounts_storage: Arc<dyn TokenAccountStorageInterface>,
processed_storage: ProcessedAccountStore,
confirmed_slot: Arc<AtomicU64>,
@ -344,7 +347,10 @@ impl InMemoryTokenStorage {
}
// add account index in account_by_owner_pubkey map
match self.account_by_owner_pubkey.entry(token_account_owner) {
match self
.account_by_owner_pubkey
.entry(token_account_owner.into())
{
dashmap::mapref::entry::Entry::Occupied(mut occ) => {
occ.get_mut().push(token_index);
}
@ -441,7 +447,7 @@ impl InMemoryTokenStorage {
if let Some(account_filters) = account_filters {
let (owner, mint) = get_spl_token_owner_mint_filter(&program_pubkey, &account_filters);
if let Some(owner) = owner {
match self.account_by_owner_pubkey.get(&owner) {
match self.account_by_owner_pubkey.get(&owner.into()) {
Some(token_acc_indexes) => {
let indexes = token_acc_indexes
.value()
@ -458,9 +464,9 @@ impl InMemoryTokenStorage {
.filter(|acc| {
// filter by mint if necessary
if let Some(mint) = mint {
acc.mint == mint
acc.mint == mint && acc.owner == owner
} else {
true
acc.owner == owner
}
})
.filter_map(|token_account| {
@ -814,10 +820,11 @@ impl AccountStorageInterface for InMemoryTokenStorage {
.save_or_update(token_account)
.await;
if is_added {
match self.account_by_owner_pubkey.get_mut(&owner) {
match self.account_by_owner_pubkey.get_mut(&owner.into()) {
Some(mut accs) => accs.push(index),
None => {
self.account_by_owner_pubkey.insert(owner, vec![index]);
self.account_by_owner_pubkey
.insert(owner.into(), vec![index]);
}
}