use solana_sdk::pubkey::Pubkey;
use solana_sdk::signature::Signature;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use dashmap::DashMap;
use std::collections::BTreeSet;
const MAX_SIGNATURES: usize = 1000;
const CLEANUP_BATCH_SIZE: usize = 100;
#[derive(Default)]
struct SignatureAddresses {
dev_addresses: BTreeSet<Pubkey>,
bonk_dev_addresses: BTreeSet<Pubkey>,
}
pub struct GlobalState {
signature_data: DashMap<Signature, SignatureAddresses>,
signature_count: AtomicUsize,
generation: AtomicU64,
}
impl GlobalState {
pub fn new() -> Self {
Self {
signature_data: DashMap::new(),
signature_count: AtomicUsize::new(0),
generation: AtomicU64::new(0),
}
}
fn maybe_cleanup(&self) {
let current_count = self.signature_count.load(Ordering::Relaxed);
if current_count <= MAX_SIGNATURES {
return;
}
let gen = self.generation.load(Ordering::Relaxed);
if self.generation.compare_exchange_weak(gen, gen + 1, Ordering::Acquire, Ordering::Relaxed).is_err() {
return; }
let mut signatures_to_remove: Vec<Signature> = self.signature_data.iter()
.map(|entry| *entry.key())
.collect();
if signatures_to_remove.len() <= MAX_SIGNATURES {
return; }
signatures_to_remove.truncate(CLEANUP_BATCH_SIZE);
for signature in signatures_to_remove {
self.signature_data.remove(&signature);
self.signature_count.fetch_sub(1, Ordering::Relaxed);
}
}
pub fn add_dev_address(&self, signature: &Signature, address: Pubkey) {
self.maybe_cleanup();
self.signature_data.entry(*signature)
.and_modify(|addresses| {
addresses.dev_addresses.insert(address);
})
.or_insert_with(|| {
self.signature_count.fetch_add(1, Ordering::Relaxed);
let mut sig_addr = SignatureAddresses::default();
sig_addr.dev_addresses.insert(address);
sig_addr
});
}
pub fn add_bonk_dev_address(&self, signature: &Signature, address: Pubkey) {
self.maybe_cleanup();
self.signature_data.entry(*signature)
.and_modify(|addresses| {
addresses.bonk_dev_addresses.insert(address);
})
.or_insert_with(|| {
self.signature_count.fetch_add(1, Ordering::Relaxed);
let mut sig_addr = SignatureAddresses::default();
sig_addr.bonk_dev_addresses.insert(address);
sig_addr
});
}
pub fn is_dev_address_in_signature(&self, signature: &Signature, address: &Pubkey) -> bool {
self.signature_data.get(signature)
.map(|entry| entry.dev_addresses.contains(address))
.unwrap_or(false)
}
pub fn is_bonk_dev_address_in_signature(&self, signature: &Signature, address: &Pubkey) -> bool {
self.signature_data.get(signature)
.map(|entry| entry.bonk_dev_addresses.contains(address))
.unwrap_or(false)
}
pub fn is_dev_address(&self, address: &Pubkey) -> bool {
self.signature_data.iter().any(|entry| entry.dev_addresses.contains(address))
}
pub fn is_bonk_dev_address(&self, address: &Pubkey) -> bool {
self.signature_data.iter().any(|entry| entry.bonk_dev_addresses.contains(address))
}
pub fn get_dev_addresses(&self) -> Vec<Pubkey> {
let mut all_addresses = BTreeSet::new();
for entry in self.signature_data.iter() {
for addr in &entry.dev_addresses {
all_addresses.insert(*addr);
}
}
all_addresses.into_iter().collect()
}
pub fn get_bonk_dev_addresses(&self) -> Vec<Pubkey> {
let mut all_addresses = BTreeSet::new();
for entry in self.signature_data.iter() {
for addr in &entry.bonk_dev_addresses {
all_addresses.insert(*addr);
}
}
all_addresses.into_iter().collect()
}
pub fn get_dev_addresses_for_signature(&self, signature: &Signature) -> Vec<Pubkey> {
self.signature_data.get(signature)
.map(|entry| entry.dev_addresses.iter().copied().collect())
.unwrap_or_default()
}
pub fn get_bonk_dev_addresses_for_signature(&self, signature: &Signature) -> Vec<Pubkey> {
self.signature_data.get(signature)
.map(|entry| entry.bonk_dev_addresses.iter().copied().collect())
.unwrap_or_default()
}
pub fn get_signature_count(&self) -> usize {
self.signature_count.load(Ordering::Relaxed)
}
pub fn clear_all_data(&self) {
self.signature_data.clear();
self.signature_count.store(0, Ordering::Relaxed);
self.generation.store(0, Ordering::Relaxed);
}
}
impl Default for GlobalState {
fn default() -> Self {
Self::new()
}
}
static GLOBAL_STATE: once_cell::sync::Lazy<GlobalState> =
once_cell::sync::Lazy::new(GlobalState::new);
pub fn get_global_state() -> &'static GlobalState {
&GLOBAL_STATE
}
pub fn add_dev_address(signature: &Signature, address: Pubkey) {
get_global_state().add_dev_address(signature, address);
}
pub fn is_dev_address(address: &Pubkey) -> bool {
get_global_state().is_dev_address(address)
}
pub fn add_bonk_dev_address(signature: &Signature, address: Pubkey) {
get_global_state().add_bonk_dev_address(signature, address);
}
pub fn is_bonk_dev_address(address: &Pubkey) -> bool {
get_global_state().is_bonk_dev_address(address)
}
pub fn get_dev_addresses() -> Vec<Pubkey> {
get_global_state().get_dev_addresses()
}
pub fn get_bonk_dev_addresses() -> Vec<Pubkey> {
get_global_state().get_bonk_dev_addresses()
}
pub fn get_dev_addresses_for_signature(signature: &Signature) -> Vec<Pubkey> {
get_global_state().get_dev_addresses_for_signature(signature)
}
pub fn get_bonk_dev_addresses_for_signature(signature: &Signature) -> Vec<Pubkey> {
get_global_state().get_bonk_dev_addresses_for_signature(signature)
}
pub fn get_signature_count() -> usize {
get_global_state().get_signature_count()
}
pub fn is_dev_address_in_signature(signature: &Signature, address: &Pubkey) -> bool {
get_global_state().is_dev_address_in_signature(signature, address)
}
pub fn is_bonk_dev_address_in_signature(signature: &Signature, address: &Pubkey) -> bool {
get_global_state().is_bonk_dev_address_in_signature(signature, address)
}