pub mod cache;
pub mod error;
pub mod fetcher;
pub mod provider;
use std::borrow::Cow;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use chia_protocol::{Bytes32, CoinStateFilters, SpendBundle};
use dig_chainsource_interface::{ProviderId, ProviderInfo, ProviderKind};
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use crate::peer::connect::PeerOrigin;
use crate::peer::frames::{FrameSource, FrameSubscription, PoolFrame, SourcedFrame};
use crate::peer::PeerBackend;
use cache::CoinStateCache;
use error::LightClientError;
use fetcher::{CoinStateFetcher, PooledFetcher};
pub use provider::LightClientProvider;
pub const DEFAULT_PROVIDER_PRIORITY: i32 = 20;
const FRAME_BUFFER: usize = 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SubmitOutcome {
Accepted,
Pending,
Failed,
Unknown(u8),
}
impl SubmitOutcome {
fn from_status(status: u8) -> Self {
match status {
1 => SubmitOutcome::Accepted,
2 => SubmitOutcome::Pending,
3 => SubmitOutcome::Failed,
other => SubmitOutcome::Unknown(other),
}
}
pub fn is_accepted(self) -> bool {
matches!(self, SubmitOutcome::Accepted | SubmitOutcome::Pending)
}
}
pub struct ChiaLightClient {
fetcher: PooledFetcher,
cache: Arc<RwLock<CoinStateCache>>,
rearm_needed: Arc<AtomicBool>,
drive: Option<JoinHandle<()>>,
}
impl ChiaLightClient {
pub async fn new(backend: Arc<PeerBackend>, request_timeout: Duration) -> Self {
let cache = Arc::new(RwLock::new(CoinStateCache::new()));
let fetcher = PooledFetcher::new(backend.clone(), request_timeout);
let rearm_needed = Arc::new(AtomicBool::new(false));
let subscription = backend.subscribe_frames(FRAME_BUFFER).await;
let drive = spawn_drive_loop(
subscription,
cache.clone(),
fetcher.clone(),
rearm_needed.clone(),
);
Self {
fetcher,
cache,
rearm_needed,
drive: Some(drive),
}
}
pub async fn subscribe_coins(&self, coin_ids: Vec<Bytes32>) -> Result<(), LightClientError> {
let states = self.fetcher.coin_states(coin_ids.clone(), true).await?;
let mut cache = self.cache.write().await;
cache.track_coins(coin_ids);
cache.seed(states);
Ok(())
}
pub async fn subscribe_puzzle_hashes(
&self,
puzzle_hashes: Vec<Bytes32>,
filters: CoinStateFilters,
) -> Result<(), LightClientError> {
let states = self
.fetcher
.puzzle_states(puzzle_hashes.clone(), filters, true)
.await?;
let mut cache = self.cache.write().await;
cache.track_puzzle_hashes(puzzle_hashes);
cache.seed(states);
Ok(())
}
pub async fn submit_spend(
&self,
bundle: SpendBundle,
) -> Result<SubmitOutcome, LightClientError> {
let status = self.fetcher.send_transaction(bundle).await?;
Ok(SubmitOutcome::from_status(status))
}
pub async fn peak(&self) -> Option<(u32, Bytes32)> {
self.cache.read().await.peak()
}
pub async fn unsubscribe_coins(&self, coin_ids: Vec<Bytes32>) -> Result<(), LightClientError> {
self.fetcher
.remove_coin_subscriptions(coin_ids.clone())
.await?;
self.cache.write().await.untrack_coins(&coin_ids);
Ok(())
}
pub fn needs_rearm(&self) -> bool {
self.rearm_needed.load(Ordering::Acquire)
}
pub async fn reconnect(&self) -> Result<(), LightClientError> {
let (coins, puzzle_hashes) = {
let cache = self.cache.read().await;
(cache.subscribed_coins(), cache.subscribed_puzzle_hashes())
};
if !coins.is_empty() {
self.subscribe_coins(coins).await?;
}
if !puzzle_hashes.is_empty() {
self.subscribe_puzzle_hashes(puzzle_hashes, all_coin_states())
.await?;
}
self.rearm_needed.store(false, Ordering::Release);
Ok(())
}
pub async fn as_chain_source_provider(
&self,
handle: tokio::runtime::Handle,
) -> LightClientProvider {
LightClientProvider::new(
Arc::new(self.fetcher.clone()),
self.cache.clone(),
handle,
self.provider_info().await,
)
}
pub async fn provider_info(&self) -> ProviderInfo {
let kind = match self.fetcher.current_anchor().await.map(|a| a.origin) {
Some(PeerOrigin::Priority) => ProviderKind::LocalNode,
Some(PeerOrigin::Discovered) | None => ProviderKind::Custom,
};
ProviderInfo {
id: ProviderId(Cow::Borrowed("chia-query-light-client")),
kind,
priority: DEFAULT_PROVIDER_PRIORITY,
trustless: false,
}
}
}
impl Drop for ChiaLightClient {
fn drop(&mut self) {
if let Some(handle) = self.drive.take() {
handle.abort();
}
}
}
fn all_coin_states() -> CoinStateFilters {
CoinStateFilters {
include_spent: true,
include_unspent: true,
include_hinted: true,
min_amount: 0,
}
}
fn spawn_drive_loop(
mut subscription: FrameSubscription,
cache: Arc<RwLock<CoinStateCache>>,
fetcher: PooledFetcher,
rearm_needed: Arc<AtomicBool>,
) -> JoinHandle<()> {
tokio::spawn(async move {
while let Some(sourced) = subscription.recv().await {
if !follows(fetcher.anchor_address().await, sourced.source) {
continue;
}
let address = sourced.source.address;
if let AfterFrame::Resubscribe = apply_frame(&cache, sourced).await {
fetcher.release_anchor(address).await;
rearm_needed.store(true, Ordering::Release);
}
}
rearm_needed.store(true, Ordering::Release);
})
}
fn follows(anchor: Option<SocketAddr>, source: FrameSource) -> bool {
anchor.is_some_and(|address| address == source.address)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AfterFrame {
Continue,
Resubscribe,
}
async fn apply_frame(cache: &RwLock<CoinStateCache>, sourced: SourcedFrame) -> AfterFrame {
match sourced.frame {
PoolFrame::Reset => AfterFrame::Resubscribe,
PoolFrame::Peak {
height,
header_hash,
} => {
cache.write().await.set_peak(height, header_hash);
AfterFrame::Continue
}
PoolFrame::CoinStates {
height,
fork_height,
peak_hash,
items,
} => {
let spent: Vec<Bytes32> = items
.iter()
.filter(|state| state.spent_height.is_some())
.map(|state| state.coin.coin_id())
.collect();
let mut cache = cache.write().await;
cache.apply_update(&items, height, fork_height, peak_hash);
cache.untrack_coins(&spent);
AfterFrame::Continue
}
PoolFrame::SessionEnded { reason } => {
log::debug!(
"light-client session {:?} ended: {reason:?}",
sourced.source.session
);
AfterFrame::Resubscribe
}
}
}
#[cfg(test)]
mod tests;