Skip to main content

linera_views/backends/
dual.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Implements [`crate::store::KeyValueStore`] by combining two existing stores.
5
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9#[cfg(with_testing)]
10use crate::store::TestKeyValueDatabase;
11use crate::{
12    batch::Batch,
13    store::{
14        KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore, WithError,
15        WritableKeyValueStore,
16    },
17};
18
19/// A dual database.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct DualDatabase<D1, D2, A> {
22    /// The first database.
23    pub first_database: D1,
24    /// The second database.
25    pub second_database: D2,
26    /// Marker for the static root key assignment.
27    _marker: std::marker::PhantomData<A>,
28}
29
30/// The initial configuration of the system.
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct DualStoreConfig<C1, C2> {
33    /// The first config.
34    pub first_config: C1,
35    /// The second config.
36    pub second_config: C2,
37}
38
39/// The store in use.
40#[derive(Clone, Copy, Debug)]
41pub enum StoreInUse {
42    /// The first store.
43    First,
44    /// The second store.
45    Second,
46}
47
48/// The trait for a (static) root key assignment.
49pub trait DualStoreRootKeyAssignment {
50    /// Obtains the store assigned to this root key.
51    fn assigned_store(root_key: &[u8]) -> Result<StoreInUse, bcs::Error>;
52}
53
54/// A partition opened in one of the two databases.
55#[derive(Clone)]
56pub enum DualStore<S1, S2> {
57    /// The first store.
58    First(S1),
59    /// The second store.
60    Second(S2),
61}
62
63impl<D1, D2, A> WithError for DualDatabase<D1, D2, A>
64where
65    D1: WithError,
66    D2: WithError,
67{
68    type Error = DualStoreError<D1::Error, D2::Error>;
69}
70
71impl<S1, S2> WithError for DualStore<S1, S2>
72where
73    S1: WithError,
74    S2: WithError,
75{
76    type Error = DualStoreError<S1::Error, S2::Error>;
77}
78
79impl<S1, S2> ReadableKeyValueStore for DualStore<S1, S2>
80where
81    S1: ReadableKeyValueStore,
82    S2: ReadableKeyValueStore,
83{
84    // TODO(#2524): consider changing MAX_KEY_SIZE into a function.
85    const MAX_KEY_SIZE: usize = if S1::MAX_KEY_SIZE < S2::MAX_KEY_SIZE {
86        S1::MAX_KEY_SIZE
87    } else {
88        S2::MAX_KEY_SIZE
89    };
90
91    fn root_key(&self) -> Result<Vec<u8>, Self::Error> {
92        Ok(match self {
93            Self::First(store) => store.root_key().map_err(DualStoreError::First)?,
94            Self::Second(store) => store.root_key().map_err(DualStoreError::Second)?,
95        })
96    }
97
98    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
99        let result = match self {
100            Self::First(store) => store
101                .read_value_bytes(key)
102                .await
103                .map_err(DualStoreError::First)?,
104            Self::Second(store) => store
105                .read_value_bytes(key)
106                .await
107                .map_err(DualStoreError::Second)?,
108        };
109        Ok(result)
110    }
111
112    async fn contains_key(&self, key: &[u8]) -> Result<bool, Self::Error> {
113        let result = match self {
114            Self::First(store) => store
115                .contains_key(key)
116                .await
117                .map_err(DualStoreError::First)?,
118            Self::Second(store) => store
119                .contains_key(key)
120                .await
121                .map_err(DualStoreError::Second)?,
122        };
123        Ok(result)
124    }
125
126    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, Self::Error> {
127        let result = match self {
128            Self::First(store) => store
129                .contains_keys(keys)
130                .await
131                .map_err(DualStoreError::First)?,
132            Self::Second(store) => store
133                .contains_keys(keys)
134                .await
135                .map_err(DualStoreError::Second)?,
136        };
137        Ok(result)
138    }
139
140    async fn read_multi_values_bytes(
141        &self,
142        keys: &[Vec<u8>],
143    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
144        let result = match self {
145            Self::First(store) => store
146                .read_multi_values_bytes(keys)
147                .await
148                .map_err(DualStoreError::First)?,
149            Self::Second(store) => store
150                .read_multi_values_bytes(keys)
151                .await
152                .map_err(DualStoreError::Second)?,
153        };
154        Ok(result)
155    }
156
157    async fn find_keys_by_prefix(&self, key_prefix: &[u8]) -> Result<Vec<Vec<u8>>, Self::Error> {
158        let result = match self {
159            Self::First(store) => store
160                .find_keys_by_prefix(key_prefix)
161                .await
162                .map_err(DualStoreError::First)?,
163            Self::Second(store) => store
164                .find_keys_by_prefix(key_prefix)
165                .await
166                .map_err(DualStoreError::Second)?,
167        };
168        Ok(result)
169    }
170
171    async fn find_key_values_by_prefix(
172        &self,
173        key_prefix: &[u8],
174    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Self::Error> {
175        let result = match self {
176            Self::First(store) => store
177                .find_key_values_by_prefix(key_prefix)
178                .await
179                .map_err(DualStoreError::First)?,
180            Self::Second(store) => store
181                .find_key_values_by_prefix(key_prefix)
182                .await
183                .map_err(DualStoreError::Second)?,
184        };
185        Ok(result)
186    }
187}
188
189impl<S1, S2> WritableKeyValueStore for DualStore<S1, S2>
190where
191    S1: WritableKeyValueStore,
192    S2: WritableKeyValueStore,
193{
194    const MAX_VALUE_SIZE: usize = usize::MAX;
195
196    async fn write_batch(&self, batch: Batch) -> Result<(), Self::Error> {
197        match self {
198            Self::First(store) => store
199                .write_batch(batch)
200                .await
201                .map_err(DualStoreError::First)?,
202            Self::Second(store) => store
203                .write_batch(batch)
204                .await
205                .map_err(DualStoreError::Second)?,
206        }
207        Ok(())
208    }
209
210    async fn clear_journal(&self) -> Result<(), Self::Error> {
211        match self {
212            Self::First(store) => store.clear_journal().await.map_err(DualStoreError::First)?,
213            Self::Second(store) => store
214                .clear_journal()
215                .await
216                .map_err(DualStoreError::Second)?,
217        }
218        Ok(())
219    }
220}
221
222impl<D1, D2, A> KeyValueDatabase for DualDatabase<D1, D2, A>
223where
224    D1: KeyValueDatabase,
225    D2: KeyValueDatabase,
226    A: DualStoreRootKeyAssignment + linera_base::util::traits::AutoTraits,
227{
228    type Config = DualStoreConfig<D1::Config, D2::Config>;
229    type Store = DualStore<D1::Store, D2::Store>;
230
231    fn get_name() -> String {
232        format!("dual {} and {}", D1::get_name(), D2::get_name())
233    }
234
235    async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, Self::Error> {
236        let first_database = D1::connect(&config.first_config, namespace)
237            .await
238            .map_err(DualStoreError::First)?;
239        let second_database = D2::connect(&config.second_config, namespace)
240            .await
241            .map_err(DualStoreError::Second)?;
242        let database = Self {
243            first_database,
244            second_database,
245            _marker: std::marker::PhantomData,
246        };
247        Ok(database)
248    }
249
250    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
251        match A::assigned_store(root_key)? {
252            StoreInUse::First => {
253                let store = self
254                    .first_database
255                    .open_shared(root_key)
256                    .map_err(DualStoreError::First)?;
257                Ok(DualStore::First(store))
258            }
259            StoreInUse::Second => {
260                let store = self
261                    .second_database
262                    .open_shared(root_key)
263                    .map_err(DualStoreError::Second)?;
264                Ok(DualStore::Second(store))
265            }
266        }
267    }
268
269    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
270        match A::assigned_store(root_key)? {
271            StoreInUse::First => {
272                let store = self
273                    .first_database
274                    .open_exclusive(root_key)
275                    .map_err(DualStoreError::First)?;
276                Ok(DualStore::First(store))
277            }
278            StoreInUse::Second => {
279                let store = self
280                    .second_database
281                    .open_exclusive(root_key)
282                    .map_err(DualStoreError::Second)?;
283                Ok(DualStore::Second(store))
284            }
285        }
286    }
287
288    async fn list_all(config: &Self::Config) -> Result<Vec<String>, Self::Error> {
289        let namespaces1 = D1::list_all(&config.first_config)
290            .await
291            .map_err(DualStoreError::First)?;
292        let mut namespaces = Vec::new();
293        for namespace in namespaces1 {
294            if D2::exists(&config.second_config, &namespace)
295                .await
296                .map_err(DualStoreError::Second)?
297            {
298                namespaces.push(namespace);
299            } else {
300                tracing::warn!("Namespace {} only exists in the first store", namespace);
301            }
302        }
303        Ok(namespaces)
304    }
305
306    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
307        let mut root_keys = self
308            .first_database
309            .list_root_keys()
310            .await
311            .map_err(DualStoreError::First)?;
312        root_keys.extend(
313            self.second_database
314                .list_root_keys()
315                .await
316                .map_err(DualStoreError::Second)?,
317        );
318        Ok(root_keys)
319    }
320
321    async fn exists(config: &Self::Config, namespace: &str) -> Result<bool, Self::Error> {
322        Ok(D1::exists(&config.first_config, namespace)
323            .await
324            .map_err(DualStoreError::First)?
325            && D2::exists(&config.second_config, namespace)
326                .await
327                .map_err(DualStoreError::Second)?)
328    }
329
330    async fn create(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
331        let exists1 = D1::exists(&config.first_config, namespace)
332            .await
333            .map_err(DualStoreError::First)?;
334        let exists2 = D2::exists(&config.second_config, namespace)
335            .await
336            .map_err(DualStoreError::Second)?;
337        if exists1 && exists2 {
338            return Err(DualStoreError::StoreAlreadyExists);
339        }
340        if !exists1 {
341            D1::create(&config.first_config, namespace)
342                .await
343                .map_err(DualStoreError::First)?;
344        }
345        if !exists2 {
346            D2::create(&config.second_config, namespace)
347                .await
348                .map_err(DualStoreError::Second)?;
349        }
350        Ok(())
351    }
352
353    async fn delete(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
354        D1::delete(&config.first_config, namespace)
355            .await
356            .map_err(DualStoreError::First)?;
357        D2::delete(&config.second_config, namespace)
358            .await
359            .map_err(DualStoreError::Second)?;
360        Ok(())
361    }
362}
363
364#[cfg(with_testing)]
365impl<D1, D2, A> TestKeyValueDatabase for DualDatabase<D1, D2, A>
366where
367    D1: TestKeyValueDatabase,
368    D2: TestKeyValueDatabase,
369    A: DualStoreRootKeyAssignment + linera_base::util::traits::AutoTraits,
370{
371    async fn new_test_config() -> Result<Self::Config, Self::Error> {
372        let first_config = D1::new_test_config().await.map_err(DualStoreError::First)?;
373        let second_config = D2::new_test_config()
374            .await
375            .map_err(DualStoreError::Second)?;
376        Ok(DualStoreConfig {
377            first_config,
378            second_config,
379        })
380    }
381}
382
383/// The error type for [`DualStore`].
384#[derive(Error, Debug)]
385pub enum DualStoreError<E1, E2> {
386    /// Store already exists during a create operation
387    #[error("Store already exists during a create operation")]
388    StoreAlreadyExists,
389
390    /// Serialization error with BCS.
391    #[error(transparent)]
392    BcsError(#[from] bcs::Error),
393
394    /// First store.
395    #[error("Error in first store: {0}")]
396    First(E1),
397
398    /// Second store.
399    #[error("Error in second store: {0}")]
400    Second(E2),
401}
402
403impl<E1, E2> KeyValueStoreError for DualStoreError<E1, E2>
404where
405    E1: KeyValueStoreError,
406    E2: KeyValueStoreError,
407{
408    const BACKEND: &'static str = "dual_store";
409
410    fn must_reload_view(&self) -> bool {
411        match self {
412            DualStoreError::First(e) => e.must_reload_view(),
413            DualStoreError::Second(e) => e.must_reload_view(),
414            _ => false,
415        }
416    }
417}