perf(storage): add append functionallity to tables (#1895)

* perf(storage): add append functionallity to tables

* CR fix

crates/papyrus_storage/src/db/mod.rs

@@ -159,6 +159,11 @@ pub enum DbError {
     /// An error that occurred when trying to open a db file that does not exist.
     #[error("The file '{0}' does not exist.")]
     FileDoesNotExist(PathBuf),
+    // TODO(dvir): consider adding more details about the error, table name, key, value and last
+    // key in the tree.
+    /// An error that occurred when trying to append a key when it is not the last.
+    #[error("Append error. The key is not the last in the table.")]
+    Append,
 }
 
 type DbResult<V> = result::Result<V, DbError>;

crates/papyrus_storage/src/db/table_types/dup_sort_tables.rs

@@ -274,6 +274,61 @@ impl<'env, K: KeyTrait + Debug, V: ValueSerde + Debug, T: DupSortTableType + Dup
         Ok(())
     }
 
+    // TODO(dvir): consider first checking if the key is equal to the last key, delete the last key,
+    // and then append instead of optimistically append.
+    fn append(
+        &'env self,
+        txn: &DbTransaction<'env, RW>,
+        key: &K,
+        value: &<V as ValueSerde>::Value,
+    ) -> DbResult<()> {
+        let main_key = T::get_main_key(key)?;
+        let sub_key_and_value = T::get_sub_key_and_value(key, value)?;
+
+        let mut cursor = txn.txn.cursor(&self.database)?;
+        match cursor.put(&main_key, &sub_key_and_value, WriteFlags::APPEND_DUP | WriteFlags::APPEND)
+        {
+            Err(libmdbx::Error::KeyMismatch) => {
+                // This case can happen if the appended sub_key_and_value is smaller than the last
+                // entry value, but the sub key itself is equal.
+                // For example: append (0,0) -> 1, old last entry: (0,0) -> 2.
+                let (last_main_key_bytes, last_key_suffix_and_value_bytes) =
+                    cursor.last::<DbKeyType<'_>, DbValueType<'_>>()?.expect(
+                        "Should have a last key. otherwise the previous put operation would \
+                         succeed.",
+                    );
+
+                // If the appended key is equal to the last key in the table, we can append it. To
+                // do that we first need to delete the old entry.
+                if last_main_key_bytes == main_key.as_slice()
+                    && last_key_suffix_and_value_bytes.starts_with(&T::get_sub_key(key)?)
+                {
+                    cursor.del(WriteFlags::empty())?;
+                    cursor.put(
+                        &main_key,
+                        &sub_key_and_value,
+                        WriteFlags::APPEND_DUP | WriteFlags::APPEND,
+                    )?;
+
+                    Ok(())
+                } else {
+                    Err(DbError::Append)
+                }
+            }
+            Ok(()) => {
+                // In the case of overriding the last key with a bigger value, we need to delete the
+                // old entry.
+                if let Some(prev) = cursor.prev_dup::<DbKeyType<'_>, DbValueType<'_>>()? {
+                    if prev.1.starts_with(&T::get_sub_key(key)?) {
+                        cursor.del(WriteFlags::empty())?;
+                    }
+                }
+                Ok(())
+            }
+            Err(err) => Err(err.into()),
+        }
+    }
+
     fn delete(&'env self, txn: &DbTransaction<'env, RW>, key: &Self::Key) -> DbResult<()> {
         let main_key = T::get_main_key(key)?;
         let first_sub_key = T::get_sub_key_lower_bound(key)?;

crates/papyrus_storage/src/db/table_types/mod.rs

@@ -14,6 +14,9 @@ pub(crate) use simple_table::SimpleTable;
 #[cfg(test)]
 pub(crate) mod test_utils;
 
+// TODO(dvir): consider adding the create_table method to the Table trait.
+// TODO(dvir): add some documentation to the Table and the cursor traits.
+// TODO(dvir): consider adding unchecked version of the those functions.
 pub(crate) trait Table<'env> {
     type Key: KeyTrait + Debug;
     type Value: ValueSerde + Debug;
@@ -46,9 +49,21 @@ pub(crate) trait Table<'env> {
         value: &<Self::Value as ValueSerde>::Value,
     ) -> DbResult<()>;
 
+    // Append a key value pair to the end of the table. The key must be bigger than or equal to
+    // the last key in the table; otherwise, an error will be returned.
+    #[allow(dead_code)]
+    fn append(
+        &'env self,
+        txn: &DbTransaction<'env, RW>,
+        key: &Self::Key,
+        value: &<Self::Value as ValueSerde>::Value,
+    ) -> DbResult<()>;
+
     fn delete(&'env self, txn: &DbTransaction<'env, RW>, key: &Self::Key) -> DbResult<()>;
 }
 
+// TODO(dvir): consider adding append functionality using a cursor. It should be more efficient for
+// more than a single append operation (also for other table types).
 pub(crate) trait DbCursorTrait {
     type Key: KeyTrait + Debug;
     type Value: ValueSerde + Debug;

crates/papyrus_storage/src/db/table_types/simple_table.rs

@@ -113,6 +113,24 @@ impl<'env, K: KeyTrait + Debug, V: ValueSerde + Debug> Table<'env>
         Ok(())
     }
 
+    #[allow(dead_code)]
+    fn append(
+        &'env self,
+        txn: &DbTransaction<'env, RW>,
+        key: &K,
+        value: &<V as ValueSerde>::Value,
+    ) -> DbResult<()> {
+        let data = V::serialize(value)?;
+        let bin_key = key.serialize()?;
+        txn.txn.put(&self.database, bin_key, data, WriteFlags::APPEND).map_err(
+            |err| match err {
+                libmdbx::Error::KeyMismatch => DbError::Append,
+                _ => err.into(),
+            },
+        )?;
+        Ok(())
+    }
+
     fn delete(&'env self, txn: &DbTransaction<'env, RW>, key: &Self::Key) -> DbResult<()> {
         let bin_key = key.serialize()?;
         txn.txn.del(&self.database, bin_key, None)?;

crates/papyrus_storage/src/db/table_types/test_utils.rs

@@ -1,3 +1,4 @@
+use assert_matches::assert_matches;
 use rand::rngs::ThreadRng;
 use rand::Rng;
 use tracing::debug;
@@ -43,6 +44,9 @@ pub(crate) fn table_test<T: TableType>(
     let upsert_test_table = create_table(&mut writer, "upsert_test").unwrap();
     upsert_test(upsert_test_table, &mut writer);
 
+    let append_test_table = create_table(&mut writer, "append_test").unwrap();
+    append_test(append_test_table, &mut writer);
+
     let delete_test_table = create_table(&mut writer, "delete_test").unwrap();
     delete_test(delete_test_table, &mut writer);
 
@@ -121,6 +125,51 @@ fn upsert_test<T: TableType>(
     assert_eq!(table.get(&txn, &(1, 1)).unwrap(), Some(0));
 }
 
+fn append_test<T: TableType>(
+    table_id: TableIdentifier<TableKey, TableValue, T>,
+    writer: &mut DbWriter,
+) where
+    for<'env> TableHandle<'env, TableKey, TableValue, T>:
+        Table<'env, Key = TableKey, Value = TableValue>,
+{
+    let txn = writer.begin_rw_txn().unwrap();
+    let table = txn.open_table(&table_id).unwrap();
+
+    // Append to an empty table.
+    table.append(&txn, &(1, 1), &11).unwrap();
+    assert_eq!(table.get(&txn, &(1, 1)).unwrap(), Some(11));
+
+    // Successful appends.
+    table.append(&txn, &(1, 1), &0).unwrap();
+    table.append(&txn, &(1, 2), &12).unwrap();
+    table.append(&txn, &(2, 0), &20).unwrap();
+    table.append(&txn, &(2, 2), &22).unwrap();
+
+    assert_eq!(table.get(&txn, &(1, 1)).unwrap(), Some(0));
+    assert_eq!(table.get(&txn, &(1, 2)).unwrap(), Some(12));
+    assert_eq!(table.get(&txn, &(2, 0)).unwrap(), Some(20));
+    assert_eq!(table.get(&txn, &(2, 2)).unwrap(), Some(22));
+
+    // Override the last key with a smaller value.
+    table.append(&txn, &(2, 2), &0).unwrap();
+    assert_eq!(table.get(&txn, &(2, 2)).unwrap(), Some(0));
+
+    // Override the last key with a bigger value.
+    table.append(&txn, &(2, 2), &100).unwrap();
+    assert_eq!(table.get(&txn, &(2, 2)).unwrap(), Some(100));
+
+    // Append key that is not the last, should fail.
+    assert_matches!(table.append(&txn, &(0, 0), &0), Err(DbError::Append));
+    assert_matches!(table.append(&txn, &(1, 3), &0), Err(DbError::Append));
+    assert_matches!(table.append(&txn, &(2, 1), &0), Err(DbError::Append));
+
+    // Check the final database.
+    assert_eq!(table.get(&txn, &(1, 1)).unwrap(), Some(0));
+    assert_eq!(table.get(&txn, &(1, 2)).unwrap(), Some(12));
+    assert_eq!(table.get(&txn, &(2, 0)).unwrap(), Some(20));
+    assert_eq!(table.get(&txn, &(2, 2)).unwrap(), Some(100));
+}
+
 fn delete_test<T: TableType>(
     table_id: TableIdentifier<TableKey, TableValue, T>,
     writer: &mut DbWriter,
@@ -261,8 +310,7 @@ pub(crate) fn random_table_test<T0: TableType, T1: TableType>(
     for iter in 0..ITERS {
         debug!("iteration: {iter:?}");
         let wtxn = writer.begin_rw_txn().unwrap();
-        // TODO(dvir): add append functionality to this test.
-        let random_op = rng.gen_range(0..3);
+        let random_op = rng.gen_range(0..4);
         let key = get_random_key(&mut rng);
         let value = rng.gen_range(0..MAX_VALUE);
 
@@ -283,6 +331,18 @@ pub(crate) fn random_table_test<T0: TableType, T1: TableType>(
             first_table.upsert(&wtxn, &key, &value).unwrap();
             second_table.upsert(&wtxn, &key, &value).unwrap();
         } else if random_op == 2 {
+            // Append
+            // TODO(dvir): consider increasing the number of successful appends (append of not the
+            // last entry will fail).
+            debug!("append: {key:?}, {value:?}");
+            let first_res = first_table.append(&wtxn, &key, &value);
+            let second_res = second_table.append(&wtxn, &key, &value);
+            assert!(
+                (first_res.is_ok() && second_res.is_ok())
+                    || (matches!(first_res.unwrap_err(), DbError::Append)
+                        && matches!(second_res.unwrap_err(), DbError::Append))
+            );
+        } else if random_op == 3 {
             // Delete
             debug!("delete: {key:?}");
             first_table.delete(&wtxn, &key).unwrap();