upsert: store errors on the heap

`DataflowError` instances are heap allocated since these values are
rarer and we don't want to be paying the cost of storing them for each
value.

The upsert operator was previosuly unpacking the general `DataflowError`
into a local `UpsertError` variant that was not boxed, leading to large
structs.

This commit re-boxes the error variant bringing the size of
`UpsertValue` down significantly.

Fixes MaterializeInc/database-issues#9567

Signed-off-by: Petros Angelatos <[email protected]>

Author: petrosagg

src/storage/src/render/sources.rs

@@ -43,7 +43,7 @@ use crate::decode::{render_decode_cdcv2, render_decode_delimited};
 use crate::healthcheck::{HealthStatusMessage, StatusNamespace};
 use crate::source::types::{DecodeResult, SourceOutput, SourceRender};
 use crate::source::{self, RawSourceCreationConfig, SourceExportCreationConfig};
-use crate::upsert::UpsertKey;
+use crate::upsert::{UpsertKey, UpsertValue};
 
 /// _Renders_ complete _differential_ [`Collection`]s
 /// that represent the final source and its errors
@@ -479,7 +479,7 @@ fn append_metadata_to_value<G: Scope, FromTime: Timestamp>(
 fn upsert_commands<G: Scope, FromTime: Timestamp>(
     input: Collection<G, DecodeResult<FromTime>, Diff>,
     upsert_envelope: UpsertEnvelope,
-) -> Collection<G, (UpsertKey, Option<Result<Row, UpsertError>>, FromTime), Diff> {
+) -> Collection<G, (UpsertKey, Option<UpsertValue>, FromTime), Diff> {
     let mut row_buf = Row::default();
     input.map(move |result| {
         let from_time = result.from_time;
@@ -493,9 +493,15 @@ fn upsert_commands<G: Scope, FromTime: Timestamp>(
         // If we have a well-formed key we can continue, otherwise we're upserting an error
         let key = match key {
             Ok(key) => key,
-            err @ Err(_) => match result.value {
-                Some(_) => return (UpsertKey::from_key(err.as_ref()), Some(err), from_time),
-                None => return (UpsertKey::from_key(err.as_ref()), None, from_time),
+            Err(err) => match result.value {
+                Some(_) => {
+                    return (
+                        UpsertKey::from_key(Err(&err)),
+                        Some(Err(Box::new(err))),
+                        from_time,
+                    );
+                }
+                None => return (UpsertKey::from_key(Err(&err)), None, from_time),
             },
         };
 
@@ -584,10 +590,10 @@ fn upsert_commands<G: Scope, FromTime: Timestamp>(
                         packer.extend_by_row(&metadata);
                         Some(Ok(row_buf.clone()))
                     }
-                    _ => Some(Err(UpsertError::Value(UpsertValueError {
+                    _ => Some(Err(Box::new(UpsertError::Value(UpsertValueError {
                         for_key: key_row,
                         inner,
-                    }))),
+                    })))),
                 }
             }
             None => None,

src/storage/src/upsert.rs

@@ -58,7 +58,7 @@ pub(crate) mod rocksdb;
 // TODO(aljoscha): Move next to upsert module, rename to upsert_types.
 pub(crate) mod types;
 
-pub type UpsertValue = Result<Row, UpsertError>;
+pub type UpsertValue = Result<Row, Box<UpsertError>>;
 
 #[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
 pub struct UpsertKey([u8; 32]);
@@ -508,7 +508,7 @@ enum DrainStyle<'a, T> {
 async fn drain_staged_input<S, G, T, FromTime, E>(
     stash: &mut Vec<(T, UpsertKey, Reverse<FromTime>, Option<UpsertValue>)>,
     commands_state: &mut indexmap::IndexMap<UpsertKey, types::UpsertValueAndSize<T, FromTime>>,
-    output_updates: &mut Vec<(Result<Row, UpsertError>, T, Diff)>,
+    output_updates: &mut Vec<(UpsertValue, T, Diff)>,
     multi_get_scratch: &mut Vec<UpsertKey>,
     drain_style: DrainStyle<'_, T>,
     error_emitter: &mut E,
@@ -691,12 +691,16 @@ where
         let value = match result {
             Ok(ok) => Ok(ok),
             Err(DataflowError::EnvelopeError(err)) => match *err {
-                EnvelopeError::Upsert(err) => Err(err),
+                EnvelopeError::Upsert(err) => Err(Box::new(err)),
                 _ => return None,
             },
             Err(_) => return None,
         };
-        Some((UpsertKey::from_value(value.as_ref(), &key_indices), value))
+        let value_ref = match value {
+            Ok(ref row) => Ok(row),
+            Err(ref err) => Err(&**err),
+        };
+        Some((UpsertKey::from_value(value_ref, &key_indices), value))
     });
     let (output_handle, output) = builder.new_output();
 
@@ -931,7 +935,7 @@ where
     (
         output.as_collection().map(|result| match result {
             Ok(ok) => Ok(ok),
-            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(err))),
+            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(*err))),
         }),
         health_stream,
         snapshot_stream,

src/storage/src/upsert/types.rs

@@ -304,15 +304,17 @@ impl<T, O> StateValue<T, O> {
     /// Other implementations may use more accurate accounting.
     #[cfg(test)]
     pub fn memory_size(&self) -> usize {
-        use std::mem::size_of_val;
+        use mz_repr::Row;
+        use std::mem::size_of;
+
         let heap_size = match self {
             Self::Consolidating(Consolidating { value_xor, .. }) => value_xor.len(),
             Self::Value(value) => {
                 let finalized_heap_size = match value.finalized {
                     Some(Ok(ref row)) => {
                         // `Row::byte_len` includes the size of `Row`, which is also in `Self`, so we
                         // subtract it.
-                        row.byte_len() - size_of_val(&row)
+                        row.byte_len() - size_of::<Row>()
                     }
                     // Assume errors are rare enough to not move the needle.
                     _ => 0,
@@ -322,7 +324,7 @@ impl<T, O> StateValue<T, O> {
                         Some(Ok(ref row)) => {
                             // `Row::byte_len` includes the size of `Row`, which is also in `Self`, so we
                             // subtract it.
-                            row.byte_len() - size_of_val(&row)
+                            row.byte_len() - size_of::<Row>()
                         }
                         // Assume errors are rare enough to not move the needle.
                         _ => 0,
@@ -332,7 +334,7 @@ impl<T, O> StateValue<T, O> {
                 finalized_heap_size + provisional_heap_size
             }
         };
-        heap_size + size_of_val(self)
+        heap_size + size_of::<Self>()
     }
 }
 
@@ -1366,23 +1368,23 @@ mod tests {
     fn test_memory_size() {
         let finalized_value: StateValue<(), ()> = StateValue::finalized_value(Ok(Row::default()));
         assert!(
-            finalized_value.memory_size() <= 144,
+            finalized_value.memory_size() <= 64,
             "memory size is {}",
             finalized_value.memory_size(),
         );
 
         let provisional_value_with_finalized_value: StateValue<(), ()> =
             finalized_value.into_provisional_value(Ok(Row::default()), (), ());
         assert!(
-            provisional_value_with_finalized_value.memory_size() <= 168,
+            provisional_value_with_finalized_value.memory_size() <= 64,
             "memory size is {}",
             provisional_value_with_finalized_value.memory_size(),
         );
 
         let provisional_value_without_finalized_value: StateValue<(), ()> =
             StateValue::new_provisional_value(Ok(Row::default()), (), ());
         assert!(
-            provisional_value_without_finalized_value.memory_size() <= 144,
+            provisional_value_without_finalized_value.memory_size() <= 64,
             "memory size is {}",
             provisional_value_without_finalized_value.memory_size(),
         );
@@ -1395,7 +1397,7 @@ mod tests {
             &mut Vec::new(),
         );
         assert!(
-            consolidating_value.memory_size() <= 146,
+            consolidating_value.memory_size() <= 66,
             "memory size is {}",
             consolidating_value.memory_size(),
         );

src/storage/src/upsert_continual_feedback.rs

@@ -20,7 +20,7 @@ use indexmap::map::Entry;
 use itertools::Itertools;
 use mz_ore::vec::VecExt;
 use mz_repr::{Diff, GlobalId, Row};
-use mz_storage_types::errors::{DataflowError, EnvelopeError, UpsertError};
+use mz_storage_types::errors::{DataflowError, EnvelopeError};
 use mz_timely_util::builder_async::{
     Event as AsyncEvent, OperatorBuilder as AsyncOperatorBuilder, PressOnDropButton,
 };
@@ -132,12 +132,16 @@ where
         let value = match result {
             Ok(ok) => Ok(ok),
             Err(DataflowError::EnvelopeError(err)) => match *err {
-                EnvelopeError::Upsert(err) => Err(err),
+                EnvelopeError::Upsert(err) => Err(Box::new(err)),
                 _ => return None,
             },
             Err(_) => return None,
         };
-        Some((UpsertKey::from_value(value.as_ref(), &key_indices), value))
+        let value_ref = match value {
+            Ok(ref row) => Ok(row),
+            Err(ref err) => Err(&**err),
+        };
+        Some((UpsertKey::from_value(value_ref, &key_indices), value))
     });
     let (output_handle, output) = builder.new_output();
 
@@ -536,7 +540,7 @@ where
     (
         output.as_collection().map(|result| match result {
             Ok(ok) => Ok(ok),
-            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(err))),
+            Err(err) => Err(DataflowError::from(EnvelopeError::Upsert(*err))),
         }),
         health_stream,
         snapshot_stream,
@@ -599,7 +603,7 @@ enum DrainStyle<'a, T> {
 async fn drain_staged_input<S, G, T, FromTime, E>(
     stash: &mut Vec<(T, UpsertKey, Reverse<FromTime>, Option<UpsertValue>)>,
     commands_state: &mut indexmap::IndexMap<UpsertKey, UpsertValueAndSize<T, FromTime>>,
-    output_updates: &mut Vec<(Result<Row, UpsertError>, T, Diff)>,
+    output_updates: &mut Vec<(UpsertValue, T, Diff)>,
     multi_get_scratch: &mut Vec<UpsertKey>,
     drain_style: DrainStyle<'_, T>,
     error_emitter: &mut E,