mz_storage/upsert/
rocksdb.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 RocksDB.
11
12use mz_rocksdb::{KeyUpdate, RocksDBInstance};
13use serde::{Serialize, de::DeserializeOwned};
14
15use super::UpsertKey;
16use super::types::{
17    GetStats, MergeStats, MergeValue, PutStats, PutValue, StateValue, UpsertStateBackend,
18    UpsertValueAndSize, ValueMetadata,
19};
20
21/// A `UpsertStateBackend` implementation backed by RocksDB.
22pub struct RocksDB<T, O> {
23    rocksdb: RocksDBInstance<UpsertKey, StateValue<T, O>>,
24}
25
26impl<T, O> RocksDB<T, O> {
27    pub fn new(rocksdb: RocksDBInstance<UpsertKey, StateValue<T, O>>) -> Self {
28        Self { rocksdb }
29    }
30}
31
32#[async_trait::async_trait(?Send)]
33impl<T, O> UpsertStateBackend<T, O> for RocksDB<T, O>
34where
35    O: Send + Sync + Serialize + DeserializeOwned + 'static,
36    T: Send + Sync + Serialize + DeserializeOwned + 'static,
37{
38    fn supports_merge(&self) -> bool {
39        self.rocksdb.supports_merges
40    }
41
42    async fn multi_put<P>(&mut self, puts: P) -> Result<PutStats, anyhow::Error>
43    where
44        P: IntoIterator<Item = (UpsertKey, PutValue<StateValue<T, O>>)>,
45    {
46        let mut p_stats = PutStats::default();
47        let stats = self
48            .rocksdb
49            .multi_update(puts.into_iter().map(
50                |(
51                    k,
52                    PutValue {
53                        value,
54                        previous_value_metadata,
55                    },
56                )| {
57                    p_stats.adjust(value.as_ref(), None, &previous_value_metadata);
58                    let value = match value {
59                        Some(v) => KeyUpdate::Put(v),
60                        None => KeyUpdate::Delete,
61                    };
62                    (k, value, None)
63                },
64            ))
65            .await?;
66        p_stats.processed_puts += stats.processed_updates;
67        let size: i64 = stats.size_written.try_into().expect("less than i64 size");
68        p_stats.size_diff += size;
69
70        Ok(p_stats)
71    }
72
73    async fn multi_merge<M>(&mut self, merges: M) -> Result<MergeStats, anyhow::Error>
74    where
75        M: IntoIterator<Item = (UpsertKey, MergeValue<StateValue<T, O>>)>,
76    {
77        let mut m_stats = MergeStats::default();
78        let stats =
79            self.rocksdb
80                .multi_update(merges.into_iter().map(|(k, MergeValue { value, diff })| {
81                    (k, KeyUpdate::Merge(value), Some(diff))
82                }))
83                .await?;
84        m_stats.written_merge_operands += stats.processed_updates;
85        m_stats.size_written += stats.size_written;
86        if let Some(diff) = stats.size_diff {
87            m_stats.size_diff += diff.into_inner();
88        }
89        Ok(m_stats)
90    }
91
92    async fn multi_get<'r, G, R>(
93        &mut self,
94        gets: G,
95        results_out: R,
96    ) -> Result<GetStats, anyhow::Error>
97    where
98        G: IntoIterator<Item = UpsertKey>,
99        R: IntoIterator<Item = &'r mut UpsertValueAndSize<T, O>>,
100    {
101        let mut g_stats = GetStats::default();
102        let stats = self
103            .rocksdb
104            .multi_get(gets, results_out, |value| {
105                value.map_or(
106                    UpsertValueAndSize {
107                        value: None,
108                        metadata: None,
109                    },
110                    |v| {
111                        let is_tombstone = v.value.is_tombstone();
112                        UpsertValueAndSize {
113                            value: Some(v.value),
114                            metadata: Some(ValueMetadata {
115                                size: v.size,
116                                is_tombstone,
117                            }),
118                        }
119                    },
120                )
121            })
122            .await?;
123
124        g_stats.processed_gets += stats.processed_gets;
125        g_stats.processed_gets_size += stats.processed_gets_size;
126        g_stats.returned_gets += stats.returned_gets;
127        Ok(g_stats)
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use std::sync::Arc;
134
135    use mz_ore::metrics::MetricsRegistry;
136    use mz_persist_client::ShardId;
137    use mz_repr::{Datum, Diff, GlobalId, Row, Timestamp};
138    use mz_rocksdb::{InstanceOptions, KeyUpdate, RocksDBConfig, RocksDBInstance, ValueIterator};
139    use mz_storage_types::sources::SourceEnvelope;
140    use mz_storage_types::sources::envelope::{KeyEnvelope, UpsertEnvelope, UpsertStyle};
141    use rocksdb::Env;
142    use timely::progress::{Antichain, Timestamp as _};
143
144    use super::RocksDB;
145    use crate::metrics::upsert::{UpsertMetricDefs, UpsertMetrics};
146    use crate::statistics::{SourceStatistics, SourceStatisticsMetricDefs};
147    use crate::upsert::UpsertKey;
148    use crate::upsert::types::{
149        BincodeOpts, StateValue, UpsertState, UpsertValueAndSize, consolidating_merge_function,
150        upsert_bincode_opts,
151    };
152
153    // The order key type. `()` suffices for a test that never exercises order keys.
154    type TestState = StateValue<Timestamp, ()>;
155
156    // Open a RocksDB instance configured like the upsert operator, with the
157    // native merge operator.
158    fn open_instance(
159        path: &std::path::Path,
160        env: Env,
161        upsert_metrics: &UpsertMetrics,
162    ) -> RocksDBInstance<UpsertKey, TestState> {
163        RocksDBInstance::new(
164            path,
165            InstanceOptions::new(
166                env,
167                2,
168                Some((
169                    "upsert_state_snapshot_merge_v1".to_string(),
170                    |a: &[u8], b: ValueIterator<BincodeOpts, TestState>| {
171                        consolidating_merge_function::<Timestamp, ()>(a.into(), b)
172                    },
173                )),
174                upsert_bincode_opts(),
175            ),
176            RocksDBConfig::new(Default::default(), None),
177            Arc::clone(&upsert_metrics.rocksdb_shared),
178            Arc::clone(&upsert_metrics.rocksdb_instance_metrics),
179        )
180        .expect("failed to open rocksdb instance")
181    }
182
183    fn test_metrics(source_id: GlobalId) -> UpsertMetrics {
184        let registry = MetricsRegistry::new();
185        let defs = UpsertMetricDefs::register_with(&registry);
186        UpsertMetrics::new(&defs, source_id, 0, None)
187    }
188
189    fn test_statistics(source_id: GlobalId) -> SourceStatistics {
190        let registry = MetricsRegistry::new();
191        let defs = SourceStatisticsMetricDefs::register_with(&registry);
192        let envelope = SourceEnvelope::Upsert(UpsertEnvelope {
193            source_arity: 2,
194            style: UpsertStyle::Default(KeyEnvelope::Flattened),
195            key_indices: vec![0],
196        });
197        SourceStatistics::new(
198            source_id,
199            0,
200            &defs,
201            source_id,
202            &ShardId::new(),
203            envelope,
204            Antichain::from_elem(Timestamp::minimum()),
205        )
206    }
207
208    /// Regression test for stale upsert state surviving an in-process dataflow
209    /// restart on replicas without a scratch directory. Those replicas keep
210    /// upsert state in an in-memory `Env`, and a restarted dataflow that
211    /// reopens the same path in the same `Env` must start empty. If a
212    /// finalized value survives, rehydration re-inserts the persist snapshot
213    /// on top of it, the merge yields `diff_sum == 2`, and `ensure_decoded`
214    /// panics with `invalid upsert state`.
215    #[mz_ore::test(tokio::test)]
216    #[cfg_attr(miri, ignore)] // rocksdb FFI is unsupported under miri
217    async fn stale_state_corrupts_rehydration() {
218        let source_id = GlobalId::User(0);
219        let upsert_metrics = test_metrics(source_id);
220
221        // One in-memory env reused by both instances, and a path that does
222        // not initially exist on the host filesystem.
223        let mem_env = Env::mem_env().unwrap();
224        let tmp = tempfile::tempdir().unwrap();
225        let path = tmp.path().join("does-not-exist-on-host").join("instance");
226
227        // A value already committed for one key, present in both local state
228        // and the persist output.
229        let key = UpsertKey::from_key(Ok(&Row::pack_slice(&[Datum::Int64(0)])));
230        let value: crate::upsert::UpsertValue =
231            Ok(Row::pack_slice(&[Datum::Int64(0), Datum::Int64(42)]));
232
233        // First instance: write the finalized value like steady-state
234        // `multi_put`, then close.
235        {
236            let mut instance = open_instance(&path, mem_env.clone(), &upsert_metrics);
237            instance
238                .multi_update([(
239                    key,
240                    KeyUpdate::Put(StateValue::finalized_value(value.clone())),
241                    None,
242                )])
243                .await
244                .unwrap();
245            instance.close().await.unwrap();
246        }
247
248        // Second instance: same `Env`, same path, mirroring an in-process
249        // `SuspendAndRestart`.
250        let mut state = UpsertState::<_, Timestamp, ()>::new(
251            RocksDB::new(open_instance(&path, mem_env.clone(), &upsert_metrics)),
252            Arc::clone(&upsert_metrics.shared),
253            &upsert_metrics,
254            test_statistics(source_id),
255            0,
256        );
257
258        // Rehydrate from the persist snapshot: re-insert the committed output
259        // as `(key, value, +1)`, merged on top of whatever is already in state.
260        state
261            .consolidate_chunk([(key, value.clone(), Diff::ONE)].into_iter(), true)
262            .await
263            .unwrap();
264
265        // Read it back, as `drain_staged_input` does. `ensure_decoded` panics
266        // on an inflated diff.
267        let mut out = vec![UpsertValueAndSize::<Timestamp, ()>::default()];
268        state.multi_get([key], out.iter_mut()).await.unwrap();
269        let mut decoded = out.into_iter().next().unwrap().value.expect("key present");
270        decoded.ensure_decoded(upsert_bincode_opts(), source_id, Some(&key));
271
272        assert_eq!(
273            decoded.into_finalized_value(),
274            Some(value),
275            "value != snapshot"
276        );
277    }
278}
mz_storage/upsert/rocksdb.rs

mz_storage/upsert/
rocksdb.rs
