studio/tuwunel
2
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Optimize get_shared_rooms()/intersection_sorted_stream2() for tighter loops.

Browse Source 
Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk
2026-02-27 01:54:24 +00:00
parent e9ca105d79
commit d959dd740f
2 changed files with 29 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
src/core/utils/set.rs
View File

@@ -1,12 +1,17 @@
use std::{
cmp::{Eq, Ord},
convert::identity,
pin::Pin,
sync::Arc,
};
use futures::{Stream, StreamExt};
use futures::{
Stream, StreamExt,
stream::{Peekable, unfold},
};
use tokio::sync::Mutex;
use crate::{is_equal_to, is_less_than};
use crate::{is_equal_to, is_less_than, utils::stream::ReadyExt};
/// Intersection of sets
///
@@ -57,30 +62,31 @@ where
/// Intersection of sets
///
/// Outputs the set of elements common to both streams. Streams must be sorted.
pub fn intersection_sorted_stream2<Item, A, B>(a: A, b: B) -> impl Stream<Item = Item> + Send
pub fn intersection_sorted_stream2<S, Item>(a: S, b: S) -> impl Stream<Item = Item> + Send
where
A: Stream<Item = Item> + Send,
B: Stream<Item = Item> + Send + Unpin,
S: Stream<Item = Item> + Send + Unpin,
Item: Eq + PartialOrd + Send + Sync,
{
use tokio::sync::Mutex;
struct State<S: Stream> {
a: S,
b: Peekable<S>,
}
let b = Arc::new(Mutex::new(b.peekable()));
a.map(move |ai| (ai, b.clone()))
.filter_map(async move |(ai, b)| {
let mut lock = b.lock().await;
while let Some(bi) = Pin::new(&mut *lock)
.next_if(|bi| *bi <= ai)
.await
.as_ref()
{
if ai == *bi {
return Some(ai);
}
unfold(State { a, b: b.peekable() }, async |mut state| {
let ai = state.a.next().await?;
while let Some(bi) = Pin::new(&mut state.b)
.next_if(|bi| *bi <= ai)
.await
.as_ref()
{
if ai == *bi {
return Some((Some(ai), state));
}
}
None
})
Some((None, state))
})
.ready_filter_map(identity)
}
/// Difference of sets
@@ -93,8 +99,6 @@ where
B: Stream<Item = Item> + Send + Unpin,
Item: Eq + PartialOrd + Send + Sync,
{
use tokio::sync::Mutex;
let b = Arc::new(Mutex::new(b.peekable()));
a.map(move |ai| (ai, b.clone()))
.filter_map(async move |(ai, b)| {

8
src/service/rooms/state_cache/mod.rs
View File

@@ -17,7 +17,7 @@ use tuwunel_core::{
result::LogErr,
trace,
utils::{
BoolExt,
self, BoolExt,
future::OptionStream,
stream::{BroadbandExt, ReadyExt, TryIgnore},
},
@@ -207,12 +207,10 @@ pub fn get_shared_rooms<'a>(
user_a: &'a UserId,
user_b: &'a UserId,
) -> impl Stream<Item = &RoomId> + Send + 'a {
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).boxed();
set::intersection_sorted_stream2(a, b)
utils::set::intersection_sorted_stream2(a, b)
}
/// Returns an iterator of all joined members of a room.