tuwunel_database/map/
get_batch.rs

1use std::sync::Arc;
2
3use futures::{Stream, StreamExt, TryStreamExt};
4use rocksdb::{DBPinnableSlice, ReadOptions};
5use tuwunel_core::{
6	Result, implement,
7	utils::{
8		IterStream,
9		stream::{WidebandExt, automatic_amplification, automatic_width},
10	},
11};
12
13use super::get::{cached_handle_from, handle_from};
14use crate::Handle;
15
16/// Extends a stream of raw keys with batched map lookup.
17///
18/// Input keys are grouped for the engine's blocking pool. The output stream
19/// yields pinned value handles or lookup errors.
20pub trait Get<'a, K, S>
21where
22	Self: Sized,
23	S: Stream<Item = K> + Send + 'a,
24	K: AsRef<[u8]> + Send + Sync + 'a,
25{
26	/// Fetches this stream's raw keys from a map.
27	///
28	/// Successful batches yield one lookup result for each input key. A
29	/// batch-level pool or channel failure appears as one stream error for that
30	/// batch. Work is split into batches sized from the server's automatic
31	/// amplification setting.
32	fn get(self, map: &'a Arc<super::Map>) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a;
33}
34
35impl<'a, K, S> Get<'a, K, S> for S
36where
37	Self: Sized,
38	S: Stream<Item = K> + Send + 'a,
39	K: AsRef<[u8]> + Send + Sync + 'a,
40{
41	#[inline]
42	fn get(self, map: &'a Arc<super::Map>) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a {
43		map.get_batch(self)
44	}
45}
46
47/// Fetches a stream of raw keys in asynchronous batches.
48///
49/// Each batch runs on the engine's blocking pool and is flattened back into
50/// individual lookup results.
51#[implement(super::Map)]
52#[tracing::instrument(skip(self, keys), level = "trace")]
53pub(crate) fn get_batch<'a, S, K>(
54	self: &'a Arc<Self>,
55	keys: S,
56) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a
57where
58	S: Stream<Item = K> + Send + 'a,
59	K: AsRef<[u8]> + Send + Sync + 'a,
60{
61	use crate::pool::Get;
62
63	keys.ready_chunks(automatic_amplification())
64		.widen_then(automatic_width(), |chunk| {
65			self.engine.pool.execute_get(Get {
66				map: self.clone(),
67				res: None,
68				key: chunk
69					.iter()
70					.map(AsRef::as_ref)
71					.map(Into::into)
72					.collect(),
73			})
74		})
75		.map_ok(|results| results.into_iter().stream())
76		.try_flatten()
77}
78
79/// Fetches an exact-size raw-key iterator from block cache.
80///
81/// Cache misses remain `Ok(None)`, while cached values and failures retain
82/// their normal result forms.
83#[implement(super::Map)]
84#[tracing::instrument(name = "batch_cached", level = "trace", skip_all)]
85pub(crate) fn _get_batch_cached<'a, I, K>(
86	&self,
87	keys: I,
88) -> impl Iterator<Item = Result<Option<Handle<'_>>>> + Send + use<'_, I, K>
89where
90	I: Iterator<Item = &'a K> + ExactSizeIterator + Send,
91	K: AsRef<[u8]> + Send + ?Sized + Sync + 'a,
92{
93	self.get_batch_blocking_opts(keys, &self.cache_read_options)
94		.map(cached_handle_from)
95}
96
97/// Fetches an exact-size raw-key iterator synchronously.
98///
99/// RocksDB performs a batched multi-get and the returned iterator classifies
100/// each point-read result.
101#[implement(super::Map)]
102#[tracing::instrument(name = "batch_blocking", level = "trace", skip_all)]
103pub(crate) fn get_batch_blocking<'a, I, K>(
104	&self,
105	keys: I,
106) -> impl Iterator<Item = Result<Handle<'_>>> + Send + use<'_, I, K>
107where
108	I: Iterator<Item = &'a K> + ExactSizeIterator + Send,
109	K: AsRef<[u8]> + Send + ?Sized + Sync + 'a,
110{
111	self.get_batch_blocking_opts(keys, &self.read_options)
112		.map(handle_from)
113}
114
115/// Performs a batched multi-get with explicit RocksDB read options.
116///
117/// Keys are treated as unsorted because callers do not promise
118/// column-comparator order.
119#[implement(super::Map)]
120fn get_batch_blocking_opts<'a, I, K>(
121	&self,
122	keys: I,
123	read_options: &ReadOptions,
124) -> impl Iterator<Item = Result<Option<DBPinnableSlice<'_>>, rocksdb::Error>> + Send + use<'_, I, K>
125where
126	I: Iterator<Item = &'a K> + ExactSizeIterator + Send,
127	K: AsRef<[u8]> + Send + ?Sized + Sync + 'a,
128{
129	// Optimization can be `true` if key vector is pre-sorted **by the column
130	// comparator**.
131	const SORTED: bool = false;
132
133	self.engine
134		.db
135		.batched_multi_get_cf_opt(&self.cf(), keys, SORTED, read_options)
136		.into_iter()
137}
tuwunel_database/map/get_batch.rs

tuwunel_database/map/
get_batch.rs
