RDMA/mlx5: Split ODP mkey search logic

Split the search for the ODP mkey when handling an rdma type page fault to
a helper function, later to be used in other page fault types.
Signed-off-by: Michael Guralnik <michaelgur@nvidia.com>
Link: https://patch.msgid.link/20240909100504.29797-6-michaelgur@nvidia.comSigned-off-by: Leon Romanovsky <leon@kernel.org>

drivers/infiniband/hw/mlx5/odp.c

@@ -819,6 +819,27 @@ static bool mkey_is_eq(struct mlx5_ib_mkey *mmkey, u32 key)
 	return mmkey->key == key;
 }
 
+static struct mlx5_ib_mkey *find_odp_mkey(struct mlx5_ib_dev *dev, u32 key)
+{
+	struct mlx5_ib_mkey *mmkey;
+
+	xa_lock(&dev->odp_mkeys);
+	mmkey = xa_load(&dev->odp_mkeys, mlx5_base_mkey(key));
+	if (!mmkey) {
+		mmkey = ERR_PTR(-ENOENT);
+		goto out;
+	}
+	if (!mkey_is_eq(mmkey, key)) {
+		mmkey = ERR_PTR(-EFAULT);
+		goto out;
+	}
+	refcount_inc(&mmkey->usecount);
+out:
+	xa_unlock(&dev->odp_mkeys);
+
+	return mmkey;
+}
+
 /*
  * Handle a single data segment in a page-fault WQE or RDMA region.
  *
@@ -846,12 +867,11 @@ static int pagefault_single_data_segment(struct mlx5_ib_dev *dev,
 
 	io_virt += *bytes_committed;
 	bcnt -= *bytes_committed;
-
 next_mr:
-	xa_lock(&dev->odp_mkeys);
-	mmkey = xa_load(&dev->odp_mkeys, mlx5_base_mkey(key));
-	if (!mmkey) {
-		xa_unlock(&dev->odp_mkeys);
+	mmkey = find_odp_mkey(dev, key);
+	if (IS_ERR(mmkey)) {
+		ret = PTR_ERR(mmkey);
+		if (ret == -ENOENT) {
 			mlx5_ib_dbg(
 				dev,
 				"skipping non ODP MR (lkey=0x%06x) in page fault handler.\n",
@@ -859,19 +879,12 @@ static int pagefault_single_data_segment(struct mlx5_ib_dev *dev,
 			if (bytes_mapped)
 				*bytes_mapped += bcnt;
 			/*
-		 * The user could specify a SGL with multiple lkeys and only
-		 * some of them are ODP. Treat the non-ODP ones as fully
-		 * faulted.
+			 * The user could specify a SGL with multiple lkeys and
+			 * only some of them are ODP. Treat the non-ODP ones as
+			 * fully faulted.
 			 */
 			ret = 0;
-		goto end;
 		}
-	refcount_inc(&mmkey->usecount);
-	xa_unlock(&dev->odp_mkeys);
-
-	if (!mkey_is_eq(mmkey, key)) {
-		mlx5_ib_dbg(dev, "failed to find mkey %x\n", key);
-		ret = -EFAULT;
 		goto end;
 	}
 
@@ -966,7 +979,7 @@ static int pagefault_single_data_segment(struct mlx5_ib_dev *dev,
 	}
 
 end:
-	if (mmkey)
+	if (!IS_ERR(mmkey))
 		mlx5r_deref_odp_mkey(mmkey);
 	while (head) {
 		frame = head;