[PATCH] nvme: Add module reference counting for multipath nvme device

From: Wen Xiong <[email protected]>

Ensure proper module reference counting for NVMe multipath head devices
on open, preventing the controller module from being unloaded while
the multipath head device is still open or in use.

This patch acquires a module reference in nvme_ns_head_open() and
releases it in nvme_ns_head_release()

Signed-off-by: Wen Xiong <[email protected]>
---
 drivers/nvme/host/multipath.c | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/drivers/nvme/host/multipath.c b/drivers/nvme/host/multipath.c
index 9b9a657fa330..74c3710bf5f7 100644
--- a/drivers/nvme/host/multipath.c
+++ b/drivers/nvme/host/multipath.c
@@ -554,14 +554,39 @@ static void nvme_ns_head_submit_bio(struct bio *bio)
 
 static int nvme_ns_head_open(struct gendisk *disk, blk_mode_t mode)
 {
-	if (!nvme_tryget_ns_head(disk->private_data))
+	struct nvme_ns_head *head = disk->private_data;
+	struct nvme_ns *ns;
+	int srcu_idx;
+
+	if (!nvme_tryget_ns_head(head))
 		return -ENXIO;
+
+	/* Get module reference from any available path */
+	srcu_idx = srcu_read_lock(&head->srcu);
+	ns = nvme_find_path(head);
+	if (ns && !try_module_get(ns->ctrl->ops->module)) {
+		srcu_read_unlock(&head->srcu, srcu_idx);
+		nvme_put_ns_head(head);
+		return -ENXIO;
+	}
+	srcu_read_unlock(&head->srcu, srcu_idx);
+
 	return 0;
 }
 
 static void nvme_ns_head_release(struct gendisk *disk)
 {
-	nvme_put_ns_head(disk->private_data);
+	struct nvme_ns_head *head = disk->private_data;
+	struct nvme_ns *ns;
+	int srcu_idx;
+
+	srcu_idx = srcu_read_lock(&head->srcu);
+	ns = nvme_find_path(head);
+	if (ns)
+		module_put(ns->ctrl->ops->module);
+	srcu_read_unlock(&head->srcu, srcu_idx);
+
+	nvme_put_ns_head(head);
 }
 
 static int nvme_ns_head_get_unique_id(struct gendisk *disk, u8 id[16],
-- 
2.52.0