100 lines
3.3 KiB
Rust
100 lines
3.3 KiB
Rust
use anyhow::{bail, Context};
|
|
use async_trait::async_trait;
|
|
use itertools::Itertools;
|
|
use solana_client::nonblocking::rpc_client::RpcClient;
|
|
use solana_lite_rpc_core::{
|
|
structures::leader_data::LeaderData, traits::leaders_fetcher_interface::LeaderFetcherInterface,
|
|
};
|
|
use std::{collections::VecDeque, sync::Arc};
|
|
use tokio::sync::RwLock;
|
|
|
|
// Stores leaders for slots from older to newer in leader schedule
|
|
// regularly removed old leaders and adds new ones
|
|
pub struct JsonRpcLeaderGetter {
|
|
rpc_client: Arc<RpcClient>,
|
|
leader_schedule: RwLock<VecDeque<LeaderData>>,
|
|
leaders_to_cache_count: u64,
|
|
old_slot_leader_cache_size: u64,
|
|
}
|
|
|
|
impl JsonRpcLeaderGetter {
|
|
pub fn new(
|
|
rpc_client: Arc<RpcClient>,
|
|
leaders_to_cache_count: u64,
|
|
old_slot_leader_cache_size: u64,
|
|
) -> Self {
|
|
Self {
|
|
rpc_client,
|
|
leader_schedule: RwLock::new(VecDeque::new()),
|
|
leaders_to_cache_count,
|
|
old_slot_leader_cache_size,
|
|
}
|
|
}
|
|
|
|
pub async fn update_leader_schedule(&self, slot: u64) -> anyhow::Result<()> {
|
|
// remove old elements
|
|
let mut leader_queue = self.leader_schedule.write().await;
|
|
// remove old leaders
|
|
while leader_queue.front().map_or(slot, |x| x.leader_slot)
|
|
< slot.saturating_sub(self.old_slot_leader_cache_size)
|
|
{
|
|
leader_queue.pop_front();
|
|
}
|
|
|
|
let last_slot_needed = slot + self.leaders_to_cache_count;
|
|
let first_slot_to_fetch = leader_queue.back().map_or(slot, |x| x.leader_slot + 1);
|
|
|
|
if last_slot_needed >= first_slot_to_fetch {
|
|
let leaders = self
|
|
.rpc_client
|
|
.get_slot_leaders(
|
|
first_slot_to_fetch,
|
|
last_slot_needed.saturating_sub(first_slot_to_fetch),
|
|
)
|
|
.await
|
|
.context("failed to get slot leaders")?;
|
|
|
|
for leader_slot in first_slot_to_fetch..last_slot_needed {
|
|
let current_leader = (leader_slot.saturating_sub(first_slot_to_fetch)) as usize;
|
|
let pubkey = leaders[current_leader];
|
|
leader_queue.push_back(LeaderData {
|
|
leader_slot,
|
|
pubkey,
|
|
});
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl LeaderFetcherInterface for JsonRpcLeaderGetter {
|
|
async fn get_slot_leaders(
|
|
&self,
|
|
from: solana_sdk::slot_history::Slot,
|
|
to: solana_sdk::slot_history::Slot,
|
|
) -> anyhow::Result<Vec<LeaderData>> {
|
|
if to <= from || to.saturating_sub(from) > self.leaders_to_cache_count {
|
|
bail!("invalid arguments for get_slot_leaders");
|
|
}
|
|
let schedule = self.leader_schedule.read().await;
|
|
|
|
let schedule = if schedule.is_empty()
|
|
|| schedule.front().unwrap().leader_slot > from
|
|
|| schedule.back().unwrap().leader_slot < to
|
|
{
|
|
// rebuilding the leader schedule
|
|
drop(schedule);
|
|
self.update_leader_schedule(from).await?;
|
|
self.leader_schedule.read().await
|
|
} else {
|
|
schedule
|
|
};
|
|
Ok(schedule
|
|
.iter()
|
|
.filter(|x| x.leader_slot >= from && x.leader_slot <= to)
|
|
.cloned()
|
|
.collect_vec())
|
|
}
|
|
}
|