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/count.rs

@@ -7,7 +7,9 @@ use serde::Serialize;
 /// Count the total number of entries in the map.
 #[implement(super::Map)]
 #[inline]
-pub fn count(&self) -> impl Future<Output = usize> + Send + '_ { self.raw_keys().count() }
+pub fn count(self: &Arc<Self>) -> impl Future<Output = usize> + Send + '_ {
+	self.raw_keys().count()
+}
 
 /// Count the number of entries in the map starting from a lower-bound.
 ///

src/database/map/keys.rs

@@ -1,11 +1,16 @@
-use conduwuit::{implement, Result};
-use futures::{Stream, StreamExt};
-use serde::Deserialize;
+use std::sync::Arc;
 
-use crate::{keyval, keyval::Key, stream, stream::Cursor};
+use conduwuit::{implement, Result};
+use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use rocksdb::Direction;
+use serde::Deserialize;
+use tokio::task;
+
+use super::stream::is_cached;
+use crate::{keyval, keyval::Key, stream};
 
 #[implement(super::Map)]
-pub fn keys<'a, K>(&'a self) -> impl Stream<Item = Result<Key<'_, K>>> + Send
+pub fn keys<'a, K>(self: &'a Arc<Self>) -> impl Stream<Item = Result<Key<'_, K>>> + Send
 where
 	K: Deserialize<'a> + Send,
 {
@@ -14,7 +19,33 @@ where
 
 #[implement(super::Map)]
 #[tracing::instrument(skip(self), fields(%self), level = "trace")]
-pub fn raw_keys(&self) -> impl Stream<Item = Result<Key<'_>>> + Send {
+pub fn raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + Send {
+	use crate::pool::Seek;
+
 	let opts = super::iter_options_default();
-	stream::Keys::new(&self.db, &self.cf, opts).init(None)
+	let state = stream::State::new(&self.db, &self.cf, opts);
+	if is_cached(self) {
+		let state = state.init_fwd(None);
+		return task::consume_budget()
+			.map(move |()| stream::Keys::<'_>::from(state))
+			.into_stream()
+			.flatten()
+			.boxed();
+	}
+
+	let seek = Seek {
+		map: self.clone(),
+		dir: Direction::Forward,
+		state: crate::pool::into_send_seek(state),
+		key: None,
+		res: None,
+	};
+
+	self.db
+		.pool
+		.execute_iter(seek)
+		.ok_into::<stream::Keys<'_>>()
+		.into_stream()
+		.try_flatten()
+		.boxed()
 }

src/database/map/keys_prefix.rs

@@ -1,11 +1,7 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	future,
-	stream::{Stream, StreamExt},
-	TryStreamExt,
-};
+use futures::{future, Stream, StreamExt, TryStreamExt};
 use serde::{Deserialize, Serialize};
 
 use crate::keyval::{result_deserialize_key, serialize_key, Key};

src/database/map/rev_keys.rs

@@ -1,11 +1,16 @@
-use conduwuit::{implement, Result};
-use futures::{Stream, StreamExt};
-use serde::Deserialize;
+use std::sync::Arc;
 
-use crate::{keyval, keyval::Key, stream, stream::Cursor};
+use conduwuit::{implement, Result};
+use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use rocksdb::Direction;
+use serde::Deserialize;
+use tokio::task;
+
+use super::rev_stream::is_cached;
+use crate::{keyval, keyval::Key, stream};
 
 #[implement(super::Map)]
-pub fn rev_keys<'a, K>(&'a self) -> impl Stream<Item = Result<Key<'_, K>>> + Send
+pub fn rev_keys<'a, K>(self: &'a Arc<Self>) -> impl Stream<Item = Result<Key<'_, K>>> + Send
 where
 	K: Deserialize<'a> + Send,
 {
@@ -14,7 +19,33 @@ where
 
 #[implement(super::Map)]
 #[tracing::instrument(skip(self), fields(%self), level = "trace")]
-pub fn rev_raw_keys(&self) -> impl Stream<Item = Result<Key<'_>>> + Send {
+pub fn rev_raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + Send {
+	use crate::pool::Seek;
+
 	let opts = super::iter_options_default();
-	stream::KeysRev::new(&self.db, &self.cf, opts).init(None)
+	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::KeysRev::<'_>::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::KeysRev<'_>>()
+		.into_stream()
+		.try_flatten()
+		.boxed()
 }

src/database/map/rev_keys_prefix.rs

@@ -1,11 +1,7 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	future,
-	stream::{Stream, StreamExt},
-	TryStreamExt,
-};
+use futures::{future, Stream, StreamExt, TryStreamExt};
 use serde::{Deserialize, Serialize};
 
 use crate::keyval::{result_deserialize_key, serialize_key, Key};

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()
 }

src/database/map/rev_stream_from.rs

@@ -1,12 +1,10 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	stream::{Stream, StreamExt},
-	FutureExt, TryFutureExt, TryStreamExt,
-};
+use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use rocksdb::Direction;
 use serde::{Deserialize, Serialize};
+use tokio::task;
 
 use crate::{
 	keyval::{result_deserialize, serialize_key, KeyVal},
@@ -85,7 +83,12 @@ where
 	let opts = super::iter_options_default();
 	let state = stream::State::new(&self.db, &self.cf, opts);
 	if is_cached(self, from) {
-		return stream::ItemsRev::<'_>::from(state.init_rev(from.as_ref().into())).boxed();
+		let state = state.init_rev(from.as_ref().into());
+		return task::consume_budget()
+			.map(move |()| stream::ItemsRev::<'_>::from(state))
+			.into_stream()
+			.flatten()
+			.boxed();
 	};
 
 	let seek = Seek {

src/database/map/rev_stream_prefix.rs

@@ -1,11 +1,7 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	future,
-	stream::{Stream, StreamExt},
-	TryStreamExt,
-};
+use futures::{future, Stream, StreamExt, TryStreamExt};
 use serde::{Deserialize, Serialize};
 
 use crate::keyval::{result_deserialize, serialize_key, KeyVal};

src/database/map/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 beginning.
 ///
 /// - Result is deserialized
 #[implement(super::Map)]
-pub fn stream<'a, K, V>(&'a self) -> impl Stream<Item = Result<KeyVal<'_, K, V>>> + Send
+pub fn 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,
@@ -21,9 +27,35 @@ where
 /// - Result is raw
 #[implement(super::Map)]
 #[tracing::instrument(skip(self), fields(%self), level = "trace")]
-pub fn raw_stream(&self) -> impl Stream<Item = Result<KeyVal<'_>>> + Send {
-	let opts = super::iter_options_default();
-	stream::Items::new(&self.db, &self.cf, opts).init(None)
+pub fn 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_fwd(None);
+		return task::consume_budget()
+			.map(move |()| stream::Items::<'_>::from(state))
+			.into_stream()
+			.flatten()
+			.boxed();
+	};
+
+	let seek = Seek {
+		map: self.clone(),
+		dir: Direction::Forward,
+		state: crate::pool::into_send_seek(state),
+		key: None,
+		res: None,
+	};
+
+	self.db
+		.pool
+		.execute_iter(seek)
+		.ok_into::<stream::Items<'_>>()
+		.into_stream()
+		.try_flatten()
+		.boxed()
 }
 
 #[tracing::instrument(
@@ -32,13 +64,9 @@ pub fn 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_fwd(None);
 
-	state.seek_fwd();
 	!state.is_incomplete()
 }

src/database/map/stream_from.rs

@@ -1,12 +1,10 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	stream::{Stream, StreamExt},
-	FutureExt, TryFutureExt, TryStreamExt,
-};
+use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
 use rocksdb::Direction;
 use serde::{Deserialize, Serialize};
+use tokio::task;
 
 use crate::{
 	keyval::{result_deserialize, serialize_key, KeyVal},
@@ -82,7 +80,12 @@ where
 	let opts = super::read_options_default();
 	let state = stream::State::new(&self.db, &self.cf, opts);
 	if is_cached(self, from) {
-		return stream::Items::<'_>::from(state.init_fwd(from.as_ref().into())).boxed();
+		let state = state.init_fwd(from.as_ref().into());
+		return task::consume_budget()
+			.map(move |()| stream::Items::<'_>::from(state))
+			.into_stream()
+			.flatten()
+			.boxed();
 	};
 
 	let seek = Seek {

src/database/map/stream_prefix.rs

@@ -1,11 +1,7 @@
 use std::{convert::AsRef, fmt::Debug, sync::Arc};
 
 use conduwuit::{implement, Result};
-use futures::{
-	future,
-	stream::{Stream, StreamExt},
-	TryStreamExt,
-};
+use futures::{future, Stream, StreamExt, TryStreamExt};
 use serde::{Deserialize, Serialize};
 
 use crate::keyval::{result_deserialize, serialize_key, KeyVal};