Reduce cloning of ResolvedDocuments in Writes (#42966)

1. do not use OrdMap in FunctionWrites, which does not need to be cheaply cloneable
2. store an Arc inside of `Writes`'s OrdMap (which I have changed to an OrdSet), because OrdMap frequently clones its contents

GitOrigin-RevId: 9a0e97384db9c83588f55ebcdda2b1a4d681c25d

Author: goffrie

crates/application/src/application_function_runner/mod.rs

@@ -402,7 +402,9 @@ impl<RT: Runtime> FunctionRouter<RT> {
                 udf_type,
                 tx.identity().clone(),
                 tx.begin_timestamp(),
-                tx.writes().as_flat()?.clone().into(),
+                FunctionWrites {
+                    updates: tx.writes().as_flat()?.coalesced_writes().cloned().collect(),
+                },
                 log_line_sender,
                 function_metadata,
                 http_action_metadata,

crates/database/src/committer.rs

@@ -73,7 +73,6 @@ use common::{
         Timestamp,
         WriteTimestamp,
     },
-    value::ResolvedDocumentId,
 };
 use errors::ErrorMetadata;
 use fastrace::prelude::*;
@@ -142,7 +141,6 @@ use crate::{
         PendingWrites,
         WriteSource,
     },
-    writes::DocumentWrite,
     ComponentRegistry,
     Snapshot,
     Transaction,
@@ -695,10 +693,10 @@ impl<RT: Runtime> Committer<RT> {
         // which is the same order they should be applied to database metadata
         // and index data structures
         let mut ordered_updates = updates;
-        ordered_updates.sort_by_key(|(id, update)| {
+        ordered_updates.sort_by_key(|update| {
             table_dependency_sort_key(
                 BootstrapTableIds::new(&transaction.table_mapping),
-                InternalDocumentId::from(**id),
+                InternalDocumentId::from(update.id),
                 update.new_document.as_ref(),
             )
         });
@@ -717,7 +715,7 @@ impl<RT: Runtime> Committer<RT> {
             commit_ts,
             ordered_updates
                 .into_iter()
-                .map(|(&id, update)| (id, PackedDocumentUpdate::pack(update)))
+                .map(|update| (update.id, PackedDocumentUpdate::pack(update)))
                 .collect(),
             write_source,
             snapshot,
@@ -735,7 +733,7 @@ impl<RT: Runtime> Committer<RT> {
     fn compute_writes(
         &self,
         commit_ts: Timestamp,
-        ordered_updates: &Vec<(&ResolvedDocumentId, &DocumentUpdateWithPrevTs)>,
+        ordered_updates: &Vec<&DocumentUpdateWithPrevTs>,
     ) -> anyhow::Result<(
         Vec<ValidatedDocumentWrite>,
         BTreeSet<(Timestamp, DatabaseIndexUpdate)>,
@@ -754,16 +752,14 @@ impl<RT: Runtime> Committer<RT> {
             .pending_writes
             .latest_snapshot()
             .unwrap_or_else(|| self.snapshot_manager.read().latest_snapshot());
-        for &(id, document_update) in ordered_updates.iter() {
+        for &document_update in ordered_updates.iter() {
             let (updates, doc_in_vector_index) =
                 latest_pending_snapshot.update(document_update, commit_ts)?;
             index_writes.extend(updates);
             document_writes.push(ValidatedDocumentWrite {
                 commit_ts,
-                id: (*id).into(),
-                write: DocumentWrite {
-                    document: document_update.new_document.clone(),
-                },
+                id: document_update.id.into(),
+                write: document_update.new_document.clone(),
                 doc_in_vector_index,
                 prev_ts: document_update.old_document.as_ref().map(|&(_, ts)| ts),
             });
@@ -934,7 +930,7 @@ impl<RT: Runtime> Committer<RT> {
                         .map(|write| DocumentLogEntry {
                             ts: write.commit_ts,
                             id: write.id,
-                            value: write.write.document,
+                            value: write.write,
                             prev_ts: write.prev_ts,
                         })
                         .collect_vec(),
@@ -1021,7 +1017,7 @@ impl<RT: Runtime> Committer<RT> {
         }
         for validated_write in document_writes {
             let ValidatedDocumentWrite {
-                write: DocumentWrite { document },
+                write: document,
                 doc_in_vector_index,
                 ..
             } = validated_write;
@@ -1090,7 +1086,7 @@ impl<RT: Runtime> Committer<RT> {
 struct ValidatedDocumentWrite {
     commit_ts: Timestamp,
     id: InternalDocumentId,
-    write: DocumentWrite,
+    write: Option<ResolvedDocument>,
     doc_in_vector_index: DocInVectorIndex,
     prev_ts: Option<Timestamp>,
 }

crates/database/src/lib.rs

@@ -93,7 +93,6 @@ pub use write_log::{
     WriteSource,
 };
 pub use writes::{
-    DocumentWrite,
     TransactionWriteSize,
     Writes,
 };

crates/database/src/tests/apply_function_runner_tx.rs

@@ -80,7 +80,12 @@ async fn test_apply_function_runner_tx_new_table(rt: TestRuntime) -> anyhow::Res
         .iter()
         .map(|(table, stats)| (*table, stats.rows_read))
         .collect();
-    let updates = function_runner_tx.writes.as_flat()?.clone().into_updates();
+    let updates = function_runner_tx
+        .writes
+        .as_flat()?
+        .coalesced_writes()
+        .cloned()
+        .collect();
     backend_tx.apply_function_runner_tx(
         *begin_timestamp,
         reads,
@@ -133,7 +138,12 @@ async fn test_apply_function_runner_tx_read_only(rt: TestRuntime) -> anyhow::Res
         .iter()
         .map(|(table, stats)| (*table, stats.rows_read))
         .collect();
-    let updates = function_runner_tx.writes.as_flat()?.clone().into_updates();
+    let updates = function_runner_tx
+        .writes
+        .as_flat()?
+        .coalesced_writes()
+        .cloned()
+        .collect();
     backend_tx.apply_function_runner_tx(
         *begin_timestamp,
         reads,
@@ -183,7 +193,12 @@ async fn test_apply_function_runner_tx_replace(rt: TestRuntime) -> anyhow::Resul
         .iter()
         .map(|(table, stats)| (*table, stats.rows_read))
         .collect();
-    let updates = function_runner_tx.writes.as_flat()?.clone().into_updates();
+    let updates = function_runner_tx
+        .writes
+        .as_flat()?
+        .coalesced_writes()
+        .cloned()
+        .collect();
     backend_tx.apply_function_runner_tx(
         *begin_timestamp,
         reads,
@@ -218,7 +233,7 @@ async fn test_apply_function_runner_tx_merge_existing_writes(
             FunctionUsageTracker::new(),
         )
         .await?;
-    let updates = backend_tx.writes().as_flat()?.clone().into_updates();
+    let updates = backend_tx.writes().as_flat()?.coalesced_writes().cloned();
     function_runner_tx.merge_writes(updates)?;
 
     // Perform writes as if in funrun
@@ -236,7 +251,12 @@ async fn test_apply_function_runner_tx_merge_existing_writes(
         .iter()
         .map(|(table, stats)| (*table, stats.rows_read))
         .collect();
-    let updates = function_runner_tx.writes.as_flat()?.clone().into_updates();
+    let updates = function_runner_tx
+        .writes
+        .as_flat()?
+        .coalesced_writes()
+        .cloned()
+        .collect();
     backend_tx.apply_function_runner_tx(
         *begin_timestamp,
         reads,
@@ -288,7 +308,12 @@ async fn test_apply_function_runner_tx_merge_existing_writes_bad(
         .iter()
         .map(|(table, stats)| (*table, stats.rows_read))
         .collect();
-    let updates = function_runner_tx.writes.as_flat()?.clone().into_updates();
+    let updates = function_runner_tx
+        .writes
+        .as_flat()?
+        .coalesced_writes()
+        .cloned()
+        .collect();
     assert!(backend_tx
         .apply_function_runner_tx(
             *begin_timestamp,

crates/database/src/transaction.rs

@@ -77,7 +77,6 @@ use common::{
     virtual_system_mapping::VirtualSystemMapping,
 };
 use errors::ErrorMetadata;
-use imbl::OrdMap;
 use indexing::backend_in_memory_indexes::RangeRequest;
 use keybroker::{
     Identity,
@@ -394,8 +393,11 @@ impl<RT: Runtime> Transaction<RT> {
         let mut biggest_document_id = None;
         let mut max_nesting = 0;
         let mut most_nested_document_id = None;
-        for (document_id, DocumentUpdateWithPrevTs { new_document, .. }) in
-            self.writes.coalesced_writes()
+        for DocumentUpdateWithPrevTs {
+            id: document_id,
+            new_document,
+            ..
+        } in self.writes.coalesced_writes()
         {
             let (size, nesting) = new_document
                 .as_ref()
@@ -440,7 +442,7 @@ impl<RT: Runtime> Transaction<RT> {
         num_intervals: usize,
         user_tx_size: crate::reads::TransactionReadSize,
         system_tx_size: crate::reads::TransactionReadSize,
-        updates: OrdMap<ResolvedDocumentId, DocumentUpdateWithPrevTs>,
+        updates: Vec<DocumentUpdateWithPrevTs>,
         rows_read_by_tablet: BTreeMap<TabletId, u64>,
     ) -> anyhow::Result<()> {
         anyhow::ensure!(
@@ -467,24 +469,29 @@ impl<RT: Runtime> Transaction<RT> {
     // In most scenarios this transaction will have no writes.
     pub fn merge_writes(
         &mut self,
-        updates: OrdMap<ResolvedDocumentId, DocumentUpdateWithPrevTs>,
+        updates: impl IntoIterator<Item = DocumentUpdateWithPrevTs>,
     ) -> anyhow::Result<()> {
         let existing_updates = self.writes().as_flat()?.clone().into_updates();
 
         let mut updates = updates.into_iter().collect::<Vec<_>>();
         let bootstrap_tables = self.bootstrap_tables();
-        updates.sort_by_cached_key(|(id, update)| {
-            table_dependency_sort_key(bootstrap_tables, (*id).into(), update.new_document.as_ref())
+        updates.sort_by_cached_key(|update| {
+            table_dependency_sort_key(
+                bootstrap_tables,
+                update.id.into(),
+                update.new_document.as_ref(),
+            )
         });
 
         let mut preserved_update_count = 0;
-        for (id, update) in updates {
+        for update in updates {
+            let id = update.id;
             // Ensure that the existing update matches, and that
             // that the merged-in writes didn't otherwise modify documents
             // already written to in this transaction.
             if let Some(existing_update) = existing_updates.get(&id) {
                 anyhow::ensure!(
-                    *existing_update == update,
+                    **existing_update == update,
                     "Conflicting updates for document {id}"
                 );
                 preserved_update_count += 1;
@@ -1197,7 +1204,7 @@ impl FinalTransaction {
             .writes
             .as_flat()?
             .coalesced_writes()
-            .map(|(id, _)| id.tablet_id)
+            .map(|update| update.id.tablet_id)
             .collect();
         Self::validate_memory_index_size(
             table_mapping,