Cleanup and refactor accountsdb index (#419)

Main changes: 
1. Use a single environment for all the accountsdb indexes
2. Introduce table abstraction to make the code more clean
3. Disable read ahead in the lmdb env configuration, which was
unnecessary due to the mostly random access nature of operations.

Due to the use of a single environment instead of previously used 3,
this is a breaking change, which requires running a migration script to
fix up the index directory.

<!-- greptile_comment -->

## Greptile Summary

Major refactoring of the AccountsDB index system to improve performance
and code organization by consolidating database environments and
optimizing LMDB operations.

- Consolidated three separate LMDB environments into a single
environment for all indices in `magicblock-accounts-db/src/index.rs`
- Added new `Table` abstraction in
`magicblock-accounts-db/src/index/table.rs` to encapsulate database
operations cleanly
- Disabled LMDB read-ahead for better performance with random access
patterns
- Changed program account retrieval to return empty vector instead of
NotFound error when no accounts exist
- **Breaking Change**: Requires migration script to fix index directory
structure due to environment consolidation



<!-- /greptile_comment -->

Author: bmuddha

magicblock-accounts-db/src/index.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/iterator.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/standalone.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 }
 }