tuwunel_database/map/
seek.rs

1use std::sync::Arc;
2
3use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
4use rocksdb::Direction;
5use tokio::task;
6use tuwunel_core::Result;
7
8use super::{Map, cache_iter_options_default, iter_options_default};
9use crate::{
10	pool::{Seek, into_send_seek},
11	stream,
12};
13
14/// Builds a forward or reverse map stream from an optional raw seek key.
15///
16/// A block-cache probe selects inline iteration when the initial seek is
17/// cached; otherwise the seek runs on the engine's blocking pool. The
18/// projection type determines whether each item contains a key alone or a
19/// key-value pair.
20pub(super) fn seek_stream<'a, C, T>(
21	map: &'a Arc<Map>,
22	dir: Direction,
23	from: Option<&[u8]>,
24) -> impl Stream<Item = Result<T>> + Send + use<'a, C, T>
25where
26	C: From<stream::State<'a>> + Stream<Item = Result<T>> + Send,
27{
28	let opts = iter_options_default(&map.engine);
29	let state = stream::State::new(map, opts);
30	if is_cached(map, dir, from) {
31		let state = init(state, dir, from);
32		return Either::Left(
33			task::consume_budget()
34				.map(move |()| C::from(state))
35				.into_stream()
36				.flatten(),
37		);
38	}
39
40	let seek = Seek {
41		map: map.clone(),
42		state: into_send_seek(state),
43		dir,
44		key: from.map(Into::into),
45		res: None,
46	};
47
48	Either::Right(
49		map.engine
50			.pool
51			.execute_iter(seek)
52			.ok_into::<C>()
53			.into_stream()
54			.try_flatten(),
55	)
56}
57
58/// Tests whether an initial seek can complete from block cache.
59///
60/// The probe uses the same direction and starting key as the real iterator
61/// without filling cache.
62#[tracing::instrument(
63    name = "cached",
64    level = "trace",
65    skip_all,
66    fields(%map),
67)]
68fn is_cached(map: &Arc<Map>, dir: Direction, from: Option<&[u8]>) -> bool {
69	let opts = cache_iter_options_default(&map.engine);
70	let state = init(stream::State::new(map, opts), dir, from);
71
72	!state.is_incomplete()
73}
74
75/// Initializes iterator state for the requested seek direction.
76///
77/// The optional raw key is interpreted as a lower bound when moving forward and
78/// an upper bound when moving backward.
79fn init<'a>(state: stream::State<'a>, dir: Direction, from: Option<&[u8]>) -> stream::State<'a> {
80	match dir {
81		| Direction::Forward => state.init_fwd(from),
82		| Direction::Reverse => state.init_rev(from),
83	}
84}
tuwunel_database/map/seek.rs

tuwunel_database/map/
seek.rs
