linera_views/backends/
value_splitting.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Adds support for large values to a given store by splitting them between several keys.
5
6use linera_base::ensure;
7use thiserror::Error;
8
9use crate::{
10    batch::{Batch, WriteOperation},
11    store::{
12        KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore, WithError,
13        WritableKeyValueStore,
14    },
15};
16#[cfg(with_testing)]
17use crate::{
18    memory::{MemoryStore, MemoryStoreError},
19    store::TestKeyValueDatabase,
20};
21
22/// A key-value database with no size limit for values.
23///
24/// It wraps a key-value store, potentially _with_ a size limit, and automatically
25/// splits up large values into smaller ones. A single logical key-value pair is
26/// stored as multiple smaller key-value pairs in the wrapped store.
27/// See the `README.md` for additional details.
28#[derive(Clone)]
29pub struct ValueSplittingDatabase<D> {
30    /// The underlying database.
31    database: D,
32}
33
34/// A key-value store with no size limit for values.
35#[derive(Clone)]
36pub struct ValueSplittingStore<S> {
37    /// The underlying store.
38    store: S,
39}
40
41/// The composed error type built from the inner error type.
42#[derive(Error, Debug)]
43pub enum ValueSplittingError<E> {
44    /// inner store error
45    #[error(transparent)]
46    InnerStoreError(#[from] E),
47
48    /// The key is of length less than 4, so we cannot extract the first byte
49    #[error("the key is of length less than 4, so we cannot extract the first byte")]
50    TooShortKey,
51
52    /// Value segment is missing from the database
53    #[error("value segment is missing from the database")]
54    MissingSegment,
55
56    /// No count of size `u32` is available in the value
57    #[error("no count of size u32 is available in the value")]
58    NoCountAvailable,
59}
60
61impl<E: KeyValueStoreError> From<bcs::Error> for ValueSplittingError<E> {
62    fn from(error: bcs::Error) -> Self {
63        let error = E::from(error);
64        ValueSplittingError::InnerStoreError(error)
65    }
66}
67
68impl<E: KeyValueStoreError + 'static> KeyValueStoreError for ValueSplittingError<E> {
69    const BACKEND: &'static str = "value splitting";
70
71    fn must_reload_view(&self) -> bool {
72        match self {
73            ValueSplittingError::InnerStoreError(e) => e.must_reload_view(),
74            _ => false,
75        }
76    }
77}
78
79impl<S> WithError for ValueSplittingDatabase<S>
80where
81    S: WithError,
82    S::Error: 'static,
83{
84    type Error = ValueSplittingError<S::Error>;
85}
86
87impl<D> WithError for ValueSplittingStore<D>
88where
89    D: WithError,
90    D::Error: 'static,
91{
92    type Error = ValueSplittingError<D::Error>;
93}
94
95impl<S> ReadableKeyValueStore for ValueSplittingStore<S>
96where
97    S: ReadableKeyValueStore,
98    S::Error: 'static,
99{
100    const MAX_KEY_SIZE: usize = S::MAX_KEY_SIZE - 4;
101
102    fn root_key(&self) -> Result<Vec<u8>, Self::Error> {
103        Ok(self.store.root_key()?)
104    }
105
106    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
107        let mut big_key = key.to_vec();
108        big_key.extend(&[0, 0, 0, 0]);
109        let value = self.store.read_value_bytes(&big_key).await?;
110        let Some(value) = value else {
111            return Ok(None);
112        };
113        let count = Self::read_count_from_value(&value)?;
114        let mut big_value = value[4..].to_vec();
115        if count == 1 {
116            return Ok(Some(big_value));
117        }
118        let mut big_keys = Vec::new();
119        for i in 1..count {
120            let big_key_segment = Self::get_segment_key(key, i)?;
121            big_keys.push(big_key_segment);
122        }
123        let segments = self.store.read_multi_values_bytes(&big_keys).await?;
124        for segment in segments {
125            match segment {
126                None => {
127                    return Err(ValueSplittingError::MissingSegment);
128                }
129                Some(segment) => {
130                    big_value.extend(segment);
131                }
132            }
133        }
134        Ok(Some(big_value))
135    }
136
137    async fn contains_key(&self, key: &[u8]) -> Result<bool, Self::Error> {
138        let mut big_key = key.to_vec();
139        big_key.extend(&[0, 0, 0, 0]);
140        Ok(self.store.contains_key(&big_key).await?)
141    }
142
143    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, Self::Error> {
144        let big_keys = keys
145            .iter()
146            .map(|key| {
147                let mut big_key = key.clone();
148                big_key.extend(&[0, 0, 0, 0]);
149                big_key
150            })
151            .collect::<Vec<_>>();
152        Ok(self.store.contains_keys(&big_keys).await?)
153    }
154
155    async fn read_multi_values_bytes(
156        &self,
157        keys: &[Vec<u8>],
158    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
159        let mut big_keys = Vec::new();
160        for key in keys {
161            let mut big_key = key.clone();
162            big_key.extend(&[0, 0, 0, 0]);
163            big_keys.push(big_key);
164        }
165        let values = self.store.read_multi_values_bytes(&big_keys).await?;
166        let mut big_values = Vec::<Option<Vec<u8>>>::new();
167        let mut keys_add = Vec::new();
168        let mut n_blocks = Vec::new();
169        for (key, value) in keys.iter().zip(values) {
170            match value {
171                None => {
172                    n_blocks.push(0);
173                    big_values.push(None);
174                }
175                Some(value) => {
176                    let count = Self::read_count_from_value(&value)?;
177                    let big_value = value[4..].to_vec();
178                    for i in 1..count {
179                        let big_key_segment = Self::get_segment_key(key, i)?;
180                        keys_add.push(big_key_segment);
181                    }
182                    n_blocks.push(count);
183                    big_values.push(Some(big_value));
184                }
185            }
186        }
187        if !keys_add.is_empty() {
188            let mut segments = self
189                .store
190                .read_multi_values_bytes(&keys_add)
191                .await?
192                .into_iter();
193            for (big_value, count) in big_values.iter_mut().zip(&n_blocks) {
194                if let Some(value) = big_value {
195                    for _ in 1..*count {
196                        let segment = segments.next().unwrap().unwrap();
197                        value.extend(segment);
198                    }
199                }
200            }
201        }
202        Ok(big_values)
203    }
204
205    async fn find_keys_by_prefix(&self, key_prefix: &[u8]) -> Result<Vec<Vec<u8>>, Self::Error> {
206        let mut keys = Vec::new();
207        for big_key in self.store.find_keys_by_prefix(key_prefix).await? {
208            let len = big_key.len();
209            if Self::read_index_from_key(&big_key)? == 0 {
210                let key = big_key[0..len - 4].to_vec();
211                keys.push(key);
212            }
213        }
214        Ok(keys)
215    }
216
217    async fn find_key_values_by_prefix(
218        &self,
219        key_prefix: &[u8],
220    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Self::Error> {
221        let small_key_values = self.store.find_key_values_by_prefix(key_prefix).await?;
222        let mut small_kv_iterator = small_key_values.into_iter();
223        let mut key_values = Vec::new();
224        while let Some((mut big_key, value)) = small_kv_iterator.next() {
225            if Self::read_index_from_key(&big_key)? != 0 {
226                continue; // Leftover segment from an earlier value.
227            }
228            big_key.truncate(big_key.len() - 4);
229            let key = big_key;
230            let count = Self::read_count_from_value(&value)?;
231            let mut big_value = value[4..].to_vec();
232            for idx in 1..count {
233                let (big_key, value) = small_kv_iterator
234                    .next()
235                    .ok_or(ValueSplittingError::MissingSegment)?;
236                ensure!(
237                    Self::read_index_from_key(&big_key)? == idx
238                        && big_key.starts_with(&key)
239                        && big_key.len() == key.len() + 4,
240                    ValueSplittingError::MissingSegment
241                );
242                big_value.extend(value);
243            }
244            key_values.push((key, big_value));
245        }
246        Ok(key_values)
247    }
248}
249
250impl<K> WritableKeyValueStore for ValueSplittingStore<K>
251where
252    K: WritableKeyValueStore,
253    K::Error: 'static,
254{
255    const MAX_VALUE_SIZE: usize = usize::MAX;
256
257    async fn write_batch(&self, batch: Batch) -> Result<(), Self::Error> {
258        let mut batch_new = Batch::new();
259        for operation in batch.operations {
260            match operation {
261                WriteOperation::Delete { key } => {
262                    let mut big_key = key.to_vec();
263                    big_key.extend(&[0, 0, 0, 0]);
264                    batch_new.delete_key(big_key);
265                }
266                WriteOperation::Put { key, mut value } => {
267                    let big_key = Self::get_segment_key(&key, 0)?;
268                    let mut count: u32 = 1;
269                    let value_ext = if value.len() <= K::MAX_VALUE_SIZE - 4 {
270                        Self::get_initial_count_first_chunk(count, &value)?
271                    } else {
272                        let remainder = value.split_off(K::MAX_VALUE_SIZE - 4);
273                        for value_chunk in remainder.chunks(K::MAX_VALUE_SIZE) {
274                            let big_key_segment = Self::get_segment_key(&key, count)?;
275                            batch_new.put_key_value_bytes(big_key_segment, value_chunk.to_vec());
276                            count += 1;
277                        }
278                        Self::get_initial_count_first_chunk(count, &value)?
279                    };
280                    batch_new.put_key_value_bytes(big_key, value_ext);
281                }
282                WriteOperation::DeletePrefix { key_prefix } => {
283                    batch_new.delete_key_prefix(key_prefix);
284                }
285            }
286        }
287        Ok(self.store.write_batch(batch_new).await?)
288    }
289
290    async fn clear_journal(&self) -> Result<(), Self::Error> {
291        Ok(self.store.clear_journal().await?)
292    }
293}
294
295impl<D> KeyValueDatabase for ValueSplittingDatabase<D>
296where
297    D: KeyValueDatabase,
298    D::Error: 'static,
299{
300    type Config = D::Config;
301
302    type Store = ValueSplittingStore<D::Store>;
303
304    fn get_name() -> String {
305        format!("value splitting {}", D::get_name())
306    }
307
308    async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, Self::Error> {
309        let database = D::connect(config, namespace).await?;
310        Ok(Self { database })
311    }
312
313    fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
314        let store = self.database.open_shared(root_key)?;
315        Ok(ValueSplittingStore { store })
316    }
317
318    fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
319        let store = self.database.open_exclusive(root_key)?;
320        Ok(ValueSplittingStore { store })
321    }
322
323    async fn list_all(config: &Self::Config) -> Result<Vec<String>, Self::Error> {
324        Ok(D::list_all(config).await?)
325    }
326
327    async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
328        Ok(self.database.list_root_keys().await?)
329    }
330
331    async fn delete_all(config: &Self::Config) -> Result<(), Self::Error> {
332        Ok(D::delete_all(config).await?)
333    }
334
335    async fn exists(config: &Self::Config, namespace: &str) -> Result<bool, Self::Error> {
336        Ok(D::exists(config, namespace).await?)
337    }
338
339    async fn create(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
340        Ok(D::create(config, namespace).await?)
341    }
342
343    async fn delete(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
344        Ok(D::delete(config, namespace).await?)
345    }
346}
347
348#[cfg(with_testing)]
349impl<D> TestKeyValueDatabase for ValueSplittingDatabase<D>
350where
351    D: TestKeyValueDatabase,
352    D::Error: 'static,
353{
354    async fn new_test_config() -> Result<D::Config, Self::Error> {
355        Ok(D::new_test_config().await?)
356    }
357}
358
359#[cfg(with_testing)]
360impl<D: crate::backends::DatabaseBackup> crate::backends::DatabaseBackup
361    for ValueSplittingDatabase<D>
362{
363    fn backup_to(&self, dir: &std::path::Path) -> anyhow::Result<()> {
364        self.database.backup_to(dir)
365    }
366}
367
368impl<D> ValueSplittingStore<D>
369where
370    D: WithError,
371{
372    /// Creates a new store that deals with big values from one that does not.
373    pub fn new(store: D) -> Self {
374        ValueSplittingStore { store }
375    }
376
377    fn get_segment_key(key: &[u8], index: u32) -> Result<Vec<u8>, ValueSplittingError<D::Error>> {
378        let mut big_key_segment = key.to_vec();
379        let mut bytes = bcs::to_bytes(&index)?;
380        bytes.reverse();
381        big_key_segment.extend(bytes);
382        Ok(big_key_segment)
383    }
384
385    fn get_initial_count_first_chunk(
386        count: u32,
387        first_chunk: &[u8],
388    ) -> Result<Vec<u8>, ValueSplittingError<D::Error>> {
389        let mut bytes = bcs::to_bytes(&count)?;
390        bytes.reverse();
391        let mut value_ext = Vec::new();
392        value_ext.extend(bytes);
393        value_ext.extend(first_chunk);
394        Ok(value_ext)
395    }
396
397    fn read_count_from_value(value: &[u8]) -> Result<u32, ValueSplittingError<D::Error>> {
398        if value.len() < 4 {
399            return Err(ValueSplittingError::NoCountAvailable);
400        }
401        let mut bytes = value[0..4].to_vec();
402        bytes.reverse();
403        Ok(bcs::from_bytes::<u32>(&bytes)?)
404    }
405
406    fn read_index_from_key(key: &[u8]) -> Result<u32, ValueSplittingError<D::Error>> {
407        let len = key.len();
408        if len < 4 {
409            return Err(ValueSplittingError::TooShortKey);
410        }
411        let mut bytes = key[len - 4..len].to_vec();
412        bytes.reverse();
413        Ok(bcs::from_bytes::<u32>(&bytes)?)
414    }
415}
416
417/// A memory store for which the values are limited to 100 bytes and can be used for tests.
418#[derive(Clone)]
419#[cfg(with_testing)]
420pub struct LimitedTestMemoryStore {
421    inner: MemoryStore,
422}
423
424#[cfg(with_testing)]
425impl Default for LimitedTestMemoryStore {
426    fn default() -> Self {
427        Self::new()
428    }
429}
430
431#[cfg(with_testing)]
432impl WithError for LimitedTestMemoryStore {
433    type Error = MemoryStoreError;
434}
435
436#[cfg(with_testing)]
437impl ReadableKeyValueStore for LimitedTestMemoryStore {
438    const MAX_KEY_SIZE: usize = usize::MAX;
439
440    fn root_key(&self) -> Result<Vec<u8>, MemoryStoreError> {
441        self.inner.root_key()
442    }
443
444    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, MemoryStoreError> {
445        self.inner.read_value_bytes(key).await
446    }
447
448    async fn contains_key(&self, key: &[u8]) -> Result<bool, MemoryStoreError> {
449        self.inner.contains_key(key).await
450    }
451
452    async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, MemoryStoreError> {
453        self.inner.contains_keys(keys).await
454    }
455
456    async fn read_multi_values_bytes(
457        &self,
458        keys: &[Vec<u8>],
459    ) -> Result<Vec<Option<Vec<u8>>>, MemoryStoreError> {
460        self.inner.read_multi_values_bytes(keys).await
461    }
462
463    async fn find_keys_by_prefix(
464        &self,
465        key_prefix: &[u8],
466    ) -> Result<Vec<Vec<u8>>, MemoryStoreError> {
467        self.inner.find_keys_by_prefix(key_prefix).await
468    }
469
470    async fn find_key_values_by_prefix(
471        &self,
472        key_prefix: &[u8],
473    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, MemoryStoreError> {
474        self.inner.find_key_values_by_prefix(key_prefix).await
475    }
476}
477
478#[cfg(with_testing)]
479impl WritableKeyValueStore for LimitedTestMemoryStore {
480    // We set up the MAX_VALUE_SIZE to the artificially low value of 100
481    // purely for testing purposes.
482    const MAX_VALUE_SIZE: usize = 100;
483
484    async fn write_batch(&self, batch: Batch) -> Result<(), MemoryStoreError> {
485        assert!(
486            batch.check_value_size(Self::MAX_VALUE_SIZE),
487            "The batch size is not adequate for this test"
488        );
489        self.inner.write_batch(batch).await
490    }
491
492    async fn clear_journal(&self) -> Result<(), MemoryStoreError> {
493        self.inner.clear_journal().await
494    }
495}
496
497#[cfg(with_testing)]
498impl LimitedTestMemoryStore {
499    /// Creates a `LimitedTestMemoryStore`
500    pub fn new() -> Self {
501        let inner = MemoryStore::new_for_testing();
502        LimitedTestMemoryStore { inner }
503    }
504}
505
506/// Provides a `LimitedTestMemoryStore<()>` that can be used for tests.
507#[cfg(with_testing)]
508pub fn create_value_splitting_memory_store() -> ValueSplittingStore<LimitedTestMemoryStore> {
509    ValueSplittingStore::new(LimitedTestMemoryStore::new())
510}
511
512#[cfg(test)]
513mod tests {
514    use linera_views::{
515        batch::Batch,
516        store::{ReadableKeyValueStore, WritableKeyValueStore},
517        value_splitting::{LimitedTestMemoryStore, ValueSplittingStore},
518    };
519    use rand::Rng;
520
521    // The key splitting means that when a key is overwritten
522    // some previous segments may still be present.
523    #[tokio::test]
524    async fn test_value_splitting1_testing_leftovers() {
525        let store = LimitedTestMemoryStore::new();
526        const MAX_LEN: usize = LimitedTestMemoryStore::MAX_VALUE_SIZE;
527        const _: () = assert!(MAX_LEN > 10);
528        let big_store = ValueSplittingStore::new(store.clone());
529        let key = vec![0, 0];
530        // Write a key with a long value
531        let mut batch = Batch::new();
532        let value = Vec::from([0; MAX_LEN + 1]);
533        batch.put_key_value_bytes(key.clone(), value.clone());
534        big_store.write_batch(batch).await.unwrap();
535        let value_read = big_store.read_value_bytes(&key).await.unwrap();
536        assert_eq!(value_read, Some(value));
537        // Write a key with a smaller value
538        let mut batch = Batch::new();
539        let value = Vec::from([0, 1]);
540        batch.put_key_value_bytes(key.clone(), value.clone());
541        big_store.write_batch(batch).await.unwrap();
542        let value_read = big_store.read_value_bytes(&key).await.unwrap();
543        assert_eq!(value_read, Some(value));
544        // Two segments are present even though only one is used
545        let keys = store.find_keys_by_prefix(&[0]).await.unwrap();
546        assert_eq!(keys, vec![vec![0, 0, 0, 0, 0], vec![0, 0, 0, 0, 1]]);
547    }
548
549    #[tokio::test]
550    async fn test_value_splitting2_testing_splitting() {
551        let store = LimitedTestMemoryStore::new();
552        const MAX_LEN: usize = LimitedTestMemoryStore::MAX_VALUE_SIZE;
553        let big_store = ValueSplittingStore::new(store.clone());
554        let key = vec![0, 0];
555        // Writing a big value
556        let mut batch = Batch::new();
557        let mut value = Vec::new();
558        let mut rng = crate::random::make_deterministic_rng();
559        for _ in 0..2 * MAX_LEN - 4 {
560            value.push(rng.gen::<u8>());
561        }
562        batch.put_key_value_bytes(key.clone(), value.clone());
563        big_store.write_batch(batch).await.unwrap();
564        let value_read = big_store.read_value_bytes(&key).await.unwrap();
565        assert_eq!(value_read, Some(value.clone()));
566        // Reading the segments and checking
567        let mut value_concat = Vec::<u8>::new();
568        for index in 0..2 {
569            let mut segment_key = key.clone();
570            let mut bytes = bcs::to_bytes(&index).unwrap();
571            bytes.reverse();
572            segment_key.extend(bytes);
573            let value_read = store.read_value_bytes(&segment_key).await.unwrap();
574            let Some(value_read) = value_read else {
575                unreachable!(
576                    "value_splitting test: segment key not found in underlying store right after a multi-segment write"
577                )
578            };
579            if index == 0 {
580                value_concat.extend(&value_read[4..]);
581            } else {
582                value_concat.extend(&value_read);
583            }
584        }
585        assert_eq!(value, value_concat);
586    }
587
588    #[tokio::test]
589    async fn test_value_splitting3_write_and_delete() {
590        let store = LimitedTestMemoryStore::new();
591        const MAX_LEN: usize = LimitedTestMemoryStore::MAX_VALUE_SIZE;
592        let big_store = ValueSplittingStore::new(store.clone());
593        let key = vec![0, 0];
594        // writing a big key
595        let mut batch = Batch::new();
596        let mut value = Vec::new();
597        let mut rng = crate::random::make_deterministic_rng();
598        for _ in 0..3 * MAX_LEN - 4 {
599            value.push(rng.gen::<u8>());
600        }
601        batch.put_key_value_bytes(key.clone(), value.clone());
602        big_store.write_batch(batch).await.unwrap();
603        // deleting it
604        let mut batch = Batch::new();
605        batch.delete_key(key.clone());
606        big_store.write_batch(batch).await.unwrap();
607        // reading everything (there are leftover keys)
608        let key_values = big_store.find_key_values_by_prefix(&[0]).await.unwrap();
609        assert_eq!(key_values.len(), 0);
610        // Two segments remain
611        let keys = store.find_keys_by_prefix(&[0]).await.unwrap();
612        assert_eq!(keys, vec![vec![0, 0, 0, 0, 1], vec![0, 0, 0, 0, 2]]);
613    }
614}
linera_views/backends/value_splitting.rs

linera_views/backends/
value_splitting.rs
