Unbox and pin database streams.

Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk
2025-11-24 09:12:14 +00:00
parent 98affbdeaf
commit 71f3ccf140
26 changed files with 217 additions and 166 deletions

View File

@@ -1,11 +1,11 @@
stack-size-threshold = 393216
future-size-threshold = 24576
array-size-threshold = 4096 array-size-threshold = 4096
cognitive-complexity-threshold = 100 # TODO reduce me ALARA
excessive-nesting-threshold = 8
future-size-threshold = 8192
stack-size-threshold = 196608 # TODO reduce me ALARA
too-many-lines-threshold = 780 # TODO reduce me to <= 100
type-complexity-threshold = 250 # reduce me to ~200
large-error-threshold = 256 # TODO reduce me ALARA large-error-threshold = 256 # TODO reduce me ALARA
too-many-lines-threshold = 780 # TODO reduce me to <= 100
excessive-nesting-threshold = 8
type-complexity-threshold = 250 # reduce me to ~200
cognitive-complexity-threshold = 100 # TODO reduce me ALARA
#disallowed-macros = [ #disallowed-macros = [
# { path = "log::error", reason = "use tuwunel_core::error" }, # { path = "log::error", reason = "use tuwunel_core::error" },

View File

@@ -1,5 +1,9 @@
use axum::extract::State; use axum::extract::State;
use futures::{FutureExt, StreamExt, TryFutureExt, future::OptionFuture, pin_mut}; use futures::{
FutureExt, StreamExt, TryFutureExt,
future::{Either, OptionFuture},
pin_mut,
};
use ruma::{ use ruma::{
RoomId, UserId, RoomId, UserId,
api::{ api::{
@@ -105,17 +109,18 @@ pub(crate) async fn get_message_events_route(
} }
let it = match body.dir { let it = match body.dir {
| Direction::Forward => services | Direction::Forward => Either::Left(
.timeline services
.pdus(Some(sender_user), room_id, Some(from)) .timeline
.ignore_err() .pdus(Some(sender_user), room_id, Some(from))
.boxed(), .ignore_err(),
),
| Direction::Backward => services | Direction::Backward => Either::Right(
.timeline services
.pdus_rev(Some(sender_user), room_id, Some(from)) .timeline
.ignore_err() .pdus_rev(Some(sender_user), room_id, Some(from))
.boxed(), .ignore_err(),
),
}; };
let events: Vec<_> = it let events: Vec<_> = it

View File

@@ -1182,6 +1182,7 @@ async fn calculate_state_changes<'a>(
.state_accessor .state_accessor
.state_full_shortids(horizon_shortstatehash) .state_full_shortids(horizon_shortstatehash)
.expect_ok() .expect_ok()
.boxed()
.into_future() .into_future()
}) })
.into(); .into();

View File

@@ -170,6 +170,7 @@ async fn collect_room(
.ready_filter(|&user_id| user_id != sender_user) .ready_filter(|&user_id| user_id != sender_user)
.map(ToOwned::to_owned) .map(ToOwned::to_owned)
.map(|user_id| (MembershipState::Join, user_id)) .map(|user_id| (MembershipState::Join, user_id))
.boxed()
.into_future() .into_future()
}) })
.into(); .into();

View File

@@ -34,7 +34,7 @@ pub(crate) async fn search_users_route(
.min(LIMIT_MAX); .min(LIMIT_MAX);
let search_term = body.search_term.to_lowercase(); let search_term = body.search_term.to_lowercase();
let mut users = services let users = services
.users .users
.stream() .stream()
.ready_filter(|&user_id| user_id != sender_user) .ready_filter(|&user_id| user_id != sender_user)
@@ -83,6 +83,7 @@ pub(crate) async fn search_users_route(
}) })
}); });
pin_mut!(users);
let results = users.by_ref().take(limit).collect().await; let results = users.by_ref().take(limit).collect().await;
let limited = users.next().await.is_some(); let limited = users.next().await.is_some();

View File

@@ -1,6 +1,9 @@
use std::{convert::AsRef, fmt::Debug, sync::Arc}; use std::{convert::AsRef, fmt::Debug, sync::Arc};
use futures::{Future, FutureExt, TryFutureExt, future::ready}; use futures::{
Future, FutureExt, TryFutureExt,
future::{Either, ready},
};
use rocksdb::{DBPinnableSlice, ReadOptions}; use rocksdb::{DBPinnableSlice, ReadOptions};
use tokio::task; use tokio::task;
use tuwunel_core::{Err, Result, err, implement, utils::result::MapExpect}; use tuwunel_core::{Err, Result, err, implement, utils::result::MapExpect};
@@ -25,9 +28,9 @@ where
let cached = self.get_cached(key); let cached = self.get_cached(key);
if matches!(cached, Err(_) | Ok(Some(_))) { if matches!(cached, Err(_) | Ok(Some(_))) {
return task::consume_budget() return Either::Left(
.map(move |()| cached.map_expect("data found in cache")) task::consume_budget().map(move |()| cached.map_expect("data found in cache")),
.boxed(); );
} }
debug_assert!(matches!(cached, Ok(None)), "expected status Incomplete"); debug_assert!(matches!(cached, Ok(None)), "expected status Incomplete");
@@ -37,11 +40,12 @@ where
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_get(cmd) .pool
.and_then(|mut res| ready(res.remove(0))) .execute_get(cmd)
.boxed() .and_then(|mut res| ready(res.remove(0))),
)
} }
/// Fetch a value from the cache without I/O. /// Fetch a value from the cache without I/O.

View File

@@ -50,12 +50,12 @@ where
.widen_then(automatic_width(), |chunk| { .widen_then(automatic_width(), |chunk| {
self.engine.pool.execute_get(Get { self.engine.pool.execute_get(Get {
map: self.clone(), map: self.clone(),
res: None,
key: chunk key: chunk
.iter() .iter()
.map(AsRef::as_ref) .map(AsRef::as_ref)
.map(Into::into) .map(Into::into)
.collect(), .collect(),
res: None,
}) })
}) })
.map_ok(|results| results.into_iter().stream()) .map_ok(|results| results.into_iter().stream())

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::Deserialize; use serde::Deserialize;
use tokio::task; use tokio::task;
@@ -27,11 +27,12 @@ pub fn raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + Send
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self) { if is_cached(self) {
let state = state.init_fwd(None); let state = state.init_fwd(None);
return task::consume_budget() return Either::Left(
.map(move |()| stream::Keys::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::Keys::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -42,11 +43,12 @@ pub fn raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + Send
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::Keys<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::Keys<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }

View File

@@ -1,8 +1,9 @@
use std::{convert::AsRef, fmt::Debug, sync::Arc}; use std::{convert::AsRef, fmt::Debug, sync::Arc};
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::task;
use tuwunel_core::{Result, implement}; use tuwunel_core::{Result, implement};
use super::stream_from::is_cached; use super::stream_from::is_cached;
@@ -64,7 +65,13 @@ where
let opts = super::iter_options_default(&self.engine); let opts = super::iter_options_default(&self.engine);
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self, from) { if is_cached(self, from) {
return stream::Keys::<'_>::from(state.init_fwd(from.as_ref().into())).boxed(); let state = state.init_fwd(from.as_ref().into());
return Either::Left(
task::consume_budget()
.map(move |()| stream::Keys::<'_>::from(state))
.into_stream()
.flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -75,11 +82,12 @@ where
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::Keys<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::Keys<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::Deserialize; use serde::Deserialize;
use tokio::task; use tokio::task;
@@ -27,11 +27,12 @@ pub fn rev_raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + S
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self) { if is_cached(self) {
let state = state.init_rev(None); let state = state.init_rev(None);
return task::consume_budget() return Either::Left(
.map(move |()| stream::KeysRev::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::KeysRev::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -42,11 +43,12 @@ pub fn rev_raw_keys(self: &Arc<Self>) -> impl Stream<Item = Result<Key<'_>>> + S
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::KeysRev<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::KeysRev<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }

View File

@@ -1,8 +1,9 @@
use std::{convert::AsRef, fmt::Debug, sync::Arc}; use std::{convert::AsRef, fmt::Debug, sync::Arc};
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::task;
use tuwunel_core::{Result, implement}; use tuwunel_core::{Result, implement};
use super::rev_stream_from::is_cached; use super::rev_stream_from::is_cached;
@@ -64,7 +65,13 @@ where
let opts = super::iter_options_default(&self.engine); let opts = super::iter_options_default(&self.engine);
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self, from) { if is_cached(self, from) {
return stream::KeysRev::<'_>::from(state.init_rev(from.as_ref().into())).boxed(); let state = state.init_rev(from.as_ref().into());
return Either::Left(
task::consume_budget()
.map(move |()| stream::KeysRev::<'_>::from(state))
.into_stream()
.flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -75,11 +82,12 @@ where
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::KeysRev<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::KeysRev<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::Deserialize; use serde::Deserialize;
use tokio::task; use tokio::task;
@@ -35,11 +35,12 @@ pub fn rev_raw_stream(self: &Arc<Self>) -> impl Stream<Item = Result<KeyVal<'_>>
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self) { if is_cached(self) {
let state = state.init_rev(None); let state = state.init_rev(None);
return task::consume_budget() return Either::Left(
.map(move |()| stream::ItemsRev::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::ItemsRev::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -50,13 +51,14 @@ pub fn rev_raw_stream(self: &Arc<Self>) -> impl Stream<Item = Result<KeyVal<'_>>
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::ItemsRev<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::ItemsRev<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }
#[tracing::instrument( #[tracing::instrument(

View File

@@ -1,6 +1,6 @@
use std::{convert::AsRef, fmt::Debug, sync::Arc}; use std::{convert::AsRef, fmt::Debug, sync::Arc};
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::task; use tokio::task;
@@ -84,11 +84,12 @@ where
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self, from) { if is_cached(self, from) {
let state = state.init_rev(from.as_ref().into()); let state = state.init_rev(from.as_ref().into());
return task::consume_budget() return Either::Left(
.map(move |()| stream::ItemsRev::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::ItemsRev::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -99,13 +100,14 @@ where
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::ItemsRev<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::ItemsRev<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }
#[tracing::instrument( #[tracing::instrument(

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::Deserialize; use serde::Deserialize;
use tokio::task; use tokio::task;
@@ -35,11 +35,12 @@ pub fn raw_stream(self: &Arc<Self>) -> impl Stream<Item = Result<KeyVal<'_>>> +
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self) { if is_cached(self) {
let state = state.init_fwd(None); let state = state.init_fwd(None);
return task::consume_budget() return Either::Left(
.map(move |()| stream::Items::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::Items::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -50,13 +51,14 @@ pub fn raw_stream(self: &Arc<Self>) -> impl Stream<Item = Result<KeyVal<'_>>> +
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::Items<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::Items<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }
#[tracing::instrument( #[tracing::instrument(

View File

@@ -1,6 +1,6 @@
use std::{convert::AsRef, fmt::Debug, sync::Arc}; use std::{convert::AsRef, fmt::Debug, sync::Arc};
use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt}; use futures::{FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt, future::Either};
use rocksdb::Direction; use rocksdb::Direction;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use tokio::task; use tokio::task;
@@ -83,11 +83,12 @@ where
let state = stream::State::new(self, opts); let state = stream::State::new(self, opts);
if is_cached(self, from) { if is_cached(self, from) {
let state = state.init_fwd(from.as_ref().into()); let state = state.init_fwd(from.as_ref().into());
return task::consume_budget() return Either::Left(
.map(move |()| stream::Items::<'_>::from(state)) task::consume_budget()
.into_stream() .map(move |()| stream::Items::<'_>::from(state))
.flatten() .into_stream()
.boxed(); .flatten(),
);
} }
let seek = Seek { let seek = Seek {
@@ -98,13 +99,14 @@ where
res: None, res: None,
}; };
self.engine Either::Right(
.pool self.engine
.execute_iter(seek) .pool
.ok_into::<stream::Items<'_>>() .execute_iter(seek)
.into_stream() .ok_into::<stream::Items<'_>>()
.try_flatten() .into_stream()
.boxed() .try_flatten(),
)
} }
#[tracing::instrument( #[tracing::instrument(

View File

@@ -22,7 +22,7 @@ pub(crate) struct State<'a> {
init: bool, init: bool,
} }
pub(crate) trait Cursor<'a, T> { pub(crate) trait Cursor<'a, T>: Send {
fn state(&self) -> &State<'a>; fn state(&self) -> &State<'a>;
fn fetch(&self) -> Option<T>; fn fetch(&self) -> Option<T>;
@@ -50,12 +50,12 @@ impl<'a> State<'a> {
#[inline] #[inline]
pub(super) fn new(map: &'a Arc<Map>, opts: ReadOptions) -> Self { pub(super) fn new(map: &'a Arc<Map>, opts: ReadOptions) -> Self {
Self { Self {
init: true,
seek: false,
inner: map inner: map
.engine() .engine()
.db .db
.raw_iterator_cf_opt(&map.cf(), opts), .raw_iterator_cf_opt(&map.cf(), opts),
init: true,
seek: false,
} }
} }

View File

@@ -3,7 +3,7 @@ mod room_tags;
use std::sync::Arc; use std::sync::Arc;
use futures::{Stream, StreamExt, TryFutureExt}; use futures::{Stream, StreamExt, TryFutureExt, pin_mut};
use ruma::{ use ruma::{
RoomId, UserId, RoomId, UserId,
events::{ events::{
@@ -175,15 +175,18 @@ pub async fn last_count<'a>(
let upper = upper.unwrap_or(u64::MAX); let upper = upper.unwrap_or(u64::MAX);
let key = (room_id, user_id, upper, Interfix); let key = (room_id, user_id, upper, Interfix);
self.db let keys = self
.db
.roomuserdataid_accountdata .roomuserdataid_accountdata
.rev_keys_from(&key) .rev_keys_from(&key)
.ignore_err() .ignore_err()
.ready_take_while(move |(room_id_, user_id_, ..): &Key<'_>| { .ready_take_while(move |(room_id_, user_id_, ..): &Key<'_>| {
room_id == *room_id_ && user_id == *user_id_ room_id == *room_id_ && user_id == *user_id_
}) })
.map(at!(2)) .map(at!(2));
.next()
pin_mut!(keys);
keys.next()
.await .await
.ok_or_else(|| err!(Request(NotFound("No account data found.")))) .ok_or_else(|| err!(Request(NotFound("No account data found."))))
} }

View File

@@ -1,6 +1,6 @@
use std::{sync::Arc, time::Duration}; use std::{sync::Arc, time::Duration};
use futures::StreamExt; use futures::{StreamExt, pin_mut};
use ruma::{Mxc, OwnedMxcUri, UserId, http_headers::ContentDisposition}; use ruma::{Mxc, OwnedMxcUri, UserId, http_headers::ContentDisposition};
use tuwunel_core::{ use tuwunel_core::{
Err, Result, debug, debug_info, err, Err, Result, debug, debug_info, err,
@@ -109,11 +109,14 @@ impl Data {
let dim: &[u32] = &[dim.width, dim.height]; let dim: &[u32] = &[dim.width, dim.height];
let prefix = (mxc, dim, Interfix); let prefix = (mxc, dim, Interfix);
let key = self let keys = self
.mediaid_file .mediaid_file
.keys_prefix_raw(&prefix) .keys_prefix_raw(&prefix)
.ignore_err() .ignore_err()
.map(ToOwned::to_owned) .map(ToOwned::to_owned);
pin_mut!(keys);
let key = keys
.next() .next()
.await .await
.ok_or_else(|| err!(Request(NotFound("Media not found"))))?; .ok_or_else(|| err!(Request(NotFound("Media not found"))))?;

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{FutureExt, StreamExt, pin_mut}; use futures::{FutureExt, StreamExt};
use ruma::RoomId; use ruma::RoomId;
use tuwunel_core::{ use tuwunel_core::{
Result, debug, Result, debug,
@@ -35,6 +35,7 @@ impl Service {
.services .services
.state_cache .state_cache
.local_users_in_room(room_id) .local_users_in_room(room_id)
.boxed()
.into_future() .into_future()
.map(|(next, ..)| next.as_ref().is_some()); .map(|(next, ..)| next.as_ref().is_some());
@@ -42,10 +43,10 @@ impl Service {
.services .services
.state_cache .state_cache
.local_users_invited_to_room(room_id) .local_users_invited_to_room(room_id)
.boxed()
.into_future() .into_future()
.map(|(next, ..)| next.as_ref().is_some()); .map(|(next, ..)| next.as_ref().is_some());
pin_mut!(has_local_users, has_local_invites);
if has_local_users.or(has_local_invites).await { if has_local_users.or(has_local_invites).await {
trace!(?room_id, "Not deleting with local joined or invited"); trace!(?room_id, "Not deleting with local joined or invited");
return; return;

View File

@@ -46,13 +46,14 @@ pub async fn exists(&self, room_id: &RoomId) -> bool {
}; };
// Look for PDUs in that room. // Look for PDUs in that room.
self.db let keys = self
.db
.pduid_pdu .pduid_pdu
.keys_prefix_raw(&prefix) .keys_prefix_raw(&prefix)
.ignore_err() .ignore_err();
.next()
.await pin_mut!(keys);
.is_some() keys.next().await.is_some()
} }
#[implement(Service)] #[implement(Service)]

View File

@@ -1,6 +1,6 @@
use std::sync::Arc; use std::sync::Arc;
use futures::{Stream, StreamExt, TryFutureExt}; use futures::{Stream, StreamExt, TryFutureExt, future::Either};
use ruma::{EventId, RoomId, UserId, api::Direction}; use ruma::{EventId, RoomId, UserId, api::Direction};
use tuwunel_core::{ use tuwunel_core::{
PduId, Result, PduId, Result,
@@ -86,16 +86,8 @@ pub fn get_relations<'a>(
}; };
match dir { match dir {
| Direction::Backward => self | Direction::Backward => Either::Left(self.db.tofrom_relation.rev_raw_keys_from(start)),
.db | Direction::Forward => Either::Right(self.db.tofrom_relation.raw_keys_from(start)),
.tofrom_relation
.rev_raw_keys_from(start)
.boxed(),
| Direction::Forward => self
.db
.tofrom_relation
.raw_keys_from(start)
.boxed(),
} }
.ignore_err() .ignore_err()
.ready_take_while(move |key| key.starts_with(&target)) .ready_take_while(move |key| key.starts_with(&target))

View File

@@ -1,6 +1,6 @@
use std::{borrow::Borrow, fmt::Debug, mem::size_of_val, sync::Arc}; use std::{borrow::Borrow, fmt::Debug, mem::size_of_val, sync::Arc};
use futures::{FutureExt, Stream, StreamExt}; use futures::{FutureExt, Stream, StreamExt, pin_mut};
use ruma::{EventId, OwnedRoomId, RoomId, events::StateEventType}; use ruma::{EventId, OwnedRoomId, RoomId, events::StateEventType};
use serde::Deserialize; use serde::Deserialize;
pub use tuwunel_core::matrix::{ShortEventId, ShortId, ShortRoomId, ShortStateKey}; pub use tuwunel_core::matrix::{ShortEventId, ShortId, ShortRoomId, ShortStateKey};
@@ -245,10 +245,14 @@ pub async fn get_shortroomid(&self, room_id: &RoomId) -> Result<ShortRoomId> {
#[implement(Service)] #[implement(Service)]
pub async fn get_roomid_from_short(&self, shortroomid_: ShortRoomId) -> Result<OwnedRoomId> { pub async fn get_roomid_from_short(&self, shortroomid_: ShortRoomId) -> Result<OwnedRoomId> {
self.db let stream = self
.db
.roomid_shortroomid .roomid_shortroomid
.stream() .stream()
.ready_filter_map(Result::ok) .ready_filter_map(Result::ok);
pin_mut!(stream);
stream
.ready_find(|&(_, shortroomid)| shortroomid == shortroomid_) .ready_find(|&(_, shortroomid)| shortroomid == shortroomid_)
.map(|found| found.map(|(room_id, _): (&RoomId, ShortRoomId)| room_id.to_owned())) .map(|found| found.map(|(room_id, _): (&RoomId, ShortRoomId)| room_id.to_owned()))
.await .await

View File

@@ -108,6 +108,7 @@ pub fn room_state_keys_with_ids<'a>(
.map_ok(|shortstatehash| { .map_ok(|shortstatehash| {
self.state_keys_with_ids(shortstatehash, event_type) self.state_keys_with_ids(shortstatehash, event_type)
.map(Ok) .map(Ok)
.boxed()
}) })
.map_err(move |e| err!(Database("Missing state for {room_id:?}: {e:?}"))) .map_err(move |e| err!(Database("Missing state for {room_id:?}: {e:?}")))
.try_flatten_stream() .try_flatten_stream()
@@ -127,6 +128,7 @@ pub fn room_state_keys<'a>(
.map_ok(|shortstatehash| { .map_ok(|shortstatehash| {
self.state_keys(shortstatehash, event_type) self.state_keys(shortstatehash, event_type)
.map(Ok) .map(Ok)
.boxed()
}) })
.map_err(move |e| err!(Database("Missing state for {room_id:?}: {e:?}"))) .map_err(move |e| err!(Database("Missing state for {room_id:?}: {e:?}")))
.try_flatten_stream() .try_flatten_stream()

View File

@@ -249,7 +249,6 @@ pub fn state_keys_with_shortids<'a>(
.ignore_err() .ignore_err()
.unzip() .unzip()
.map(|(ssks, sids): (Vec<u64>, Vec<u64>)| (ssks, sids)) .map(|(ssks, sids): (Vec<u64>, Vec<u64>)| (ssks, sids))
.boxed()
.shared(); .shared();
let shortstatekeys = short_ids let shortstatekeys = short_ids
@@ -410,7 +409,6 @@ pub fn state_full_shortids(
.map_ok(Vec::into_iter) .map_ok(Vec::into_iter)
.map_ok(IterStream::try_stream) .map_ok(IterStream::try_stream)
.try_flatten_stream() .try_flatten_stream()
.boxed()
} }
#[implement(super::Service)] #[implement(super::Service)]

View File

@@ -212,8 +212,9 @@ pub fn get_shared_rooms<'a>(
) -> impl Stream<Item = &RoomId> + Send + 'a { ) -> impl Stream<Item = &RoomId> + Send + 'a {
use tuwunel_core::utils::set; use tuwunel_core::utils::set;
let a = self.rooms_joined(user_a); let a = self.rooms_joined(user_a).boxed();
let b = self.rooms_joined(user_b); let b = self.rooms_joined(user_b).boxed();
set::intersection_sorted_stream2(a, b) set::intersection_sorted_stream2(a, b)
} }
@@ -415,6 +416,7 @@ pub fn user_memberships<'a>(
.then(|| { .then(|| {
self.rooms_joined(user_id) self.rooms_joined(user_id)
.map(|room_id| (Join, room_id)) .map(|room_id| (Join, room_id))
.boxed()
.into_future() .into_future()
}) })
.into(); .into();
@@ -424,6 +426,7 @@ pub fn user_memberships<'a>(
.then(|| { .then(|| {
self.rooms_invited(user_id) self.rooms_invited(user_id)
.map(|room_id| (Invite, room_id)) .map(|room_id| (Invite, room_id))
.boxed()
.into_future() .into_future()
}) })
.into(); .into();
@@ -433,6 +436,7 @@ pub fn user_memberships<'a>(
.then(|| { .then(|| {
self.rooms_knocked(user_id) self.rooms_knocked(user_id)
.map(|room_id| (Knock, room_id)) .map(|room_id| (Knock, room_id))
.boxed()
.into_future() .into_future()
}) })
.into(); .into();
@@ -442,6 +446,7 @@ pub fn user_memberships<'a>(
.then(|| { .then(|| {
self.rooms_left(user_id) self.rooms_left(user_id)
.map(|room_id| (Leave, room_id)) .map(|room_id| (Leave, room_id))
.boxed()
.into_future() .into_future()
}) })
.into(); .into();

View File

@@ -1,6 +1,6 @@
use std::{collections::BTreeMap, mem}; use std::{collections::BTreeMap, mem};
use futures::{Stream, StreamExt, TryFutureExt}; use futures::{Stream, StreamExt, TryFutureExt, pin_mut};
use ruma::{ use ruma::{
DeviceId, KeyId, OneTimeKeyAlgorithm, OneTimeKeyId, OneTimeKeyName, OwnedKeyId, RoomId, UInt, DeviceId, KeyId, OneTimeKeyAlgorithm, OneTimeKeyId, OneTimeKeyName, OwnedKeyId, RoomId, UInt,
UserId, UserId,
@@ -113,7 +113,7 @@ pub async fn take_one_time_key(
prefix.extend_from_slice(key_algorithm.as_ref().as_bytes()); prefix.extend_from_slice(key_algorithm.as_ref().as_bytes());
prefix.push(b':'); prefix.push(b':');
let one_time_key = self let one_time_keys = self
.db .db
.onetimekeyid_onetimekeys .onetimekeyid_onetimekeys
.raw_stream_prefix(&prefix) .raw_stream_prefix(&prefix)
@@ -136,11 +136,13 @@ pub async fn take_one_time_key(
.unwrap(); .unwrap();
(key, val) (key, val)
}) });
.next()
.await;
one_time_key.ok_or_else(|| err!(Request(NotFound("No one-time-key found")))) pin_mut!(one_time_keys);
one_time_keys
.next()
.await
.ok_or_else(|| err!(Request(NotFound("No one-time-key found"))))
} }
#[implement(super::Service)] #[implement(super::Service)]