iota_faucet/faucet/
write_ahead_log.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::path::Path;
6
7use iota_sdk_types::{Address, ObjectId, Transaction};
8use serde::{Deserialize, Serialize};
9use tracing::info;
10use typed_store::{DBMapUtils, Map, TypedStoreError, rocks::DBMap};
11use uuid::Uuid;
12
13/// Persistent log of transactions paying out iota from the faucet, keyed by the
14/// coin serving the request.  Transactions are expected to be written to the
15/// log before they are sent to full-node, and removed after receiving a
16/// response back, before the coin becomes available for subsequent writes.
17///
18/// This allows the faucet to go down and back up, and not forget which requests
19/// were in-flight that it needs to confirm succeeded or failed.
20#[derive(DBMapUtils, Clone)]
21pub struct WriteAheadLog {
22    pub log: DBMap<ObjectId, Entry>,
23}
24
25#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
26pub struct Entry {
27    pub uuid: uuid::Bytes,
28    // TODO (jian): remove recipient
29    pub recipient: Address,
30    pub tx: Transaction,
31    pub retry_count: u64,
32    pub in_flight: bool,
33}
34
35impl WriteAheadLog {
36    pub(crate) fn open(path: &Path) -> Self {
37        Self::open_tables_read_write(
38            path.to_path_buf(),
39            typed_store::rocks::MetricConf::new("faucet_write_ahead_log"),
40            None,
41            None,
42        )
43    }
44
45    /// Mark `coin` as reserved for transaction `tx` sending coin to
46    /// `recipient`. Fails if `coin` is already in the WAL pointing to an
47    /// existing transaction.
48    pub(crate) fn reserve(
49        &mut self,
50        uuid: Uuid,
51        coin: ObjectId,
52        recipient: Address,
53        tx: Transaction,
54    ) -> Result<(), TypedStoreError> {
55        if self.log.contains_key(&coin)? {
56            // Don't permit multiple writes against the same coin
57            // TODO: Use a better error type than `TypedStoreError`.
58            return Err(TypedStoreError::Serialization(format!(
59                "Duplicate WAL entry for coin {coin:?}",
60            )));
61        }
62
63        let uuid = *uuid.as_bytes();
64        self.log.insert(
65            &coin,
66            &Entry {
67                uuid,
68                recipient,
69                tx,
70                retry_count: 0,
71                in_flight: true,
72            },
73        )
74    }
75
76    /// Check whether `coin` has a pending transaction in the WAL.  Returns
77    /// `Ok(Some(entry))` if a pending transaction exists, `Ok(None)` if
78    /// not, and `Err(_)` if there was an internal error accessing the WAL.
79    pub(crate) fn reclaim(&self, coin: ObjectId) -> Result<Option<Entry>, TypedStoreError> {
80        match self.log.get(&coin) {
81            Ok(entry) => Ok(entry),
82            Err(TypedStoreError::Serialization(_)) => {
83                // Remove bad log from the store, so we don't crash on start up, this can happen
84                // if we update the WAL Entry and have some leftover Entry from
85                // the WAL.
86                self.log
87                    .remove(&coin)
88                    .unwrap_or_else(|_| panic!("coin: {coin:?} unable to be removed from log."));
89                Ok(None)
90            }
91            Err(err) => Err(err),
92        }
93    }
94
95    /// Indicate that the transaction in flight for `coin` has landed, and the
96    /// entry in the WAL can be removed.
97    pub(crate) fn commit(&mut self, coin: ObjectId) -> Result<(), TypedStoreError> {
98        self.log.remove(&coin)
99    }
100
101    pub(crate) fn increment_retry_count(&mut self, coin: ObjectId) -> Result<(), TypedStoreError> {
102        if let Some(mut entry) = self.log.get(&coin)? {
103            entry.retry_count += 1;
104            self.log.insert(&coin, &entry)?;
105        }
106        Ok(())
107    }
108
109    pub(crate) fn set_in_flight(
110        &mut self,
111        coin: ObjectId,
112        bool: bool,
113    ) -> Result<(), TypedStoreError> {
114        if let Some(mut entry) = self.log.get(&coin)? {
115            entry.in_flight = bool;
116            self.log.insert(&coin, &entry)?;
117        } else {
118            info!(
119                ?coin,
120                "Attempted to set inflight a coin that was not in the WAL."
121            );
122
123            return Err(TypedStoreError::RocksDB(format!(
124                "Coin object {coin:?} not found in WAL."
125            )));
126        }
127        Ok(())
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use iota_sdk_types::ObjectReference;
134    use iota_types::{
135        base_types::random_object_ref,
136        transaction::{TEST_ONLY_GAS_UNIT_FOR_TRANSFER, TransactionAPI},
137    };
138
139    use super::*;
140
141    #[tokio::test]
142    async fn reserve_reclaim_reclaim() {
143        let tmp_dir = iota_common::tempdir();
144        let mut wal = WriteAheadLog::open(&tmp_dir.path().join("wal"));
145
146        let uuid = Uuid::new_v4();
147        let coin = random_object_ref();
148        let (recv, tx) = random_request(coin);
149
150        assert!(wal.reserve(uuid, coin.object_id, recv, tx.clone()).is_ok());
151
152        // Reclaim once
153        let Some(entry) = wal.reclaim(coin.object_id).unwrap() else {
154            panic!("entry not found for {}", coin.object_id);
155        };
156
157        assert_eq!(uuid, Uuid::from_bytes(entry.uuid));
158        assert_eq!(recv, entry.recipient);
159        assert_eq!(tx, entry.tx);
160
161        // Reclaim again, should still be there.
162        let Some(entry) = wal.reclaim(coin.object_id).unwrap() else {
163            panic!("entry not found for {}", coin.object_id);
164        };
165
166        assert_eq!(uuid, Uuid::from_bytes(entry.uuid));
167        assert_eq!(recv, entry.recipient);
168        assert_eq!(tx, entry.tx);
169    }
170
171    #[tokio::test]
172    async fn test_increment_wal() {
173        let tmp_dir = iota_common::tempdir();
174        let mut wal = WriteAheadLog::open(&tmp_dir.path().join("wal"));
175        let uuid = Uuid::new_v4();
176        let coin = random_object_ref();
177        let (recv0, tx0) = random_request(coin);
178
179        // First write goes through
180        wal.reserve(uuid, coin.object_id, recv0, tx0).unwrap();
181        wal.increment_retry_count(coin.object_id).unwrap();
182
183        let entry = wal.reclaim(coin.object_id).unwrap().unwrap();
184        assert_eq!(entry.retry_count, 1);
185    }
186
187    #[tokio::test]
188    async fn reserve_reserve() {
189        let tmp_dir = iota_common::tempdir();
190        let mut wal = WriteAheadLog::open(&tmp_dir.path().join("wal"));
191
192        let uuid = Uuid::new_v4();
193        let coin = random_object_ref();
194        let (recv0, tx0) = random_request(coin);
195        let (recv1, tx1) = random_request(coin);
196
197        // First write goes through
198        wal.reserve(uuid, coin.object_id, recv0, tx0).unwrap();
199
200        // Second write fails because it tries to write to the same coin
201        assert!(matches!(
202            wal.reserve(uuid, coin.object_id, recv1, tx1),
203            Err(TypedStoreError::Serialization(_)),
204        ));
205    }
206
207    #[tokio::test]
208    async fn reserve_reclaim_commit_reclaim() {
209        let tmp_dir = iota_common::tempdir();
210        let mut wal = WriteAheadLog::open(&tmp_dir.path().join("wal"));
211
212        let uuid = Uuid::new_v4();
213        let coin = random_object_ref();
214        let (recv, tx) = random_request(coin);
215
216        wal.reserve(uuid, coin.object_id, recv, tx.clone()).unwrap();
217
218        // Reclaim to show that the entry is there
219        let Some(entry) = wal.reclaim(coin.object_id).unwrap() else {
220            panic!("entry not found for {}", coin.object_id);
221        };
222
223        assert_eq!(uuid, Uuid::from_bytes(entry.uuid));
224        assert_eq!(recv, entry.recipient);
225        assert_eq!(tx, entry.tx);
226
227        // Commit the transaction, which removes it from the log.
228        wal.commit(coin.object_id).unwrap();
229
230        // Expect it to now be gone
231        assert_eq!(Ok(None), wal.reclaim(coin.object_id));
232    }
233
234    #[tokio::test]
235    async fn reserve_commit_reserve() {
236        let tmp_dir = iota_common::tempdir();
237        let mut wal = WriteAheadLog::open(&tmp_dir.path().join("wal"));
238
239        let uuid = Uuid::new_v4();
240        let coin = random_object_ref();
241        let (recv0, tx0) = random_request(coin);
242        let (recv1, tx1) = random_request(coin);
243
244        // Write the transaction
245        wal.reserve(uuid, coin.object_id, recv0, tx0).unwrap();
246
247        // Commit the transaction, which removes it from the log.
248        wal.commit(coin.object_id).unwrap();
249
250        // Write a fresh transaction, which should now pass
251        wal.reserve(uuid, coin.object_id, recv1, tx1).unwrap();
252    }
253
254    fn random_request(coin: ObjectReference) -> (Address, Transaction) {
255        let gas_price = 1;
256        let send = Address::random();
257        let recv = Address::random();
258        (
259            recv,
260            Transaction::new_pay_iota(
261                send,
262                vec![coin],
263                vec![recv],
264                vec![1000],
265                coin,
266                gas_price * TEST_ONLY_GAS_UNIT_FOR_TRANSFER,
267                gas_price,
268            )
269            .unwrap(),
270        )
271    }
272}
iota_faucet/faucet/write_ahead_log.rs

iota_faucet/faucet/
write_ahead_log.rs
