mz_storage/upsert/
memory.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! An `UpsertStateBackend` that stores values in memory.
11
12// Allow usage of `std::collections::HashMap`.
13// We need to iterate through all the values in the map, so we can't use `mz_ore` wrapper.
14// Also, we don't need any ordering for the values fetched, so using std HashMap.
15#![allow(clippy::disallowed_types)]
16
17use std::collections::HashMap;
18
19use itertools::Itertools;
20
21use super::UpsertKey;
22use super::types::{
23    GetStats, MergeStats, MergeValue, PutStats, PutValue, StateValue, UpsertStateBackend,
24    UpsertValueAndSize, ValueMetadata,
25};
26
27/// A `HashMap` tracking its total size
28pub struct InMemoryHashMap<T, O> {
29    state: HashMap<UpsertKey, StateValue<T, O>>,
30    total_size: i64,
31}
32
33impl<T, O> Default for InMemoryHashMap<T, O> {
34    fn default() -> Self {
35        Self {
36            state: HashMap::new(),
37            total_size: 0,
38        }
39    }
40}
41
42#[async_trait::async_trait(?Send)]
43impl<T, O> UpsertStateBackend<T, O> for InMemoryHashMap<T, O>
44where
45    O: Clone + 'static,
46    T: Clone + 'static,
47{
48    fn supports_merge(&self) -> bool {
49        false
50    }
51
52    async fn multi_put<P>(&mut self, puts: P) -> Result<PutStats, anyhow::Error>
53    where
54        P: IntoIterator<Item = (UpsertKey, PutValue<StateValue<T, O>>)>,
55    {
56        let mut stats = PutStats::default();
57        for (key, p_value) in puts {
58            stats.processed_puts += 1;
59            match p_value.value {
60                Some(value) => {
61                    let size: i64 = value.memory_size().try_into().expect("less than i64 size");
62                    stats.adjust(Some(&value), Some(size), &p_value.previous_value_metadata);
63                    self.state.insert(key, value);
64                }
65                None => {
66                    stats.adjust::<T, O>(None, None, &p_value.previous_value_metadata);
67                    self.state.remove(&key);
68                }
69            }
70        }
71        self.total_size += stats.size_diff;
72        Ok(stats)
73    }
74
75    async fn multi_merge<M>(&mut self, _merges: M) -> Result<MergeStats, anyhow::Error>
76    where
77        M: IntoIterator<Item = (UpsertKey, MergeValue<StateValue<T, O>>)>,
78    {
79        anyhow::bail!("InMemoryHashMap does not support merging");
80    }
81
82    async fn multi_get<'r, G, R>(
83        &mut self,
84        gets: G,
85        results_out: R,
86    ) -> Result<GetStats, anyhow::Error>
87    where
88        G: IntoIterator<Item = UpsertKey>,
89        R: IntoIterator<Item = &'r mut UpsertValueAndSize<T, O>>,
90    {
91        let mut stats = GetStats::default();
92        for (key, result_out) in gets.into_iter().zip_eq(results_out) {
93            stats.processed_gets += 1;
94            let value = self.state.get(&key).cloned();
95            let metadata = value.as_ref().map(|v| ValueMetadata {
96                size: v.memory_size(),
97                is_tombstone: v.is_tombstone(),
98            });
99            stats.processed_gets_size += metadata.map_or(0, |m| m.size);
100            stats.returned_gets += metadata.map_or(0, |_| 1);
101            *result_out = UpsertValueAndSize { value, metadata };
102        }
103        Ok(stats)
104    }
105}
mz_storage/upsert/memory.rs

mz_storage/upsert/
memory.rs
