cfxcore/pos/mempool/core_mempool/
transaction_store.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8use crate::pos::mempool::{
9    core_mempool::{
10        index::{
11            AccountTransactionIter, AccountTransactions, TTLIndex,
12            TimelineIndex,
13        },
14        transaction::{MempoolTransaction, TimelineState},
15    },
16    logging::{LogEntry, LogEvent, LogSchema, TxnsLog},
17};
18use diem_config::config::MempoolConfig;
19use diem_crypto::{hash::CryptoHash, HashValue};
20use diem_logger::prelude::*;
21use diem_types::{
22    account_address::AccountAddress,
23    mempool_status::{MempoolStatus, MempoolStatusCode},
24    transaction::{SignedTransaction, TransactionPayload},
25};
26use std::{
27    collections::{hash_map::Values, HashMap, HashSet},
28    time::Duration,
29};
30
31/// TransactionStore is in-memory storage for all transactions in mempool.
32pub struct TransactionStore {
33    // normal transactions
34    transactions: AccountTransactions,
35    // pivot decision helper structure
36    pivot_decisions: HashMap<HashValue, HashSet<(AccountAddress, HashValue)>>,
37
38    // Evicts txns after `system_transaction_timeout` so stalled commit
39    // callbacks cannot clog the mempool indefinitely.
40    system_ttl_index: TTLIndex,
41    timeline_index: TimelineIndex,
42
43    // Per-sender cap against Byzantine spam. Invariant: every removal
44    // of `self.transactions` must route through `index_remove`.
45    per_sender_count: HashMap<AccountAddress, usize>,
46    capacity_per_sender: usize,
47}
48
49pub type PivotDecisionIter<'a> =
50    Values<'a, HashValue, HashSet<(AccountAddress, HashValue)>>;
51
52impl TransactionStore {
53    pub(crate) fn new(config: &MempoolConfig) -> Self {
54        assert!(
55            config.capacity_per_sender > 0,
56            "mempool.capacity_per_sender must be > 0",
57        );
58        Self {
59            // main DS
60            transactions: AccountTransactions::new(),
61            pivot_decisions: HashMap::new(),
62
63            // various indexes
64            system_ttl_index: TTLIndex::new(Box::new(
65                |t: &MempoolTransaction| t.expiration_time,
66            )),
67            timeline_index: TimelineIndex::new(),
68
69            per_sender_count: HashMap::new(),
70            capacity_per_sender: config.capacity_per_sender,
71        }
72    }
73
74    /// Fetch transaction by account address + hash.
75    pub(crate) fn get(&self, hash: &HashValue) -> Option<SignedTransaction> {
76        if let Some(txn) = self.transactions.get(hash) {
77            return Some(txn.txn.clone());
78        }
79        None
80    }
81
82    /// Fetch pivot decisions by pivot hash.
83    pub(crate) fn get_pivot_decisions(
84        &self, hash: &HashValue,
85    ) -> Vec<HashValue> {
86        if let Some(decisions) = self.pivot_decisions.get(hash) {
87            decisions
88                .iter()
89                .map(|(_, tx_hash)| tx_hash.clone())
90                .collect::<_>()
91        } else {
92            vec![]
93        }
94    }
95
96    /// Insert transaction into TransactionStore. Performs validation checks and
97    /// updates indexes.
98    pub(crate) fn insert(
99        &mut self, mut txn: MempoolTransaction,
100    ) -> MempoolStatus {
101        let address = txn.get_sender();
102        let hash = txn.get_hash();
103        let has_tx = self.get(&hash).is_some();
104
105        if has_tx {
106            return MempoolStatus::new(MempoolStatusCode::Accepted);
107        }
108
109        let sender_entry = self.per_sender_count.entry(address).or_insert(0);
110        if *sender_entry >= self.capacity_per_sender {
111            let sender_count = *sender_entry;
112            // Rate-limited so a sustained attack doesn't flood logs.
113            diem_sample!(
114                SampleRate::Duration(Duration::from_secs(60)),
115                diem_warn!(
116                    sender = %address,
117                    sender_count = sender_count,
118                    cap = self.capacity_per_sender,
119                    "mempool: per-sender capacity reached, rejecting txn",
120                )
121            );
122            return MempoolStatus::new(MempoolStatusCode::TooManyTransactions)
123                .with_message(format!(
124                    "sender {} already has {} transactions (cap {})",
125                    address, sender_count, self.capacity_per_sender,
126                ));
127        }
128        *sender_entry += 1;
129
130        self.timeline_index.insert(&mut txn);
131        self.system_ttl_index.insert(&txn);
132
133        if let TransactionPayload::PivotDecision(pivot_decision) =
134            txn.txn.payload()
135        {
136            let pivot_decision_hash = pivot_decision.hash();
137            self.pivot_decisions
138                .entry(pivot_decision_hash)
139                .or_insert_with(HashSet::new);
140            if let Some(account_decision) =
141                self.pivot_decisions.get_mut(&pivot_decision_hash)
142            {
143                diem_debug!("txpool::insert pivot {:?}", hash);
144                account_decision.insert((address, hash));
145            }
146            self.transactions.insert(hash, txn, true);
147        } else {
148            self.transactions.insert(hash, txn, false);
149        }
150        diem_debug!(
151            LogSchema::new(LogEntry::AddTxn)
152                .txns(TxnsLog::new_txn(address, hash)),
153            hash = hash,
154            has_tx = has_tx
155        );
156
157        MempoolStatus::new(MempoolStatusCode::Accepted)
158    }
159
160    /// Handles transaction commit: deletes the transaction and cleans up
161    /// its entries in the timeline and TTL indexes.
162    pub(crate) fn commit_transaction(&mut self, hash: HashValue) {
163        let mut txns_log = TxnsLog::new();
164        if let Some(transaction) = self.transactions.remove(&hash) {
165            txns_log.add(transaction.get_sender(), transaction.get_hash());
166            self.index_remove(&transaction);
167            // handle pivot decision
168            let payload = transaction.txn.into_raw_transaction().into_payload();
169            if let TransactionPayload::PivotDecision(pivot_decision) = payload {
170                let pivot_decision_hash = pivot_decision.hash();
171                if let Some(indices) =
172                    self.pivot_decisions.remove(&pivot_decision_hash)
173                {
174                    for (_, hash) in indices {
175                        if let Some(txn) = self.transactions.remove(&hash) {
176                            txns_log.add(txn.get_sender(), txn.get_hash());
177                            self.index_remove(&txn);
178                        }
179                    }
180                }
181            }
182        }
183        diem_debug!(LogSchema::new(LogEntry::CleanCommittedTxn).txns(txns_log));
184    }
185
186    /// Removes transaction from all indexes.
187    fn index_remove(&mut self, txn: &MempoolTransaction) {
188        self.system_ttl_index.remove(&txn);
189        self.timeline_index.remove(&txn);
190        let sender = txn.get_sender();
191        debug_assert!(
192            self.per_sender_count.contains_key(&sender),
193            "per_sender_count missing entry for {} at index_remove",
194            sender,
195        );
196        if let Some(count) = self.per_sender_count.get_mut(&sender) {
197            *count -= 1;
198            if *count == 0 {
199                self.per_sender_count.remove(&sender);
200            }
201        }
202    }
203
204    /// Read `count` transactions from timeline since `timeline_id`.
205    /// Returns block of transactions and new last_timeline_id.
206    pub(crate) fn read_timeline(
207        &mut self, timeline_id: u64, count: usize,
208    ) -> (Vec<SignedTransaction>, u64) {
209        let mut batch = vec![];
210        let mut last_timeline_id = timeline_id;
211        for (_, hash) in self.timeline_index.read_timeline(timeline_id, count) {
212            if let Some(txn) = self.transactions.get(&hash) {
213                batch.push(txn.txn.clone());
214                if let TimelineState::Ready(timeline_id) = txn.timeline_state {
215                    last_timeline_id = timeline_id;
216                }
217            }
218        }
219        (batch, last_timeline_id)
220    }
221
222    pub(crate) fn timeline_range(
223        &mut self, start_id: u64, end_id: u64,
224    ) -> Vec<SignedTransaction> {
225        self.timeline_index
226            .timeline_range(start_id, end_id)
227            .iter()
228            .filter_map(|(_, hash)| {
229                self.transactions.get(hash).map(|txn| txn.txn.clone())
230            })
231            .collect()
232    }
233
234    /// Garbage collect old transactions by system TTL.
235    pub(crate) fn gc_by_system_ttl(&mut self) {
236        let now = std::time::SystemTime::now()
237            .duration_since(std::time::UNIX_EPOCH)
238            .expect("System time is before UNIX_EPOCH");
239
240        let mut gc_txns = self.system_ttl_index.gc(now);
241        gc_txns.sort_by_key(|key| (key.address, key.hash));
242
243        let mut gc_txns_log = TxnsLog::new();
244        for key in gc_txns.iter() {
245            if let Some(txn) = self.transactions.remove(&key.hash) {
246                gc_txns_log.add(txn.get_sender(), txn.get_hash());
247                self.index_remove(&txn);
248                if let TransactionPayload::PivotDecision(pivot_decision) =
249                    txn.txn.into_raw_transaction().into_payload()
250                {
251                    self.pivot_decisions.remove(&pivot_decision.hash());
252                }
253            }
254        }
255
256        diem_debug!(LogSchema::event_log(
257            LogEntry::GCRemoveTxns,
258            LogEvent::SystemTTLExpiration
259        )
260        .txns(gc_txns_log));
261    }
262
263    pub(crate) fn iter(&self) -> AccountTransactionIter<'_> {
264        self.transactions.iter()
265    }
266
267    pub(crate) fn iter_pivot_decision(&self) -> PivotDecisionIter<'_> {
268        self.pivot_decisions.values()
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275    use diem_crypto::{
276        bls::{BLSPrivateKey, BLSPublicKey},
277        PrivateKey, SigningKey, Uniform,
278    };
279    use diem_types::{
280        chain_id::ChainId,
281        transaction::{RawTransaction, RetirePayload, TransactionPayload},
282    };
283    use std::time::Duration;
284
285    fn store_with_cap(cap: usize) -> TransactionStore {
286        let mut cfg = MempoolConfig::default();
287        cfg.capacity_per_sender = cap;
288        TransactionStore::new(&cfg)
289    }
290
291    // Address is arbitrary — `per_sender_count` keys on address only.
292    fn new_sender() -> (BLSPrivateKey, BLSPublicKey, AccountAddress) {
293        let sk = BLSPrivateKey::generate_for_testing();
294        let pk = sk.public_key();
295        (sk, pk, AccountAddress::random())
296    }
297
298    fn mk_txn(
299        sk: &BLSPrivateKey, pk: &BLSPublicKey, sender: AccountAddress,
300        nonce: u64,
301    ) -> MempoolTransaction {
302        let payload = TransactionPayload::Retire(RetirePayload {
303            node_id: sender,
304            votes: nonce,
305        });
306        let raw =
307            RawTransaction::new(sender, payload, u64::MAX, ChainId::test());
308        let sig = sk.sign(&raw);
309        MempoolTransaction::new(
310            SignedTransaction::new(raw, pk.clone(), sig),
311            Duration::from_secs(3600),
312            TimelineState::NotReady,
313        )
314    }
315
316    #[test]
317    fn per_sender_count_insert_and_commit_lifecycle() {
318        let mut store = store_with_cap(3);
319        let (sk, pk, sender) = new_sender();
320        assert!(!store.per_sender_count.contains_key(&sender));
321
322        let mut hashes = Vec::new();
323        for n in 0..3 {
324            let txn = mk_txn(&sk, &pk, sender, n);
325            hashes.push(txn.get_hash());
326            assert_eq!(store.insert(txn).code, MempoolStatusCode::Accepted);
327        }
328        assert_eq!(store.per_sender_count[&sender], 3);
329
330        for (i, h) in hashes.iter().enumerate() {
331            store.commit_transaction(*h);
332            let remaining = 3 - (i + 1);
333            if remaining == 0 {
334                assert!(!store.per_sender_count.contains_key(&sender));
335            } else {
336                assert_eq!(store.per_sender_count[&sender], remaining);
337            }
338        }
339    }
340
341    #[test]
342    fn per_sender_count_cap_rejects_without_growth() {
343        let mut store = store_with_cap(2);
344        let (sk, pk, sender) = new_sender();
345
346        for n in 0..2 {
347            assert_eq!(
348                store.insert(mk_txn(&sk, &pk, sender, n)).code,
349                MempoolStatusCode::Accepted
350            );
351        }
352        assert_eq!(store.per_sender_count[&sender], 2);
353
354        for n in 2..6 {
355            assert_eq!(
356                store.insert(mk_txn(&sk, &pk, sender, n)).code,
357                MempoolStatusCode::TooManyTransactions
358            );
359            assert_eq!(store.per_sender_count[&sender], 2);
360        }
361    }
362
363    #[test]
364    fn per_sender_count_duplicate_hash_no_double_count() {
365        let mut store = store_with_cap(8);
366        let (sk, pk, sender) = new_sender();
367        let txn = mk_txn(&sk, &pk, sender, 0);
368        let dup = MempoolTransaction::new(
369            txn.txn.clone(),
370            txn.expiration_time,
371            txn.timeline_state,
372        );
373
374        assert_eq!(store.insert(txn).code, MempoolStatusCode::Accepted);
375        assert_eq!(store.per_sender_count[&sender], 1);
376
377        assert_eq!(store.insert(dup).code, MempoolStatusCode::Accepted);
378        assert_eq!(store.per_sender_count[&sender], 1);
379    }
380
381    #[test]
382    #[should_panic(expected = "mempool.capacity_per_sender must be > 0")]
383    fn capacity_per_sender_zero_panics_on_construction() {
384        let _ = store_with_cap(0);
385    }
386}
cfxcore/pos/mempool/core_mempool/transaction_store.rs

cfxcore/pos/mempool/core_mempool/
transaction_store.rs
