ferogram_fsm/
storage.rs

1/*
2 * Copyright (c) 2026 Ankit Chaubey <ankitchaubey.dev@gmail.com>
3 * https://github.com/ankit-chaubey
4 *
5 * Project: ferogram
6 * Website: https://ferogram.dev
7 *
8 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
9 * https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
10 * <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your option.
11 * This file may not be copied, modified, or distributed except according
12 * to those terms.
13 */
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use async_trait::async_trait;
19use dashmap::DashMap;
20
21use crate::error::StorageError;
22use crate::key::StateKey;
23
24/// Persistent storage backend for FSM state.
25///
26/// All methods are async and return `Result<_, StorageError>`.
27/// Implement this trait to add custom backends (database, Redis, etc.).
28///
29/// Built-in implementations:
30/// - [`MemoryStorage`] - in-process `DashMap`, zero setup, no persistence.
31#[async_trait]
32pub trait StateStorage: Send + Sync + 'static {
33    /// Return the current state key for this slot, or `None` if no state is set.
34    async fn get_state(&self, key: StateKey) -> Result<Option<String>, StorageError>;
35
36    /// Persist a new state. Overwrites any previously set state.
37    async fn set_state(&self, key: StateKey, state: String) -> Result<(), StorageError>;
38
39    /// Clear the state for this slot. Data is NOT cleared.
40    async fn clear_state(&self, key: StateKey) -> Result<(), StorageError>;
41
42    /// Retrieve a single data field as a raw JSON value.
43    async fn get_data(
44        &self,
45        key: StateKey,
46        field: &str,
47    ) -> Result<Option<serde_json::Value>, StorageError>;
48
49    /// Persist a single data field as a raw JSON value.
50    async fn set_data(
51        &self,
52        key: StateKey,
53        field: &str,
54        value: serde_json::Value,
55    ) -> Result<(), StorageError>;
56
57    /// Return all data fields stored for this slot.
58    async fn get_all_data(
59        &self,
60        key: StateKey,
61    ) -> Result<HashMap<String, serde_json::Value>, StorageError>;
62
63    /// Remove all data fields for this slot. State is NOT cleared.
64    async fn clear_data(&self, key: StateKey) -> Result<(), StorageError>;
65
66    /// Clear both state and all data for this slot (full reset).
67    async fn clear_all(&self, key: StateKey) -> Result<(), StorageError>;
68}
69
70/// An in-process, non-persistent [`StateStorage`] backed by `DashMap`.
71///
72/// State is lost on process restart. Suitable for development and bots that
73/// do not need persistence.
74///
75/// `MemoryStorage` is `Send + Sync + Clone` - each clone shares the same
76/// underlying map, so you can hold an `Arc<MemoryStorage>` or clone freely.
77#[derive(Clone, Default)]
78pub struct MemoryStorage {
79    entries: Arc<DashMap<StateKey, StorageEntry>>,
80}
81
82#[derive(Clone, Default)]
83struct StorageEntry {
84    state: Option<String>,
85    data: HashMap<String, serde_json::Value>,
86}
87
88impl MemoryStorage {
89    /// Create a new, empty in-memory storage.
90    pub fn new() -> Self {
91        Self::default()
92    }
93
94    /// Returns the number of active conversation slots.
95    pub fn len(&self) -> usize {
96        self.entries.len()
97    }
98
99    /// Returns `true` if no slots are currently active.
100    pub fn is_empty(&self) -> bool {
101        self.entries.is_empty()
102    }
103}
104
105#[async_trait]
106impl StateStorage for MemoryStorage {
107    async fn get_state(&self, key: StateKey) -> Result<Option<String>, StorageError> {
108        Ok(self.entries.get(&key).and_then(|e| e.state.clone()))
109    }
110
111    async fn set_state(&self, key: StateKey, state: String) -> Result<(), StorageError> {
112        self.entries.entry(key).or_default().state = Some(state);
113        Ok(())
114    }
115
116    async fn clear_state(&self, key: StateKey) -> Result<(), StorageError> {
117        if let Some(mut entry) = self.entries.get_mut(&key) {
118            entry.state = None;
119            if entry.data.is_empty() {
120                drop(entry);
121                self.entries.remove(&key);
122            }
123        }
124        Ok(())
125    }
126
127    async fn get_data(
128        &self,
129        key: StateKey,
130        field: &str,
131    ) -> Result<Option<serde_json::Value>, StorageError> {
132        Ok(self
133            .entries
134            .get(&key)
135            .and_then(|e| e.data.get(field).cloned()))
136    }
137
138    async fn set_data(
139        &self,
140        key: StateKey,
141        field: &str,
142        value: serde_json::Value,
143    ) -> Result<(), StorageError> {
144        self.entries
145            .entry(key)
146            .or_default()
147            .data
148            .insert(field.to_string(), value);
149        Ok(())
150    }
151
152    async fn get_all_data(
153        &self,
154        key: StateKey,
155    ) -> Result<HashMap<String, serde_json::Value>, StorageError> {
156        Ok(self
157            .entries
158            .get(&key)
159            .map(|e| e.data.clone())
160            .unwrap_or_default())
161    }
162
163    async fn clear_data(&self, key: StateKey) -> Result<(), StorageError> {
164        if let Some(mut entry) = self.entries.get_mut(&key) {
165            entry.data.clear();
166            if entry.state.is_none() {
167                drop(entry);
168                self.entries.remove(&key);
169            }
170        }
171        Ok(())
172    }
173
174    async fn clear_all(&self, key: StateKey) -> Result<(), StorageError> {
175        self.entries.remove(&key);
176        Ok(())
177    }
178}
ferogram_fsm/storage.rs

ferogram_fsm/
storage.rs
