Skip to content

Commit

Permalink
feat(electrum)!: use new sync/full-scan structs for ElectrumExt
Browse files Browse the repository at this point in the history 
`ElectrumResultExt` trait is also introduced that adds methods which can
convert the `Anchor` type for the update `TxGraph`.

Examples and tests are updated to use the new `ElectrumExt` API.
  • Loading branch information
@evanlinjin
evanlinjin committed May 4, 2024
1 parent 7217a88 commit a39a7e9
Show file tree
Hide file tree
Showing 5 changed files with 218 additions and 294 deletions.
2 changes: 1 addition & 1 deletion crates/electrum/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ readme = "README.md"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
bdk_chain = { path = "../chain", version = "0.13.0", default-features = false }
bdk_chain = { path = "../chain", version = "0.13.0" }
electrum-client = { version = "0.19" }
#rustls = { version = "=0.21.1", optional = true, features = ["dangerous_configuration"] }

Expand Down
239 changes: 111 additions & 128 deletions crates/electrum/src/electrum_ext.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,87 +2,17 @@ use bdk_chain::{
bitcoin::{OutPoint, ScriptBuf, Transaction, Txid},
collections::{BTreeMap, HashMap, HashSet},
local_chain::CheckPoint,
spk_client::{FullScanRequest, FullScanResult, SyncRequest, SyncResult, TxCache},
tx_graph::TxGraph,
BlockId, ConfirmationHeightAnchor, ConfirmationTimeHeightAnchor,
};
use core::{fmt::Debug, str::FromStr};
use core::str::FromStr;
use electrum_client::{ElectrumApi, Error, HeaderNotification};
use std::sync::Arc;

/// We include a chain suffix of a certain length for the purpose of robustness.
const CHAIN_SUFFIX_LENGTH: u32 = 8;

/// Type that maintains a cache of [`Arc`]-wrapped transactions.
pub type TxCache = HashMap<Txid, Arc<Transaction>>;

/// Combination of chain and transactions updates from electrum
///
/// We have to update the chain and the txids at the same time since we anchor the txids to
/// the same chain tip that we check before and after we gather the txids.
#[derive(Debug)]
pub struct ElectrumUpdate<A = ConfirmationHeightAnchor> {
/// Chain update
pub chain_update: CheckPoint,
/// Tracks electrum updates in TxGraph
pub graph_update: TxGraph<A>,
}

impl<A: Clone + Ord> ElectrumUpdate<A> {
/// Transform the [`ElectrumUpdate`] to have [`bdk_chain::Anchor`]s of another type.
///
/// Refer to [`TxGraph::map_anchors`].
pub fn map_anchors<A2: Clone + Ord, F>(self, f: F) -> ElectrumUpdate<A2>
where
F: FnMut(A) -> A2,
{
ElectrumUpdate {
chain_update: self.chain_update,
graph_update: self.graph_update.map_anchors(f),
}
}
}

impl ElectrumUpdate {
/// Transforms the [`TxGraph`]'s [`bdk_chain::Anchor`] type to [`ConfirmationTimeHeightAnchor`].
pub fn into_confirmation_time_update(
self,
client: &impl ElectrumApi,
) -> Result<ElectrumUpdate<ConfirmationTimeHeightAnchor>, Error> {
let relevant_heights = self
.graph_update
.all_anchors()
.iter()
.map(|(a, _)| a.confirmation_height)
.collect::<HashSet<_>>();

let height_to_time = relevant_heights
.clone()
.into_iter()
.zip(
client
.batch_block_header(relevant_heights)?
.into_iter()
.map(|bh| bh.time as u64),
)
.collect::<HashMap<u32, u64>>();

let chain_update = self.chain_update;
let graph_update =
self.graph_update
.clone()
.map_anchors(|a| ConfirmationTimeHeightAnchor {
anchor_block: a.anchor_block,
confirmation_height: a.confirmation_height,
confirmation_time: height_to_time[&a.confirmation_height],
});

Ok(ElectrumUpdate {
chain_update,
graph_update,
})
}
}

/// Trait to extend [`electrum_client::Client`] functionality.
pub trait ElectrumExt {
/// Full scan the keychain scripts specified with the blockchain (via an Electrum client) and
Expand All @@ -97,12 +27,10 @@ pub trait ElectrumExt {
/// single batch request.
fn full_scan<K: Ord + Clone>(
&self,
tx_cache: &mut TxCache,
prev_tip: CheckPoint,
keychain_spks: BTreeMap<K, impl IntoIterator<Item = (u32, ScriptBuf)>>,
request: FullScanRequest<K>,
stop_gap: usize,
batch_size: usize,
) -> Result<(ElectrumUpdate, BTreeMap<K, u32>), Error>;
) -> Result<FullScanResult<K, ConfirmationHeightAnchor>, Error>;

/// Sync a set of scripts with the blockchain (via an Electrum client) for the data specified
/// and returns updates for [`bdk_chain`] data structures.
Expand All @@ -123,28 +51,19 @@ pub trait ElectrumExt {
/// [`full_scan`]: ElectrumExt::full_scan
fn sync(
&self,
tx_cache: &mut TxCache,
prev_tip: CheckPoint,
misc_spks: impl IntoIterator<Item = ScriptBuf>,
txids: impl IntoIterator<Item = Txid>,
outpoints: impl IntoIterator<Item = OutPoint>,
request: SyncRequest,
batch_size: usize,
) -> Result<ElectrumUpdate, Error>;
) -> Result<SyncResult<ConfirmationHeightAnchor>, Error>;
}

impl<E: ElectrumApi> ElectrumExt for E {
fn full_scan<K: Ord + Clone>(
&self,
tx_cache: &mut TxCache,
prev_tip: CheckPoint,
keychain_spks: BTreeMap<K, impl IntoIterator<Item = (u32, ScriptBuf)>>,
mut request: FullScanRequest<K>,
stop_gap: usize,
batch_size: usize,
) -> Result<(ElectrumUpdate, BTreeMap<K, u32>), Error> {
let mut request_spks = keychain_spks
.into_iter()
.map(|(k, s)| (k, s.into_iter()))
.collect::<BTreeMap<K, _>>();
) -> Result<FullScanResult<K, ConfirmationHeightAnchor>, Error> {
let mut request_spks = request.spks_by_keychain;

// We keep track of already-scanned spks just in case a reorg happens and we need to do a
// rescan. We need to keep track of this as iterators in `keychain_spks` are "unbounded" so
Expand All @@ -154,8 +73,8 @@ impl<E: ElectrumApi> ElectrumExt for E {
// * val: (script_pubkey, has_tx_history).
let mut scanned_spks = BTreeMap::<(K, u32), (ScriptBuf, bool)>::new();

let (electrum_update, keychain_update) = loop {
let (tip, _) = construct_update_tip(self, prev_tip.clone())?;
let update = loop {
let (tip, _) = construct_update_tip(self, request.chain_tip.clone())?;
let mut graph_update = TxGraph::<ConfirmationHeightAnchor>::default();
let cps = tip
.iter()
Expand All @@ -168,7 +87,7 @@ impl<E: ElectrumApi> ElectrumExt for E {
scanned_spks.append(&mut populate_with_spks(
self,
&cps,
tx_cache,
&mut request.tx_cache,
&mut graph_update,
&mut scanned_spks
.iter()
Expand All @@ -182,7 +101,7 @@ impl<E: ElectrumApi> ElectrumExt for E {
populate_with_spks(
self,
&cps,
tx_cache,
&mut request.tx_cache,
&mut graph_update,
keychain_spks,
stop_gap,
Expand Down Expand Up @@ -213,55 +132,118 @@ impl<E: ElectrumApi> ElectrumExt for E {
})
.collect::<BTreeMap<_, _>>();

break (
ElectrumUpdate {
chain_update,
graph_update,
},
keychain_update,
);
break FullScanResult {
graph_update,
chain_update,
last_active_indices: keychain_update,
};
};

Ok((electrum_update, keychain_update))
Ok(update)
}

fn sync(
&self,
tx_cache: &mut TxCache,
prev_tip: CheckPoint,
misc_spks: impl IntoIterator<Item = ScriptBuf>,
txids: impl IntoIterator<Item = Txid>,
outpoints: impl IntoIterator<Item = OutPoint>,
request: SyncRequest,
batch_size: usize,
) -> Result<ElectrumUpdate, Error> {
let spk_iter = misc_spks
.into_iter()
.enumerate()
.map(|(i, spk)| (i as u32, spk));

let (electrum_update, _) = self.full_scan(
tx_cache,
prev_tip.clone(),
[((), spk_iter)].into(),
usize::MAX,
batch_size,
)?;

let (tip, _) = construct_update_tip(self, prev_tip)?;
) -> Result<SyncResult<ConfirmationHeightAnchor>, Error> {
let mut tx_cache = request.tx_cache.clone();

let full_scan_req = FullScanRequest::from_chain_tip(request.chain_tip.clone())
.cache_txs(request.tx_cache)
.set_spks_for_keychain((), request.spks.enumerate().map(|(i, spk)| (i as u32, spk)));
let full_scan_res = self.full_scan(full_scan_req, usize::MAX, batch_size)?;

let (tip, _) = construct_update_tip(self, request.chain_tip)?;
let cps = tip
.iter()
.take(10)
.map(|cp| (cp.height(), cp))
.collect::<BTreeMap<u32, CheckPoint>>();

let mut tx_graph = TxGraph::<ConfirmationHeightAnchor>::default();
populate_with_txids(self, &cps, tx_cache, &mut tx_graph, txids)?;
populate_with_outpoints(self, &cps, &mut tx_graph, outpoints)?;
populate_with_txids(self, &cps, &mut tx_cache, &mut tx_graph, request.txids)?;
populate_with_outpoints(self, &cps, &mut tx_cache, &mut tx_graph, request.outpoints)?;

Ok(SyncResult {
chain_update: full_scan_res.chain_update,
graph_update: full_scan_res.graph_update,
})
}
}

/// Trait that extends [`SyncResult`] and [`FullScanResult`] functionality.
///
/// Currently, only a single method exists that converts the update [`TxGraph`] to have an anchor
/// type of [`ConfirmationTimeHeightAnchor`].
pub trait ElectrumResultExt {
/// New result type with a [`TxGraph`] that contains the [`ConfirmationTimeHeightAnchor`].
type NewResult;

/// Convert result type to have an update [`TxGraph`] that contains the [`ConfirmationTimeHeightAnchor`] .
fn try_into_confirmation_time_result(
self,
client: &impl ElectrumApi,
) -> Result<Self::NewResult, Error>;
}

impl<K> ElectrumResultExt for FullScanResult<K, ConfirmationHeightAnchor> {
type NewResult = FullScanResult<K, ConfirmationTimeHeightAnchor>;

fn try_into_confirmation_time_result(
self,
client: &impl ElectrumApi,
) -> Result<Self::NewResult, Error> {
Ok(FullScanResult::<K, ConfirmationTimeHeightAnchor> {
graph_update: try_into_confirmation_time_result(self.graph_update, client)?,
chain_update: self.chain_update,
last_active_indices: self.last_active_indices,
})
}
}

impl ElectrumResultExt for SyncResult<ConfirmationHeightAnchor> {
type NewResult = SyncResult<ConfirmationTimeHeightAnchor>;

Ok(electrum_update)
fn try_into_confirmation_time_result(
self,
client: &impl ElectrumApi,
) -> Result<Self::NewResult, Error> {
Ok(SyncResult {
graph_update: try_into_confirmation_time_result(self.graph_update, client)?,
chain_update: self.chain_update,
})
}
}

fn try_into_confirmation_time_result(
graph_update: TxGraph<ConfirmationHeightAnchor>,
client: &impl ElectrumApi,
) -> Result<TxGraph<ConfirmationTimeHeightAnchor>, Error> {
let relevant_heights = graph_update
.all_anchors()
.iter()
.map(|(a, _)| a.confirmation_height)
.collect::<HashSet<_>>();

let height_to_time = relevant_heights
.clone()
.into_iter()
.zip(
client
.batch_block_header(relevant_heights)?
.into_iter()
.map(|bh| bh.time as u64),
)
.collect::<HashMap<u32, u64>>();

Ok(graph_update.map_anchors(|a| ConfirmationTimeHeightAnchor {
anchor_block: a.anchor_block,
confirmation_height: a.confirmation_height,
confirmation_time: height_to_time[&a.confirmation_height],
}))
}

/// Return a [`CheckPoint`] of the latest tip, that connects with `prev_tip`.
fn construct_update_tip(
client: &impl ElectrumApi,
Expand Down Expand Up @@ -380,6 +362,7 @@ fn determine_tx_anchor(
fn populate_with_outpoints(
client: &impl ElectrumApi,
cps: &BTreeMap<u32, CheckPoint>,
tx_cache: &mut TxCache,
tx_graph: &mut TxGraph<ConfirmationHeightAnchor>,
outpoints: impl IntoIterator<Item = OutPoint>,
) -> Result<(), Error> {
Expand Down Expand Up @@ -415,9 +398,9 @@ fn populate_with_outpoints(
let res_tx = match tx_graph.get_tx(res.tx_hash) {
Some(tx) => tx,
None => {
let res_tx = client.transaction_get(&res.tx_hash)?;
let _ = tx_graph.insert_tx(res_tx);
tx_graph.get_tx(res.tx_hash).expect("just inserted")
let res_tx = fetch_tx(client, tx_cache, res.tx_hash)?;
let _ = tx_graph.insert_tx(Arc::clone(&res_tx));
res_tx
}
};
has_spending = res_tx
Expand Down
Loading

0 comments on commit a39a7e9

Please sign in to comment.