mz_storage/upsert/
types.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//! # State-management for UPSERT.
11//!
12//! This module and provide structures for use within an UPSERT
13//! operator implementation.
14//!
15//! UPSERT is a effectively a process which transforms a `Stream<(Key, Option<Data>)>`
16//! into a differential collection, by indexing the data based on the key.
17//!
18//! _This module does not implement this transformation, instead exposing APIs designed
19//! for use within an UPSERT operator. There is one exception to this: `consolidate_chunk`
20//! implements an efficient upsert-like transformation to re-index a collection using the
21//! _output collection_ of an upsert transformation. More on this below.
22//!
23//! ## `UpsertState`
24//!
25//! Its primary export is `UpsertState`, which wraps an `UpsertStateBackend` and provides 3 APIs:
26//!
27//! ### `multi_get`
28//! `multi_get` returns the current value for a (unique) set of keys. To keep implementations
29//! efficient, the set of keys is an iterator, and results are written back into another parallel
30//! iterator. In addition to returning the current values, implementations must also return the
31//! _size_ of those values _as they are stored within the implementation_. Implementations are
32//! required to chunk large iterators if they need to operate over smaller batches.
33//!
34//! `multi_get` is implemented directly with `UpsertStateBackend::multi_get`.
35//!
36//! ### `multi_put`
37//! Update or delete values for a set of keys. To keep implementations efficient, the set
38//! of updates is an iterator. Implementations are also required to return the difference
39//! in values and total size after processing the updates. To simplify this (and because
40//! in the `upsert` usecase we have this data readily available), the updates are input
41//! with the size of the current value (if any) that was returned from a previous `multi_get`.
42//! Implementations are required to chunk large iterators if they need to operate over smaller
43//! batches.
44//!
45//! `multi_put` is implemented directly with `UpsertStateBackend::multi_put`.
46//!
47//! ### `consolidate_chunk`
48//!
49//! `consolidate_chunk` re-indexes an UPSERT collection based on its _output collection_ (as
50//! opposed to its _input `Stream`_. Please see the docs on `consolidate_chunk` and `StateValue`
51//! for more information.
52//!
53//! `consolidate_chunk` is implemented with both `UpsertStateBackend::multi_put` and
54//! `UpsertStateBackend::multi_get`
55//!
56//! ## Order Keys
57//!
58//! In practice, the input stream for UPSERT collections includes an _order key_. This is used to
59//! sort data with the same key occurring in the same timestamp. This module provides support
60//! for serializing and deserializing order keys with their associated data. Being able to ingest
61//! data on non-frontier boundaries requires this support.
62//!
63//! A consequence of this is that tombstones with an order key can be stored within the state.
64//! There is currently no support for cleaning these tombstones up, as they are considered rare and
65//! small enough.
66//!
67//! Because `consolidate_chunk` handles data that consolidates correctly, it does not handle
68//! order keys.
69//!
70//!
71//! ## A note on state size
72//!
73//! The `UpsertStateBackend` trait requires implementations report _relatively accurate_ information about
74//! how the state size changes over time. Note that it does NOT ask the implementations to give
75//! accurate information about actual resource consumption (like disk space including space
76//! amplification), and instead is just asking about the size of the values, after they have been
77//! encoded. For implementations like `RocksDB`, these may be highly accurate (it literally
78//! reports the encoded size as written to the RocksDB API, and for others like the
79//! `InMemoryHashMap`, they may be rough estimates of actual memory usage. See
80//! `StateValue::memory_size` for more information.
81//!
82//! Note also that after consolidation, additional space may be used if `StateValue` is
83//! used.
84//!
85
86use std::fmt;
87use std::num::Wrapping;
88use std::sync::Arc;
89use std::time::Instant;
90
91use bincode::Options;
92use itertools::Itertools;
93use mz_ore::error::ErrorExt;
94use mz_repr::{Diff, GlobalId};
95use serde::{Serialize, de::DeserializeOwned};
96
97use crate::metrics::upsert::{UpsertMetrics, UpsertSharedMetrics};
98use crate::statistics::SourceStatistics;
99use crate::upsert::{UpsertKey, UpsertValue};
100
101/// The default set of `bincode` options used for consolidating
102/// upsert updates (and writing values to RocksDB).
103pub type BincodeOpts = bincode::config::DefaultOptions;
104
105/// Build the default `BincodeOpts`.
106pub fn upsert_bincode_opts() -> BincodeOpts {
107    // We don't allow trailing bytes, for now,
108    // and use varint encoding for space saving.
109    bincode::DefaultOptions::new()
110}
111
112/// The result type for `multi_get`.
113/// The value and size are stored in individual `Option`s so callees
114/// can reuse this value as they overwrite this value, keeping
115/// track of the previous metadata. Additionally, values
116/// may be `None` for tombstones.
117#[derive(Clone)]
118pub struct UpsertValueAndSize<T, O> {
119    /// The value, if there was one.
120    pub value: Option<StateValue<T, O>>,
121    /// The size of original`value` as persisted,
122    /// Useful for users keeping track of statistics.
123    pub metadata: Option<ValueMetadata<u64>>,
124}
125
126impl<T, O> std::fmt::Debug for UpsertValueAndSize<T, O> {
127    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128        f.debug_struct("UpsertValueAndSize")
129            .field("value", &self.value)
130            .field("metadata", &self.metadata)
131            .finish()
132    }
133}
134
135impl<T, O> Default for UpsertValueAndSize<T, O> {
136    fn default() -> Self {
137        Self {
138            value: None,
139            metadata: None,
140        }
141    }
142}
143
144/// Metadata about an existing value in the upsert state backend, as returned
145/// by `multi_get`.
146#[derive(Copy, Clone, Debug)]
147pub struct ValueMetadata<S> {
148    /// The size of the value.
149    pub size: S,
150    /// If the value is a tombstone.
151    pub is_tombstone: bool,
152}
153
154/// A value to put in with `multi_put`.
155#[derive(Clone, Debug)]
156pub struct PutValue<V> {
157    /// The new value, or a `None` to indicate a delete.
158    pub value: Option<V>,
159    /// The value of the previous value for this key, if known.
160    /// Passed into efficiently calculate statistics.
161    pub previous_value_metadata: Option<ValueMetadata<i64>>,
162}
163
164/// A value to put in with a `multi_merge`.
165pub struct MergeValue<V> {
166    /// The value of the merge operand to write to the backend.
167    pub value: V,
168    /// The 'diff' of this merge operand value, used to estimate the overall size diff
169    /// of the working set after this merge operand is merged by the backend.
170    pub diff: Diff,
171}
172
173/// `UpsertState` has 2 modes:
174/// - Normal operation
175/// - Consolidation.
176///
177/// This struct and its substructs are helpers to simplify the logic that
178/// individual `UpsertState` implementations need to do to manage these 2 modes.
179///
180/// Normal operation is simple, we just store an ordinary `UpsertValue`, and allow the implementer
181/// to store it any way they want. During consolidation, the logic is more complex.
182/// See the docs on `StateValue::merge_update` for more information.
183///
184/// Note also that this type is designed to support _partial updates_. All values are
185/// associated with an _order key_ `O` that can be used to determine if a value existing in the
186/// `UpsertStateBackend` occurred before or after a value being considered for insertion.
187///
188/// `O` typically required to be `: Default`, with the default value sorting below all others.
189/// During consolidation, values consolidate correctly (as they are actual
190/// differential updates with diffs), so order keys are not required.
191#[derive(Clone, serde::Serialize, serde::Deserialize)]
192pub enum StateValue<T, O> {
193    Consolidating(Consolidating),
194    Value(Value<T, O>),
195}
196
197impl<T, O> std::fmt::Debug for StateValue<T, O> {
198    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
199        match self {
200            StateValue::Consolidating(c) => std::fmt::Display::fmt(c, f),
201            StateValue::Value(_) => write!(f, "Value"),
202        }
203    }
204}
205
206/// A totally consolidated value stored within the `UpsertStateBackend`.
207///
208/// This type contains support for _tombstones_, that contain an _order key_,
209/// and provisional values.
210///
211/// What is considered finalized and provisional depends on the implementation
212/// of the UPSERT operator: it might consider everything that it writes to its
213/// state finalized, and assume that what it emits will be written down in the
214/// output exactly as presented. Or it might consider everything it writes down
215/// provisional, and only consider updates that it _knows_ to be persisted as
216/// finalized.
217///
218/// Provisional values should only be considered while still "working off"
219/// updates with the same timestamp at which the provisional update was
220/// recorded.
221#[derive(Clone, serde::Serialize, serde::Deserialize, Debug)]
222pub struct Value<T, O> {
223    /// The finalized value of a key is the value we know to be correct for the last complete
224    /// timestamp that got processed. A finalized value of None means that the key has been deleted
225    /// and acts as a tombstone.
226    pub finalized: Option<UpsertValue>,
227    /// When `Some(_)` it contains the upsert value has been processed for a yet incomplete
228    /// timestamp. When None, no provisional update has been emitted yet.
229    pub provisional: Option<ProvisionalValue<T, O>>,
230}
231
232/// A provisional value emitted for a timestamp. This struct contains enough information to
233#[derive(Clone, serde::Serialize, serde::Deserialize, Debug)]
234pub struct ProvisionalValue<T, O> {
235    /// The timestamp at which this provisional value occured at
236    pub timestamp: T,
237    /// The order of this upsert command *within* the timestamp. Commands that happen at the same
238    /// timestamp with lower order get ignored. Commands with higher order override this one. If
239    /// there a case of equal order then the value itself is used as a tie breaker.
240    pub order: O,
241    /// The provisional value. A provisional value of None means that the key has been deleted and
242    /// acts as a tombstone.
243    pub value: Option<UpsertValue>,
244}
245
246/// A value as produced during consolidation.
247#[derive(Clone, Default, serde::Serialize, serde::Deserialize, Debug)]
248pub struct Consolidating {
249    #[serde(with = "serde_bytes")]
250    value_xor: Vec<u8>,
251    len_sum: Wrapping<i64>,
252    checksum_sum: Wrapping<i64>,
253    diff_sum: Wrapping<i64>,
254}
255
256impl fmt::Display for Consolidating {
257    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
258        f.debug_struct("Consolidating")
259            .field("len_sum", &self.len_sum)
260            .field("checksum_sum", &self.checksum_sum)
261            .field("diff_sum", &self.diff_sum)
262            .finish_non_exhaustive()
263    }
264}
265
266impl<T, O> StateValue<T, O> {
267    /// A finalized, that is (assumed) persistent, value.
268    pub fn finalized_value(value: UpsertValue) -> Self {
269        Self::Value(Value {
270            finalized: Some(value),
271            provisional: None,
272        })
273    }
274
275    #[allow(unused)]
276    /// A tombstoned value.
277    pub fn tombstone() -> Self {
278        Self::Value(Value {
279            finalized: None,
280            provisional: None,
281        })
282    }
283
284    /// Whether the value is a tombstone.
285    pub fn is_tombstone(&self) -> bool {
286        match self {
287            Self::Value(value) => value.finalized.is_none(),
288            _ => false,
289        }
290    }
291
292    /// Pull out the `Value` value for a `StateValue`, after `ensure_decoded` has been called.
293    pub fn into_decoded(self) -> Value<T, O> {
294        match self {
295            Self::Value(value) => value,
296            _ => panic!("called `into_decoded without calling `ensure_decoded`"),
297        }
298    }
299
300    /// The size of a `StateValue`, in memory. This is:
301    /// 1. only used in the `InMemoryHashMap` implementation.
302    /// 2. An estimate (it only looks at value sizes, and not errors)
303    ///
304    /// Other implementations may use more accurate accounting.
305    #[cfg(test)]
306    pub fn memory_size(&self) -> usize {
307        use mz_repr::Row;
308        use std::mem::size_of;
309
310        let heap_size = match self {
311            Self::Consolidating(Consolidating { value_xor, .. }) => value_xor.len(),
312            Self::Value(value) => {
313                let finalized_heap_size = match value.finalized {
314                    Some(Ok(ref row)) => {
315                        // `Row::byte_len` includes the size of `Row`, which is also in `Self`, so we
316                        // subtract it.
317                        row.byte_len() - size_of::<Row>()
318                    }
319                    // Assume errors are rare enough to not move the needle.
320                    _ => 0,
321                };
322                let provisional_heap_size = match value.provisional {
323                    Some(ref provisional) => match provisional.value {
324                        Some(Ok(ref row)) => {
325                            // `Row::byte_len` includes the size of `Row`, which is also in `Self`, so we
326                            // subtract it.
327                            row.byte_len() - size_of::<Row>()
328                        }
329                        // Assume errors are rare enough to not move the needle.
330                        _ => 0,
331                    },
332                    None => 0,
333                };
334                finalized_heap_size + provisional_heap_size
335            }
336        };
337        heap_size + size_of::<Self>()
338    }
339}
340
341impl<T: Eq, O> StateValue<T, O> {
342    /// Creates a new provisional value, occurring at some order key, observed
343    /// at the given timestamp.
344    pub fn new_provisional_value(value: UpsertValue, timestamp: T, order: O) -> Self {
345        Self::Value(Value {
346            finalized: None,
347            provisional: Some(ProvisionalValue {
348                value: Some(value),
349                timestamp,
350                order,
351            }),
352        })
353    }
354
355    /// Creates a provisional value, that retains the finalized value in this `StateValue`.
356    pub fn into_provisional_value(self, value: UpsertValue, timestamp: T, order: O) -> Self {
357        match self {
358            StateValue::Value(finalized) => Self::Value(Value {
359                finalized: finalized.finalized,
360                provisional: Some(ProvisionalValue {
361                    value: Some(value),
362                    timestamp,
363                    order,
364                }),
365            }),
366            StateValue::Consolidating(_) => {
367                panic!("called `into_provisional_value` without calling `ensure_decoded`")
368            }
369        }
370    }
371
372    /// Creates a new provisional tombstone occurring at some order key,
373    /// observed at the given timestamp.
374    pub fn new_provisional_tombstone(timestamp: T, order: O) -> Self {
375        Self::Value(Value {
376            finalized: None,
377            provisional: Some(ProvisionalValue {
378                value: None,
379                timestamp,
380                order,
381            }),
382        })
383    }
384
385    /// Creates a provisional tombstone, that retains the finalized value in this `StateValue`.
386    ///
387    /// We record the current finalized value, so that we can present it when
388    /// needed or when trying to read a provisional value at a different
389    /// timestamp.
390    pub fn into_provisional_tombstone(self, timestamp: T, order: O) -> Self {
391        match self {
392            StateValue::Value(finalized) => Self::Value(Value {
393                finalized: finalized.finalized,
394                provisional: Some(ProvisionalValue {
395                    value: None,
396                    timestamp,
397                    order,
398                }),
399            }),
400            StateValue::Consolidating(_) => {
401                panic!("called `into_provisional_tombstone` without calling `ensure_decoded`")
402            }
403        }
404    }
405
406    /// Returns the order of a provisional value at the given timestamp, if any.
407    pub fn provisional_order(&self, ts: &T) -> Option<&O> {
408        match self {
409            Self::Value(value) => match &value.provisional {
410                Some(p) if &p.timestamp == ts => Some(&p.order),
411                _ => None,
412            },
413            Self::Consolidating(_) => {
414                panic!("called `provisional_order` without calling `ensure_decoded`")
415            }
416        }
417    }
418
419    /// Returns the provisional value, if one is present at the given timestamp.
420    /// Falls back to the finalized value, or `None` if there is neither.
421    pub fn provisional_value_ref(&self, ts: &T) -> Option<&UpsertValue> {
422        match self {
423            Self::Value(value) => match &value.provisional {
424                Some(p) if &p.timestamp == ts => p.value.as_ref(),
425                _ => value.finalized.as_ref(),
426            },
427            Self::Consolidating(_) => {
428                panic!("called `provisional_value_ref` without calling `ensure_decoded`")
429            }
430        }
431    }
432
433    /// Returns the the finalized value, if one is present.
434    pub fn into_finalized_value(self) -> Option<UpsertValue> {
435        match self {
436            Self::Value(v) => v.finalized,
437            _ => panic!("called `into_finalized_value` without calling `ensure_decoded`"),
438        }
439    }
440}
441
442impl<T: Eq, O> StateValue<T, O> {
443    /// We use a XOR trick in order to accumulate the values without having to store the full
444    /// unconsolidated history in memory. For all (value, diff) updates of a key we track:
445    /// - diff_sum = SUM(diff)
446    /// - checksum_sum = SUM(checksum(bincode(value)) * diff)
447    /// - len_sum = SUM(len(bincode(value)) * diff)
448    /// - value_xor = XOR(bincode(value))
449    ///
450    /// ## Return value
451    /// Returns a `bool` indicating whether or not the current merged value is able to be deleted.
452    ///
453    /// ## Correctness
454    ///
455    /// The method is correct because a well formed upsert collection at a given
456    /// timestamp will have for each key:
457    /// - Zero or one updates of the form (cur_value, +1)
458    /// - Zero or more pairs of updates of the form (prev_value, +1), (prev_value, -1)
459    ///
460    /// We are interested in extracting the cur_value of each key and discard all prev_values
461    /// that might be included in the stream. Since the history of prev_values always comes in
462    /// pairs, computing the XOR of those is always going to cancel their effects out. Also,
463    /// since XOR is commutative this property is true independent of the order. The same is
464    /// true for the summations of the length and checksum since the sum will contain the
465    /// unrelated values zero times.
466    ///
467    /// Therefore the accumulators will end up precisely in one of two states:
468    /// 1. diff == 0, checksum == 0, value == [0..] => the key is not present
469    /// 2. diff == 1, checksum == checksum(cur_value) value == cur_value => the key is present
470    ///
471    /// ## Robustness
472    ///
473    /// In the absense of bugs, accumulating the diff and checksum is not required since we know
474    /// that a well formed collection always satisfies XOR(bincode(values)) == bincode(cur_value).
475    /// However bugs may happen and so storing 16 more bytes per key to have a very high
476    /// guarantee that we're not decoding garbage is more than worth it.
477    /// The main key->value used to store previous values.
478    #[allow(clippy::as_conversions)]
479    pub fn merge_update(
480        &mut self,
481        value: UpsertValue,
482        diff: mz_repr::Diff,
483        bincode_opts: BincodeOpts,
484        bincode_buffer: &mut Vec<u8>,
485    ) -> bool {
486        match self {
487            Self::Consolidating(Consolidating {
488                value_xor,
489                len_sum,
490                checksum_sum,
491                diff_sum,
492            }) => {
493                bincode_buffer.clear();
494                bincode_opts
495                    .serialize_into(&mut *bincode_buffer, &value)
496                    .unwrap();
497                let len = i64::try_from(bincode_buffer.len()).unwrap();
498
499                *diff_sum += diff.into_inner();
500                *len_sum += len.wrapping_mul(diff.into_inner());
501                // Truncation is fine (using `as`) as this is just a checksum
502                *checksum_sum +=
503                    (seahash::hash(&*bincode_buffer) as i64).wrapping_mul(diff.into_inner());
504
505                // XOR of even diffs cancel out, so we only do it if diff is odd
506                if diff.abs() % Diff::from(2) == Diff::ONE {
507                    if value_xor.len() < bincode_buffer.len() {
508                        value_xor.resize(bincode_buffer.len(), 0);
509                    }
510                    // Note that if the new value is _smaller_ than the `value_xor`, and
511                    // the values at the end are zeroed out, we can shrink the buffer. This
512                    // is extremely sensitive code, so we don't (yet) do that.
513                    for (acc, val) in value_xor.iter_mut().zip(bincode_buffer.drain(..)) {
514                        *acc ^= val;
515                    }
516                }
517
518                // Returns whether or not the value can be deleted. This allows
519                // us to delete values in `UpsertState::consolidate_chunk` (even
520                // if they come back later), to minimize space usage.
521                diff_sum.0 == 0 && checksum_sum.0 == 0 && value_xor.iter().all(|&x| x == 0)
522            }
523            StateValue::Value(_value) => {
524                // We can turn a Value back into a Consolidating state:
525                // `std::mem::take` will leave behind a default value, which
526                // happens to be a default `Consolidating` `StateValue`.
527                let this = std::mem::take(self);
528
529                let finalized_value = this.into_finalized_value();
530                if let Some(finalized_value) = finalized_value {
531                    // If we had a value before, merge it into the
532                    // now-consolidating state first.
533                    let _ =
534                        self.merge_update(finalized_value, Diff::ONE, bincode_opts, bincode_buffer);
535
536                    // Then merge the new value in.
537                    self.merge_update(value, diff, bincode_opts, bincode_buffer)
538                } else {
539                    // We didn't have a value before, might have been a
540                    // tombstone. So just merge in the new value.
541                    self.merge_update(value, diff, bincode_opts, bincode_buffer)
542                }
543            }
544        }
545    }
546
547    /// Merge an existing StateValue into this one, using the same method described in `merge_update`.
548    /// See the docstring above for more information on correctness and robustness.
549    pub fn merge_update_state(&mut self, other: &Self) {
550        match (self, other) {
551            (
552                Self::Consolidating(Consolidating {
553                    value_xor,
554                    len_sum,
555                    checksum_sum,
556                    diff_sum,
557                }),
558                Self::Consolidating(other_consolidating),
559            ) => {
560                *diff_sum += other_consolidating.diff_sum;
561                *len_sum += other_consolidating.len_sum;
562                *checksum_sum += other_consolidating.checksum_sum;
563                if other_consolidating.value_xor.len() > value_xor.len() {
564                    value_xor.resize(other_consolidating.value_xor.len(), 0);
565                }
566                for (acc, val) in value_xor
567                    .iter_mut()
568                    .zip(other_consolidating.value_xor.iter())
569                {
570                    *acc ^= val;
571                }
572            }
573            _ => panic!("`merge_update_state` called with non-consolidating state"),
574        }
575    }
576
577    /// During and after consolidation, we assume that values in the `UpsertStateBackend` implementation
578    /// can be `Self::Consolidating`, with a `diff_sum` of 1 (or 0, if they have been deleted).
579    /// Afterwards, if we need to retract one of these values, we need to assert that its in this correct state,
580    /// then mutate it to its `Value` state, so the `upsert` operator can use it.
581    #[allow(clippy::as_conversions)]
582    pub fn ensure_decoded(&mut self, bincode_opts: BincodeOpts, source_id: GlobalId) {
583        match self {
584            StateValue::Consolidating(consolidating) => {
585                match consolidating.diff_sum.0 {
586                    1 => {
587                        let len = usize::try_from(consolidating.len_sum.0)
588                            .map_err(|_| {
589                                format!(
590                                    "len_sum can't be made into a usize, state: {}, {}",
591                                    consolidating, source_id,
592                                )
593                            })
594                            .expect("invalid upsert state");
595                        let value = &consolidating
596                            .value_xor
597                            .get(..len)
598                            .ok_or_else(|| {
599                                format!(
600                                    "value_xor is not the same length ({}) as len ({}), state: {}, {}",
601                                    consolidating.value_xor.len(),
602                                    len,
603                                    consolidating,
604                                    source_id,
605                                )
606                            })
607                            .expect("invalid upsert state");
608                        // Truncation is fine (using `as`) as this is just a checksum
609                        assert_eq!(
610                            consolidating.checksum_sum.0,
611                            // Hash the value, not the full buffer, which may have extra 0's
612                            seahash::hash(value) as i64,
613                            "invalid upsert state: checksum_sum does not match, state: {}, {}",
614                            consolidating,
615                            source_id,
616                        );
617                        *self = Self::finalized_value(bincode_opts.deserialize(value).unwrap());
618                    }
619                    0 => {
620                        assert_eq!(
621                            consolidating.len_sum.0, 0,
622                            "invalid upsert state: len_sum is non-0, state: {}, {}",
623                            consolidating, source_id,
624                        );
625                        assert_eq!(
626                            consolidating.checksum_sum.0, 0,
627                            "invalid upsert state: checksum_sum is non-0, state: {}, {}",
628                            consolidating, source_id,
629                        );
630                        assert!(
631                            consolidating.value_xor.iter().all(|&x| x == 0),
632                            "invalid upsert state: value_xor not all 0s with 0 diff. \
633                            Non-zero positions: {:?}, state: {}, {}",
634                            consolidating
635                                .value_xor
636                                .iter()
637                                .positions(|&x| x != 0)
638                                .collect::<Vec<_>>(),
639                            consolidating,
640                            source_id,
641                        );
642                        *self = Self::tombstone();
643                    }
644                    other => panic!(
645                        "invalid upsert state: non 0/1 diff_sum: {}, state: {}, {}",
646                        other, consolidating, source_id
647                    ),
648                }
649            }
650            _ => {}
651        }
652    }
653}
654
655impl<T, O> Default for StateValue<T, O> {
656    fn default() -> Self {
657        Self::Consolidating(Consolidating::default())
658    }
659}
660
661/// Statistics for a single call to `consolidate_chunk`.
662#[derive(Clone, Default, Debug)]
663pub struct SnapshotStats {
664    /// The number of updates processed.
665    pub updates: u64,
666    /// The aggregated number of values inserted or deleted into `state`.
667    pub values_diff: Diff,
668    /// The total aggregated size of values inserted, deleted, or updated in `state`.
669    /// If the current call to `consolidate_chunk` deletes a lot of values,
670    /// or updates values to smaller ones, this can be negative!
671    pub size_diff: i64,
672    /// The number of inserts i.e. +1 diff
673    pub inserts: u64,
674    /// The number of deletes i.e. -1 diffs
675    pub deletes: u64,
676}
677
678impl std::ops::AddAssign for SnapshotStats {
679    fn add_assign(&mut self, rhs: Self) {
680        self.updates += rhs.updates;
681        self.values_diff += rhs.values_diff;
682        self.size_diff += rhs.size_diff;
683        self.inserts += rhs.inserts;
684        self.deletes += rhs.deletes;
685    }
686}
687
688/// Statistics for a single call to `multi_merge`.
689#[derive(Clone, Default, Debug)]
690pub struct MergeStats {
691    /// The number of updates written as merge operands to the backend, for the backend
692    /// to process async in the `consolidating_merge_function`.
693    /// Should be equal to number of inserts + deletes
694    pub written_merge_operands: u64,
695    /// The total size of values provided to `multi_merge`. The backend will write these
696    /// down and then later merge them in the `consolidating_merge_function`.
697    pub size_written: u64,
698    /// The estimated diff of the total size of the working set after the merge operands
699    /// are merged by the backend. This is an estimate since it can't account for the
700    /// size overhead of `StateValue` for values that consolidate to 0 (tombstoned-values).
701    pub size_diff: i64,
702}
703
704/// Statistics for a single call to `multi_put`.
705#[derive(Clone, Default, Debug)]
706pub struct PutStats {
707    /// The number of puts/deletes processed
708    /// Should be equal to number of inserts + updates + deletes
709    pub processed_puts: u64,
710    /// The aggregated number of non-tombstone values inserted or deleted into `state`.
711    pub values_diff: i64,
712    /// The aggregated number of tombstones inserted or deleted into `state`
713    pub tombstones_diff: i64,
714    /// The total aggregated size of values inserted, deleted, or updated in `state`.
715    /// If the current call to `multi_put` deletes a lot of values,
716    /// or updates values to smaller ones, this can be negative!
717    pub size_diff: i64,
718    /// The number of inserts
719    pub inserts: u64,
720    /// The number of updates
721    pub updates: u64,
722    /// The number of deletes
723    pub deletes: u64,
724}
725
726impl PutStats {
727    /// Adjust the `PutStats` based on the new value and the previous metadata.
728    ///
729    /// The size parameter is separate as its value is backend-dependent. Its optional
730    /// as some backends increase the total size after an entire batch is processed.
731    ///
732    /// This method is provided for implementors of `UpsertStateBackend::multi_put`.
733    pub fn adjust<T, O>(
734        &mut self,
735        new_value: Option<&StateValue<T, O>>,
736        new_size: Option<i64>,
737        previous_metdata: &Option<ValueMetadata<i64>>,
738    ) {
739        self.adjust_size(new_value, new_size, previous_metdata);
740        self.adjust_values(new_value, previous_metdata);
741        self.adjust_tombstone(new_value, previous_metdata);
742    }
743
744    fn adjust_size<T, O>(
745        &mut self,
746        new_value: Option<&StateValue<T, O>>,
747        new_size: Option<i64>,
748        previous_metdata: &Option<ValueMetadata<i64>>,
749    ) {
750        match (&new_value, previous_metdata.as_ref()) {
751            (Some(_), Some(ps)) => {
752                self.size_diff -= ps.size;
753                if let Some(new_size) = new_size {
754                    self.size_diff += new_size;
755                }
756            }
757            (None, Some(ps)) => {
758                self.size_diff -= ps.size;
759            }
760            (Some(_), None) => {
761                if let Some(new_size) = new_size {
762                    self.size_diff += new_size;
763                }
764            }
765            (None, None) => {}
766        }
767    }
768
769    fn adjust_values<T, O>(
770        &mut self,
771        new_value: Option<&StateValue<T, O>>,
772        previous_metdata: &Option<ValueMetadata<i64>>,
773    ) {
774        let truly_new_value = new_value.map_or(false, |v| !v.is_tombstone());
775        let truly_old_value = previous_metdata.map_or(false, |v| !v.is_tombstone);
776
777        match (truly_new_value, truly_old_value) {
778            (false, true) => {
779                self.values_diff -= 1;
780            }
781            (true, false) => {
782                self.values_diff += 1;
783            }
784            _ => {}
785        }
786    }
787
788    fn adjust_tombstone<T, O>(
789        &mut self,
790        new_value: Option<&StateValue<T, O>>,
791        previous_metdata: &Option<ValueMetadata<i64>>,
792    ) {
793        let new_tombstone = new_value.map_or(false, |v| v.is_tombstone());
794        let old_tombstone = previous_metdata.map_or(false, |v| v.is_tombstone);
795
796        match (new_tombstone, old_tombstone) {
797            (false, true) => {
798                self.tombstones_diff -= 1;
799            }
800            (true, false) => {
801                self.tombstones_diff += 1;
802            }
803            _ => {}
804        }
805    }
806}
807
808/// Statistics for a single call to `multi_get`.
809#[derive(Clone, Default, Debug)]
810pub struct GetStats {
811    /// The number of gets processed
812    pub processed_gets: u64,
813    /// The total size in bytes returned
814    pub processed_gets_size: u64,
815    /// The number of non-empty records returned
816    pub returned_gets: u64,
817}
818
819/// A trait that defines the fundamental primitives required by a state-backing of
820/// `UpsertState`.
821///
822/// Implementors of this trait are blind maps that associate keys and values. They need
823/// not understand the semantics of `StateValue`, tombstones, or anything else related
824/// to a correct `upsert` implementation. The singular exception to this is that they
825/// **must** produce accurate `PutStats` and `GetStats`. The reasoning for this is two-fold:
826/// - efficiency: this avoids additional buffer allocation.
827/// - value sizes: only the backend implementation understands the size of values as recorded
828///
829/// This **must** is not a correctness requirement (we won't panic when emitting statistics), but
830/// rather a requirement to ensure the upsert operator is introspectable.
831#[async_trait::async_trait(?Send)]
832pub trait UpsertStateBackend<T, O>
833where
834    T: 'static,
835    O: 'static,
836{
837    /// Whether this backend supports the `multi_merge` operation.
838    fn supports_merge(&self) -> bool;
839
840    /// Insert or delete for all `puts` keys, prioritizing the last value for
841    /// repeated keys.
842    ///
843    /// The `PutValue` is _guaranteed_ to have an accurate and up-to-date
844    /// record of the metadata for existing value for the given key (if one existed),
845    /// as reported by a previous call to `multi_get`.
846    ///
847    /// `PutStats` **must** be populated correctly, according to these semantics:
848    /// - `values_diff` must record the difference in number of new non-tombstone values being
849    /// inserted into the backend.
850    /// - `tombstones_diff` must record the difference in number of tombstone values being
851    /// inserted into the backend.
852    /// - `size_diff` must record the change in size for the values being inserted/deleted/updated
853    /// in the backend, regardless of whether the values are tombstones or not.
854    async fn multi_put<P>(&mut self, puts: P) -> Result<PutStats, anyhow::Error>
855    where
856        P: IntoIterator<Item = (UpsertKey, PutValue<StateValue<T, O>>)>;
857
858    /// Get the `gets` keys, which must be unique, placing the results in `results_out`.
859    ///
860    /// Panics if `gets` and `results_out` are not the same length.
861    async fn multi_get<'r, G, R>(
862        &mut self,
863        gets: G,
864        results_out: R,
865    ) -> Result<GetStats, anyhow::Error>
866    where
867        G: IntoIterator<Item = UpsertKey>,
868        R: IntoIterator<Item = &'r mut UpsertValueAndSize<T, O>>;
869
870    /// For each key in `merges` writes a 'merge operand' to the backend. The backend stores these
871    /// merge operands and periodically calls the `consolidating_merge_function` to merge them into
872    /// any existing value for each key. The backend will merge the merge operands in the order
873    /// they are provided, and the merge function will always be run for a given key when a `get`
874    /// operation is performed on that key, or when the backend decides to run the merge based
875    /// on its own internal logic.
876    /// This allows avoiding the read-modify-write method of updating many values to
877    /// improve performance.
878    ///
879    /// The `MergeValue` should include a `diff` field that represents the update diff for the
880    /// value. This is used to estimate the overall size diff of the working set
881    /// after the merge operands are merged by the backend `sum[merges: m](m.diff * m.size)`.
882    ///
883    /// `MergeStats` **must** be populated correctly, according to these semantics:
884    /// - `written_merge_operands` must record the number of merge operands written to the backend.
885    /// - `size_written` must record the total size of values written to the backend.
886    ///     Note that the size of the post-merge values are not known, so this is the size of the
887    ///     values written to the backend as merge operands.
888    /// - `size_diff` must record the estimated diff of the total size of the working set after the
889    ///    merge operands are merged by the backend.
890    async fn multi_merge<P>(&mut self, merges: P) -> Result<MergeStats, anyhow::Error>
891    where
892        P: IntoIterator<Item = (UpsertKey, MergeValue<StateValue<T, O>>)>;
893}
894
895/// A function that merges a set of updates for a key into the existing value
896/// for the key. This is called by the backend implementation when it has
897/// accumulated a set of updates for a key, and needs to merge them into the
898/// existing value for the key.
899///
900/// The function is called with the following arguments:
901/// - The key for which the merge is being performed.
902/// - An iterator over any current value and merge operands queued for the key.
903///
904/// The function should return the new value for the key after merging all the updates.
905pub(crate) fn consolidating_merge_function<T: Eq, O>(
906    _key: UpsertKey,
907    updates: impl Iterator<Item = StateValue<T, O>>,
908) -> StateValue<T, O> {
909    let mut current: StateValue<T, O> = Default::default();
910
911    let mut bincode_buf = Vec::new();
912    for update in updates {
913        match update {
914            StateValue::Consolidating(_) => {
915                current.merge_update_state(&update);
916            }
917            StateValue::Value(_) => {
918                // This branch is more expensive, but we hopefully rarely hit
919                // it.
920                if let Some(finalized_value) = update.into_finalized_value() {
921                    let mut update = StateValue::default();
922                    update.merge_update(
923                        finalized_value,
924                        Diff::ONE,
925                        upsert_bincode_opts(),
926                        &mut bincode_buf,
927                    );
928                    current.merge_update_state(&update);
929                }
930            }
931        }
932    }
933
934    current
935}
936
937/// An `UpsertStateBackend` wrapper that supports consolidating merging, and
938/// reports basic metrics about the usage of the `UpsertStateBackend`.
939pub struct UpsertState<'metrics, S, T, O> {
940    inner: S,
941
942    // The status, start time, and stats about calls to `consolidate_chunk`.
943    pub snapshot_start: Instant,
944    snapshot_stats: SnapshotStats,
945    snapshot_completed: bool,
946
947    // Metrics shared across all workers running the `upsert` operator.
948    metrics: Arc<UpsertSharedMetrics>,
949    // Metrics for a specific worker.
950    worker_metrics: &'metrics UpsertMetrics,
951    // User-facing statistics.
952    stats: SourceStatistics,
953
954    // Bincode options and buffer used in `consolidate_chunk`.
955    bincode_opts: BincodeOpts,
956    bincode_buffer: Vec<u8>,
957
958    // We need to iterate over `updates` in `consolidate_chunk` twice, so we
959    // have a scratch vector for this.
960    consolidate_scratch: Vec<(UpsertKey, UpsertValue, mz_repr::Diff)>,
961    // "mini-upsert" map used in `consolidate_chunk`
962    consolidate_upsert_scratch: indexmap::IndexMap<UpsertKey, UpsertValueAndSize<T, O>>,
963    // a scratch vector for calling `multi_get`
964    multi_get_scratch: Vec<UpsertKey>,
965    shrink_upsert_unused_buffers_by_ratio: usize,
966}
967
968impl<'metrics, S, T, O> UpsertState<'metrics, S, T, O> {
969    pub(crate) fn new(
970        inner: S,
971        metrics: Arc<UpsertSharedMetrics>,
972        worker_metrics: &'metrics UpsertMetrics,
973        stats: SourceStatistics,
974        shrink_upsert_unused_buffers_by_ratio: usize,
975    ) -> Self {
976        Self {
977            inner,
978            snapshot_start: Instant::now(),
979            snapshot_stats: SnapshotStats::default(),
980            snapshot_completed: false,
981            metrics,
982            worker_metrics,
983            stats,
984            bincode_opts: upsert_bincode_opts(),
985            bincode_buffer: Vec::new(),
986            consolidate_scratch: Vec::new(),
987            consolidate_upsert_scratch: indexmap::IndexMap::new(),
988            multi_get_scratch: Vec::new(),
989            shrink_upsert_unused_buffers_by_ratio,
990        }
991    }
992}
993
994impl<S, T, O> UpsertState<'_, S, T, O>
995where
996    S: UpsertStateBackend<T, O>,
997    T: Eq + Clone + Send + Sync + Serialize + 'static,
998    O: Clone + Send + Sync + Serialize + DeserializeOwned + 'static,
999{
1000    /// Consolidate the following differential updates into the state. Updates
1001    /// provided to this method can be assumed to consolidate into a single
1002    /// value per-key, after all chunks of updates for a given timestamp have
1003    /// been processed,
1004    ///
1005    /// Therefore, after all updates of a given timestamp have been
1006    /// `consolidated`, all values must be in the correct state (as determined
1007    /// by `StateValue::ensure_decoded`).
1008    ///
1009    /// The `completed` boolean communicates whether or not this is the final
1010    /// chunk of updates for the initial "snapshot" from persist.
1011    ///
1012    /// If the backend supports it, this method will use `multi_merge` to
1013    /// consolidate the updates to avoid having to read the existing value for
1014    /// each key first. On some backends (like RocksDB), this can be
1015    /// significantly faster than the read-then-write consolidation strategy.
1016    ///
1017    /// Also note that we use `self.inner.multi_*`, not `self.multi_*`. This is
1018    /// to avoid erroneously changing metric and stats values.
1019    pub async fn consolidate_chunk<U>(
1020        &mut self,
1021        updates: U,
1022        completed: bool,
1023    ) -> Result<(), anyhow::Error>
1024    where
1025        U: IntoIterator<Item = (UpsertKey, UpsertValue, mz_repr::Diff)> + ExactSizeIterator,
1026    {
1027        fail::fail_point!("fail_consolidate_chunk", |_| {
1028            Err(anyhow::anyhow!("Error consolidating values"))
1029        });
1030
1031        if completed && self.snapshot_completed {
1032            panic!("attempted completion of already completed upsert snapshot")
1033        }
1034
1035        let now = Instant::now();
1036        let batch_size = updates.len();
1037
1038        self.consolidate_scratch.clear();
1039        self.consolidate_upsert_scratch.clear();
1040        self.multi_get_scratch.clear();
1041
1042        // Shrinking the scratch vectors if the capacity is significantly more than batch size
1043        if self.shrink_upsert_unused_buffers_by_ratio > 0 {
1044            let reduced_capacity =
1045                self.consolidate_scratch.capacity() / self.shrink_upsert_unused_buffers_by_ratio;
1046            if reduced_capacity > batch_size {
1047                // These vectors have already been cleared above and should be empty here
1048                self.consolidate_scratch.shrink_to(reduced_capacity);
1049                self.consolidate_upsert_scratch.shrink_to(reduced_capacity);
1050                self.multi_get_scratch.shrink_to(reduced_capacity);
1051            }
1052        }
1053
1054        // Depending on if the backend supports multi_merge, call the appropriate method.
1055        let stats = if self.inner.supports_merge() {
1056            self.consolidate_merge_inner(updates).await?
1057        } else {
1058            self.consolidate_read_write_inner(updates).await?
1059        };
1060
1061        // NOTE: These metrics use the term `merge` to refer to the consolidation of values.
1062        // This is because they were introduced before we the `multi_merge` operation was added.
1063        self.metrics
1064            .merge_snapshot_latency
1065            .observe(now.elapsed().as_secs_f64());
1066        self.worker_metrics
1067            .merge_snapshot_updates
1068            .inc_by(stats.updates);
1069        self.worker_metrics
1070            .merge_snapshot_inserts
1071            .inc_by(stats.inserts);
1072        self.worker_metrics
1073            .merge_snapshot_deletes
1074            .inc_by(stats.deletes);
1075
1076        self.stats.update_bytes_indexed_by(stats.size_diff);
1077        self.stats
1078            .update_records_indexed_by(stats.values_diff.into_inner());
1079
1080        self.snapshot_stats += stats;
1081
1082        if !self.snapshot_completed {
1083            // Updating the metrics
1084            self.worker_metrics.rehydration_total.set(
1085                self.snapshot_stats
1086                    .values_diff
1087                    .into_inner()
1088                    .try_into()
1089                    .unwrap_or_else(|e: std::num::TryFromIntError| {
1090                        tracing::warn!(
1091                            "rehydration_total metric overflowed or is negative \
1092                        and is innacurate: {}. Defaulting to 0",
1093                            e.display_with_causes(),
1094                        );
1095
1096                        0
1097                    }),
1098            );
1099            self.worker_metrics
1100                .rehydration_updates
1101                .set(self.snapshot_stats.updates);
1102        }
1103
1104        if completed {
1105            if self.shrink_upsert_unused_buffers_by_ratio > 0 {
1106                // After rehydration is done, these scratch buffers should now be empty
1107                // shrinking them entirely
1108                self.consolidate_scratch.shrink_to_fit();
1109                self.consolidate_upsert_scratch.shrink_to_fit();
1110                self.multi_get_scratch.shrink_to_fit();
1111            }
1112
1113            self.worker_metrics
1114                .rehydration_latency
1115                .set(self.snapshot_start.elapsed().as_secs_f64());
1116
1117            self.snapshot_completed = true;
1118        }
1119        Ok(())
1120    }
1121
1122    /// Consolidate the updates into the state. This method requires the backend
1123    /// has support for the `multi_merge` operation, and will panic if
1124    /// `self.inner.supports_merge()` was not checked before calling this
1125    /// method. `multi_merge` will write the updates as 'merge operands' to the
1126    /// backend, and then the backend will consolidate those updates with any
1127    /// existing state using the `consolidating_merge_function`.
1128    ///
1129    /// This method can have significant performance benefits over the
1130    /// read-then-write method of `consolidate_read_write_inner`.
1131    async fn consolidate_merge_inner<U>(
1132        &mut self,
1133        updates: U,
1134    ) -> Result<SnapshotStats, anyhow::Error>
1135    where
1136        U: IntoIterator<Item = (UpsertKey, UpsertValue, mz_repr::Diff)> + ExactSizeIterator,
1137    {
1138        let mut updates = updates.into_iter().peekable();
1139
1140        let mut stats = SnapshotStats::default();
1141
1142        if updates.peek().is_some() {
1143            let m_stats = self
1144                .inner
1145                .multi_merge(updates.map(|(k, v, diff)| {
1146                    // Transform into a `StateValue<O>` that can be used by the
1147                    // `consolidating_merge_function` to merge with any existing
1148                    // value for the key.
1149                    let mut val: StateValue<T, O> = Default::default();
1150                    val.merge_update(v, diff, self.bincode_opts, &mut self.bincode_buffer);
1151
1152                    stats.updates += 1;
1153                    if diff.is_positive() {
1154                        stats.inserts += 1;
1155                    } else if diff.is_negative() {
1156                        stats.deletes += 1;
1157                    }
1158
1159                    // To keep track of the overall `values_diff` we can use the sum of diffs which
1160                    // should be equal to the number of non-tombstoned values in the backend.
1161                    // This is a bit misleading as this represents the eventual state after the
1162                    // `consolidating_merge_function` has been called to merge all the updates,
1163                    // and not the state after this `multi_merge` call.
1164                    //
1165                    // This does not accurately report values that have been consolidated to diff == 0, as tracking that
1166                    // per-key is extremely difficult.
1167                    stats.values_diff += diff;
1168
1169                    (k, MergeValue { value: val, diff })
1170                }))
1171                .await?;
1172
1173            stats.size_diff = m_stats.size_diff;
1174        }
1175
1176        Ok(stats)
1177    }
1178
1179    /// Consolidates the updates into the state. This method reads the existing
1180    /// values for each key, consolidates the updates, and writes the new values
1181    /// back to the state.
1182    async fn consolidate_read_write_inner<U>(
1183        &mut self,
1184        updates: U,
1185    ) -> Result<SnapshotStats, anyhow::Error>
1186    where
1187        U: IntoIterator<Item = (UpsertKey, UpsertValue, mz_repr::Diff)> + ExactSizeIterator,
1188    {
1189        let mut updates = updates.into_iter().peekable();
1190
1191        let mut stats = SnapshotStats::default();
1192
1193        if updates.peek().is_some() {
1194            self.consolidate_scratch.extend(updates);
1195            self.consolidate_upsert_scratch.extend(
1196                self.consolidate_scratch
1197                    .iter()
1198                    .map(|(k, _, _)| (*k, UpsertValueAndSize::default())),
1199            );
1200            self.multi_get_scratch
1201                .extend(self.consolidate_upsert_scratch.iter().map(|(k, _)| *k));
1202            self.inner
1203                .multi_get(
1204                    self.multi_get_scratch.drain(..),
1205                    self.consolidate_upsert_scratch.iter_mut().map(|(_, v)| v),
1206                )
1207                .await?;
1208
1209            for (key, value, diff) in self.consolidate_scratch.drain(..) {
1210                stats.updates += 1;
1211                if diff.is_positive() {
1212                    stats.inserts += 1;
1213                } else if diff.is_negative() {
1214                    stats.deletes += 1;
1215                }
1216
1217                // We rely on the diffs in our input instead of the result of
1218                // multi_put below. This makes sure we report the same stats as
1219                // `consolidate_merge_inner`, regardless of what values
1220                // there were in state before.
1221                stats.values_diff += diff;
1222
1223                let entry = self.consolidate_upsert_scratch.get_mut(&key).unwrap();
1224                let val = entry.value.get_or_insert_with(Default::default);
1225
1226                if val.merge_update(value, diff, self.bincode_opts, &mut self.bincode_buffer) {
1227                    entry.value = None;
1228                }
1229            }
1230
1231            // Note we do 1 `multi_get` and 1 `multi_put` while processing a _batch of updates_.
1232            // Within the batch, we effectively consolidate each key, before persisting that
1233            // consolidated value. Easy!!
1234            let p_stats = self
1235                .inner
1236                .multi_put(self.consolidate_upsert_scratch.drain(..).map(|(k, v)| {
1237                    (
1238                        k,
1239                        PutValue {
1240                            value: v.value,
1241                            previous_value_metadata: v.metadata.map(|v| ValueMetadata {
1242                                size: v.size.try_into().expect("less than i64 size"),
1243                                is_tombstone: v.is_tombstone,
1244                            }),
1245                        },
1246                    )
1247                }))
1248                .await?;
1249
1250            stats.size_diff = p_stats.size_diff;
1251        }
1252
1253        Ok(stats)
1254    }
1255
1256    /// Insert or delete for all `puts` keys, prioritizing the last value for
1257    /// repeated keys.
1258    pub async fn multi_put<P>(
1259        &mut self,
1260        update_per_record_stats: bool,
1261        puts: P,
1262    ) -> Result<(), anyhow::Error>
1263    where
1264        P: IntoIterator<Item = (UpsertKey, PutValue<Value<T, O>>)>,
1265    {
1266        fail::fail_point!("fail_state_multi_put", |_| {
1267            Err(anyhow::anyhow!("Error putting values into state"))
1268        });
1269        let now = Instant::now();
1270        let stats = self
1271            .inner
1272            .multi_put(puts.into_iter().map(|(k, pv)| {
1273                (
1274                    k,
1275                    PutValue {
1276                        value: pv.value.map(StateValue::Value),
1277                        previous_value_metadata: pv.previous_value_metadata,
1278                    },
1279                )
1280            }))
1281            .await?;
1282
1283        self.metrics
1284            .multi_put_latency
1285            .observe(now.elapsed().as_secs_f64());
1286        self.worker_metrics
1287            .multi_put_size
1288            .inc_by(stats.processed_puts);
1289
1290        if update_per_record_stats {
1291            self.worker_metrics.upsert_inserts.inc_by(stats.inserts);
1292            self.worker_metrics.upsert_updates.inc_by(stats.updates);
1293            self.worker_metrics.upsert_deletes.inc_by(stats.deletes);
1294
1295            self.stats.update_bytes_indexed_by(stats.size_diff);
1296            self.stats.update_records_indexed_by(stats.values_diff);
1297            self.stats
1298                .update_envelope_state_tombstones_by(stats.tombstones_diff);
1299        }
1300
1301        Ok(())
1302    }
1303
1304    /// Get the `gets` keys, which must be unique, placing the results in `results_out`.
1305    ///
1306    /// Panics if `gets` and `results_out` are not the same length.
1307    pub async fn multi_get<'r, G, R>(
1308        &mut self,
1309        gets: G,
1310        results_out: R,
1311    ) -> Result<(), anyhow::Error>
1312    where
1313        G: IntoIterator<Item = UpsertKey>,
1314        R: IntoIterator<Item = &'r mut UpsertValueAndSize<T, O>>,
1315        O: 'r,
1316    {
1317        fail::fail_point!("fail_state_multi_get", |_| {
1318            Err(anyhow::anyhow!("Error getting values from state"))
1319        });
1320        let now = Instant::now();
1321        let stats = self.inner.multi_get(gets, results_out).await?;
1322
1323        self.metrics
1324            .multi_get_latency
1325            .observe(now.elapsed().as_secs_f64());
1326        self.worker_metrics
1327            .multi_get_size
1328            .inc_by(stats.processed_gets);
1329        self.worker_metrics
1330            .multi_get_result_count
1331            .inc_by(stats.returned_gets);
1332        self.worker_metrics
1333            .multi_get_result_bytes
1334            .inc_by(stats.processed_gets_size);
1335
1336        Ok(())
1337    }
1338}
1339
1340#[cfg(test)]
1341mod tests {
1342    use mz_repr::Row;
1343
1344    use super::*;
1345    #[mz_ore::test]
1346    fn test_merge_update() {
1347        let mut buf = Vec::new();
1348        let opts = upsert_bincode_opts();
1349
1350        let mut s = StateValue::<(), ()>::Consolidating(Consolidating::default());
1351
1352        let small_row = Ok(Row::default());
1353        let longer_row = Ok(Row::pack([mz_repr::Datum::Null]));
1354        s.merge_update(small_row, Diff::ONE, opts, &mut buf);
1355        s.merge_update(longer_row.clone(), Diff::MINUS_ONE, opts, &mut buf);
1356        // This clears the retraction of the `longer_row`, but the
1357        // `value_xor` is the length of the `longer_row`. This tests
1358        // that we are tracking checksums correctly.
1359        s.merge_update(longer_row, Diff::ONE, opts, &mut buf);
1360
1361        // Assert that the `Consolidating` value is fully merged.
1362        s.ensure_decoded(opts, GlobalId::User(1));
1363    }
1364
1365    // We guard some of our assumptions. Increasing in-memory size of StateValue
1366    // has a direct impact on memory usage of in-memory UPSERT sources.
1367    #[mz_ore::test]
1368    fn test_memory_size() {
1369        let finalized_value: StateValue<(), ()> = StateValue::finalized_value(Ok(Row::default()));
1370        assert!(
1371            finalized_value.memory_size() <= 64,
1372            "memory size is {}",
1373            finalized_value.memory_size(),
1374        );
1375
1376        let provisional_value_with_finalized_value: StateValue<(), ()> =
1377            finalized_value.into_provisional_value(Ok(Row::default()), (), ());
1378        assert!(
1379            provisional_value_with_finalized_value.memory_size() <= 64,
1380            "memory size is {}",
1381            provisional_value_with_finalized_value.memory_size(),
1382        );
1383
1384        let provisional_value_without_finalized_value: StateValue<(), ()> =
1385            StateValue::new_provisional_value(Ok(Row::default()), (), ());
1386        assert!(
1387            provisional_value_without_finalized_value.memory_size() <= 64,
1388            "memory size is {}",
1389            provisional_value_without_finalized_value.memory_size(),
1390        );
1391
1392        let mut consolidating_value: StateValue<(), ()> = StateValue::default();
1393        consolidating_value.merge_update(
1394            Ok(Row::default()),
1395            Diff::ONE,
1396            upsert_bincode_opts(),
1397            &mut Vec::new(),
1398        );
1399        assert!(
1400            consolidating_value.memory_size() <= 66,
1401            "memory size is {}",
1402            consolidating_value.memory_size(),
1403        );
1404    }
1405
1406    #[mz_ore::test]
1407    #[should_panic(
1408        expected = "invalid upsert state: len_sum is non-0, state: Consolidating { len_sum: 1"
1409    )]
1410    fn test_merge_update_len_0_assert() {
1411        let mut buf = Vec::new();
1412        let opts = upsert_bincode_opts();
1413
1414        let mut s = StateValue::<(), ()>::Consolidating(Consolidating::default());
1415
1416        let small_row = Ok(mz_repr::Row::default());
1417        let longer_row = Ok(mz_repr::Row::pack([mz_repr::Datum::Null]));
1418        s.merge_update(longer_row.clone(), Diff::ONE, opts, &mut buf);
1419        s.merge_update(small_row.clone(), Diff::MINUS_ONE, opts, &mut buf);
1420
1421        s.ensure_decoded(opts, GlobalId::User(1));
1422    }
1423
1424    #[mz_ore::test]
1425    #[should_panic(
1426        expected = "invalid upsert state: \"value_xor is not the same length (3) as len (4), state: Consolidating { len_sum: 4"
1427    )]
1428    fn test_merge_update_len_to_long_assert() {
1429        let mut buf = Vec::new();
1430        let opts = upsert_bincode_opts();
1431
1432        let mut s = StateValue::<(), ()>::Consolidating(Consolidating::default());
1433
1434        let small_row = Ok(mz_repr::Row::default());
1435        let longer_row = Ok(mz_repr::Row::pack([mz_repr::Datum::Null]));
1436        s.merge_update(longer_row.clone(), Diff::ONE, opts, &mut buf);
1437        s.merge_update(small_row.clone(), Diff::MINUS_ONE, opts, &mut buf);
1438        s.merge_update(longer_row.clone(), Diff::ONE, opts, &mut buf);
1439
1440        s.ensure_decoded(opts, GlobalId::User(1));
1441    }
1442
1443    #[mz_ore::test]
1444    #[should_panic(expected = "invalid upsert state: checksum_sum does not match")]
1445    fn test_merge_update_checksum_doesnt_match() {
1446        let mut buf = Vec::new();
1447        let opts = upsert_bincode_opts();
1448
1449        let mut s = StateValue::<(), ()>::Consolidating(Consolidating::default());
1450
1451        let small_row = Ok(mz_repr::Row::pack([mz_repr::Datum::Int64(2)]));
1452        let longer_row = Ok(mz_repr::Row::pack([mz_repr::Datum::Int64(1)]));
1453        s.merge_update(longer_row.clone(), Diff::ONE, opts, &mut buf);
1454        s.merge_update(small_row.clone(), Diff::MINUS_ONE, opts, &mut buf);
1455        s.merge_update(longer_row.clone(), Diff::ONE, opts, &mut buf);
1456
1457        s.ensure_decoded(opts, GlobalId::User(1));
1458    }
1459}