Cleanup and refactor accountsdb index

magicblock-accounts-db/src/index/iterator.rs

@@ -1,60 +1,59 @@
-use lmdb::{Cursor, Database, RoCursor, RoTransaction, Transaction};
+use lmdb::{Cursor, RoCursor, RoTransaction};
 use log::error;
 use solana_pubkey::Pubkey;
 
-use super::lmdb_utils::MDB_GET_CURRENT_OP;
+use super::{table::Table, MDB_SET_OP};
 use crate::AdbResult;
 
 /// Iterator over pubkeys and offsets, where accounts
 /// for those pubkeys can be found in database
 ///
 /// S: Starting position operation, determines where to place cursor initially
 /// N: Next position operation, determines where to move cursor next
-pub(crate) struct OffsetPubkeyIter<'env, const S: u32, const N: u32> {
-    cursor: RoCursor<'env>,
-    terminated: bool,
+pub(crate) struct OffsetPubkeyIter<'env> {
+    iter: lmdb::Iter<'env>,
+    _cursor: RoCursor<'env>,
     _txn: RoTransaction<'env>,
 }
 
-impl<'a, const S: u32, const N: u32> OffsetPubkeyIter<'a, S, N> {
+impl<'env> OffsetPubkeyIter<'env> {
     pub(super) fn new(
-        db: Database,
-        txn: RoTransaction<'a>,
+        table: &Table,
+        txn: RoTransaction<'env>,
         pubkey: Option<&Pubkey>,
     ) -> AdbResult<Self> {
-        let cursor = txn.open_ro_cursor(db)?;
+        let cursor = table.cursor_ro(&txn)?;
         // SAFETY:
         // nasty/neat trick for lifetime erasure, but we are upholding
-        // the rust's  ownership contracts by keeping txn around as well
-        let cursor: RoCursor = unsafe { std::mem::transmute(cursor) };
-        // jump to the first entry, key might be ignored depending on OP
-        cursor.get(pubkey.map(AsRef::as_ref), None, S)?;
+        // the rust's ownership contracts by keeping txn around as well
+        let mut cursor: RoCursor = unsafe { std::mem::transmute(cursor) };
+        let iter = if let Some(pubkey) = pubkey {
+            // NOTE: there's a bug in the LMDB, which ignores NotFound error when
+            // iterating on DUPSORT databases, where it just jumps to the next key,
+            // here we check for the error explicitly to prevent this behavior
+            if let Err(lmdb::Error::NotFound) =
+                cursor.get(Some(pubkey.as_ref()), None, MDB_SET_OP)
+            {
+                lmdb::Iter::Err(lmdb::Error::NotFound)
+            } else {
+                cursor.iter_dup_of(pubkey)
+            }
+        } else {
+            cursor.iter_start()
+        };
         Ok(Self {
             _txn: txn,
-            cursor,
-            terminated: false,
+            _cursor: cursor,
+            iter,
         })
     }
 }
 
-impl<const S: u32, const N: u32> Iterator for OffsetPubkeyIter<'_, S, N> {
+impl Iterator for OffsetPubkeyIter<'_> {
     type Item = (u32, Pubkey);
     fn next(&mut self) -> Option<Self::Item> {
-        if self.terminated {
-            return None;
-        }
-
-        match self.cursor.get(None, None, MDB_GET_CURRENT_OP) {
-            Ok(entry) => {
-                // advance the cursor,
-                let advance = self.cursor.get(None, None, N);
-                // if we move past the iterable range, NotFound will be
-                // triggered by OP, and we can terminate the iteration
-                if let Err(lmdb::Error::NotFound) = advance {
-                    self.terminated = true;
-                }
-                Some(bytes!(#unpack, entry.1, u32, Pubkey))
-            }
+        match self.iter.next()? {
+            Ok(entry) => Some(bytes!(#unpack, entry.1, u32, Pubkey)),
             Err(error) => {
                 error!("error advancing offset iterator cursor: {error}");
                 None

magicblock-accounts-db/src/index/table.rs

@@ -0,0 +1,92 @@
+use lmdb::{
+    Database, DatabaseFlags, Environment, RoCursor, RwCursor, RwTransaction,
+    Transaction, WriteFlags,
+};
+
+use super::WEMPTY;
+
+pub(super) struct Table {
+    db: Database,
+}
+
+impl Table {
+    pub(super) fn new(
+        env: &Environment,
+        name: &str,
+        flags: DatabaseFlags,
+    ) -> lmdb::Result<Self> {
+        let db = env.create_db(Some(name), flags)?;
+        Ok(Self { db })
+    }
+
+    #[inline]
+    pub(super) fn get<'txn, T: Transaction, K: AsRef<[u8]>>(
+        &self,
+        txn: &'txn T,
+        key: K,
+    ) -> lmdb::Result<Option<&'txn [u8]>> {
+        match txn.get(self.db, &key) {
+            Ok(bytes) => Ok(Some(bytes)),
+            Err(lmdb::Error::NotFound) => Ok(None),
+            Err(e) => Err(e),
+        }
+    }
+
+    #[inline]
+    pub(super) fn put<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+        &self,
+        txn: &mut RwTransaction,
+        key: K,
+        value: V,
+    ) -> lmdb::Result<()> {
+        txn.put(self.db, &key, &value, WEMPTY)
+    }
+
+    #[inline]
+    pub(super) fn del<K: AsRef<[u8]>>(
+        &self,
+        txn: &mut RwTransaction,
+        key: K,
+        value: Option<&[u8]>,
+    ) -> lmdb::Result<()> {
+        match txn.del(self.db, &key, value) {
+            Ok(_) | Err(lmdb::Error::NotFound) => Ok(()),
+            Err(e) => Err(e),
+        }
+    }
+
+    #[inline]
+    pub(super) fn put_if_not_exists<K: AsRef<[u8]>, V: AsRef<[u8]>>(
+        &self,
+        txn: &mut RwTransaction,
+        key: K,
+        value: V,
+    ) -> lmdb::Result<bool> {
+        let result = txn.put(self.db, &key, &value, WriteFlags::NO_OVERWRITE);
+        match result {
+            Ok(_) => Ok(true),
+            Err(lmdb::Error::KeyExist) => Ok(false),
+            Err(err) => Err(err),
+        }
+    }
+
+    #[inline]
+    pub(super) fn cursor_ro<'txn, T: Transaction>(
+        &self,
+        txn: &'txn T,
+    ) -> lmdb::Result<RoCursor<'txn>> {
+        txn.open_ro_cursor(self.db)
+    }
+
+    #[inline]
+    pub(super) fn cursor_rw<'txn>(
+        &self,
+        txn: &'txn mut RwTransaction,
+    ) -> lmdb::Result<RwCursor<'txn>> {
+        txn.open_rw_cursor(self.db)
+    }
+
+    pub(super) fn entries<T: Transaction>(&self, txn: &T) -> usize {
+        txn.stat(self.db).map(|s| s.entries()).unwrap_or_default()
+    }
+}

magicblock-accounts-db/src/index/tests.rs

@@ -171,7 +171,7 @@ fn test_ensure_correct_owner() {
     );
     let result = tenv.get_program_accounts_iter(&owner);
     assert!(
-        matches!(result, Err(AccountsDbError::NotFound)),
+        matches!(result.map(|i| i.count()), Ok(0)),
         "programs index still has record of account after owner change"
     );
 
@@ -201,14 +201,18 @@ fn test_program_index_cleanup() {
         .env
         .begin_rw_txn()
         .expect("failed to start new RW transaction");
-    let result =
-        tenv.remove_programs_index_entry(&pubkey, &mut txn, allocation.offset);
+    let result = tenv.remove_programs_index_entry(
+        &pubkey,
+        None,
+        &mut txn,
+        allocation.offset,
+    );
     assert!(result.is_ok(), "failed to remove entry from programs index");
     txn.commit().expect("failed to commit transaction");
 
     let result = tenv.get_program_accounts_iter(&owner);
     assert!(
-        matches!(result, Err(AccountsDbError::NotFound)),
+        matches!(result.map(|i| i.count()), Ok(0)),
         "programs index still has record of account after cleanup"
     );
 }
@@ -339,7 +343,7 @@ fn setup() -> IndexTestEnv {
     let directory = tempfile::tempdir()
         .expect("failed to create temp directory for index tests")
         .keep();
-    let index = AccountsDbIndex::new(&config, &directory)
+    let index = AccountsDbIndex::new(config.index_map_size, &directory)
         .expect("failed to create accountsdb index in temp dir");
     IndexTestEnv { index, directory }
 }