[PATCH v6 06/11] rust: pci: drop driver data before remove returns

[Simon Song <xinmengs@nvidia.com>]

From: Zhi Wang <zhiw@nvidia.com>

A Rust PCI driver's private data is currently dropped by
driver::Registration::post_unbind_callback(). The driver core calls it
from device_unbind_cleanup(), after the PCI bus remove callback,
pci_device_remove(), has returned.

VF drivers will store a vfio::pci::Registration in their private data,
which calls vfio_pci_core_register_device() when created and
vfio_pci_core_unregister_device() when destroyed.

vfio_pci_core_unregister_device() expects to be called from
pci_device_remove(), and hence can't be deferred beyond it. This also
makes it incompatible with devres, i.e. a
devm_vfio_pci_core_register_device() helper wouldn't work either, since
devres_release_all() also runs from device_unbind_cleanup().

Hence, drop the private data in pci::Adapter::remove_callback(), so it
is dropped before pci_device_remove() returns.

This is a special case for VFIO-PCI. In general, a driver's private
data must not rely on being dropped from the bus remove callback; any
point after the driver's remove callback and before devres_release_all()
in device_unbind_cleanup() is sufficient.

Suggested-by: Danilo Krummrich <dakr@kernel.org>
Signed-off-by: Zhi Wang <zhiw@nvidia.com>
Signed-off-by: Simon Song <xinmengs@nvidia.com>
---
 rust/kernel/pci.rs | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/rust/kernel/pci.rs b/rust/kernel/pci.rs
index 9e31b9bb4d0d..75d764975a4e 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -149,6 +149,19 @@ extern "C" fn remove_callback(pdev: *mut bindings::pci_dev) {
         let data = unsafe { pdev.as_ref().drvdata_borrow::<T::Data<'_>>() };
 
         T::unbind(pdev, data);
+
+        // SAFETY: The driver's unbind callback has returned, and no callbacks retain a borrow.
+        let data = unsafe { pdev.as_ref().drvdata_obtain::<T::Data<'_>>() };
+
+        // Drop the private data here, ahead of `driver::Registration::post_unbind_callback()`,
+        // which would otherwise drop it after `pci_device_remove()` returns. `drvdata_obtain()`
+        // has cleared the driver data pointer, so that callback finds nothing left to drop.
+        //
+        // This is a special case for VFIO-PCI: `vfio_pci_core_unregister_device()` must be
+        // called from `pci_device_remove()`. In general, private data must not rely on being
+        // dropped from the bus remove callback; any point after the driver's remove callback and
+        // before `devres_release_all()` in `device_unbind_cleanup()` is sufficient.
+        drop(data);
     }
 }
 
-- 
2.43.7