steel_core/command/
storage.rs

1//! Persistent command storage isolated by Steel domain.
2
3use std::{
4    collections::BTreeMap,
5    io::{self, Cursor},
6    sync::atomic::{AtomicU64, Ordering},
7};
8
9use rustc_hash::FxHashMap;
10use serde::{Deserialize, Serialize};
11use simdnbt::owned::{NbtCompound, read_compound};
12use steel_utils::{
13    Identifier,
14    locks::{AsyncMutex, SyncRwLock},
15};
16
17use crate::{server::worlds::WorldMap, world::World};
18use steel_utils::saved_data::names as saved_data_names;
19
20#[derive(Default, Deserialize, Serialize)]
21struct PersistentCommandStorage {
22    entries: BTreeMap<String, Vec<u8>>,
23}
24
25struct CommandStorageSaveSnapshot {
26    revision: u64,
27    state: PersistentCommandStorage,
28}
29
30/// Vanilla command storage for one Steel domain.
31pub(crate) struct CommandStorage {
32    entries: SyncRwLock<FxHashMap<Identifier, NbtCompound>>,
33    revision: AtomicU64,
34    saved_revision: AtomicU64,
35}
36
37impl CommandStorage {
38    /// Creates an empty, clean command storage.
39    #[must_use]
40    pub(crate) fn new() -> Self {
41        Self {
42            entries: SyncRwLock::new(FxHashMap::default()),
43            revision: AtomicU64::new(0),
44            saved_revision: AtomicU64::new(0),
45        }
46    }
47
48    fn from_persistent(persistent: PersistentCommandStorage) -> io::Result<Self> {
49        let mut entries = FxHashMap::default();
50        for (raw_key, bytes) in persistent.entries {
51            let key = raw_key.parse::<Identifier>().map_err(|error| {
52                io::Error::new(
53                    io::ErrorKind::InvalidData,
54                    format!("invalid command storage key '{raw_key}': {error}"),
55                )
56            })?;
57            let mut cursor = Cursor::new(bytes.as_slice());
58            let compound = read_compound(&mut cursor).map_err(|error| {
59                io::Error::new(
60                    io::ErrorKind::InvalidData,
61                    format!("invalid NBT for command storage key '{key}': {error:?}"),
62                )
63            })?;
64            if cursor.position() != bytes.len() as u64 {
65                return Err(io::Error::new(
66                    io::ErrorKind::InvalidData,
67                    format!("trailing NBT data for command storage key '{key}'"),
68                ));
69            }
70            if !compound.is_empty() {
71                entries.insert(key, compound);
72            }
73        }
74        Ok(Self {
75            entries: SyncRwLock::new(entries),
76            revision: AtomicU64::new(0),
77            saved_revision: AtomicU64::new(0),
78        })
79    }
80
81    /// Returns the compound stored at `id`, or an empty compound when absent.
82    #[must_use]
83    pub(crate) fn get(&self, id: &Identifier) -> NbtCompound {
84        self.entries
85            .read()
86            .get(id)
87            .cloned()
88            .unwrap_or_else(NbtCompound::new)
89    }
90
91    /// Stores a compound, removing the key when the compound is empty.
92    pub(crate) fn set(&self, id: Identifier, contents: NbtCompound) {
93        let mut entries = self.entries.write();
94        if contents.is_empty() {
95            entries.remove(&id);
96        } else {
97            entries.insert(id, contents);
98        }
99        self.revision.fetch_add(1, Ordering::Release);
100    }
101
102    /// Returns stored keys in stable resource-location order.
103    #[must_use]
104    pub(crate) fn keys(&self) -> Vec<Identifier> {
105        let mut keys = self.entries.read().keys().cloned().collect::<Vec<_>>();
106        keys.sort_by_cached_key(ToString::to_string);
107        keys
108    }
109
110    fn pending_save(&self) -> Option<CommandStorageSaveSnapshot> {
111        let entries = self.entries.read();
112        let revision = self.revision.load(Ordering::Acquire);
113        if revision == self.saved_revision.load(Ordering::Acquire) {
114            return None;
115        }
116
117        let entries = entries
118            .iter()
119            .map(|(key, compound)| {
120                let mut bytes = Vec::new();
121                compound.write(&mut bytes);
122                (key.to_string(), bytes)
123            })
124            .collect();
125        Some(CommandStorageSaveSnapshot {
126            revision,
127            state: PersistentCommandStorage { entries },
128        })
129    }
130
131    fn mark_saved(&self, revision: u64) {
132        self.saved_revision.fetch_max(revision, Ordering::Release);
133    }
134}
135
136impl Default for CommandStorage {
137    fn default() -> Self {
138        Self::new()
139    }
140}
141
142/// Loaded command storages keyed by Steel domain.
143pub(crate) struct DomainCommandStorage {
144    storages: BTreeMap<String, CommandStorage>,
145    save_lock: AsyncMutex<()>,
146}
147
148impl DomainCommandStorage {
149    /// Loads one command storage through each domain's default world.
150    pub(crate) async fn load(worlds: &WorldMap) -> io::Result<Self> {
151        let mut domains = worlds.domain_names().collect::<Vec<_>>();
152        domains.sort_unstable();
153        let mut storages = BTreeMap::new();
154        for domain in domains {
155            let world = domain_default_world(worlds, domain)?;
156            let persistent: PersistentCommandStorage = world
157                .saved_data
158                .load_or_default(saved_data_names::COMMAND_STORAGE)
159                .await
160                .map_err(|error| storage_io_error(domain, error))?;
161            let storage = CommandStorage::from_persistent(persistent)
162                .map_err(|error| storage_io_error(domain, error))?;
163            storages.insert(domain.to_owned(), storage);
164        }
165        Ok(Self {
166            storages,
167            save_lock: AsyncMutex::new(()),
168        })
169    }
170
171    /// Returns command storage for a domain.
172    #[must_use]
173    pub(crate) fn get(&self, domain: &str) -> Option<&CommandStorage> {
174        self.storages.get(domain)
175    }
176
177    /// Saves every dirty domain storage and returns the number written.
178    pub(crate) async fn save(&self, worlds: &WorldMap) -> io::Result<usize> {
179        let _save_guard = self.save_lock.lock().await;
180        let mut saved = 0;
181        for (domain, storage) in &self.storages {
182            let Some(snapshot) = storage.pending_save() else {
183                continue;
184            };
185            let world = domain_default_world(worlds, domain)?;
186            world
187                .saved_data
188                .save(saved_data_names::COMMAND_STORAGE, &snapshot.state)
189                .await
190                .map_err(|error| storage_io_error(domain, error))?;
191            storage.mark_saved(snapshot.revision);
192            saved += 1;
193        }
194        Ok(saved)
195    }
196}
197
198fn domain_default_world<'a>(worlds: &'a WorldMap, domain: &str) -> io::Result<&'a World> {
199    worlds
200        .default_world(domain)
201        .map(AsRef::as_ref)
202        .ok_or_else(|| {
203            io::Error::new(
204                io::ErrorKind::NotFound,
205                format!("domain '{domain}' has no loaded default world"),
206            )
207        })
208}
209
210fn storage_io_error(domain: &str, error: io::Error) -> io::Error {
211    io::Error::new(
212        error.kind(),
213        format!("command storage I/O failed for domain '{domain}': {error}"),
214    )
215}
216
217#[cfg(test)]
218mod tests {
219    use std::{
220        env::temp_dir,
221        slice,
222        time::{SystemTime, UNIX_EPOCH},
223    };
224
225    use simdnbt::owned::{NbtCompound, NbtTag};
226    use steel_utils::Identifier;
227    use tokio::fs;
228
229    use steel_utils::saved_data::SavedDataManager;
230
231    use super::*;
232
233    #[test]
234    fn missing_and_empty_values_match_vanilla_storage_semantics() {
235        let storage = CommandStorage::new();
236        let key = Identifier::from_steel("data");
237
238        assert!(storage.get(&key).is_empty());
239        let mut value = NbtCompound::new();
240        value.insert("value", 3);
241        storage.set(key.clone(), value);
242        assert_eq!(storage.keys(), slice::from_ref(&key));
243
244        storage.set(key.clone(), NbtCompound::new());
245        assert!(storage.get(&key).is_empty());
246        assert!(storage.keys().is_empty());
247    }
248
249    #[tokio::test]
250    async fn persistent_storage_round_trips_binary_nbt() {
251        let unique = SystemTime::now()
252            .duration_since(UNIX_EPOCH)
253            .expect("system time should be after Unix epoch")
254            .as_nanos();
255        let path = temp_dir().join(format!("steel-command-storage-{unique}"));
256        let manager = SavedDataManager::new(Some(&path));
257        let storage = CommandStorage::new();
258        let key = Identifier::from_steel("nested");
259        let mut nested = NbtCompound::new();
260        nested.insert("bytes", NbtTag::ByteArray(vec![1, 2, 255]));
261        let mut value = NbtCompound::new();
262        value.insert("nested", nested);
263        storage.set(key.clone(), value.clone());
264
265        let snapshot = storage.pending_save().expect("storage should become dirty");
266        manager
267            .save(saved_data_names::COMMAND_STORAGE, &snapshot.state)
268            .await
269            .expect("command storage should save");
270        storage.mark_saved(snapshot.revision);
271        assert!(storage.pending_save().is_none());
272
273        let persistent: PersistentCommandStorage = manager
274            .load_or_default(saved_data_names::COMMAND_STORAGE)
275            .await
276            .expect("command storage should load");
277        let restored =
278            CommandStorage::from_persistent(persistent).expect("stored NBT should validate");
279        assert_eq!(restored.get(&key), value);
280
281        fs::remove_dir_all(path)
282            .await
283            .expect("temporary command storage directory should be removed");
284    }
285
286    #[test]
287    fn mutation_after_snapshot_remains_dirty() {
288        let storage = CommandStorage::new();
289        let first = Identifier::from_steel("first");
290        let second = Identifier::from_steel("second");
291        let mut value = NbtCompound::new();
292        value.insert("value", 1);
293        storage.set(first, value.clone());
294        let snapshot = storage.pending_save().expect("storage should become dirty");
295
296        storage.set(second.clone(), value);
297        storage.mark_saved(snapshot.revision);
298
299        let pending = storage
300            .pending_save()
301            .expect("newer mutation should remain dirty");
302        assert!(pending.revision > snapshot.revision);
303        assert!(pending.state.entries.contains_key(&second.to_string()));
304    }
305
306    #[test]
307    fn domains_keep_independent_command_storage() {
308        let storages = DomainCommandStorage {
309            storages: [
310                ("alpha".to_owned(), CommandStorage::new()),
311                ("beta".to_owned(), CommandStorage::new()),
312            ]
313            .into_iter()
314            .collect(),
315            save_lock: AsyncMutex::new(()),
316        };
317        let key = Identifier::from_steel("data");
318        let mut value = NbtCompound::new();
319        value.insert("value", 1);
320        storages
321            .get("alpha")
322            .expect("alpha storage should exist")
323            .set(key.clone(), value);
324
325        assert!(
326            storages
327                .get("beta")
328                .expect("beta storage should exist")
329                .get(&key)
330                .is_empty()
331        );
332    }
333
334    #[test]
335    fn invalid_persisted_nbt_is_rejected() {
336        let persistent = PersistentCommandStorage {
337            entries: [("steel:data".to_owned(), vec![10])].into_iter().collect(),
338        };
339
340        let error = CommandStorage::from_persistent(persistent)
341            .err()
342            .expect("invalid NBT should fail loading");
343        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
344    }
345}
steel_core/command/storage.rs

steel_core/command/
storage.rs
