rabia_kvstore_example/
smr_impl.rs

1//! # KVStore SMR Implementation
2//!
3//! This module implements the StateMachine trait for the KVStore,
4//! making it compatible with the Rabia consensus protocol.
5
6use crate::operations::{KVOperation, KVResult, StoreError};
7use crate::store::{KVStore, KVStoreConfig, ValueEntry};
8use async_trait::async_trait;
9use rabia_core::smr::StateMachine;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12
13/// KVStore state that can be serialized/deserialized
14#[derive(Debug, Clone, Serialize, Deserialize, Default)]
15pub struct KVStoreState {
16    pub data: HashMap<String, ValueEntry>,
17    pub version: u64,
18}
19
20/// SMR wrapper for KVStore that implements the StateMachine trait
21pub struct KVStoreSMR {
22    store: KVStore,
23}
24
25impl Clone for KVStoreSMR {
26    fn clone(&self) -> Self {
27        // For state machine replication, we need to create a new instance
28        // with the same configuration and current state
29        let config = self.store.config.clone();
30
31        // Create a new store instance with the same configuration
32        let new_store = tokio::task::block_in_place(|| {
33            tokio::runtime::Handle::current().block_on(async { KVStore::new(config).await })
34        })
35        .expect("Failed to create new KVStore instance");
36
37        // Copy the current state
38        let current_data = self.store.get_all_data();
39        let current_version = self.store.current_version();
40
41        new_store.set_all_data(current_data);
42        new_store.set_version(current_version);
43
44        Self { store: new_store }
45    }
46}
47
48impl KVStoreSMR {
49    /// Create a new KVStoreSMR instance
50    pub async fn new(config: KVStoreConfig) -> Result<Self, StoreError> {
51        let store = KVStore::new(config).await?;
52        Ok(Self { store })
53    }
54
55    /// Create a new KVStoreSMR instance with default configuration
56    pub async fn new_default() -> Result<Self, StoreError> {
57        Self::new(KVStoreConfig::default()).await
58    }
59
60    /// Get access to the underlying store for advanced operations
61    pub fn store(&self) -> &KVStore {
62        &self.store
63    }
64}
65
66#[async_trait]
67impl StateMachine for KVStoreSMR {
68    type Command = KVOperation;
69    type Response = KVResult;
70    type State = KVStoreState;
71
72    async fn apply_command(&mut self, command: Self::Command) -> Self::Response {
73        match command {
74            KVOperation::Set { key, value } => match self.store.set(&key, &value).await {
75                Ok(result) => result,
76                Err(e) => KVResult::Error(e.to_string()),
77            },
78            KVOperation::Get { key } => match self.store.get(&key).await {
79                Ok(Some(_)) => KVResult::Success,
80                Ok(None) => KVResult::NotFound,
81                Err(e) => KVResult::Error(e.to_string()),
82            },
83            KVOperation::Delete { key } => match self.store.delete(&key).await {
84                Ok(result) => result,
85                Err(e) => KVResult::Error(e.to_string()),
86            },
87            KVOperation::Exists { key } => match self.store.exists(&key).await {
88                Ok(true) => KVResult::Success,
89                Ok(false) => KVResult::NotFound,
90                Err(e) => KVResult::Error(e.to_string()),
91            },
92        }
93    }
94
95    fn get_state(&self) -> Self::State {
96        // Create a state snapshot from the current store data
97        let data = self.store.get_all_data();
98        let version = self.store.current_version();
99
100        KVStoreState { data, version }
101    }
102
103    fn set_state(&mut self, state: Self::State) {
104        // Clear current data and restore from state
105        self.store.set_all_data(state.data);
106        self.store.set_version(state.version);
107    }
108
109    fn serialize_state(&self) -> Vec<u8> {
110        let state = self.get_state();
111        bincode::serialize(&state).unwrap_or_default()
112    }
113
114    fn deserialize_state(&mut self, data: &[u8]) -> Result<(), Box<dyn std::error::Error>> {
115        let state: KVStoreState = bincode::deserialize(data)?;
116        self.set_state(state);
117        Ok(())
118    }
119
120    async fn apply_commands(&mut self, commands: Vec<Self::Command>) -> Vec<Self::Response> {
121        let mut responses = Vec::with_capacity(commands.len());
122        for command in commands {
123            responses.push(self.apply_command(command).await);
124        }
125        responses
126    }
127
128    fn is_deterministic(&self) -> bool {
129        true
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[tokio::test]
138    async fn test_kvstore_smr_basic_operations() {
139        let mut smr = KVStoreSMR::new_default().await.unwrap();
140
141        // Test SET command
142        let set_cmd = KVOperation::Set {
143            key: "test_key".to_string(),
144            value: "test_value".to_string(),
145        };
146        let result = smr.apply_command(set_cmd).await;
147        assert!(result.is_success());
148
149        // Test GET command
150        let get_cmd = KVOperation::Get {
151            key: "test_key".to_string(),
152        };
153        let result = smr.apply_command(get_cmd).await;
154        assert!(result.is_success());
155
156        // Test EXISTS command
157        let exists_cmd = KVOperation::Exists {
158            key: "test_key".to_string(),
159        };
160        let result = smr.apply_command(exists_cmd).await;
161        assert!(result.is_success());
162
163        // Test DELETE command
164        let delete_cmd = KVOperation::Delete {
165            key: "test_key".to_string(),
166        };
167        let result = smr.apply_command(delete_cmd).await;
168        assert!(result.is_success());
169
170        // Test GET after DELETE
171        let get_cmd = KVOperation::Get {
172            key: "test_key".to_string(),
173        };
174        let result = smr.apply_command(get_cmd).await;
175        assert!(result.is_not_found());
176    }
177
178    #[tokio::test]
179    async fn test_kvstore_smr_state_serialization() {
180        let mut smr = KVStoreSMR::new_default().await.unwrap();
181
182        // Add some data
183        let set_cmd1 = KVOperation::Set {
184            key: "key1".to_string(),
185            value: "value1".to_string(),
186        };
187        let set_cmd2 = KVOperation::Set {
188            key: "key2".to_string(),
189            value: "value2".to_string(),
190        };
191        smr.apply_command(set_cmd1).await;
192        smr.apply_command(set_cmd2).await;
193
194        // Serialize state
195        let serialized = smr.serialize_state();
196        assert!(!serialized.is_empty());
197
198        // Create new SMR instance and deserialize
199        let mut new_smr = KVStoreSMR::new_default().await.unwrap();
200        new_smr.deserialize_state(&serialized).unwrap();
201
202        // Verify state was restored
203        let state = new_smr.get_state();
204        assert_eq!(state.data.len(), 2);
205        assert!(state.data.contains_key("key1"));
206        assert!(state.data.contains_key("key2"));
207        assert_eq!(state.data["key1"].value, "value1");
208        assert_eq!(state.data["key2"].value, "value2");
209    }
210
211    #[tokio::test]
212    async fn test_kvstore_smr_multiple_commands() {
213        let mut smr = KVStoreSMR::new_default().await.unwrap();
214
215        let commands = vec![
216            KVOperation::Set {
217                key: "key1".to_string(),
218                value: "value1".to_string(),
219            },
220            KVOperation::Set {
221                key: "key2".to_string(),
222                value: "value2".to_string(),
223            },
224            KVOperation::Get {
225                key: "key1".to_string(),
226            },
227            KVOperation::Delete {
228                key: "key2".to_string(),
229            },
230            KVOperation::Get {
231                key: "key2".to_string(),
232            },
233        ];
234
235        let responses = smr.apply_commands(commands).await;
236        assert_eq!(responses.len(), 5);
237        assert!(responses[0].is_success()); // SET key1
238        assert!(responses[1].is_success()); // SET key2
239        assert!(responses[2].is_success()); // GET key1
240        assert!(responses[3].is_success()); // DELETE key2
241        assert!(responses[4].is_not_found()); // GET key2 (after delete)
242    }
243}
rabia_kvstore_example/smr_impl.rs

rabia_kvstore_example/
smr_impl.rs
