fix: Record device list changes before syncs can read their count

Signed-off-by: Erwan Leboucher <erwanleboucher@gmail.com>

changelog.d/+device-list-count-race.bugfix.md

@@ -0,0 +1 @@
+Fix new devices sometimes never reaching other users' device lists. Contributed by @eleboucher.

src/service/globals/data.rs

@@ -21,7 +21,9 @@ impl Data {
 		}
 	}
 
-	pub fn next_count(&self) -> Result<u64> {
+	pub fn next_count(&self) -> Result<u64> { self.next_count_with(|_| ()) }
+
+	pub fn next_count_with(&self, write: impl FnOnce(u64)) -> Result<u64> {
 		let _cork = self.db.cork();
 		let mut lock = self.counter.write();
 		let counter: &mut u64 = &mut lock;
@@ -30,7 +32,9 @@ impl Data {
 			"counter mismatch"
 		);
 
-		*counter = counter.checked_add(1).unwrap_or(*counter);
+		let next = counter.checked_add(1).unwrap_or(*counter);
+		write(next);
+		*counter = next;
 
 		self.global.insert(COUNTER, counter.to_be_bytes());
 

src/service/globals/mod.rs

@@ -80,6 +80,11 @@ impl Service {
 	#[inline]
 	pub fn next_count(&self) -> Result<u64> { self.db.next_count() }
 
+	#[inline]
+	pub fn next_count_with(&self, write: impl FnOnce(u64)) -> Result<u64> {
+		self.db.next_count_with(write)
+	}
+
 	#[inline]
 	pub fn current_count(&self) -> Result<u64> { Ok(self.db.current_count()) }
 

src/service/users/keys.rs

@@ -442,9 +442,8 @@ impl super::Service {
 	/// Marks that a user's device keys have been updated, so that other users
 	/// can be notified of the change.
 	pub async fn mark_device_key_update(&self, user_id: &UserId) {
-		let count = self.services.globals.next_count().unwrap();
-
-		self.services
+		let encrypted_rooms: Vec<_> = self
+			.services
 			.state_cache
 			.rooms_joined(user_id)
 			// Don't send key updates to unencrypted rooms
@@ -460,16 +459,26 @@ impl super::Service {
 					None
 				}
 			})
-			.for_each(async |room_id| {
-				let key = (&room_id, count);
-				self.db.keychangeid_userid.put_raw(key, user_id);
-
-				self.services.sync.wake_all_joined(&room_id).await;
-			})
+			.collect()
 			.await;
 
-		let key = (user_id, count);
-		self.db.keychangeid_userid.put_raw(key, user_id);
+		self.services
+			.globals
+			.next_count_with(|count| {
+				self.db
+					.keychangeid_userid
+					.put_raw((user_id, count), user_id);
+				for room_id in &encrypted_rooms {
+					self.db
+						.keychangeid_userid
+						.put_raw((room_id, count), user_id);
+				}
+			})
+			.unwrap();
+
+		for room_id in &encrypted_rooms {
+			self.services.sync.wake_all_joined(room_id).await;
+		}
 	}
 
 	/// Returns the device identity keys for a given device.