implemented files batching

Author: dean-starkware

crates/apollo_storage/src/body/mod.rs

@@ -473,34 +473,51 @@ fn write_transactions<'env>(
     events_table: &'env EventsTable<'env>,
     block_number: BlockNumber,
 ) -> StorageResult<()> {
-    for (index, ((tx, tx_output), tx_hash)) in block_body
-        .transactions
+    if block_body.transactions.is_empty() {
+        return Ok(());
+    }
+
+    // Batch append all transactions and transaction outputs at once
+    let transactions_refs: Vec<&Transaction> = block_body.transactions.iter().collect();
+    let tx_outputs_refs: Vec<&TransactionOutput> = block_body.transaction_outputs.iter().collect();
+    
+    let tx_locations = file_handlers.append_transactions_batch(&transactions_refs);
+    let tx_output_locations = file_handlers.append_transaction_outputs_batch(&tx_outputs_refs);
+
+    // Write metadata and events to MDBX tables
+    for (index, ((tx_output, tx_hash), (tx_location, tx_output_location))) in block_body
+        .transaction_outputs
         .iter()
-        .zip(block_body.transaction_outputs.iter())
         .zip(block_body.transaction_hashes.iter())
+        .zip(tx_locations.iter().zip(tx_output_locations.iter()))
         .enumerate()
     {
         let tx_offset_in_block = TransactionOffsetInBlock(index);
         let transaction_index = TransactionIndex(block_number, tx_offset_in_block);
-        let tx_location = file_handlers.append_transaction(tx);
-        let tx_output_location = file_handlers.append_transaction_output(tx_output);
+        
         write_events(tx_output, txn, events_table, transaction_index)?;
         transaction_hash_to_idx_table.insert(txn, tx_hash, &transaction_index)?;
         transaction_metadata_table.append(
             txn,
             &transaction_index,
-            &TransactionMetadata { tx_location, tx_output_location, tx_hash: *tx_hash },
+            &TransactionMetadata { 
+                tx_location: *tx_location, 
+                tx_output_location: *tx_output_location, 
+                tx_hash: *tx_hash 
+            },
         )?;
+    }
 
-        // If this is the last iteration, update the file offset table.
-        if index == block_body.transactions.len() - 1 {
-            file_offset_table.upsert(txn, &OffsetKind::Transaction, &tx_location.next_offset())?;
-            file_offset_table.upsert(
-                txn,
-                &OffsetKind::TransactionOutput,
-                &tx_output_location.next_offset(),
-            )?;
-        }
+    // Update file offsets with the last locations
+    if let (Some(last_tx_location), Some(last_tx_output_location)) = 
+        (tx_locations.last(), tx_output_locations.last()) 
+    {
+        file_offset_table.upsert(txn, &OffsetKind::Transaction, &last_tx_location.next_offset())?;
+        file_offset_table.upsert(
+            txn,
+            &OffsetKind::TransactionOutput,
+            &last_tx_output_location.next_offset(),
+        )?;
     }
 
     Ok(())

crates/apollo_storage/src/class.rs

@@ -207,11 +207,27 @@ fn write_classes<'env>(
     file_handlers: &FileHandlers<RW>,
     file_offset_table: &'env FileOffsetTable<'env>,
 ) -> StorageResult<()> {
-    for (class_hash, contract_class) in classes {
-        let location = file_handlers.append_contract_class(contract_class);
-        declared_classes_table.insert(txn, class_hash, &location)?;
-        file_offset_table.upsert(txn, &OffsetKind::ContractClass, &location.next_offset())?;
+    if classes.is_empty() {
+        return Ok(());
     }
+
+    // Extract contract classes for batch append
+    let contract_classes: Vec<&SierraContractClass> =
+        classes.iter().map(|(_, class)| *class).collect();
+
+    // Batch append all contract classes at once
+    let locations = file_handlers.append_contract_classes_batch(&contract_classes);
+
+    // Write to MDBX tables
+    for ((class_hash, _), location) in classes.iter().zip(locations.iter()) {
+        declared_classes_table.insert(txn, class_hash, location)?;
+    }
+
+    // Update file offset with the last location
+    if let Some(last_location) = locations.last() {
+        file_offset_table.upsert(txn, &OffsetKind::ContractClass, &last_location.next_offset())?;
+    }
+
     Ok(())
 }
 
@@ -223,18 +239,42 @@ fn write_deprecated_classes<'env>(
     file_handlers: &FileHandlers<RW>,
     file_offset_table: &'env FileOffsetTable<'env>,
 ) -> StorageResult<()> {
+    if deprecated_classes.is_empty() {
+        return Ok(());
+    }
+
+    // Filter out classes that already exist and collect the ones to write
+    let mut classes_to_write = Vec::new();
+    let mut hashes_to_write = Vec::new();
+    
     for (class_hash, deprecated_contract_class) in deprecated_classes {
-        if deprecated_declared_classes_table.get(txn, class_hash)?.is_some() {
-            continue;
+        if deprecated_declared_classes_table.get(txn, class_hash)?.is_none() {
+            classes_to_write.push(*deprecated_contract_class);
+            hashes_to_write.push(*class_hash);
         }
-        let location = file_handlers.append_deprecated_contract_class(deprecated_contract_class);
-        let value = IndexedDeprecatedContractClass { block_number, location_in_file: location };
+    }
+
+    if classes_to_write.is_empty() {
+        return Ok(());
+    }
+
+    // Batch append all deprecated contract classes at once
+    let locations = file_handlers.append_deprecated_contract_classes_batch(&classes_to_write);
+
+    // Write to MDBX tables
+    for (class_hash, location) in hashes_to_write.iter().zip(locations.iter()) {
+        let value = IndexedDeprecatedContractClass { block_number, location_in_file: *location };
+        deprecated_declared_classes_table.insert(txn, class_hash, &value)?;
+    }
+
+    // Update file offset with the last location
+    if let Some(last_location) = locations.last() {
         file_offset_table.upsert(
             txn,
             &OffsetKind::DeprecatedContractClass,
-            &location.next_offset(),
+            &last_location.next_offset(),
         )?;
-        deprecated_declared_classes_table.insert(txn, class_hash, &value)?;
     }
+
     Ok(())
 }

crates/apollo_storage/src/lib.rs

@@ -727,11 +727,24 @@ impl FileHandlers<RW> {
         self.clone().contract_class.append(contract_class)
     }
 
+    // Appends multiple contract classes in batch and returns their locations.
+    fn append_contract_classes_batch(
+        &self,
+        contract_classes: &[&SierraContractClass],
+    ) -> Vec<LocationInFile> {
+        self.clone().contract_class.append_batch(contract_classes)
+    }
+
     // Appends a CASM to the corresponding file and returns its location.
     fn append_casm(&self, casm: &CasmContractClass) -> LocationInFile {
         self.clone().casm.append(casm)
     }
 
+    // Appends multiple CASMs in batch and returns their locations.
+    fn append_casms_batch(&self, casms: &[&CasmContractClass]) -> Vec<LocationInFile> {
+        self.clone().casm.append_batch(casms)
+    }
+
     // Appends a deprecated contract class to the corresponding file and returns its location.
     fn append_deprecated_contract_class(
         &self,
@@ -740,16 +753,37 @@ impl FileHandlers<RW> {
         self.clone().deprecated_contract_class.append(deprecated_contract_class)
     }
 
+    // Appends multiple deprecated contract classes in batch and returns their locations.
+    fn append_deprecated_contract_classes_batch(
+        &self,
+        deprecated_contract_classes: &[&DeprecatedContractClass],
+    ) -> Vec<LocationInFile> {
+        self.clone().deprecated_contract_class.append_batch(deprecated_contract_classes)
+    }
+
     // Appends a thin transaction output to the corresponding file and returns its location.
     fn append_transaction_output(&self, transaction_output: &TransactionOutput) -> LocationInFile {
         self.clone().transaction_output.append(transaction_output)
     }
 
+    // Appends multiple transaction outputs in batch and returns their locations.
+    fn append_transaction_outputs_batch(
+        &self,
+        transaction_outputs: &[&TransactionOutput],
+    ) -> Vec<LocationInFile> {
+        self.clone().transaction_output.append_batch(transaction_outputs)
+    }
+
     // Appends a transaction to the corresponding file and returns its location.
     fn append_transaction(&self, transaction: &Transaction) -> LocationInFile {
         self.clone().transaction.append(transaction)
     }
 
+    // Appends multiple transactions in batch and returns their locations.
+    fn append_transactions_batch(&self, transactions: &[&Transaction]) -> Vec<LocationInFile> {
+        self.clone().transaction.append_batch(transactions)
+    }
+
     // TODO(dan): Consider 1. flushing only the relevant files, 2. flushing concurrently.
     #[latency_histogram("storage_file_handler_flush_latency_seconds", false)]
     fn flush(&self) {

crates/apollo_storage/src/mmap_file/mod.rs

@@ -107,6 +107,13 @@ pub(crate) trait Writer<V: ValueSerde> {
     /// Inserts an object to the file, returns the [`LocationInFile`] of the object.
     fn append(&mut self, val: &V::Value) -> LocationInFile;
 
+    /// Inserts multiple objects to the file in a batch, returns their [`LocationInFile`]s.
+    /// This is more efficient than calling `append` multiple times as it:
+    /// - Serializes all objects first
+    /// - Writes them all at once
+    /// - Performs a single flush operation for all objects
+    fn append_batch(&mut self, vals: &[&V::Value]) -> Vec<LocationInFile>;
+
     /// Flushes the mmap to the file.
     fn flush(&self);
 }
@@ -247,6 +254,58 @@ impl<V: ValueSerde + Debug> Writer<V> for FileHandler<V, RW> {
         location
     }
 
+    fn append_batch(&mut self, vals: &[&V::Value]) -> Vec<LocationInFile> {
+        if vals.is_empty() {
+            return Vec::new();
+        }
+
+        trace!("Batch inserting {} objects", vals.len());
+
+        // First, serialize all values
+        let mut serialized_vals = Vec::with_capacity(vals.len());
+        let mut total_len = 0;
+        for val in vals {
+            let serialized = V::serialize(val).expect("Should be able to serialize");
+            total_len += serialized.len();
+            serialized_vals.push(serialized);
+        }
+
+        // Then, write all at once and record locations
+        let mut locations = Vec::with_capacity(vals.len());
+        let start_offset;
+        {
+            let mut mmap_file = self.mmap_file.lock().expect("Lock should not be poisoned");
+            start_offset = mmap_file.offset;
+            let mut current_offset = start_offset;
+
+            trace!("Batch inserting at starting offset: {}, total size: {}", start_offset, total_len);
+
+            for serialized in &serialized_vals {
+                let len = serialized.len();
+                let mmap_slice = &mut mmap_file.mmap[current_offset..];
+                mmap_slice[..len].copy_from_slice(serialized);
+                locations.push(LocationInFile { offset: current_offset, len });
+                current_offset += len;
+            }
+
+            // Single flush for all objects
+            mmap_file
+                .mmap
+                .flush_async_range(start_offset, total_len)
+                .expect("Failed to asynchronously flush the mmap after batch insert");
+            
+            mmap_file.offset = current_offset;
+            mmap_file.should_flush = true;
+        }
+
+        // Check if file needs to grow after all writes
+        let final_offset = start_offset + total_len;
+        self.grow_file_if_needed(final_offset);
+        
+        trace!("Batch insert complete. Wrote {} objects totaling {} bytes", vals.len(), total_len);
+        locations
+    }
+
     fn flush(&self) {
         let mut mmap_file = self.mmap_file.lock().expect("Lock should not be poisoned");
         if mmap_file.should_flush {