diff --git a/drivers/iommu/iommufd/main.c b/drivers/iommu/iommufd/main.c
index 0a96cc8f2..d2683ad82 100644
--- a/drivers/iommu/iommufd/main.c
+++ b/drivers/iommu/iommufd/main.c
@@ -122,6 +122,7 @@ int iommufd_object_remove(struct iommufd_ctx *ictx,
 {
 	struct iommufd_object *obj;
 	XA_STATE(xas, &ictx->objects, id);
+	DEFINE_MUTEX(remove_mutex); // Mutex for thread safety
 	bool zerod_shortterm = false;
 	int ret;
 
@@ -139,11 +140,14 @@ int iommufd_object_remove(struct iommufd_ctx *ictx,
 			 * defer cleaning this object until close.
 			 */
 			refcount_dec(&to_destroy->users);
+			mutex_unlock(&remove_mutex);
 			return ret;
 		}
 		zerod_shortterm = true;
 	}
 
+	mutex_lock(&remove_mutex);
+
 	xa_lock(&ictx->objects);
 	obj = xas_load(&xas);
 	if (to_destroy) {
@@ -176,6 +180,9 @@ int iommufd_object_remove(struct iommufd_ctx *ictx,
 	 * Since users is zero any positive users_shortterm must be racing
 	 * iommufd_put_object(), or we have a bug.
 	 */
+
+	mutex_unlock(&remove_mutex);
+
 	if (!zerod_shortterm) {
 		ret = iommufd_object_dec_wait_shortterm(ictx, obj);
 		if (WARN_ON(ret))
@@ -187,11 +194,13 @@ int iommufd_object_remove(struct iommufd_ctx *ictx,
 	return 0;
 
 err_xa:
+	xa_unlock(&ictx->objects);
+	mutex_unlock(&remove_mutex);
+
 	if (zerod_shortterm) {
 		/* Restore the xarray owned reference */
 		refcount_set(&obj->shortterm_users, 1);
 	}
-	xa_unlock(&ictx->objects);
 
 	/* The returned object reference count is zero */
 	return ret;