3 files changed, 13 insertions, 4 deletions
diff --git a/synapse/handlers/device.py b/synapse/handlers/device.py
index 0a74e4d266..68a0c8ccb4 100644
--- a/synapse/handlers/device.py
+++ b/synapse/handlers/device.py
@@ -1212,9 +1212,18 @@ class DeviceListUpdater(DeviceListWorkerUpdater):
             raise InvalidAPICallError(f"Only one origin permitted, got {origins!r}")
 
         result = {}
+        failed = set()
         # TODO(Perf): Actually batch these up
         for user_id in user_ids:
-            result[user_id] = await self.user_device_resync(user_id)
+            user_result, user_failed = await self._user_device_resync_returning_failed(
+                user_id
+            )
+            result[user_id] = user_result
+            if user_failed:
+                failed.add(user_id)
+
+        if mark_failed_as_stale:
+            await self.store.mark_remote_users_device_caches_as_stale(failed)
 
         return result
 
@@ -1226,7 +1235,7 @@ class DeviceListUpdater(DeviceListWorkerUpdater):
         if failed and mark_failed_as_stale:
             # Mark the remote user's device list as stale so we know we need to retry
             # it later.
-            await self.store.mark_remote_user_device_cache_as_stale(user_id)
+            await self.store.mark_remote_users_device_caches_as_stale((user_id,))
 
         return result
 
diff --git a/synapse/handlers/devicemessage.py b/synapse/handlers/devicemessage.py
index 75e89850f5..00c403db49 100644
--- a/synapse/handlers/devicemessage.py
+++ b/synapse/handlers/devicemessage.py
@@ -195,7 +195,7 @@ class DeviceMessageHandler:
                 sender_user_id,
                 unknown_devices,
             )
-            await self.store.mark_remote_user_device_cache_as_stale(sender_user_id)
+            await self.store.mark_remote_users_device_caches_as_stale((sender_user_id,))
 
             # Immediately attempt a resync in the background
             run_in_background(self._user_device_resync, user_id=sender_user_id)
diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index 31df7f55cc..6df000faaf 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -1423,7 +1423,7 @@ class FederationEventHandler:
         """
 
         try:
-            await self._store.mark_remote_user_device_cache_as_stale(sender)
+            await self._store.mark_remote_users_device_caches_as_stale((sender,))
 
             # Immediately attempt a resync in the background
             if self._config.worker.worker_app: