diff mbox series

[02/13] vfio: Introduce a vfio_uninit_group_dev() API call

Message ID 2-v1-eaf3ccbba33c+1add0-vfio_reflck_jgg@nvidia.com (mailing list archive)
State New, archived
Headers show
Series Provide core infrastructure for managing open/release | expand

Commit Message

Jason Gunthorpe July 15, 2021, 12:20 a.m. UTC 
From: Max Gurtovoy <mgurtovoy@nvidia.com>

This pairs with vfio_init_group_dev() and allows undoing any state that is
stored in the vfio_device unrelated to registration. Add appropriately
placed calls to all the drivers.

The following patch will use this to add pre-registration state for the
device set.

Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
---
 Documentation/driver-api/vfio.rst            |  4 ++-
 drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
 drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
 drivers/vfio/pci/vfio_pci.c                  |  6 +++--
 drivers/vfio/platform/vfio_platform_common.c |  7 +++--
 drivers/vfio/vfio.c                          |  5 ++++
 include/linux/vfio.h                         |  1 +
 samples/vfio-mdev/mbochs.c                   |  2 ++
 samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
 samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
 10 files changed, 64 insertions(+), 32 deletions(-)

Comments

Leon Romanovsky July 15, 2021, 3:49 a.m. UTC | #1 
On Wed, Jul 14, 2021 at 09:20:31PM -0300, Jason Gunthorpe wrote:
> From: Max Gurtovoy <mgurtovoy@nvidia.com>
> 
> This pairs with vfio_init_group_dev() and allows undoing any state that is
> stored in the vfio_device unrelated to registration. Add appropriately
> placed calls to all the drivers.
> 
> The following patch will use this to add pre-registration state for the
> device set.
> 
> Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
> Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
> ---
>  Documentation/driver-api/vfio.rst            |  4 ++-
>  drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
>  drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
>  drivers/vfio/pci/vfio_pci.c                  |  6 +++--
>  drivers/vfio/platform/vfio_platform_common.c |  7 +++--
>  drivers/vfio/vfio.c                          |  5 ++++
>  include/linux/vfio.h                         |  1 +
>  samples/vfio-mdev/mbochs.c                   |  2 ++
>  samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
>  samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
>  10 files changed, 64 insertions(+), 32 deletions(-)

<...>

> @@ -674,6 +675,7 @@ static int vfio_fsl_mc_remove(struct fsl_mc_device *mc_dev)
>  
>  	dprc_remove_devices(mc_dev, NULL, 0);
>  	vfio_fsl_uninit_device(vdev);
> +	vfio_uninit_group_dev(&vdev->vdev);

This is wrong place, the _uninit_ should be after vfio_fsl_mc_reflck_put().

>  	vfio_fsl_mc_reflck_put(vdev->reflck);
>  
>  	kfree(vdev);
Jason Gunthorpe July 15, 2021, 12:45 p.m. UTC | #2 
On Thu, Jul 15, 2021 at 06:49:05AM +0300, Leon Romanovsky wrote:
> On Wed, Jul 14, 2021 at 09:20:31PM -0300, Jason Gunthorpe wrote:
> > From: Max Gurtovoy <mgurtovoy@nvidia.com>
> > 
> > This pairs with vfio_init_group_dev() and allows undoing any state that is
> > stored in the vfio_device unrelated to registration. Add appropriately
> > placed calls to all the drivers.
> > 
> > The following patch will use this to add pre-registration state for the
> > device set.
> > 
> > Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
> > Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
> >  Documentation/driver-api/vfio.rst            |  4 ++-
> >  drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
> >  drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
> >  drivers/vfio/pci/vfio_pci.c                  |  6 +++--
> >  drivers/vfio/platform/vfio_platform_common.c |  7 +++--
> >  drivers/vfio/vfio.c                          |  5 ++++
> >  include/linux/vfio.h                         |  1 +
> >  samples/vfio-mdev/mbochs.c                   |  2 ++
> >  samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
> >  samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
> >  10 files changed, 64 insertions(+), 32 deletions(-)
> 
> <...>
> 
> > @@ -674,6 +675,7 @@ static int vfio_fsl_mc_remove(struct fsl_mc_device *mc_dev)
> >  
> >  	dprc_remove_devices(mc_dev, NULL, 0);
> >  	vfio_fsl_uninit_device(vdev);
> > +	vfio_uninit_group_dev(&vdev->vdev);
> 
> This is wrong place, the _uninit_ should be after vfio_fsl_mc_reflck_put().

Well, maybe, but it doesn't matter, the uninit doesn't effect the
reflck and the next fsl patch deletes the line below. I can switch it
if there is a v2

> >  	vfio_fsl_mc_reflck_put(vdev->reflck);

Jason
Cornelia Huck July 19, 2021, 12:11 p.m. UTC | #3 
On Wed, Jul 14 2021, Jason Gunthorpe <jgg@nvidia.com> wrote:

> From: Max Gurtovoy <mgurtovoy@nvidia.com>
>
> This pairs with vfio_init_group_dev() and allows undoing any state that is
> stored in the vfio_device unrelated to registration. Add appropriately
> placed calls to all the drivers.
>
> The following patch will use this to add pre-registration state for the
> device set.
>
> Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
> Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
> ---
>  Documentation/driver-api/vfio.rst            |  4 ++-
>  drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
>  drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
>  drivers/vfio/pci/vfio_pci.c                  |  6 +++--
>  drivers/vfio/platform/vfio_platform_common.c |  7 +++--
>  drivers/vfio/vfio.c                          |  5 ++++
>  include/linux/vfio.h                         |  1 +
>  samples/vfio-mdev/mbochs.c                   |  2 ++
>  samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
>  samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
>  10 files changed, 64 insertions(+), 32 deletions(-)

(...)

> diff --git a/samples/vfio-mdev/mbochs.c b/samples/vfio-mdev/mbochs.c
> index e81b875b4d87b4..cf264d0bf11053 100644
> --- a/samples/vfio-mdev/mbochs.c
> +++ b/samples/vfio-mdev/mbochs.c
> @@ -558,6 +558,7 @@ static int mbochs_probe(struct mdev_device *mdev)
>  	return 0;
>  
>  err_mem:
> +	vfio_uninit_group_dev(&mdev_state->vdev);
>  	kfree(mdev_state->vconfig);
>  	kfree(mdev_state);
>  	return ret;
> @@ -571,6 +572,7 @@ static void mbochs_remove(struct mdev_device *mdev)
>  	vfio_unregister_group_dev(&mdev_state->vdev);
>  	kfree(mdev_state->pages);
>  	kfree(mdev_state->vconfig);
> +	vfio_uninit_group_dev(&mdev_state->vdev);

Does the location of the uninit vs the kfree matter? Even if not, it
might be good to keep it consistent.

>  	kfree(mdev_state);
>  }
>
Jason Gunthorpe July 19, 2021, 12:17 p.m. UTC | #4 
On Mon, Jul 19, 2021 at 02:11:38PM +0200, Cornelia Huck wrote:
> On Wed, Jul 14 2021, Jason Gunthorpe <jgg@nvidia.com> wrote:
> 
> > From: Max Gurtovoy <mgurtovoy@nvidia.com>
> >
> > This pairs with vfio_init_group_dev() and allows undoing any state that is
> > stored in the vfio_device unrelated to registration. Add appropriately
> > placed calls to all the drivers.
> >
> > The following patch will use this to add pre-registration state for the
> > device set.
> >
> > Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
> > Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
> >  Documentation/driver-api/vfio.rst            |  4 ++-
> >  drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
> >  drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
> >  drivers/vfio/pci/vfio_pci.c                  |  6 +++--
> >  drivers/vfio/platform/vfio_platform_common.c |  7 +++--
> >  drivers/vfio/vfio.c                          |  5 ++++
> >  include/linux/vfio.h                         |  1 +
> >  samples/vfio-mdev/mbochs.c                   |  2 ++
> >  samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
> >  samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
> >  10 files changed, 64 insertions(+), 32 deletions(-)
> 
> (...)
> 
> > diff --git a/samples/vfio-mdev/mbochs.c b/samples/vfio-mdev/mbochs.c
> > index e81b875b4d87b4..cf264d0bf11053 100644
> > +++ b/samples/vfio-mdev/mbochs.c
> > @@ -558,6 +558,7 @@ static int mbochs_probe(struct mdev_device *mdev)
> >  	return 0;
> >  
> >  err_mem:
> > +	vfio_uninit_group_dev(&mdev_state->vdev);
> >  	kfree(mdev_state->vconfig);
> >  	kfree(mdev_state);
> >  	return ret;

Doesn't this leak pages? Sigh.

> > @@ -571,6 +572,7 @@ static void mbochs_remove(struct mdev_device *mdev)
> >  	vfio_unregister_group_dev(&mdev_state->vdev);
> >  	kfree(mdev_state->pages);
> >  	kfree(mdev_state->vconfig);
> > +	vfio_uninit_group_dev(&mdev_state->vdev);
> 
> Does the location of the uninit vs the kfree matter? Even if not, it
> might be good to keep it consistent.

It does not, but I will reorder it anyhow

Jason
Cornelia Huck July 19, 2021, 12:43 p.m. UTC | #5 
On Mon, Jul 19 2021, Jason Gunthorpe <jgg@nvidia.com> wrote:

> On Mon, Jul 19, 2021 at 02:11:38PM +0200, Cornelia Huck wrote:
>> On Wed, Jul 14 2021, Jason Gunthorpe <jgg@nvidia.com> wrote:
>> 
>> > From: Max Gurtovoy <mgurtovoy@nvidia.com>
>> >
>> > This pairs with vfio_init_group_dev() and allows undoing any state that is
>> > stored in the vfio_device unrelated to registration. Add appropriately
>> > placed calls to all the drivers.
>> >
>> > The following patch will use this to add pre-registration state for the
>> > device set.
>> >
>> > Signed-off-by: Max Gurtovoy <mgurtovoy@nvidia.com>
>> > Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
>> >  Documentation/driver-api/vfio.rst            |  4 ++-
>> >  drivers/vfio/fsl-mc/vfio_fsl_mc.c            |  6 +++--
>> >  drivers/vfio/mdev/vfio_mdev.c                | 13 +++++++---
>> >  drivers/vfio/pci/vfio_pci.c                  |  6 +++--
>> >  drivers/vfio/platform/vfio_platform_common.c |  7 +++--
>> >  drivers/vfio/vfio.c                          |  5 ++++
>> >  include/linux/vfio.h                         |  1 +
>> >  samples/vfio-mdev/mbochs.c                   |  2 ++
>> >  samples/vfio-mdev/mdpy.c                     | 25 ++++++++++--------
>> >  samples/vfio-mdev/mtty.c                     | 27 ++++++++++++--------
>> >  10 files changed, 64 insertions(+), 32 deletions(-)
>> 
>> (...)
>> 
>> > diff --git a/samples/vfio-mdev/mbochs.c b/samples/vfio-mdev/mbochs.c
>> > index e81b875b4d87b4..cf264d0bf11053 100644
>> > +++ b/samples/vfio-mdev/mbochs.c
>> > @@ -558,6 +558,7 @@ static int mbochs_probe(struct mdev_device *mdev)
>> >  	return 0;
>> >  
>> >  err_mem:
>> > +	vfio_uninit_group_dev(&mdev_state->vdev);
>> >  	kfree(mdev_state->vconfig);
>> >  	kfree(mdev_state);
>> >  	return ret;
>
> Doesn't this leak pages? Sigh.

I think it also fails to decrease mbochs_used_mbytes; looks like there
need to be two cleanup labels.

>
>> > @@ -571,6 +572,7 @@ static void mbochs_remove(struct mdev_device *mdev)
>> >  	vfio_unregister_group_dev(&mdev_state->vdev);
>> >  	kfree(mdev_state->pages);
>> >  	kfree(mdev_state->vconfig);
>> > +	vfio_uninit_group_dev(&mdev_state->vdev);
>> 
>> Does the location of the uninit vs the kfree matter? Even if not, it
>> might be good to keep it consistent.
>
> It does not, but I will reorder it anyhow
>
> Jason
diff mbox series

Patch

diff --git a/Documentation/driver-api/vfio.rst b/Documentation/driver-api/vfio.rst
index 606eed8823ceab..c663b6f978255b 100644
--- a/Documentation/driver-api/vfio.rst
+++ b/Documentation/driver-api/vfio.rst
@@ -255,11 +255,13 @@  vfio_unregister_group_dev() respectively::
 	void vfio_init_group_dev(struct vfio_device *device,
 				struct device *dev,
 				const struct vfio_device_ops *ops);
+	void vfio_uninit_group_dev(struct vfio_device *device);
 	int vfio_register_group_dev(struct vfio_device *device);
 	void vfio_unregister_group_dev(struct vfio_device *device);
 
 The driver should embed the vfio_device in its own structure and call
-vfio_init_group_dev() to pre-configure it before going to registration.
+vfio_init_group_dev() to pre-configure it before going to registration
+and call vfio_uninit_group_dev() after completing the un-registration.
 vfio_register_group_dev() indicates to the core to begin tracking the
 iommu_group of the specified dev and register the dev as owned by a VFIO bus
 driver. Once vfio_register_group_dev() returns it is possible for userspace to
diff --git a/drivers/vfio/fsl-mc/vfio_fsl_mc.c b/drivers/vfio/fsl-mc/vfio_fsl_mc.c
index 90cad109583b80..3d2be06e1bc146 100644
--- a/drivers/vfio/fsl-mc/vfio_fsl_mc.c
+++ b/drivers/vfio/fsl-mc/vfio_fsl_mc.c
@@ -627,7 +627,7 @@  static int vfio_fsl_mc_probe(struct fsl_mc_device *mc_dev)
 
 	ret = vfio_fsl_mc_reflck_attach(vdev);
 	if (ret)
-		goto out_kfree;
+		goto out_uninit;
 
 	ret = vfio_fsl_mc_init_device(vdev);
 	if (ret)
@@ -657,7 +657,8 @@  static int vfio_fsl_mc_probe(struct fsl_mc_device *mc_dev)
 	vfio_fsl_uninit_device(vdev);
 out_reflck:
 	vfio_fsl_mc_reflck_put(vdev->reflck);
-out_kfree:
+out_uninit:
+	vfio_uninit_group_dev(&vdev->vdev);
 	kfree(vdev);
 out_group_put:
 	vfio_iommu_group_put(group, dev);
@@ -674,6 +675,7 @@  static int vfio_fsl_mc_remove(struct fsl_mc_device *mc_dev)
 
 	dprc_remove_devices(mc_dev, NULL, 0);
 	vfio_fsl_uninit_device(vdev);
+	vfio_uninit_group_dev(&vdev->vdev);
 	vfio_fsl_mc_reflck_put(vdev->reflck);
 
 	kfree(vdev);
diff --git a/drivers/vfio/mdev/vfio_mdev.c b/drivers/vfio/mdev/vfio_mdev.c
index 39ef7489fe4719..a5c77ccb24f70a 100644
--- a/drivers/vfio/mdev/vfio_mdev.c
+++ b/drivers/vfio/mdev/vfio_mdev.c
@@ -120,12 +120,16 @@  static int vfio_mdev_probe(struct mdev_device *mdev)
 
 	vfio_init_group_dev(vdev, &mdev->dev, &vfio_mdev_dev_ops);
 	ret = vfio_register_group_dev(vdev);
-	if (ret) {
-		kfree(vdev);
-		return ret;
-	}
+	if (ret)
+		goto out_uninit;
+
 	dev_set_drvdata(&mdev->dev, vdev);
 	return 0;
+
+out_uninit:
+	vfio_uninit_group_dev(vdev);
+	kfree(vdev);
+	return ret;
 }
 
 static void vfio_mdev_remove(struct mdev_device *mdev)
@@ -133,6 +137,7 @@  static void vfio_mdev_remove(struct mdev_device *mdev)
 	struct vfio_device *vdev = dev_get_drvdata(&mdev->dev);
 
 	vfio_unregister_group_dev(vdev);
+	vfio_uninit_group_dev(vdev);
 	kfree(vdev);
 }
 
diff --git a/drivers/vfio/pci/vfio_pci.c b/drivers/vfio/pci/vfio_pci.c
index 318864d5283782..fab3715d60d4ba 100644
--- a/drivers/vfio/pci/vfio_pci.c
+++ b/drivers/vfio/pci/vfio_pci.c
@@ -2022,7 +2022,7 @@  static int vfio_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 
 	ret = vfio_pci_reflck_attach(vdev);
 	if (ret)
-		goto out_free;
+		goto out_uninit;
 	ret = vfio_pci_vf_init(vdev);
 	if (ret)
 		goto out_reflck;
@@ -2059,7 +2059,8 @@  static int vfio_pci_probe(struct pci_dev *pdev, const struct pci_device_id *id)
 	vfio_pci_vf_uninit(vdev);
 out_reflck:
 	vfio_pci_reflck_put(vdev->reflck);
-out_free:
+out_uninit:
+	vfio_uninit_group_dev(&vdev->vdev);
 	kfree(vdev->pm_save);
 	kfree(vdev);
 out_group_put:
@@ -2077,6 +2078,7 @@  static void vfio_pci_remove(struct pci_dev *pdev)
 
 	vfio_pci_vf_uninit(vdev);
 	vfio_pci_reflck_put(vdev->reflck);
+	vfio_uninit_group_dev(&vdev->vdev);
 	vfio_pci_vga_uninit(vdev);
 
 	vfio_iommu_group_put(pdev->dev.iommu_group, &pdev->dev);
diff --git a/drivers/vfio/platform/vfio_platform_common.c b/drivers/vfio/platform/vfio_platform_common.c
index 703164df7637db..bdde8605178cd2 100644
--- a/drivers/vfio/platform/vfio_platform_common.c
+++ b/drivers/vfio/platform/vfio_platform_common.c
@@ -667,7 +667,7 @@  int vfio_platform_probe_common(struct vfio_platform_device *vdev,
 		ret = vfio_platform_of_probe(vdev, dev);
 
 	if (ret)
-		return ret;
+		goto out_uninit;
 
 	vdev->device = dev;
 
@@ -675,7 +675,7 @@  int vfio_platform_probe_common(struct vfio_platform_device *vdev,
 	if (ret && vdev->reset_required) {
 		dev_err(dev, "No reset function found for device %s\n",
 			vdev->name);
-		return ret;
+		goto out_uninit;
 	}
 
 	group = vfio_iommu_group_get(dev);
@@ -698,6 +698,8 @@  int vfio_platform_probe_common(struct vfio_platform_device *vdev,
 	vfio_iommu_group_put(group, dev);
 put_reset:
 	vfio_platform_put_reset(vdev);
+out_uninit:
+	vfio_uninit_group_dev(&vdev->vdev);
 	return ret;
 }
 EXPORT_SYMBOL_GPL(vfio_platform_probe_common);
@@ -708,6 +710,7 @@  void vfio_platform_remove_common(struct vfio_platform_device *vdev)
 
 	pm_runtime_disable(vdev->device);
 	vfio_platform_put_reset(vdev);
+	vfio_uninit_group_dev(&vdev->vdev);
 	vfio_iommu_group_put(vdev->vdev.dev->iommu_group, vdev->vdev.dev);
 }
 EXPORT_SYMBOL_GPL(vfio_platform_remove_common);
diff --git a/drivers/vfio/vfio.c b/drivers/vfio/vfio.c
index 02cc51ce6891a9..cc375df0fd5dda 100644
--- a/drivers/vfio/vfio.c
+++ b/drivers/vfio/vfio.c
@@ -749,6 +749,11 @@  void vfio_init_group_dev(struct vfio_device *device, struct device *dev,
 }
 EXPORT_SYMBOL_GPL(vfio_init_group_dev);
 
+void vfio_uninit_group_dev(struct vfio_device *device)
+{
+}
+EXPORT_SYMBOL_GPL(vfio_uninit_group_dev);
+
 int vfio_register_group_dev(struct vfio_device *device)
 {
 	struct vfio_device *existing_device;
diff --git a/include/linux/vfio.h b/include/linux/vfio.h
index a2c5b30e1763ba..b0875cf8e496bb 100644
--- a/include/linux/vfio.h
+++ b/include/linux/vfio.h
@@ -61,6 +61,7 @@  extern void vfio_iommu_group_put(struct iommu_group *group, struct device *dev);
 
 void vfio_init_group_dev(struct vfio_device *device, struct device *dev,
 			 const struct vfio_device_ops *ops);
+void vfio_uninit_group_dev(struct vfio_device *device);
 int vfio_register_group_dev(struct vfio_device *device);
 void vfio_unregister_group_dev(struct vfio_device *device);
 extern struct vfio_device *vfio_device_get_from_dev(struct device *dev);
diff --git a/samples/vfio-mdev/mbochs.c b/samples/vfio-mdev/mbochs.c
index e81b875b4d87b4..cf264d0bf11053 100644
--- a/samples/vfio-mdev/mbochs.c
+++ b/samples/vfio-mdev/mbochs.c
@@ -558,6 +558,7 @@  static int mbochs_probe(struct mdev_device *mdev)
 	return 0;
 
 err_mem:
+	vfio_uninit_group_dev(&mdev_state->vdev);
 	kfree(mdev_state->vconfig);
 	kfree(mdev_state);
 	return ret;
@@ -571,6 +572,7 @@  static void mbochs_remove(struct mdev_device *mdev)
 	vfio_unregister_group_dev(&mdev_state->vdev);
 	kfree(mdev_state->pages);
 	kfree(mdev_state->vconfig);
+	vfio_uninit_group_dev(&mdev_state->vdev);
 	kfree(mdev_state);
 }
 
diff --git a/samples/vfio-mdev/mdpy.c b/samples/vfio-mdev/mdpy.c
index a7d4ed28d66411..57334034cde6dd 100644
--- a/samples/vfio-mdev/mdpy.c
+++ b/samples/vfio-mdev/mdpy.c
@@ -235,17 +235,16 @@  static int mdpy_probe(struct mdev_device *mdev)
 
 	mdev_state->vconfig = kzalloc(MDPY_CONFIG_SPACE_SIZE, GFP_KERNEL);
 	if (mdev_state->vconfig == NULL) {
-		kfree(mdev_state);
-		return -ENOMEM;
+		ret = -ENOMEM;
+		goto err_state;
 	}
 
 	fbsize = roundup_pow_of_two(type->width * type->height * type->bytepp);
 
 	mdev_state->memblk = vmalloc_user(fbsize);
 	if (!mdev_state->memblk) {
-		kfree(mdev_state->vconfig);
-		kfree(mdev_state);
-		return -ENOMEM;
+		ret = -ENOMEM;
+		goto err_vconfig;
 	}
 	dev_info(dev, "%s: %s (%dx%d)\n", __func__, type->name, type->width,
 		 type->height);
@@ -260,13 +259,18 @@  static int mdpy_probe(struct mdev_device *mdev)
 	mdpy_count++;
 
 	ret = vfio_register_group_dev(&mdev_state->vdev);
-	if (ret) {
-		kfree(mdev_state->vconfig);
-		kfree(mdev_state);
-		return ret;
-	}
+	if (ret)
+		goto err_mem;
 	dev_set_drvdata(&mdev->dev, mdev_state);
 	return 0;
+err_mem:
+	vfree(mdev_state->memblk);
+err_vconfig:
+	kfree(mdev_state->vconfig);
+err_state:
+	vfio_uninit_group_dev(&mdev_state->vdev);
+	kfree(mdev_state);
+	return ret;
 }
 
 static void mdpy_remove(struct mdev_device *mdev)
@@ -278,6 +282,7 @@  static void mdpy_remove(struct mdev_device *mdev)
 	vfio_unregister_group_dev(&mdev_state->vdev);
 	vfree(mdev_state->memblk);
 	kfree(mdev_state->vconfig);
+	vfio_uninit_group_dev(&mdev_state->vdev);
 	kfree(mdev_state);
 
 	mdpy_count--;
diff --git a/samples/vfio-mdev/mtty.c b/samples/vfio-mdev/mtty.c
index 8b26fecc4afedd..37cc9067e1601d 100644
--- a/samples/vfio-mdev/mtty.c
+++ b/samples/vfio-mdev/mtty.c
@@ -718,8 +718,8 @@  static int mtty_probe(struct mdev_device *mdev)
 
 	mdev_state = kzalloc(sizeof(struct mdev_state), GFP_KERNEL);
 	if (mdev_state == NULL) {
-		atomic_add(nr_ports, &mdev_avail_ports);
-		return -ENOMEM;
+		ret = -ENOMEM;
+		goto err_nr_ports;
 	}
 
 	vfio_init_group_dev(&mdev_state->vdev, &mdev->dev, &mtty_dev_ops);
@@ -732,9 +732,8 @@  static int mtty_probe(struct mdev_device *mdev)
 	mdev_state->vconfig = kzalloc(MTTY_CONFIG_SPACE_SIZE, GFP_KERNEL);
 
 	if (mdev_state->vconfig == NULL) {
-		kfree(mdev_state);
-		atomic_add(nr_ports, &mdev_avail_ports);
-		return -ENOMEM;
+		ret = -ENOMEM;
+		goto err_state;
 	}
 
 	mutex_init(&mdev_state->ops_lock);
@@ -743,14 +742,19 @@  static int mtty_probe(struct mdev_device *mdev)
 	mtty_create_config_space(mdev_state);
 
 	ret = vfio_register_group_dev(&mdev_state->vdev);
-	if (ret) {
-		kfree(mdev_state);
-		atomic_add(nr_ports, &mdev_avail_ports);
-		return ret;
-	}
-
+	if (ret)
+		goto err_vconfig;
 	dev_set_drvdata(&mdev->dev, mdev_state);
 	return 0;
+
+err_vconfig:
+	kfree(mdev_state->vconfig);
+err_state:
+	vfio_uninit_group_dev(&mdev_state->vdev);
+	kfree(mdev_state);
+err_nr_ports:
+	atomic_add(nr_ports, &mdev_avail_ports);
+	return ret;
 }
 
 static void mtty_remove(struct mdev_device *mdev)
@@ -761,6 +765,7 @@  static void mtty_remove(struct mdev_device *mdev)
 	vfio_unregister_group_dev(&mdev_state->vdev);
 
 	kfree(mdev_state->vconfig);
+	vfio_uninit_group_dev(&mdev_state->vdev);
 	kfree(mdev_state);
 	atomic_add(nr_ports, &mdev_avail_ports);
 }