[PATCH v6 13/17] iommu/arm-smmu-v3: Add streams_lock for atomic-context SID->master lookup

[Nicolin Chen <nicolinc@nvidia.com>]

A subsequent change will look up arm_smmu_master entries by SID from inside
arm_smmu_cmdq_batch_retry(), which runs with invs->rwlock read_lock held in
IRQ-disabled context, and so cannot take the sleeping streams_mutex.

Add a spinlock_t streams_lock that protects rb_root mutations alongside the
existing streams_mutex:
 - atomic-context readers will hold the spinlock alone
 - writers (insert/remove paths) take both

A reader under the streams_lock uses all the streams of the found master,
so the insertion has to be all-or-nothing: make arm_smmu_insert_master()
initialize all the L2 strtabs first and then insert all the stream nodes in
one critical section, making a master found via any single SID always fully
initialized.

Update the lockdep assertion in arm_smmu_find_master() to accept either of
the locks so the helper is callable from both contexts.

Suggested-by: Jason Gunthorpe <jgg@nvidia.com>
Assisted-by: LLM
Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>
---
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h |  2 +
 drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c | 95 ++++++++++++++-------
 2 files changed, 65 insertions(+), 32 deletions(-)

diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
index ca8f52142b07f..2f8d22cc404ac 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h
@@ -972,6 +972,8 @@ struct arm_smmu_device {
 
 	struct rb_root			streams;
 	struct mutex			streams_mutex;
+	/* Held during rb_root updates; allows atomic-context lookups */
+	spinlock_t			streams_lock;
 };
 
 struct arm_smmu_stream {
diff --git a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
index c3c31d4010a0a..6f217e916ab0c 100644
--- a/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c
@@ -2126,7 +2126,8 @@ arm_smmu_find_master(struct arm_smmu_device *smmu, u32 sid)
 {
 	struct rb_node *node;
 
-	lockdep_assert_held(&smmu->streams_mutex);
+	lockdep_assert(lockdep_is_held(&smmu->streams_mutex) ||
+		       lockdep_is_held(&smmu->streams_lock));
 
 	node = rb_find(&sid, &smmu->streams, arm_smmu_streams_cmp_key);
 	if (!node)
@@ -4191,6 +4192,51 @@ static int arm_smmu_stream_id_cmp(const void *_l, const void *_r)
 	return cmp_int(*l, *r);
 }
 
+/* Caller must hold the streams_mutex. Publishes all the nodes, or none */
+static int arm_smmu_insert_streams(struct arm_smmu_device *smmu,
+				   struct arm_smmu_master *master)
+{
+	struct arm_smmu_master *existing_master = NULL;
+	u32 existing_sid = 0;
+	unsigned long flags;
+	int ret = 0;
+	int i;
+
+	spin_lock_irqsave(&smmu->streams_lock, flags);
+	for (i = 0; i < master->num_streams; i++) {
+		struct rb_node *existing;
+
+		existing = rb_find_add(&master->streams[i].node,
+				       &smmu->streams,
+				       arm_smmu_streams_cmp_node);
+		if (!existing)
+			continue;
+
+		existing_master = rb_entry(existing, struct arm_smmu_stream,
+					   node)->master;
+
+		/* Bridged PCI devices may end up with duplicated IDs */
+		if (existing_master == master)
+			continue;
+
+		existing_sid = master->streams[i].id;
+		ret = -ENODEV;
+		break;
+	}
+	if (ret)
+		for (i--; i >= 0; i--)
+			if (!RB_EMPTY_NODE(&master->streams[i].node))
+				rb_erase(&master->streams[i].node,
+					 &smmu->streams);
+	spin_unlock_irqrestore(&smmu->streams_lock, flags);
+
+	if (ret)
+		dev_warn(master->dev,
+			 "Aliasing StreamID 0x%x (from %s) unsupported, expect DMA to be broken\n",
+			 existing_sid, dev_name(existing_master->dev));
+	return ret;
+}
+
 static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
 				  struct arm_smmu_master *master)
 {
@@ -4237,40 +4283,22 @@ static int arm_smmu_insert_master(struct arm_smmu_device *smmu,
 		RB_CLEAR_NODE(&master->streams[i].node);
 
 	mutex_lock(&smmu->streams_mutex);
-	for (i = 0; i < fwspec->num_ids; i++) {
-		struct arm_smmu_stream *new_stream = &master->streams[i];
-		struct rb_node *existing;
-		u32 sid = new_stream->id;
 
-		ret = arm_smmu_init_sid_strtab(smmu, sid);
+	/*
+	 * Initialize the L2 strtabs before publishing any stream node, and
+	 * insert all the nodes in one critical section, so an atomic reader
+	 * never sees a partially initialized master.
+	 */
+	for (i = 0; i < fwspec->num_ids; i++) {
+		ret = arm_smmu_init_sid_strtab(smmu, master->streams[i].id);
 		if (ret)
 			break;
-
-		/* Insert into SID tree */
-		existing = rb_find_add(&new_stream->node, &smmu->streams,
-				       arm_smmu_streams_cmp_node);
-		if (existing) {
-			struct arm_smmu_master *existing_master =
-				rb_entry(existing, struct arm_smmu_stream, node)
-					->master;
-
-			/* Bridged PCI devices may end up with duplicated IDs */
-			if (existing_master == master)
-				continue;
-
-			dev_warn(master->dev,
-				 "Aliasing StreamID 0x%x (from %s) unsupported, expect DMA to be broken\n",
-				 sid, dev_name(existing_master->dev));
-			ret = -ENODEV;
-			break;
-		}
 	}
 
+	if (!ret)
+		ret = arm_smmu_insert_streams(smmu, master);
+
 	if (ret) {
-		for (i--; i >= 0; i--)
-			if (!RB_EMPTY_NODE(&master->streams[i].node))
-				rb_erase(&master->streams[i].node,
-					 &smmu->streams);
 		kfree(master->streams);
 		kfree(master->build_invs);
 	}
@@ -4289,9 +4317,11 @@ static void arm_smmu_remove_master(struct arm_smmu_master *master)
 		return;
 
 	mutex_lock(&smmu->streams_mutex);
-	for (i = 0; i < fwspec->num_ids; i++)
-		if (!RB_EMPTY_NODE(&master->streams[i].node))
-			rb_erase(&master->streams[i].node, &smmu->streams);
+	scoped_guard(spinlock_irqsave, &smmu->streams_lock)
+		for (i = 0; i < fwspec->num_ids; i++)
+			if (!RB_EMPTY_NODE(&master->streams[i].node))
+				rb_erase(&master->streams[i].node,
+					 &smmu->streams);
 	mutex_unlock(&smmu->streams_mutex);
 
 	kfree(master->streams);
@@ -4726,6 +4756,7 @@ static int arm_smmu_init_structures(struct arm_smmu_device *smmu)
 	int ret;
 
 	mutex_init(&smmu->streams_mutex);
+	spin_lock_init(&smmu->streams_lock);
 	smmu->streams = RB_ROOT;
 
 	ret = arm_smmu_init_queues(smmu);
-- 
2.43.0