* Copyright (C) 2019 Mellanox Technologies. All rights reserved
  */
 
+#include <linux/completion.h>
 #include <linux/device.h>
 #include <linux/idr.h>
 #include <linux/kernel.h>
 #include <linux/list.h>
 #include <linux/mutex.h>
+#include <linux/refcount.h>
 #include <linux/slab.h>
 #include <linux/sysfs.h>
 
 static LIST_HEAD(nsim_bus_dev_list);
 static DEFINE_MUTEX(nsim_bus_dev_list_lock);
 static bool nsim_bus_enable;
+static refcount_t nsim_bus_devs; /* Including the bus itself. */
+static DECLARE_COMPLETION(nsim_bus_devs_released);
 
 static struct nsim_bus_dev *to_nsim_bus_dev(struct device *dev)
 {
 
        nsim_bus_dev = container_of(dev, struct nsim_bus_dev, dev);
        kfree(nsim_bus_dev);
+       if (refcount_dec_and_test(&nsim_bus_devs))
+               complete(&nsim_bus_devs_released);
 }
 
 static struct device_type nsim_bus_dev_type = {
                goto err;
        }
 
+       refcount_inc(&nsim_bus_devs);
        /* Allow using nsim_bus_dev */
        smp_store_release(&nsim_bus_dev->init, true);
 
        err = driver_register(&nsim_driver);
        if (err)
                goto err_bus_unregister;
+       refcount_set(&nsim_bus_devs, 1);
        /* Allow using resources */
        smp_store_release(&nsim_bus_enable, true);
        return 0;
 
        /* Disallow using resources */
        smp_store_release(&nsim_bus_enable, false);
+       if (refcount_dec_and_test(&nsim_bus_devs))
+               complete(&nsim_bus_devs_released);
 
        mutex_lock(&nsim_bus_dev_list_lock);
        list_for_each_entry_safe(nsim_bus_dev, tmp, &nsim_bus_dev_list, list) {
        }
        mutex_unlock(&nsim_bus_dev_list_lock);
 
+       wait_for_completion(&nsim_bus_devs_released);
+
        driver_unregister(&nsim_driver);
        bus_unregister(&nsim_bus);
 }