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}