From 446ffcd1d17ccf098e8ed3d041faf15d2abaff14 Mon Sep 17 00:00:00 2001
From: Ivan Shumkov <ivan@shumkov.ru>
Date: Wed, 27 Mar 2024 01:17:44 +0700
Subject: [PATCH] pref(rocksdb): enable async io for iterators

---
 .../rocksdb_storage/storage_context/context_immediate.rs   | 7 +++++--
 .../src/rocksdb_storage/storage_context/context_no_tx.rs   | 7 +++++--
 storage/src/rocksdb_storage/storage_context/context_tx.rs  | 7 +++++--
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/storage/src/rocksdb_storage/storage_context/context_immediate.rs b/storage/src/rocksdb_storage/storage_context/context_immediate.rs
index f5785b0c..2dec5f4b 100644
--- a/storage/src/rocksdb_storage/storage_context/context_immediate.rs
+++ b/storage/src/rocksdb_storage/storage_context/context_immediate.rs
@@ -33,7 +33,7 @@ use grovedb_costs::{
     storage_cost::key_value_cost::KeyValueStorageCost, ChildrenSizesWithIsSumTree, CostResult,
     CostsExt,
 };
-use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode, WriteBatchWithTransaction};
+use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode, ReadOptions, WriteBatchWithTransaction};
 
 use super::{make_prefixed_key, PrefixedRocksDbBatch, PrefixedRocksDbRawIterator};
 use crate::{
@@ -232,9 +232,12 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbImmediateStorageContext<'db> {
     }
 
     fn raw_iter(&self) -> Self::RawIterator {
+        let mut opts = ReadOptions::default();
+        opts.set_async_io(true);
+
         PrefixedRocksDbRawIterator {
             prefix: self.prefix,
-            raw_iterator: self.transaction.raw_iterator(),
+            raw_iterator: self.transaction.raw_iterator_opt(opts),
         }
     }
 }
diff --git a/storage/src/rocksdb_storage/storage_context/context_no_tx.rs b/storage/src/rocksdb_storage/storage_context/context_no_tx.rs
index fd639a5a..097e3705 100644
--- a/storage/src/rocksdb_storage/storage_context/context_no_tx.rs
+++ b/storage/src/rocksdb_storage/storage_context/context_no_tx.rs
@@ -33,7 +33,7 @@ use grovedb_costs::{
     storage_cost::key_value_cost::KeyValueStorageCost, ChildrenSizesWithIsSumTree, CostResult,
     CostsExt, OperationCost,
 };
-use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode};
+use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode, ReadOptions};
 
 use super::{batch::PrefixedMultiContextBatchPart, make_prefixed_key, PrefixedRocksDbRawIterator};
 use crate::{
@@ -278,9 +278,12 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbStorageContext<'db> {
     }
 
     fn raw_iter(&self) -> Self::RawIterator {
+        let mut opts = ReadOptions::default();
+        opts.set_async_io(true);
+
         PrefixedRocksDbRawIterator {
             prefix: self.prefix,
-            raw_iterator: self.storage.raw_iterator(),
+            raw_iterator: self.storage.raw_iterator_opt(opts),
         }
     }
 }
diff --git a/storage/src/rocksdb_storage/storage_context/context_tx.rs b/storage/src/rocksdb_storage/storage_context/context_tx.rs
index d5a480c3..32b28031 100644
--- a/storage/src/rocksdb_storage/storage_context/context_tx.rs
+++ b/storage/src/rocksdb_storage/storage_context/context_tx.rs
@@ -33,7 +33,7 @@ use grovedb_costs::{
     cost_return_on_error, storage_cost::key_value_cost::KeyValueStorageCost,
     ChildrenSizesWithIsSumTree, CostResult, CostsExt, OperationCost,
 };
-use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode};
+use rocksdb::{ColumnFamily, DBRawIteratorWithThreadMode, ReadOptions};
 
 use super::{batch::PrefixedMultiContextBatchPart, make_prefixed_key, PrefixedRocksDbRawIterator};
 use crate::{
@@ -310,9 +310,12 @@ impl<'db> StorageContext<'db> for PrefixedRocksDbTransactionContext<'db> {
     }
 
     fn raw_iter(&self) -> Self::RawIterator {
+        let mut opts = ReadOptions::default();
+        opts.set_async_io(true);
+
         PrefixedRocksDbRawIterator {
             prefix: self.prefix,
-            raw_iterator: self.transaction.raw_iterator(),
+            raw_iterator: self.transaction.raw_iterator_opt(opts),
         }
     }
 }