mirror of
https://forgejo.ellis.link/continuwuation/continuwuity/
synced 2026-08-28 09:25:12 +00:00
style: Document & refactor rooms/search service
This commit is contained in:
timedout
+147
-144
@@ -3,7 +3,7 @@
|
||||
use conduwuit::{
|
||||
PduCount, PduEvent, Result,
|
||||
arrayvec::ArrayVec,
|
||||
debug_warn, implement,
|
||||
debug_warn,
|
||||
matrix::event::{Event, Matches},
|
||||
utils::{
|
||||
ArrayVecExt, IterStream, ReadyExt, set,
|
||||
@@ -70,158 +70,161 @@ fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
|
||||
fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
|
||||
}
|
||||
|
||||
#[implement(Service)]
|
||||
pub fn index_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) {
|
||||
let batch = tokenize(message_body)
|
||||
.map(|word| {
|
||||
impl Service {
|
||||
/// Adds a single PDU to the search index.
|
||||
pub fn index_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) {
|
||||
let batch = tokenize(message_body)
|
||||
.map(|word| {
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(word.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here
|
||||
key
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
self.db
|
||||
.tokenids
|
||||
.insert_batch(batch.iter().map(|k| (k.as_slice(), &[])));
|
||||
}
|
||||
|
||||
/// Removes a single PDU from the search index.
|
||||
pub fn deindex_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) {
|
||||
let batch = tokenize(message_body).map(|word| {
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(word.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here
|
||||
key
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
self.db
|
||||
.tokenids
|
||||
.insert_batch(batch.iter().map(|k| (k.as_slice(), &[])));
|
||||
}
|
||||
|
||||
#[implement(Service)]
|
||||
pub fn deindex_pdu(&self, shortroomid: ShortRoomId, pdu_id: &RawPduId, message_body: &str) {
|
||||
let batch = tokenize(message_body).map(|word| {
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(word.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(pdu_id.as_ref()); // TODO: currently we save the room id a second time here
|
||||
key
|
||||
});
|
||||
|
||||
for token in batch {
|
||||
self.db.tokenids.remove(&token);
|
||||
}
|
||||
}
|
||||
|
||||
#[implement(Service)]
|
||||
pub async fn search_pdus<'a>(
|
||||
&'a self,
|
||||
query: &'a RoomQuery<'a>,
|
||||
sender_user: &'a UserId,
|
||||
) -> Result<(usize, impl Stream<Item = PduEvent> + Send + 'a)> {
|
||||
let pdu_ids: Vec<_> = self.search_pdu_ids(query).await?.collect().await;
|
||||
|
||||
let filter = &query.criteria.filter;
|
||||
let count = pdu_ids.len();
|
||||
let pdus = pdu_ids
|
||||
.into_iter()
|
||||
.stream()
|
||||
.wide_filter_map(move |result_pdu_id: RawPduId| async move {
|
||||
self.services
|
||||
.timeline
|
||||
.get_pdu_from_id(&result_pdu_id)
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.ready_filter(|pdu| !pdu.is_redacted())
|
||||
.ready_filter(move |pdu| filter.matches(pdu))
|
||||
.wide_filter_map(move |pdu| async move {
|
||||
self.services
|
||||
.state_accessor
|
||||
.user_can_see_event(query.user_id?, pdu.room_id().unwrap(), pdu.event_id())
|
||||
.await
|
||||
.then_some(pdu)
|
||||
})
|
||||
.skip(query.skip)
|
||||
.take(query.limit)
|
||||
.map(move |mut pdu| {
|
||||
pdu.set_unsigned(query.user_id);
|
||||
|
||||
pdu
|
||||
})
|
||||
.then(async move |mut pdu| {
|
||||
if let Err(e) = self
|
||||
.services
|
||||
.pdu_metadata
|
||||
.add_bundled_aggregations_to_pdu(sender_user, &mut pdu)
|
||||
.await
|
||||
{
|
||||
debug_warn!("Failed to add bundled aggregations: {e}");
|
||||
}
|
||||
pdu
|
||||
});
|
||||
|
||||
Ok((count, pdus))
|
||||
}
|
||||
|
||||
// result is modeled as a stream such that callers don't have to be refactored
|
||||
// though an additional async/wrap still exists for now
|
||||
#[implement(Service)]
|
||||
pub async fn search_pdu_ids<'a>(
|
||||
&'a self,
|
||||
query: &'a RoomQuery<'_>,
|
||||
) -> Result<impl Stream<Item = RawPduId> + Send + 'a + use<'a>> {
|
||||
let shortroomid = self.services.short.get_shortroomid(query.room_id).await?;
|
||||
|
||||
let pdu_ids = self.search_pdu_ids_query_room(query, shortroomid).await;
|
||||
|
||||
let iters = pdu_ids.into_iter().map(IntoIterator::into_iter);
|
||||
|
||||
Ok(set::intersection(iters).stream())
|
||||
}
|
||||
|
||||
#[implement(Service)]
|
||||
async fn search_pdu_ids_query_room(
|
||||
&self,
|
||||
query: &RoomQuery<'_>,
|
||||
shortroomid: ShortRoomId,
|
||||
) -> Vec<Vec<RawPduId>> {
|
||||
tokenize(&query.criteria.search_term)
|
||||
.stream()
|
||||
.wide_then(|word| async move {
|
||||
self.search_pdu_ids_query_words(shortroomid, &word)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
}
|
||||
|
||||
/// Iterate over PduId's containing a word
|
||||
#[implement(Service)]
|
||||
fn search_pdu_ids_query_words<'a>(
|
||||
&'a self,
|
||||
shortroomid: ShortRoomId,
|
||||
word: &'a str,
|
||||
) -> impl Stream<Item = RawPduId> + Send + 'a {
|
||||
self.search_pdu_ids_query_word(shortroomid, word)
|
||||
.map(move |key| -> RawPduId {
|
||||
let key = &key[prefix_len(word)..];
|
||||
key.into()
|
||||
})
|
||||
}
|
||||
|
||||
/// Iterate over raw database results for a word
|
||||
#[implement(Service)]
|
||||
fn search_pdu_ids_query_word<'a>(
|
||||
&'a self,
|
||||
shortroomid: ShortRoomId,
|
||||
word: &'a str,
|
||||
) -> impl Stream<Item = Val<'a>> + Send + 'a + use<'a> {
|
||||
// rustc says const'ing this not yet stable
|
||||
let end_id: RawPduId = PduId {
|
||||
shortroomid,
|
||||
shorteventid: PduCount::max(),
|
||||
for token in batch {
|
||||
self.db.tokenids.remove(&token);
|
||||
}
|
||||
}
|
||||
.into();
|
||||
|
||||
// Newest pdus first
|
||||
let end = make_tokenid(shortroomid, word, &end_id);
|
||||
let prefix = make_prefix(shortroomid, word);
|
||||
self.db
|
||||
.tokenids
|
||||
.rev_raw_keys_from(&end)
|
||||
.ignore_err()
|
||||
.ready_take_while(move |key| key.starts_with(&prefix))
|
||||
/// Searches through the search index to find PDUs matching the query.
|
||||
/// Filters history visibility based on what the sender can see, and bundles
|
||||
/// aggregations before returning.
|
||||
pub async fn search_pdus<'a>(
|
||||
&'a self,
|
||||
query: &'a RoomQuery<'a>,
|
||||
sender_user: &'a UserId,
|
||||
) -> Result<(usize, impl Stream<Item = PduEvent> + Send + 'a)> {
|
||||
let pdu_ids: Vec<_> = self.search_pdu_ids(query).await?.collect().await;
|
||||
|
||||
let filter = &query.criteria.filter;
|
||||
let count = pdu_ids.len();
|
||||
let pdus = pdu_ids
|
||||
.into_iter()
|
||||
.stream()
|
||||
.wide_filter_map(move |result_pdu_id: RawPduId| async move {
|
||||
self.services
|
||||
.timeline
|
||||
.get_pdu_from_id(&result_pdu_id)
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.ready_filter(|pdu| !pdu.is_redacted())
|
||||
.ready_filter(move |pdu| filter.matches(pdu))
|
||||
.wide_filter_map(move |pdu| async move {
|
||||
self.services
|
||||
.state_accessor
|
||||
.user_can_see_event(query.user_id?, pdu.room_id().unwrap(), pdu.event_id())
|
||||
.await
|
||||
.then_some(pdu)
|
||||
})
|
||||
.skip(query.skip)
|
||||
.take(query.limit)
|
||||
.map(move |mut pdu| {
|
||||
pdu.set_unsigned(query.user_id);
|
||||
|
||||
pdu
|
||||
})
|
||||
.then(async move |mut pdu| {
|
||||
if let Err(e) = self
|
||||
.services
|
||||
.pdu_metadata
|
||||
.add_bundled_aggregations_to_pdu(sender_user, &mut pdu)
|
||||
.await
|
||||
{
|
||||
debug_warn!("Failed to add bundled aggregations: {e}");
|
||||
}
|
||||
pdu
|
||||
});
|
||||
|
||||
Ok((count, pdus))
|
||||
}
|
||||
|
||||
// TODO: The next 4 functions can be, theoretically, squashed into one.
|
||||
/// Same as search_pdus (searches the search index for matching PDUs), but
|
||||
/// only returns their short IDs with no filtering or bundled aggregations.
|
||||
pub async fn search_pdu_ids<'a>(
|
||||
&'a self,
|
||||
query: &'a RoomQuery<'_>,
|
||||
) -> Result<impl Stream<Item = RawPduId> + Send + 'a + use<'a>> {
|
||||
// result is modeled as a stream such that callers don't have to be refactored
|
||||
// though an additional async/wrap still exists for now
|
||||
let shortroomid = self.services.short.get_shortroomid(query.room_id).await?;
|
||||
|
||||
let pdu_ids = self.search_pdu_ids_query_room(query, shortroomid).await;
|
||||
|
||||
let iters = pdu_ids.into_iter().map(IntoIterator::into_iter);
|
||||
|
||||
Ok(set::intersection(iters).stream())
|
||||
}
|
||||
|
||||
async fn search_pdu_ids_query_room(
|
||||
&self,
|
||||
query: &RoomQuery<'_>,
|
||||
shortroomid: ShortRoomId,
|
||||
) -> Vec<Vec<RawPduId>> {
|
||||
tokenize(&query.criteria.search_term)
|
||||
.stream()
|
||||
.wide_then(|word| async move {
|
||||
self.search_pdu_ids_query_words(shortroomid, &word)
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
}
|
||||
|
||||
/// Iterate over PduId's containing a word
|
||||
fn search_pdu_ids_query_words<'a>(
|
||||
&'a self,
|
||||
shortroomid: ShortRoomId,
|
||||
word: &'a str,
|
||||
) -> impl Stream<Item = RawPduId> + Send + 'a {
|
||||
self.search_pdu_ids_query_word(shortroomid, word)
|
||||
.map(move |key| -> RawPduId {
|
||||
let key = &key[prefix_len(word)..];
|
||||
key.into()
|
||||
})
|
||||
}
|
||||
|
||||
/// Iterate over raw database results for a word
|
||||
fn search_pdu_ids_query_word<'a>(
|
||||
&'a self,
|
||||
shortroomid: ShortRoomId,
|
||||
word: &'a str,
|
||||
) -> impl Stream<Item = Val<'a>> + Send + 'a + use<'a> {
|
||||
// rustc says const'ing this not yet stable
|
||||
let end_id: RawPduId = PduId {
|
||||
shortroomid,
|
||||
shorteventid: PduCount::max(),
|
||||
}
|
||||
.into();
|
||||
|
||||
// Newest pdus first
|
||||
let end = make_tokenid(shortroomid, word, &end_id);
|
||||
let prefix = make_prefix(shortroomid, word);
|
||||
self.db
|
||||
.tokenids
|
||||
.rev_raw_keys_from(&end)
|
||||
.ignore_err()
|
||||
.ready_take_while(move |key| key.starts_with(&prefix))
|
||||
}
|
||||
}
|
||||
|
||||
/// Splits a string into tokens used as keys in the search inverted index
|
||||
|
||||
Reference in New Issue
Block a user