Skip to main content

cuprate_txpool/service/
write.rs

1use std::{
2    collections::{hash_map::Entry, HashSet},
3    sync::Arc,
4    task::{Context, Poll},
5};
6
7use futures::channel::oneshot;
8use monero_oxide::transaction::Input;
9use rayon::ThreadPool;
10use tower::Service;
11
12use cuprate_helper::asynch::InfallibleOneshotReceiver;
13use cuprate_types::TransactionVerificationData;
14
15use crate::{
16    error::TxPoolError,
17    ops::{self, TxPoolWriteError},
18    service::interface::{TxpoolWriteRequest, TxpoolWriteResponse},
19    txpool::TxpoolDatabase,
20    types::{KeyImage, TransactionHash, TransactionInfo, TxStateFlags},
21};
22
23/// The txpool [`Service`] write handle.
24#[derive(Clone)]
25pub struct TxpoolWriteHandle {
26    /// Handle to the custom `rayon` DB reader thread-pool.
27    ///
28    /// Requests are [`rayon::ThreadPool::spawn`]ed in this thread-pool,
29    /// and responses are returned via a channel we (the caller) provide.
30    pub pool: Arc<ThreadPool>,
31
32    pub txpool: Arc<TxpoolDatabase>,
33}
34
35impl Service<TxpoolWriteRequest> for TxpoolWriteHandle {
36    type Response = TxpoolWriteResponse;
37    type Error = TxPoolError;
38    type Future = InfallibleOneshotReceiver<Result<Self::Response, Self::Error>>;
39
40    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
41        Poll::Ready(Ok(()))
42    }
43
44    fn call(&mut self, req: TxpoolWriteRequest) -> Self::Future {
45        let (tx, rx) = oneshot::channel();
46
47        let db = Arc::clone(&self.txpool);
48        self.pool.spawn(move || {
49            let res = handle_txpool_request(&db, req);
50
51            drop(tx.send(res));
52        });
53
54        InfallibleOneshotReceiver::from(rx)
55    }
56}
57
58//---------------------------------------------------------------------------------------------------- handle_txpool_request
59/// Handle an incoming [`TxpoolWriteRequest`], returning a [`TxpoolWriteResponse`].
60fn handle_txpool_request(
61    env: &TxpoolDatabase,
62    req: TxpoolWriteRequest,
63) -> Result<TxpoolWriteResponse, TxPoolError> {
64    match req {
65        TxpoolWriteRequest::AddTransaction { tx, state_stem } => {
66            add_transaction(env, &tx, state_stem)
67        }
68        TxpoolWriteRequest::RemoveTransaction(tx_hash) => remove_transaction(env, &tx_hash),
69        TxpoolWriteRequest::Promote(tx_hash) => promote(env, &tx_hash),
70        TxpoolWriteRequest::NewBlock { spent_key_images } => new_block(env, &spent_key_images),
71    }
72}
73
74//---------------------------------------------------------------------------------------------------- Handler functions
75// These are the actual functions that do stuff according to the incoming [`TxpoolWriteRequest`].
76//
77// Each function name is a 1-1 mapping (from CamelCase -> snake_case) to
78// the enum variant name, e.g: `BlockExtendedHeader` -> `block_extended_header`.
79//
80// Each function will return the [`Response`] that we
81// should send back to the caller in [`map_request()`].
82
83/// [`TxpoolWriteRequest::AddTransaction`]
84fn add_transaction(
85    db: &TxpoolDatabase,
86    tx: &TransactionVerificationData,
87    state_stem: bool,
88) -> Result<TxpoolWriteResponse, TxPoolError> {
89    struct KiDropGuard<'a>(Vec<[u8; 32]>, &'a TxpoolDatabase);
90
91    impl Drop for KiDropGuard<'_> {
92        fn drop(&mut self) {
93            for ki in &self.0 {
94                self.1.in_progress_key_images.lock().unwrap().remove(ki);
95            }
96        }
97    }
98
99    let mut guard = KiDropGuard(Vec::with_capacity(tx.tx.prefix().inputs.len()), db);
100
101    let mut in_progress_key_images = db.in_progress_key_images.lock().unwrap();
102    for ki in tx.tx.prefix().inputs.iter().map(|i| match i {
103        Input::ToKey { key_image, .. } => key_image,
104        Input::Gen(_) => unreachable!(),
105    }) {
106        let e = in_progress_key_images.entry(ki.to_bytes());
107
108        match e {
109            Entry::Occupied(o) => return Ok(TxpoolWriteResponse::AddTransaction(Some(*o.get()))),
110            Entry::Vacant(v) => {
111                v.insert(tx.tx_hash);
112            }
113        }
114
115        guard.0.push(ki.to_bytes());
116    }
117    drop(in_progress_key_images);
118
119    let mut writer = db.fjall_database.batch();
120
121    if let Err(e) = ops::add_transaction(tx, state_stem, &mut writer, db) {
122        // error adding the tx, abort the DB transaction.
123        drop(writer);
124
125        return match e {
126            TxPoolWriteError::DoubleSpend(tx_hash) => {
127                // If we couldn't add the tx due to a double spend still return ok, but include the tx
128                // this double spent.
129                // TODO: mark the double spent tx?
130                Ok(TxpoolWriteResponse::AddTransaction(Some(tx_hash)))
131            }
132            TxPoolWriteError::TxPool(e) => Err(e),
133        };
134    }
135
136    // The tx was added to the pool successfully.
137    writer.commit()?;
138
139    Ok(TxpoolWriteResponse::AddTransaction(None))
140}
141
142/// [`TxpoolWriteRequest::RemoveTransaction`]
143fn remove_transaction(
144    db: &TxpoolDatabase,
145    tx_hash: &TransactionHash,
146) -> Result<TxpoolWriteResponse, TxPoolError> {
147    let mut writer = db.fjall_database.batch();
148
149    ops::remove_transaction(tx_hash, &mut writer, db)?;
150
151    writer.commit()?;
152
153    Ok(TxpoolWriteResponse::Ok)
154}
155
156/// [`TxpoolWriteRequest::Promote`]
157fn promote(
158    db: &TxpoolDatabase,
159    tx_hash: &TransactionHash,
160) -> Result<TxpoolWriteResponse, TxPoolError> {
161    let tx_info = db.tx_infos.get(tx_hash)?.ok_or(TxPoolError::NotFound)?;
162    let mut tx_info: TransactionInfo = bytemuck::pod_read_unaligned(tx_info.as_ref());
163
164    if !tx_info.flags.private() {
165        return Ok(TxpoolWriteResponse::Ok);
166    }
167
168    tx_info.flags.remove(TxStateFlags::STATE_STEM);
169
170    db.tx_infos.insert(tx_hash, bytemuck::bytes_of(&tx_info))?;
171
172    if !db.tx_blobs.contains_key(tx_hash)? {
173        db.tx_infos.remove(tx_hash)?;
174    }
175
176    Ok(TxpoolWriteResponse::Ok)
177}
178
179/// [`TxpoolWriteRequest::NewBlock`]
180fn new_block(
181    db: &TxpoolDatabase,
182    spent_key_images: &[KeyImage],
183) -> Result<TxpoolWriteResponse, TxPoolError> {
184    let mut txs_removed = HashSet::new();
185
186    let mut writer = db.fjall_database.batch();
187
188    // Remove all txs which spend key images that were spent in the new block.
189    for key_image in spent_key_images {
190        if let Some(tx_hash) = db.spent_key_images.get(key_image)? {
191            let tx_hash = tx_hash.as_ref().try_into().unwrap();
192
193            if txs_removed.insert(tx_hash) {
194                ops::remove_transaction(&tx_hash, &mut writer, db)?;
195            }
196        }
197    }
198
199    writer.commit()?;
200    Ok(TxpoolWriteResponse::NewBlock(
201        txs_removed.into_iter().collect(),
202    ))
203}