[PATCH] nvme-tcp: unregister lockdep keys at socket destruction

[Shin'ichiro Kawasaki]

Commit 19bdb70c77d3 ("nvme-tcp: lockdep: use dynamic lockdep keys per
socket instance") introduced dynamic lockdep keys for nvme-tcp socket
instances. The lockdep keys are unregistered in nvme_tcp_free_queue(),
just after __fput_sync(queue->sock->file) call. However, at this point
still in-flight skbs are there. When the skbs are freed, the socket and
the unregistered lockdep keys can be referenced, which resutls in WARNs
[1].

To avoid the WARN, keep the lockdep keys alive until the socket is
destroyed. Allocate the keys separately from struct nvme_tcp_queue and
replace the socket's sk_destruct callback with an NVMe/TCP wrapper. The
wrapper invokes the original destructor and queues work to unregister
and free the keys in process context, since socket destruction can run
in softirq context.

Hold an explicit module reference until the deferred work completes so
that both the destructor and work callback remain valid.

[1] https://lore.kernel.org/netdev/CANn89i+wnTLC==UnXCpjsS4YxvEfhe5oK0N7fttbqr1zKyqdug@xxxxxxxxxxxxxx/

Fixes: 19bdb70c77d3 ("nvme-tcp: lockdep: use dynamic lockdep keys per socket instance")
Signed-off-by: Shin'ichiro Kawasaki <shinichiro.kawasaki@xxxxxxx>
---
 drivers/nvme/host/tcp.c | 117 +++++++++++++++++++++++++++++-----------
 1 file changed, 86 insertions(+), 31 deletions(-)

diff --git a/drivers/nvme/host/tcp.c b/drivers/nvme/host/tcp.c
index 5fda9661bdb7..645913edf1f4 100644
--- a/drivers/nvme/host/tcp.c
+++ b/drivers/nvme/host/tcp.c
@@ -144,11 +144,6 @@ struct nvme_tcp_queue {
 	void (*state_change)(struct sock *);
 	void (*data_ready)(struct sock *);
 	void (*write_space)(struct sock *);
-
-#ifdef CONFIG_DEBUG_LOCK_ALLOC
-	struct lock_class_key nvme_tcp_sk_key;
-	struct lock_class_key nvme_tcp_slock_key;
-#endif
 };
 
 static DEFINE_MUTEX(nvme_tcp_ctrl_mutex);
@@ -179,35 +174,93 @@ static const struct blk_mq_ops nvme_tcp_admin_mq_ops;
 static int nvme_tcp_try_send(struct nvme_tcp_queue *queue);
 
 #ifdef CONFIG_DEBUG_LOCK_ALLOC
+struct nvme_tcp_lockdep_keys {
+	struct lock_class_key sk_key;
+	struct lock_class_key slock_key;
+	void (*sk_destruct)(struct sock *sk);
+	struct work_struct free_work;
+};
+
+static inline struct nvme_tcp_lockdep_keys *nvme_tcp_sock_to_lockdep_keys(struct sock *sk)
+{
+	struct nvme_tcp_lockdep_keys *keys = container_of(
+		sk->sk_lock.dep_map.key, struct nvme_tcp_lockdep_keys, sk_key);
+
+	return keys;
+}
+
+static void nvme_tcp_free_lockdep_keys(struct work_struct *work)
+{
+	struct nvme_tcp_lockdep_keys *keys = container_of(work,
+			struct nvme_tcp_lockdep_keys, free_work);
+
+	lockdep_unregister_key(&keys->sk_key);
+	lockdep_unregister_key(&keys->slock_key);
+	kfree(keys);
+	module_put(THIS_MODULE);
+}
+
+static void nvme_tcp_sk_destruct(struct sock *sk)
+{
+	struct nvme_tcp_lockdep_keys *keys = nvme_tcp_sock_to_lockdep_keys(sk);
+
+	if (keys->sk_destruct)
+		keys->sk_destruct(sk);
+
+	/*
+	 * sk_destruct may run in softirq context. Do cleanup in process
+	 * context.
+	 */
+	queue_work(nvme_tcp_wq, &keys->free_work);
+}
+
+static void nvme_tcp_set_sk_destruct(struct sock *sk)
+{
+	struct nvme_tcp_lockdep_keys *keys = nvme_tcp_sock_to_lockdep_keys(sk);
+
+	keys->sk_destruct = sk->sk_destruct;
+
+	/* keep nvme_tcp loaded until the lockdep key cleanup work completes */
+	__module_get(THIS_MODULE);
+	sk->sk_destruct = nvme_tcp_sk_destruct;
+}
+
 /* lockdep can detect a circular dependency of the form
  *   sk_lock -> mmap_lock (page fault) -> fs locks -> sk_lock
  * because dependencies are tracked for both nvme-tcp and user contexts. Using
  * a separate class prevents lockdep from conflating nvme-tcp socket use with
  * user-space socket API use.
  */
-static void nvme_tcp_reclassify_socket(struct nvme_tcp_queue *queue)
+static int nvme_tcp_reclassify_socket(struct nvme_tcp_queue *queue)
 {
+	struct nvme_tcp_lockdep_keys *keys;
 	struct sock *sk = queue->sock->sk;
 
 	if (WARN_ON_ONCE(!sock_allow_reclassification(sk)))
-		return;
+		return -EINVAL;
+	if (WARN_ON_ONCE(sk->sk_family != AF_INET && sk->sk_family != AF_INET6))
+		return -EAFNOSUPPORT;
+
+	keys = kzalloc_obj(*keys);
+	if (!keys)
+		return -ENOMEM;
+
+	lockdep_register_key(&keys->sk_key);
+	lockdep_register_key(&keys->slock_key);
+	INIT_WORK(&keys->free_work, nvme_tcp_free_lockdep_keys);
 
-	switch (sk->sk_family) {
-	case AF_INET:
+	if (sk->sk_family == AF_INET)
 		sock_lock_init_class_and_name(sk, "slock-AF_INET-NVME",
-					      &queue->nvme_tcp_slock_key,
+					      &keys->slock_key,
 					      "sk_lock-AF_INET-NVME",
-					      &queue->nvme_tcp_sk_key);
-		break;
-	case AF_INET6:
+					      &keys->sk_key);
+	else
 		sock_lock_init_class_and_name(sk, "slock-AF_INET6-NVME",
-					      &queue->nvme_tcp_slock_key,
+					      &keys->slock_key,
 					      "sk_lock-AF_INET6-NVME",
-					      &queue->nvme_tcp_sk_key);
-		break;
-	default:
-		WARN_ON_ONCE(1);
-	}
+					      &keys->sk_key);
+
+	return 0;
 }
 #endif
 
@@ -1511,11 +1564,6 @@ static void nvme_tcp_free_queue(struct nvme_ctrl *nctrl, int qid)
 	mutex_destroy(&queue->send_mutex);
 	mutex_destroy(&queue->queue_lock);
 	mutex_destroy(&queue->pf_cache_lock);
-
-#ifdef CONFIG_DEBUG_LOCK_ALLOC
-	lockdep_unregister_key(&queue->nvme_tcp_sk_key);
-	lockdep_unregister_key(&queue->nvme_tcp_slock_key);
-#endif
 }
 
 static int nvme_tcp_init_connection(struct nvme_tcp_queue *queue)
@@ -1831,6 +1879,9 @@ static int nvme_tcp_alloc_queue(struct nvme_ctrl *nctrl, int qid,
 	struct nvme_tcp_queue *queue = &ctrl->queues[qid];
 	int ret, rcv_pdu_size;
 	struct file *sock_file;
+#ifdef CONFIG_DEBUG_LOCK_ALLOC
+	bool reclassified = false;
+#endif
 
 	mutex_init(&queue->queue_lock);
 	queue->ctrl = ctrl;
@@ -1864,9 +1915,10 @@ static int nvme_tcp_alloc_queue(struct nvme_ctrl *nctrl, int qid,
 	sk_net_refcnt_upgrade(queue->sock->sk);
 
 #ifdef CONFIG_DEBUG_LOCK_ALLOC
-	lockdep_register_key(&queue->nvme_tcp_sk_key);
-	lockdep_register_key(&queue->nvme_tcp_slock_key);
-	nvme_tcp_reclassify_socket(queue);
+	ret = nvme_tcp_reclassify_socket(queue);
+	if (ret)
+		goto err_sock;
+	reclassified = true;
 #endif
 
 	/* Single syn retry */
@@ -1960,6 +2012,9 @@ static int nvme_tcp_alloc_queue(struct nvme_ctrl *nctrl, int qid,
 	if (ret)
 		goto err_init_connect;
 
+#ifdef CONFIG_DEBUG_LOCK_ALLOC
+	nvme_tcp_set_sk_destruct(queue->sock->sk);
+#endif
 	set_bit(NVME_TCP_Q_ALLOCATED, &queue->flags);
 
 	return 0;
@@ -1969,13 +2024,13 @@ static int nvme_tcp_alloc_queue(struct nvme_ctrl *nctrl, int qid,
 err_rcv_pdu:
 	kfree(queue->pdu);
 err_sock:
+#ifdef CONFIG_DEBUG_LOCK_ALLOC
+	if (reclassified)
+		nvme_tcp_set_sk_destruct(queue->sock->sk);
+#endif
 	/* Use sync variant - see nvme_tcp_free_queue() for explanation */
 	__fput_sync(queue->sock->file);
 	queue->sock = NULL;
-#ifdef CONFIG_DEBUG_LOCK_ALLOC
-	lockdep_unregister_key(&queue->nvme_tcp_sk_key);
-	lockdep_unregister_key(&queue->nvme_tcp_slock_key);
-#endif
 err_destroy_mutex:
 	mutex_destroy(&queue->send_mutex);
 	mutex_destroy(&queue->queue_lock);
-- 
2.54.0


--btsgphlxdv6zmmxj--