Refactor counter increment sites for TwoPhaseCounter.

Signed-off-by: Jason Volk <jason@zemos.net>

src/service/rooms/alias/mod.rs

@@ -71,6 +71,8 @@ impl Service {
 			return Err!(Request(Forbidden("Only the server user can set this alias")));
 		}
 
+		let count = self.services.globals.next_count();
+
 		// Comes first as we don't want a stuck alias
 		self.db
 			.alias_userid
@@ -82,7 +84,8 @@ impl Service {
 
 		let mut aliasid = room_id.as_bytes().to_vec();
 		aliasid.push(0xFF);
-		aliasid.extend_from_slice(&self.services.globals.next_count()?.to_be_bytes());
+		aliasid.extend_from_slice(&count.to_be_bytes());
+
 		self.db
 			.aliasid_alias
 			.insert(&aliasid, alias.as_bytes());

src/service/rooms/read_receipt/data.rs

@@ -56,8 +56,8 @@ impl Data {
 			.ready_for_each(|key| self.readreceiptid_readreceipt.del(key))
 			.await;
 
-		let count = self.services.globals.next_count().unwrap();
-		let latest_id = (room_id, count, user_id);
+		let count = self.services.globals.next_count();
+		let latest_id = (room_id, *count, user_id);
 		self.readreceiptid_readreceipt
 			.put(latest_id, Json(event));
 	}
@@ -89,11 +89,11 @@ impl Data {
 
 	pub(super) fn private_read_set(&self, room_id: &RoomId, user_id: &UserId, pdu_count: u64) {
 		let key = (room_id, user_id);
-		let next_count = self.services.globals.next_count().unwrap();
+		let next_count = self.services.globals.next_count();
 
 		self.roomuserid_privateread.put(key, pdu_count);
 		self.roomuserid_lastprivatereadupdate
-			.put(key, next_count);
+			.put(key, *next_count);
 	}
 
 	pub(super) async fn private_read_get_count(

src/service/rooms/short/mod.rs

@@ -81,18 +81,18 @@ where
 fn create_shorteventid(&self, event_id: &EventId) -> ShortEventId {
 	const BUFSIZE: usize = size_of::<ShortEventId>();
 
-	let short = self.services.globals.next_count().unwrap();
-	debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed");
+	let short = self.services.globals.next_count();
+	debug_assert!(size_of_val(&*short) == BUFSIZE, "buffer requirement changed");
 
 	self.db
 		.eventid_shorteventid
-		.raw_aput::<BUFSIZE, _, _>(event_id, short);
+		.raw_aput::<BUFSIZE, _, _>(event_id, *short);
 
 	self.db
 		.shorteventid_eventid
-		.aput_raw::<BUFSIZE, _, _>(short, event_id);
+		.aput_raw::<BUFSIZE, _, _>(*short, event_id);
 
-	short
+	*short
 }
 
 #[implement(Service)]
@@ -120,18 +120,19 @@ pub async fn get_or_create_shortstatekey(
 	}
 
 	let key = (event_type, state_key);
-	let shortstatekey = self.services.globals.next_count().unwrap();
-	debug_assert!(size_of_val(&shortstatekey) == BUFSIZE, "buffer requirement changed");
+	let shortstatekey = self.services.globals.next_count();
+
+	debug_assert!(size_of_val(&*shortstatekey) == BUFSIZE, "buffer requirement changed");
 
 	self.db
 		.statekey_shortstatekey
-		.put_aput::<BUFSIZE, _, _>(key, shortstatekey);
+		.put_aput::<BUFSIZE, _, _>(key, *shortstatekey);
 
 	self.db
 		.shortstatekey_statekey
-		.aput_put::<BUFSIZE, _, _>(shortstatekey, key);
+		.aput_put::<BUFSIZE, _, _>(*shortstatekey, key);
 
-	shortstatekey
+	*shortstatekey
 }
 
 #[implement(Service)]
@@ -226,14 +227,14 @@ pub async fn get_or_create_shortstatehash(&self, state_hash: &[u8]) -> (ShortSta
 		return (shortstatehash, true);
 	}
 
-	let shortstatehash = self.services.globals.next_count().unwrap();
-	debug_assert!(size_of_val(&shortstatehash) == BUFSIZE, "buffer requirement changed");
+	let shortstatehash = self.services.globals.next_count();
+	debug_assert!(size_of_val(&*shortstatehash) == BUFSIZE, "buffer requirement changed");
 
 	self.db
 		.statehash_shortstatehash
-		.raw_aput::<BUFSIZE, _, _>(state_hash, shortstatehash);
+		.raw_aput::<BUFSIZE, _, _>(state_hash, *shortstatehash);
 
-	(shortstatehash, false)
+	(*shortstatehash, false)
 }
 
 #[implement(Service)]
@@ -255,13 +256,13 @@ pub async fn get_or_create_shortroomid(&self, room_id: &RoomId) -> ShortRoomId {
 		.unwrap_or_else(|_| {
 			const BUFSIZE: usize = size_of::<ShortRoomId>();
 
-			let short = self.services.globals.next_count().unwrap();
-			debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed");
+			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);
+				.raw_aput::<BUFSIZE, _, _>(room_id, *short);
 
-			short
+			*short
 		})
 }

src/service/rooms/state/mod.rs

@@ -305,7 +305,7 @@ impl Service {
 				}
 
 				// TODO: statehash with deterministic inputs
-				let shortstatehash = self.services.globals.next_count()?;
+				let shortstatehash = self.services.globals.next_count();
 
 				let mut statediffnew = CompressedState::new();
 				statediffnew.insert(new);
@@ -318,14 +318,14 @@ impl Service {
 				self.services
 					.state_compressor
 					.save_state_from_diff(
-						shortstatehash,
+						*shortstatehash,
 						Arc::new(statediffnew),
 						Arc::new(statediffremoved),
 						2,
 						states_parents,
 					)?;
 
-				Ok(shortstatehash)
+				Ok(*shortstatehash)
 			},
 			| _ =>
 				Ok(previous_shortstatehash.expect("first event in room must be a state event")),

src/service/rooms/state_cache/update.rs

@@ -272,6 +272,8 @@ pub fn mark_as_joined(&self, user_id: &UserId, room_id: &RoomId) {
 #[implement(super::Service)]
 #[tracing::instrument(skip(self), level = "debug")]
 pub fn mark_as_left(&self, user_id: &UserId, room_id: &RoomId) {
+	let count = self.services.globals.next_count();
+
 	let userroom_id = (user_id, room_id);
 	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
 
@@ -286,7 +288,7 @@ pub fn mark_as_left(&self, user_id: &UserId, room_id: &RoomId) {
 		.raw_put(&userroom_id, Json(leftstate));
 	self.db
 		.roomuserid_leftcount
-		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+		.raw_aput::<8, _, _>(&roomuser_id, *count);
 
 	self.db.userroomid_joined.remove(&userroom_id);
 	self.db.roomuserid_joined.remove(&roomuser_id);
@@ -319,6 +321,8 @@ pub fn mark_as_knocked(
 	room_id: &RoomId,
 	knocked_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
 ) {
+	let count = self.services.globals.next_count();
+
 	let userroom_id = (user_id, room_id);
 	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
 
@@ -330,7 +334,7 @@ pub fn mark_as_knocked(
 		.raw_put(&userroom_id, Json(knocked_state.unwrap_or_default()));
 	self.db
 		.roomuserid_knockedcount
-		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+		.raw_aput::<8, _, _>(&roomuser_id, *count);
 
 	self.db.userroomid_joined.remove(&userroom_id);
 	self.db.roomuserid_joined.remove(&roomuser_id);
@@ -375,6 +379,8 @@ pub async fn mark_as_invited(
 	last_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
 	invite_via: Option<Vec<OwnedServerName>>,
 ) {
+	let count = self.services.globals.next_count();
+
 	let roomuser_id = (room_id, user_id);
 	let roomuser_id = serialize_key(roomuser_id).expect("failed to serialize roomuser_id");
 
@@ -386,7 +392,7 @@ pub async fn mark_as_invited(
 		.raw_put(&userroom_id, Json(last_state.unwrap_or_default()));
 	self.db
 		.roomuserid_invitecount
-		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+		.raw_aput::<8, _, _>(&roomuser_id, *count);
 
 	self.db.userroomid_joined.remove(&userroom_id);
 	self.db.roomuserid_joined.remove(&roomuser_id);

src/service/rooms/timeline/append.rs

@@ -159,20 +159,20 @@ where
 		.await;
 
 	let insert_lock = self.mutex_insert.lock(pdu.room_id()).await;
-
-	let count1 = self.services.globals.next_count().unwrap();
+	let count1 = self.services.globals.next_count();
+	let count2 = self.services.globals.next_count();
 
 	// Mark as read first so the sending client doesn't get a notification even if
 	// appending fails
 	self.services
 		.read_receipt
-		.private_read_set(pdu.room_id(), pdu.sender(), count1);
+		.private_read_set(pdu.room_id(), pdu.sender(), *count1);
 
 	self.services
 		.user
 		.reset_notification_counts(pdu.sender(), pdu.room_id());
 
-	let count2 = PduCount::Normal(self.services.globals.next_count().unwrap());
+	let count2 = PduCount::Normal(*count2);
 	let pdu_id: RawPduId = PduId { shortroomid, shorteventid: count2 }.into();
 
 	// Insert pdu

src/service/rooms/timeline/backfill.rs

@@ -181,13 +181,9 @@ pub async fn backfill_pdu(&self, origin: &ServerName, pdu: Box<RawJsonValue>) ->
 
 	let insert_lock = self.mutex_insert.lock(&room_id).await;
 
-	let count: i64 = self
-		.services
-		.globals
-		.next_count()
-		.unwrap()
-		.try_into()?;
+	let count = self.services.globals.next_count();
 
+	let count: i64 = (*count).try_into()?;
 	let pdu_id: RawPduId = PduId {
 		shortroomid,
 		shorteventid: PduCount::Backfilled(validated!(0 - count)),

src/service/rooms/typing/mod.rs

@@ -53,6 +53,7 @@ impl Service {
 	/// roomtyping_remove is called.
 	pub async fn typing_add(&self, user_id: &UserId, room_id: &RoomId, timeout: u64) -> Result {
 		debug_info!("typing started {user_id:?} in {room_id:?} timeout:{timeout:?}");
+
 		// update clients
 		self.typing
 			.write()
@@ -61,10 +62,11 @@ impl Service {
 			.or_default()
 			.insert(user_id.to_owned(), timeout);
 
+		let count = self.services.globals.next_count();
 		self.last_typing_update
 			.write()
 			.await
-			.insert(room_id.to_owned(), self.services.globals.next_count()?);
+			.insert(room_id.to_owned(), *count);
 
 		if self
 			.typing_update_sender
@@ -86,6 +88,7 @@ impl Service {
 	/// Removes a user from typing before the timeout is reached.
 	pub async fn typing_remove(&self, user_id: &UserId, room_id: &RoomId) -> Result {
 		debug_info!("typing stopped {user_id:?} in {room_id:?}");
+
 		// update clients
 		self.typing
 			.write()
@@ -94,10 +97,11 @@ impl Service {
 			.or_default()
 			.remove(user_id);
 
+		let count = self.services.globals.next_count();
 		self.last_typing_update
 			.write()
 			.await
-			.insert(room_id.to_owned(), self.services.globals.next_count()?);
+			.insert(room_id.to_owned(), *count);
 
 		if self
 			.typing_update_sender
@@ -146,16 +150,18 @@ impl Service {
 		if !removable.is_empty() {
 			let typing = &mut self.typing.write().await;
 			let room = typing.entry(room_id.to_owned()).or_default();
+
 			for user in &removable {
 				debug_info!("typing timeout {user:?} in {room_id:?}");
 				room.remove(user);
 			}
 
 			// update clients
+			let count = self.services.globals.next_count();
 			self.last_typing_update
 				.write()
 				.await
-				.insert(room_id.to_owned(), self.services.globals.next_count()?);
+				.insert(room_id.to_owned(), *count);
 
 			if self
 				.typing_update_sender

src/service/rooms/user/mod.rs

@@ -47,6 +47,8 @@ impl crate::Service for Service {
 
 #[implement(Service)]
 pub fn reset_notification_counts(&self, user_id: &UserId, room_id: &RoomId) {
+	let count = self.services.globals.next_count();
+
 	let userroom_id = (user_id, room_id);
 	self.db
 		.userroomid_highlightcount
@@ -56,10 +58,9 @@ pub fn reset_notification_counts(&self, user_id: &UserId, room_id: &RoomId) {
 		.put(userroom_id, 0_u64);
 
 	let roomuser_id = (room_id, user_id);
-	let count = self.services.globals.next_count().unwrap();
 	self.db
 		.roomuserid_lastnotificationread
-		.put(roomuser_id, count);
+		.put(roomuser_id, *count);
 }
 
 #[implement(Service)]