offload remaining db iterator initial seeks on cache miss

consume task budget on cache hit

Signed-off-by: Jason Volk <jason@zemos.net>

src/database/map/rev_stream.rs

@@ -1,14 +1,20 @@
-use conduwuit::{implement, Result};
-use futures::stream::{Stream, StreamExt};
-use serde::Deserialize;
+use std::sync::Arc;
 
-use crate::{keyval, keyval::KeyVal, stream, stream::Cursor};
+use conduwuit::{implement, Result};
+use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use rocksdb::Direction;
+use serde::Deserialize;
+use tokio::task;
+
+use crate::{keyval, keyval::KeyVal, stream};
 
 /// Iterate key-value entries in the map from the end.
 ///
 /// - Result is deserialized
 #[implement(super::Map)]
-pub fn rev_stream<'a, K, V>(&'a self) -> impl Stream<Item = Result<KeyVal<'_, K, V>>> + Send
+pub fn rev_stream<'a, K, V>(
+	self: &'a Arc<Self>,
+) -> impl Stream<Item = Result<KeyVal<'_, K, V>>> + Send
 where
 	K: Deserialize<'a> + Send,
 	V: Deserialize<'a> + Send,
@@ -22,9 +28,35 @@ where
 /// - Result is raw
 #[implement(super::Map)]
 #[tracing::instrument(skip(self), fields(%self), level = "trace")]
-pub fn rev_raw_stream(&self) -> impl Stream<Item = Result<KeyVal<'_>>> + Send {
-	let opts = super::iter_options_default();
-	stream::ItemsRev::new(&self.db, &self.cf, opts).init(None)
+pub fn rev_raw_stream(self: &Arc<Self>) -> impl Stream<Item = Result<KeyVal<'_>>> + Send {
+	use crate::pool::Seek;
+
+	let opts = super::read_options_default();
+	let state = stream::State::new(&self.db, &self.cf, opts);
+	if is_cached(self) {
+		let state = state.init_rev(None);
+		return task::consume_budget()
+			.map(move |()| stream::ItemsRev::<'_>::from(state))
+			.into_stream()
+			.flatten()
+			.boxed();
+	};
+
+	let seek = Seek {
+		map: self.clone(),
+		dir: Direction::Reverse,
+		state: crate::pool::into_send_seek(state),
+		key: None,
+		res: None,
+	};
+
+	self.db
+		.pool
+		.execute_iter(seek)
+		.ok_into::<stream::ItemsRev<'_>>()
+		.into_stream()
+		.try_flatten()
+		.boxed()
 }
 
 #[tracing::instrument(
@@ -33,13 +65,9 @@ pub fn rev_raw_stream(&self) -> impl Stream<Item = Result<KeyVal<'_>>> + Send {
     skip_all,
     fields(%map),
 )]
-pub(super) fn _is_cached<P>(map: &super::Map) -> bool
-where
-	P: AsRef<[u8]> + ?Sized,
-{
+pub(super) fn is_cached(map: &super::Map) -> bool {
 	let opts = super::cache_read_options_default();
-	let mut state = stream::State::new(&map.db, &map.cf, opts);
+	let state = stream::State::new(&map.db, &map.cf, opts).init_rev(None);
 
-	state.seek_rev();
 	!state.is_incomplete()
 }