tuwunel_service/rooms/short/

mod.rs

use std::{borrow::Borrow, sync::Arc};

use futures::{FutureExt, Stream, StreamExt, pin_mut};
use ruma::{EventId, OwnedEventId, OwnedRoomId, RoomId, events::StateEventType};
use serde::Deserialize;
pub use tuwunel_core::matrix::{ShortEventId, ShortId, ShortRoomId, ShortStateKey};
use tuwunel_core::{
	Err, Result, err, implement,
	matrix::StateKey,
	utils,
	utils::{
		IterStream, MutexMap,
		hash::sha256::Digest,
		stream::{ReadyExt, WidebandExt},
	},
};
use tuwunel_database::{Deserialized, Get, Map, Qry, Txn};

pub struct Service {
	db: Data,
	creating: Creating,
	services: Arc<crate::services::OnceServices>,
}

struct Data {
	eventid_shorteventid: Arc<Map>,
	shorteventid_eventid: Arc<Map>,
	statekey_shortstatekey: Arc<Map>,
	shortstatekey_statekey: Arc<Map>,
	roomid_shortroomid: Arc<Map>,
	statehash_shortstatehash: Arc<Map>,
}

/// Serializes concurrent allocations so one identity maps to one short id.
///
/// A guard is held across both the re-read that detects a competing
/// allocation and the writes that publish this one. Entries exist only while
/// a caller holds or awaits one, so an uncontended allocation leaves nothing
/// behind.
#[derive(Default)]
struct Creating {
	shorteventid: MutexMap<OwnedEventId, ()>,
	shortstatekey: MutexMap<(StateEventType, StateKey), ()>,
	shortstatehash: MutexMap<Digest, ()>,
	shortroomid: MutexMap<OwnedRoomId, ()>,
}

pub type ShortStateHash = ShortId;

impl crate::Service for Service {
	fn build(args: &crate::Args<'_>) -> Result<Arc<Self>> {
		Ok(Arc::new(Self {
			db: Data {
				eventid_shorteventid: args.db["eventid_shorteventid"].clone(),
				shorteventid_eventid: args.db["shorteventid_eventid"].clone(),
				statekey_shortstatekey: args.db["statekey_shortstatekey"].clone(),
				shortstatekey_statekey: args.db["shortstatekey_statekey"].clone(),
				roomid_shortroomid: args.db["roomid_shortroomid"].clone(),
				statehash_shortstatehash: args.db["statehash_shortstatehash"].clone(),
			},
			creating: Creating::default(),
			services: args.services.clone(),
		}))
	}

	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
}

#[implement(Service)]
pub async fn get_or_create_shorteventid(&self, event_id: &EventId) -> ShortEventId {
	if let Ok(shorteventid) = self.get_shorteventid(event_id).await {
		return shorteventid;
	}

	self.create_shorteventid(event_id).await
}

/// Resolves each event id to its short id, allocating any that are absent.
///
/// Allocation runs ahead of consumer demand, so a caller that stops early
/// still allocates for the events already buffered. Today's callers drain the
/// stream in full.
#[implement(Service)]
pub fn multi_get_or_create_shorteventid<'a, I>(
	&'a self,
	event_ids: I,
) -> impl Stream<Item = ShortEventId> + Send + '_
where
	I: Iterator<Item = &'a EventId> + Clone + Send + 'a,
{
	event_ids
		.clone()
		.stream()
		.get(&self.db.eventid_shorteventid)
		.zip(event_ids.into_iter().stream())
		.wide_then(async |(result, event_id)| match result {
			| Ok(ref short) => utils::u64_from_u8(short),
			| Err(_) => self.create_shorteventid(event_id).await,
		})
}

#[implement(Service)]
async fn create_shorteventid(&self, event_id: &EventId) -> ShortEventId {
	let _lock = self.creating.shorteventid.lock(event_id).await;

	if let Ok(shorteventid) = self.get_shorteventid(event_id).await {
		return shorteventid;
	}

	let short = self.services.globals.next_count();
	let mut txn = self.services.db.txn();

	txn.insert_raw(&self.db.shorteventid_eventid, (*short).to_be_bytes(), event_id);
	txn.insert_raw(&self.db.eventid_shorteventid, event_id, (*short).to_be_bytes());
	txn.execute();

	*short
}

#[implement(Service)]
pub async fn get_shorteventid(&self, event_id: &EventId) -> Result<ShortEventId> {
	self.db
		.eventid_shorteventid
		.get(event_id)
		.await
		.deserialized()
}

#[implement(Service)]
pub async fn get_or_create_shortstatekey(
	&self,
	event_type: &StateEventType,
	state_key: &str,
) -> ShortStateKey {
	if let Ok(shortstatekey) = self
		.get_shortstatekey(event_type, state_key)
		.await
	{
		return shortstatekey;
	}

	self.create_shortstatekey(event_type, state_key)
		.await
}

#[implement(Service)]
async fn create_shortstatekey(
	&self,
	event_type: &StateEventType,
	state_key: &str,
) -> ShortStateKey {
	let owned_key = (event_type.clone(), StateKey::from_str(state_key));
	let _lock = self.creating.shortstatekey.lock(&owned_key).await;

	if let Ok(shortstatekey) = self
		.get_shortstatekey(event_type, state_key)
		.await
	{
		return shortstatekey;
	}

	let key = (event_type, state_key);
	let shortstatekey = self.services.globals.next_count();
	let mut txn = self.services.db.txn();

	txn.put(&self.db.shortstatekey_statekey, *shortstatekey, key);
	txn.put(&self.db.statekey_shortstatekey, key, *shortstatekey);
	txn.execute();

	*shortstatekey
}

#[implement(Service)]
pub async fn get_shortstatekey(
	&self,
	event_type: &StateEventType,
	state_key: &str,
) -> Result<ShortStateKey> {
	let key = (event_type, state_key);
	self.db
		.statekey_shortstatekey
		.qry(&key)
		.await
		.deserialized()
}

#[implement(Service)]
pub async fn get_eventid_from_short<Id>(&self, shorteventid: ShortEventId) -> Result<Id>
where
	Id: for<'de> Deserialize<'de> + Send + Sized + ToOwned,
	<Id as ToOwned>::Owned: Borrow<EventId>,
{
	const BUFSIZE: usize = size_of::<ShortEventId>();

	self.db
		.shorteventid_eventid
		.aqry::<BUFSIZE, _>(&shorteventid)
		.await
		.deserialized()
		.map_err(|e| err!(Database("Failed to find EventId from short {shorteventid:?}: {e:?}")))
}

#[implement(Service)]
pub fn multi_get_eventid_from_short<'a, Id, S>(
	&'a self,
	shorteventid: S,
) -> impl Stream<Item = Result<Id>> + Send + 'a
where
	S: Stream<Item = ShortEventId> + Send + 'a,
	Id: for<'de> Deserialize<'de> + Send + Sized + ToOwned + 'a,
	<Id as ToOwned>::Owned: Borrow<EventId>,
{
	shorteventid
		.qry(&self.db.shorteventid_eventid)
		.map(Deserialized::deserialized)
}

#[implement(Service)]
pub async fn get_statekey_from_short(
	&self,
	shortstatekey: ShortStateKey,
) -> Result<(StateEventType, StateKey)> {
	const BUFSIZE: usize = size_of::<ShortStateKey>();

	self.db
		.shortstatekey_statekey
		.aqry::<BUFSIZE, _>(&shortstatekey)
		.await
		.deserialized()
		.map_err(|e| {
			err!(Database(
				"Failed to find (StateEventType, state_key) from short {shortstatekey:?}: {e:?}"
			))
		})
}

#[implement(Service)]
pub fn multi_get_statekey_from_short<'a, S>(
	&'a self,
	shortstatekey: S,
) -> impl Stream<Item = Result<(StateEventType, StateKey)>> + Send + 'a
where
	S: Stream<Item = ShortStateKey> + Send + 'a,
{
	shortstatekey
		.qry(&self.db.shortstatekey_statekey)
		.map(Deserialized::deserialized)
}

/// Returns (shortstatehash, already_existed)
#[implement(Service)]
pub async fn get_or_create_shortstatehash<F>(
	&self,
	state_hash: &Digest,
	write_statediff: F,
) -> Result<(ShortStateHash, bool)>
where
	F: FnOnce(&mut Txn, ShortStateHash) -> Result,
{
	if let Ok(shortstatehash) = self.get_shortstatehash(state_hash).await {
		return Ok((shortstatehash, true));
	}

	self.create_shortstatehash(state_hash, write_statediff)
		.await
}

#[implement(Service)]
async fn create_shortstatehash<F>(
	&self,
	state_hash: &Digest,
	write_statediff: F,
) -> Result<(ShortStateHash, bool)>
where
	F: FnOnce(&mut Txn, ShortStateHash) -> Result,
{
	let _lock = self
		.creating
		.shortstatehash
		.lock(state_hash)
		.await;

	if let Ok(shortstatehash) = self.get_shortstatehash(state_hash).await {
		return Ok((shortstatehash, true));
	}

	let shortstatehash = self.services.globals.next_count();
	let mut txn = self.services.db.txn();

	txn.insert_raw(
		&self.db.statehash_shortstatehash,
		state_hash,
		(*shortstatehash).to_be_bytes(),
	);
	write_statediff(&mut txn, *shortstatehash)?;
	txn.execute();

	Ok((*shortstatehash, false))
}

#[implement(Service)]
pub async fn get_shortstatehash(&self, state_hash: &Digest) -> Result<ShortStateHash> {
	self.db
		.statehash_shortstatehash
		.get(state_hash)
		.await
		.deserialized()
}

#[implement(Service)]
pub async fn get_shortroomid(&self, room_id: &RoomId) -> Result<ShortRoomId> {
	self.db
		.roomid_shortroomid
		.get(room_id)
		.await
		.deserialized()
}

#[implement(Service)]
pub async fn get_roomid_from_short(&self, shortroomid_: ShortRoomId) -> Result<OwnedRoomId> {
	let stream = self
		.db
		.roomid_shortroomid
		.stream()
		.ready_filter_map(Result::ok);

	pin_mut!(stream);
	stream
		.ready_find(|&(_, shortroomid)| shortroomid == shortroomid_)
		.map(|found| found.map(|(room_id, _): (&RoomId, ShortRoomId)| room_id.to_owned()))
		.await
		.ok_or_else(|| err!(Database("Failed to find RoomId from {shortroomid_:?}")))
}

#[implement(Service)]
pub async fn get_or_create_shortroomid(&self, room_id: &RoomId) -> ShortRoomId {
	if let Ok(shortroomid) = self.get_shortroomid(room_id).await {
		return shortroomid;
	}

	self.create_shortroomid(room_id).await
}

#[implement(Service)]
async fn create_shortroomid(&self, room_id: &RoomId) -> ShortRoomId {
	const BUFSIZE: usize = size_of::<ShortRoomId>();

	let _lock = self.creating.shortroomid.lock(room_id).await;

	if let Ok(shortroomid) = self.get_shortroomid(room_id).await {
		return shortroomid;
	}

	let short = self.services.globals.next_count();

	debug_assert!(size_of_val(&*short) == BUFSIZE, "buffer requirement changed");

	self.db
		.roomid_shortroomid
		.raw_aput::<BUFSIZE, _, _>(room_id, *short);

	*short
}

#[implement(Service)]
pub async fn delete_shortroomid(&self, room_id: &RoomId) -> Result {
	if self
		.db
		.roomid_shortroomid
		.exists(room_id)
		.await
		.is_ok()
	{
		self.db.roomid_shortroomid.remove(room_id);
		Ok(())
	} else {
		Err!(Database("not found"))
	}
}