refactor(electrum)!: put the tx cache in electrum

Previously there was a tx cache that you passed in as part of the sync
request. This seems bad and the example show'd that you should copy all
your transactions from the transaction graph into the sync request every
time you sync'd. If you forgot to do this then you would always download everything.

Instead just do a plain old simple cache inside the electrum client.
This way at least you only download transactions once. You can
pre-populate the cache with a method also and I did this in the examples.
This commit is contained in:
LLFourn
2024-05-31 13:52:49 +10:00
committed by 志宇
parent 4a8452f9b8
commit 53fa35096f
8 changed files with 302 additions and 353 deletions

View File

@@ -0,0 +1,582 @@
use bdk_chain::{
bitcoin::{OutPoint, ScriptBuf, Transaction, Txid},
collections::{BTreeMap, HashMap, HashSet},
local_chain::CheckPoint,
spk_client::{FullScanRequest, FullScanResult, SyncRequest, SyncResult},
tx_graph::TxGraph,
BlockId, ConfirmationHeightAnchor, ConfirmationTimeHeightAnchor,
};
use core::str::FromStr;
use electrum_client::{ElectrumApi, Error, HeaderNotification};
use std::sync::{Arc, Mutex};
/// We include a chain suffix of a certain length for the purpose of robustness.
const CHAIN_SUFFIX_LENGTH: u32 = 8;
/// Wrapper around an [`electrum_client::ElectrumApi`] which includes an internal in-memory
/// transaction cache to avoid re-fetching already downloaded transactions.
#[derive(Debug)]
pub struct BdkElectrumClient<E> {
/// The internal [`electrum_client::ElectrumApi`]
pub inner: E,
/// The transaction cache
tx_cache: Mutex<HashMap<Txid, Arc<Transaction>>>,
}
impl<E: ElectrumApi> BdkElectrumClient<E> {
/// Creates a new bdk client from a [`electrum_client::ElectrumApi`]
pub fn new(client: E) -> Self {
Self {
inner: client,
tx_cache: Default::default(),
}
}
/// Inserts transactions into the transaction cache so that the client will not fetch these
/// transactions.
pub fn populate_tx_cache<A>(&self, tx_graph: impl AsRef<TxGraph<A>>) {
let txs = tx_graph
.as_ref()
.full_txs()
.map(|tx_node| (tx_node.txid, tx_node.tx));
let mut tx_cache = self.tx_cache.lock().unwrap();
for (txid, tx) in txs {
tx_cache.insert(txid, tx);
}
}
/// Fetch transaction of given `txid`.
///
/// If it hits the cache it will return the cached version and avoid making the request.
pub fn fetch_tx(&self, txid: Txid) -> Result<Arc<Transaction>, Error> {
let tx_cache = self.tx_cache.lock().unwrap();
if let Some(tx) = tx_cache.get(&txid) {
return Ok(Arc::clone(tx));
}
drop(tx_cache);
let tx = Arc::new(self.inner.transaction_get(&txid)?);
self.tx_cache.lock().unwrap().insert(txid, Arc::clone(&tx));
Ok(tx)
}
/// Full scan the keychain scripts specified with the blockchain (via an Electrum client) and
/// returns updates for [`bdk_chain`] data structures.
///
/// - `request`: struct with data required to perform a spk-based blockchain client full scan,
/// see [`FullScanRequest`]
/// - `stop_gap`: the full scan for each keychain stops after a gap of script pubkeys with no
/// associated transactions
/// - `batch_size`: specifies the max number of script pubkeys to request for in a single batch
/// request
/// - `fetch_prev_txouts`: specifies whether or not we want previous `TxOut`s for fee
pub fn full_scan<K: Ord + Clone>(
&self,
request: FullScanRequest<K>,
stop_gap: usize,
batch_size: usize,
fetch_prev_txouts: bool,
) -> Result<ElectrumFullScanResult<K>, 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
// cannot be collected. In addition, we keep track of whether an spk has an active tx
// history for determining the `last_active_index`.
// * key: (keychain, spk_index) that identifies the spk.
// * val: (script_pubkey, has_tx_history).
let mut scanned_spks = BTreeMap::<(K, u32), (ScriptBuf, bool)>::new();
let update = loop {
let (tip, _) = construct_update_tip(&self.inner, request.chain_tip.clone())?;
let mut graph_update = TxGraph::<ConfirmationHeightAnchor>::default();
let cps = tip
.iter()
.take(10)
.map(|cp| (cp.height(), cp))
.collect::<BTreeMap<u32, CheckPoint>>();
if !request_spks.is_empty() {
if !scanned_spks.is_empty() {
scanned_spks.append(
&mut self.populate_with_spks(
&cps,
&mut graph_update,
&mut scanned_spks
.iter()
.map(|(i, (spk, _))| (i.clone(), spk.clone())),
stop_gap,
batch_size,
)?,
);
}
for (keychain, keychain_spks) in &mut request_spks {
scanned_spks.extend(
self.populate_with_spks(
&cps,
&mut graph_update,
keychain_spks,
stop_gap,
batch_size,
)?
.into_iter()
.map(|(spk_i, spk)| ((keychain.clone(), spk_i), spk)),
);
}
}
// check for reorgs during scan process
let server_blockhash = self.inner.block_header(tip.height() as usize)?.block_hash();
if tip.hash() != server_blockhash {
continue; // reorg
}
// Fetch previous `TxOut`s for fee calculation if flag is enabled.
if fetch_prev_txouts {
self.fetch_prev_txout(&mut graph_update)?;
}
let chain_update = tip;
let keychain_update = request_spks
.into_keys()
.filter_map(|k| {
scanned_spks
.range((k.clone(), u32::MIN)..=(k.clone(), u32::MAX))
.rev()
.find(|(_, (_, active))| *active)
.map(|((_, i), _)| (k, *i))
})
.collect::<BTreeMap<_, _>>();
break FullScanResult {
graph_update,
chain_update,
last_active_indices: keychain_update,
};
};
Ok(ElectrumFullScanResult(update))
}
/// Sync a set of scripts with the blockchain (via an Electrum client) for the data specified
/// and returns updates for [`bdk_chain`] data structures.
///
/// - `request`: struct with data required to perform a spk-based blockchain client sync,
/// see [`SyncRequest`]
/// - `batch_size`: specifies the max number of script pubkeys to request for in a single batch
/// request
/// - `fetch_prev_txouts`: specifies whether or not we want previous `TxOut`s for fee
/// calculation
///
/// If the scripts to sync are unknown, such as when restoring or importing a keychain that
/// may include scripts that have been used, use [`full_scan`] with the keychain.
///
/// [`full_scan`]: Self::full_scan
pub fn sync(
&self,
request: SyncRequest,
batch_size: usize,
fetch_prev_txouts: bool,
) -> Result<ElectrumSyncResult, Error> {
let full_scan_req = FullScanRequest::from_chain_tip(request.chain_tip.clone())
.set_spks_for_keychain((), request.spks.enumerate().map(|(i, spk)| (i as u32, spk)));
let mut full_scan_res = self
.full_scan(full_scan_req, usize::MAX, batch_size, false)?
.with_confirmation_height_anchor();
let (tip, _) = construct_update_tip(&self.inner, request.chain_tip)?;
let cps = tip
.iter()
.take(10)
.map(|cp| (cp.height(), cp))
.collect::<BTreeMap<u32, CheckPoint>>();
self.populate_with_txids(&cps, &mut full_scan_res.graph_update, request.txids)?;
self.populate_with_outpoints(&cps, &mut full_scan_res.graph_update, request.outpoints)?;
// Fetch previous `TxOut`s for fee calculation if flag is enabled.
if fetch_prev_txouts {
self.fetch_prev_txout(&mut full_scan_res.graph_update)?;
}
Ok(ElectrumSyncResult(SyncResult {
chain_update: full_scan_res.chain_update,
graph_update: full_scan_res.graph_update,
}))
}
/// Populate the `graph_update` with transactions/anchors associated with the given `spks`.
///
/// Transactions that contains an output with requested spk, or spends form an output with
/// requested spk will be added to `graph_update`. Anchors of the aforementioned transactions are
/// also included.
///
/// Checkpoints (in `cps`) are used to create anchors. The `tx_cache` is self-explanatory.
fn populate_with_spks<I: Ord + Clone>(
&self,
cps: &BTreeMap<u32, CheckPoint>,
graph_update: &mut TxGraph<ConfirmationHeightAnchor>,
spks: &mut impl Iterator<Item = (I, ScriptBuf)>,
stop_gap: usize,
batch_size: usize,
) -> Result<BTreeMap<I, (ScriptBuf, bool)>, Error> {
let mut unused_spk_count = 0_usize;
let mut scanned_spks = BTreeMap::new();
loop {
let spks = (0..batch_size)
.map_while(|_| spks.next())
.collect::<Vec<_>>();
if spks.is_empty() {
return Ok(scanned_spks);
}
let spk_histories = self
.inner
.batch_script_get_history(spks.iter().map(|(_, s)| s.as_script()))?;
for ((spk_index, spk), spk_history) in spks.into_iter().zip(spk_histories) {
if spk_history.is_empty() {
scanned_spks.insert(spk_index, (spk, false));
unused_spk_count += 1;
if unused_spk_count > stop_gap {
return Ok(scanned_spks);
}
continue;
} else {
scanned_spks.insert(spk_index, (spk, true));
unused_spk_count = 0;
}
for tx_res in spk_history {
let _ = graph_update.insert_tx(self.fetch_tx(tx_res.tx_hash)?);
if let Some(anchor) = determine_tx_anchor(cps, tx_res.height, tx_res.tx_hash) {
let _ = graph_update.insert_anchor(tx_res.tx_hash, anchor);
}
}
}
}
}
// Helper function which fetches the `TxOut`s of our relevant transactions' previous transactions,
// which we do not have by default. This data is needed to calculate the transaction fee.
fn fetch_prev_txout(
&self,
graph_update: &mut TxGraph<ConfirmationHeightAnchor>,
) -> Result<(), Error> {
let full_txs: Vec<Arc<Transaction>> =
graph_update.full_txs().map(|tx_node| tx_node.tx).collect();
for tx in full_txs {
for vin in &tx.input {
let outpoint = vin.previous_output;
let vout = outpoint.vout;
let prev_tx = self.fetch_tx(outpoint.txid)?;
let txout = prev_tx.output[vout as usize].clone();
let _ = graph_update.insert_txout(outpoint, txout);
}
}
Ok(())
}
/// Populate the `graph_update` with associated transactions/anchors of `outpoints`.
///
/// Transactions in which the outpoint resides, and transactions that spend from the outpoint are
/// included. Anchors of the aforementioned transactions are included.
///
/// Checkpoints (in `cps`) are used to create anchors. The `tx_cache` is self-explanatory.
fn populate_with_outpoints(
&self,
cps: &BTreeMap<u32, CheckPoint>,
graph_update: &mut TxGraph<ConfirmationHeightAnchor>,
outpoints: impl IntoIterator<Item = OutPoint>,
) -> Result<(), Error> {
for outpoint in outpoints {
let op_txid = outpoint.txid;
let op_tx = self.fetch_tx(op_txid)?;
let op_txout = match op_tx.output.get(outpoint.vout as usize) {
Some(txout) => txout,
None => continue,
};
debug_assert_eq!(op_tx.txid(), op_txid);
// attempt to find the following transactions (alongside their chain positions), and
// add to our sparsechain `update`:
let mut has_residing = false; // tx in which the outpoint resides
let mut has_spending = false; // tx that spends the outpoint
for res in self.inner.script_get_history(&op_txout.script_pubkey)? {
if has_residing && has_spending {
break;
}
if !has_residing && res.tx_hash == op_txid {
has_residing = true;
let _ = graph_update.insert_tx(Arc::clone(&op_tx));
if let Some(anchor) = determine_tx_anchor(cps, res.height, res.tx_hash) {
let _ = graph_update.insert_anchor(res.tx_hash, anchor);
}
}
if !has_spending && res.tx_hash != op_txid {
let res_tx = self.fetch_tx(res.tx_hash)?;
// we exclude txs/anchors that do not spend our specified outpoint(s)
has_spending = res_tx
.input
.iter()
.any(|txin| txin.previous_output == outpoint);
if !has_spending {
continue;
}
let _ = graph_update.insert_tx(Arc::clone(&res_tx));
if let Some(anchor) = determine_tx_anchor(cps, res.height, res.tx_hash) {
let _ = graph_update.insert_anchor(res.tx_hash, anchor);
}
}
}
}
Ok(())
}
/// Populate the `graph_update` with transactions/anchors of the provided `txids`.
fn populate_with_txids(
&self,
cps: &BTreeMap<u32, CheckPoint>,
graph_update: &mut TxGraph<ConfirmationHeightAnchor>,
txids: impl IntoIterator<Item = Txid>,
) -> Result<(), Error> {
for txid in txids {
let tx = match self.fetch_tx(txid) {
Ok(tx) => tx,
Err(electrum_client::Error::Protocol(_)) => continue,
Err(other_err) => return Err(other_err),
};
let spk = tx
.output
.first()
.map(|txo| &txo.script_pubkey)
.expect("tx must have an output");
// because of restrictions of the Electrum API, we have to use the `script_get_history`
// call to get confirmation status of our transaction
let anchor = match self
.inner
.script_get_history(spk)?
.into_iter()
.find(|r| r.tx_hash == txid)
{
Some(r) => determine_tx_anchor(cps, r.height, txid),
None => continue,
};
let _ = graph_update.insert_tx(tx);
if let Some(anchor) = anchor {
let _ = graph_update.insert_anchor(txid, anchor);
}
}
Ok(())
}
}
/// The result of [`BdkElectrumClient::full_scan`].
///
/// This can be transformed into a [`FullScanResult`] with either [`ConfirmationHeightAnchor`] or
/// [`ConfirmationTimeHeightAnchor`] anchor types.
pub struct ElectrumFullScanResult<K>(FullScanResult<K, ConfirmationHeightAnchor>);
impl<K> ElectrumFullScanResult<K> {
/// Return [`FullScanResult`] with [`ConfirmationHeightAnchor`].
pub fn with_confirmation_height_anchor(self) -> FullScanResult<K, ConfirmationHeightAnchor> {
self.0
}
/// Return [`FullScanResult`] with [`ConfirmationTimeHeightAnchor`].
///
/// This requires additional calls to the Electrum server.
pub fn with_confirmation_time_height_anchor(
self,
client: &BdkElectrumClient<impl ElectrumApi>,
) -> Result<FullScanResult<K, ConfirmationTimeHeightAnchor>, Error> {
let res = self.0;
Ok(FullScanResult {
graph_update: try_into_confirmation_time_result(res.graph_update, &client.inner)?,
chain_update: res.chain_update,
last_active_indices: res.last_active_indices,
})
}
}
/// The result of [`BdkElectrumClient::sync`].
///
/// This can be transformed into a [`SyncResult`] with either [`ConfirmationHeightAnchor`] or
/// [`ConfirmationTimeHeightAnchor`] anchor types.
pub struct ElectrumSyncResult(SyncResult<ConfirmationHeightAnchor>);
impl ElectrumSyncResult {
/// Return [`SyncResult`] with [`ConfirmationHeightAnchor`].
pub fn with_confirmation_height_anchor(self) -> SyncResult<ConfirmationHeightAnchor> {
self.0
}
/// Return [`SyncResult`] with [`ConfirmationTimeHeightAnchor`].
///
/// This requires additional calls to the Electrum server.
pub fn with_confirmation_time_height_anchor(
self,
client: &BdkElectrumClient<impl ElectrumApi>,
) -> Result<SyncResult<ConfirmationTimeHeightAnchor>, Error> {
let res = self.0;
Ok(SyncResult {
graph_update: try_into_confirmation_time_result(res.graph_update, &client.inner)?,
chain_update: res.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,
prev_tip: CheckPoint,
) -> Result<(CheckPoint, Option<u32>), Error> {
let HeaderNotification { height, .. } = client.block_headers_subscribe()?;
let new_tip_height = height as u32;
// If electrum returns a tip height that is lower than our previous tip, then checkpoints do
// not need updating. We just return the previous tip and use that as the point of agreement.
if new_tip_height < prev_tip.height() {
return Ok((prev_tip.clone(), Some(prev_tip.height())));
}
// Atomically fetch the latest `CHAIN_SUFFIX_LENGTH` count of blocks from Electrum. We use this
// to construct our checkpoint update.
let mut new_blocks = {
let start_height = new_tip_height.saturating_sub(CHAIN_SUFFIX_LENGTH - 1);
let hashes = client
.block_headers(start_height as _, CHAIN_SUFFIX_LENGTH as _)?
.headers
.into_iter()
.map(|h| h.block_hash());
(start_height..).zip(hashes).collect::<BTreeMap<u32, _>>()
};
// Find the "point of agreement" (if any).
let agreement_cp = {
let mut agreement_cp = Option::<CheckPoint>::None;
for cp in prev_tip.iter() {
let cp_block = cp.block_id();
let hash = match new_blocks.get(&cp_block.height) {
Some(&hash) => hash,
None => {
assert!(
new_tip_height >= cp_block.height,
"already checked that electrum's tip cannot be smaller"
);
let hash = client.block_header(cp_block.height as _)?.block_hash();
new_blocks.insert(cp_block.height, hash);
hash
}
};
if hash == cp_block.hash {
agreement_cp = Some(cp);
break;
}
}
agreement_cp
};
let agreement_height = agreement_cp.as_ref().map(CheckPoint::height);
let new_tip = new_blocks
.into_iter()
// Prune `new_blocks` to only include blocks that are actually new.
.filter(|(height, _)| Some(*height) > agreement_height)
.map(|(height, hash)| BlockId { height, hash })
.fold(agreement_cp, |prev_cp, block| {
Some(match prev_cp {
Some(cp) => cp.push(block).expect("must extend checkpoint"),
None => CheckPoint::new(block),
})
})
.expect("must have at least one checkpoint");
Ok((new_tip, agreement_height))
}
/// A [tx status] comprises of a concatenation of `tx_hash:height:`s. We transform a single one of
/// these concatenations into a [`ConfirmationHeightAnchor`] if possible.
///
/// We use the lowest possible checkpoint as the anchor block (from `cps`). If an anchor block
/// cannot be found, or the transaction is unconfirmed, [`None`] is returned.
///
/// [tx status](https://electrumx-spesmilo.readthedocs.io/en/latest/protocol-basics.html#status)
fn determine_tx_anchor(
cps: &BTreeMap<u32, CheckPoint>,
raw_height: i32,
txid: Txid,
) -> Option<ConfirmationHeightAnchor> {
// The electrum API has a weird quirk where an unconfirmed transaction is presented with a
// height of 0. To avoid invalid representation in our data structures, we manually set
// transactions residing in the genesis block to have height 0, then interpret a height of 0 as
// unconfirmed for all other transactions.
if txid
== Txid::from_str("4a5e1e4baab89f3a32518a88c31bc87f618f76673e2cc77ab2127b7afdeda33b")
.expect("must deserialize genesis coinbase txid")
{
let anchor_block = cps.values().next()?.block_id();
return Some(ConfirmationHeightAnchor {
anchor_block,
confirmation_height: 0,
});
}
match raw_height {
h if h <= 0 => {
debug_assert!(h == 0 || h == -1, "unexpected height ({}) from electrum", h);
None
}
h => {
let h = h as u32;
let anchor_block = cps.range(h..).next().map(|(_, cp)| cp.block_id())?;
if h > anchor_block.height {
None
} else {
Some(ConfirmationHeightAnchor {
anchor_block,
confirmation_height: h,
})
}
}
}
}