tuwunel_database/map/
qry_batch.rs

1use std::{fmt::Debug, sync::Arc};
2
3use futures::{Stream, StreamExt, TryStreamExt};
4use serde::Serialize;
5use tuwunel_core::{
6	Result, implement,
7	utils::{
8		IterStream,
9		stream::{WidebandExt, automatic_amplification, automatic_width},
10	},
11};
12
13use crate::{Handle, keyval::KeyBuf, ser};
14
15/// Extends a stream of structured keys with serialized batched lookup.
16///
17/// Input keys are encoded and grouped for the engine's blocking pool. The
18/// output stream yields pinned value handles or lookup errors.
19pub trait Qry<'a, K, S>
20where
21	S: Stream<Item = K> + Send + 'a,
22	K: Serialize + Debug,
23{
24	/// Fetches this stream's serialized keys from a map.
25	///
26	/// The returned stream yields lookup results from automatically sized
27	/// batches. Serialization is deferred until the stream is polled.
28	///
29	/// # Panics
30	///
31	/// Panics if an input key cannot be serialized.
32	fn qry(self, map: &'a Arc<super::Map>) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a;
33}
34
35impl<'a, K, S> Qry<'a, K, S> for S
36where
37	Self: 'a,
38	S: Stream<Item = K> + Send + 'a,
39	K: Serialize + Debug + 'a,
40{
41	#[inline]
42	fn qry(self, map: &'a Arc<super::Map>) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a {
43		map.qry_batch(self)
44	}
45}
46
47/// Fetches a stream of structured keys in serialized asynchronous batches.
48///
49/// Each batch is encoded, run on the engine's blocking pool, and flattened back
50/// into individual lookup results.
51///
52/// # Panics
53///
54/// Panics if an input key cannot be serialized.
55#[implement(super::Map)]
56#[tracing::instrument(skip(self, keys), level = "trace")]
57pub(crate) fn qry_batch<'a, S, K>(
58	self: &'a Arc<Self>,
59	keys: S,
60) -> impl Stream<Item = Result<Handle<'_>>> + Send + 'a
61where
62	S: Stream<Item = K> + Send + 'a,
63	K: Serialize + Debug + 'a,
64{
65	use crate::pool::Get;
66
67	keys.ready_chunks(automatic_amplification())
68		.widen_then(automatic_width(), |chunk| {
69			let keys = chunk
70				.iter()
71				.map(ser::serialize_to::<KeyBuf, _>)
72				.map(|result| result.expect("failed to serialize query key"))
73				.collect();
74
75			self.engine
76				.pool
77				.execute_get(Get { map: self.clone(), key: keys, res: None })
78		})
79		.map_ok(|results| results.into_iter().stream())
80		.try_flatten()
81}
tuwunel_database/map/qry_batch.rs

tuwunel_database/map/
qry_batch.rs
