char/avril

(wip) why'd you have to go and make things so complicaaaaated?

git clone https://git.t4t.associates/char/avril

Charlotte Somadd `avril:kv` module78f4748

main
4.4 KiB123 linesraw
1use std::collections::BTreeMap;
2use std::io;
3
4use super::{Decision, Store, check_keyspace_name, parse_stamp};
5
6// Bounds on one batch so a hostile or buggy guest cannot make a single
7// commit arbitrarily large. Bytes count keyspace names, keys and values.
8pub const MAX_BATCH_OPS: usize = 1_000;
9pub const MAX_BATCH_BYTES: usize = 4 << 20;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum BatchOp {
13    Check {
14        keyspace: String,
15        key: Vec<u8>,
16        expected: Option<String>,
17    },
18    Set {
19        keyspace: String,
20        key: Vec<u8>,
21        value: Vec<u8>,
22    },
23    Delete {
24        keyspace: String,
25        key: Vec<u8>,
26    },
27}
28
29impl BatchOp {
30    pub fn keyspace(&self) -> &str {
31        match self {
32            BatchOp::Check { keyspace, .. }
33            | BatchOp::Set { keyspace, .. }
34            | BatchOp::Delete { keyspace, .. } => keyspace,
35        }
36    }
37}
38
39fn invalid(message: &str) -> io::Error {
40    io::Error::new(io::ErrorKind::InvalidInput, message)
41}
42
43impl Store {
44    pub fn atomic_commit(&self, ops: &[BatchOp]) -> io::Result<(bool, Option<String>)> {
45        if ops.len() > MAX_BATCH_OPS {
46            return Err(invalid("Atomic batch exceeds operation limit"));
47        }
48        let mut bytes = 0;
49        for op in ops {
50            check_keyspace_name(op.keyspace())?;
51            bytes += op.keyspace().len()
52                + match op {
53                    BatchOp::Check { key, .. } | BatchOp::Delete { key, .. } => key.len(),
54                    BatchOp::Set { key, value, .. } => key.len() + value.len(),
55                };
56            // Malformed stamps are caller error, not a quiet check failure.
57            if let BatchOp::Check { expected: Some(stamp), .. } = op {
58                parse_stamp(stamp)?;
59            }
60        }
61        if bytes > MAX_BATCH_BYTES {
62            return Err(invalid("Atomic batch exceeds size limit"));
63        }
64        if ops.is_empty() {
65            return Ok((true, None));
66        }
67        self.write_tx(|tx| {
68            // Checks observe pre-commit state and never register keyspaces: a
69            // missing registry row means the entry is absent.
70            for op in ops {
71                if let BatchOp::Check { keyspace, key, expected } = op {
72                    let current = tx.get(keyspace, key.as_slice())?;
73                    let pass = match (&current, expected) {
74                        (None, None) => true,
75                        (Some(_), None) => false,
76                        (None, Some(_)) => false,
77                        // Both sides are canonical renderings (parse_stamp
78                        // accepted `want` above), so string equality suffices.
79                        (Some((_, have)), Some(want)) => have == want,
80                    };
81                    if !pass {
82                        // Abort drops the transaction uncommitted: a failed
83                        // batch changes nothing, not even metadata.
84                        return Ok(Decision::Abort((false, None)));
85                    }
86                }
87            }
88            let mut mutations: BTreeMap<(&str, &[u8]), Option<&[u8]>> = BTreeMap::new();
89            for op in ops {
90                match op {
91                    BatchOp::Check { .. } => {}
92                    BatchOp::Set { keyspace, key, value } => {
93                        mutations
94                            .insert((keyspace.as_str(), key.as_slice()), Some(value.as_slice()));
95                    }
96                    BatchOp::Delete { keyspace, key } => {
97                        mutations.insert((keyspace.as_str(), key.as_slice()), None);
98                    }
99                }
100            }
101            // A check-only batch commits nothing and burns no versionstamp.
102            if mutations.is_empty() {
103                return Ok(Decision::Abort((true, None)));
104            }
105            for ((keyspace, key), mutation) in &mutations {
106                match mutation {
107                    Some(value) => {
108                        tx.put(keyspace, key, value)?;
109                    }
110                    None => {
111                        tx.delete(keyspace, key)?;
112                    }
113                }
114            }
115            let stamp = match tx.stamp() {
116                Some(stamp) => stamp,
117                // Every mutation missed: the batch still consumes a commit.
118                None => tx.allocate_stamp()?,
119            };
120            Ok(Decision::Commit((true, Some(stamp))))
121        })
122    }
123}