1 files changed, 7 insertions, 12 deletions

diff --git a/drivers/vfio/vfio.c b/drivers/vfio/vfio.c
index 9591e2b509d..17830c9c7cc 100644
--- a/drivers/vfio/vfio.c
+++ b/drivers/vfio/vfio.c
@@ -264,6 +264,7 @@ static struct vfio_group *vfio_create_group(struct iommu_group *iommu_group)
 	return group;
 }
 
+/* called with vfio.group_lock held */
 static void vfio_group_release(struct kref *kref)
 {
 	struct vfio_group *group = container_of(kref, struct vfio_group, kref);
@@ -287,13 +288,7 @@ static void vfio_group_release(struct kref *kref)
 
 static void vfio_group_put(struct vfio_group *group)
 {
-	mutex_lock(&vfio.group_lock);
-	/*
-	 * Release needs to unlock to unregister the notifier, so only
-	 * unlock if not released.
-	 */
-	if (!kref_put(&group->kref, vfio_group_release))
-		mutex_unlock(&vfio.group_lock);
+	kref_put_mutex(&group->kref, vfio_group_release, &vfio.group_lock);
 }
 
 /* Assume group_lock or group reference is held */
@@ -401,7 +396,6 @@ static void vfio_device_release(struct kref *kref)
 						  struct vfio_device, kref);
 	struct vfio_group *group = device->group;
 
-	mutex_lock(&group->device_lock);
 	list_del(&device->group_next);
 	mutex_unlock(&group->device_lock);
 
@@ -416,8 +410,9 @@ static void vfio_device_release(struct kref *kref)
 /* Device reference always implies a group reference */
 static void vfio_device_put(struct vfio_device *device)
 {
-	kref_put(&device->kref, vfio_device_release);
-	vfio_group_put(device->group);
+	struct vfio_group *group = device->group;
+	kref_put_mutex(&device->kref, vfio_device_release, &group->device_lock);
+	vfio_group_put(group);
 }
 
 static void vfio_device_get(struct vfio_device *device)
@@ -1116,10 +1111,10 @@ static int vfio_group_get_device_fd(struct vfio_group *group, char *buf)
 		 */
 		filep->f_mode |= (FMODE_LSEEK | FMODE_PREAD | FMODE_PWRITE);
 
-		fd_install(ret, filep);
-
 		vfio_device_get(device);
 		atomic_inc(&group->container_users);
+
+		fd_install(ret, filep);
 		break;
 	}
 	mutex_unlock(&group->device_lock);