[PATCH rdma-next v1 6/6] RDMA/mlx5: Annotate MR allocation

[Leon Romanovsky <leon@xxxxxxxxxx>]

From: Leon Romanovsky <leonro@xxxxxxxxxxxx>

Use restrack ID ability to track MRs.

Signed-off-by: Leon Romanovsky <leonro@xxxxxxxxxxxx>
---
 drivers/infiniband/core/verbs.c |  2 ++
 drivers/infiniband/hw/mlx5/mr.c | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/drivers/infiniband/core/verbs.c b/drivers/infiniband/core/verbs.c
index 0a85bb2c3408..fb4b0ac85c09 100644
--- a/drivers/infiniband/core/verbs.c
+++ b/drivers/infiniband/core/verbs.c
@@ -287,6 +287,8 @@ struct ib_pd *__ib_alloc_pd(struct ib_device *device, unsigned int flags,
 			ib_dealloc_pd(pd);
 			return ERR_CAST(mr);
 		}
+		rdma_restrack_set_task(&mr->res, caller);
+		rdma_restrack_kadd(&mr->res);
 
 		mr->device	= pd->device;
 		mr->pd		= pd;
diff --git a/drivers/infiniband/hw/mlx5/mr.c b/drivers/infiniband/hw/mlx5/mr.c
index 8f557f90ad7f..4d374846eaf6 100644
--- a/drivers/infiniband/hw/mlx5/mr.c
+++ b/drivers/infiniband/hw/mlx5/mr.c
@@ -198,6 +198,7 @@ static int add_keys(struct mlx5_ib_dev *dev, int c, int num)
 		mr->order = ent->order;
 		mr->allocated_from_cache = 1;
 		mr->dev = dev;
+		mr->ibmr.device = &dev->ib_dev;
 
 		MLX5_SET(mkc, mkc, free, 1);
 		MLX5_SET(mkc, mkc, umr_en, 1);
@@ -812,10 +813,15 @@ struct ib_mr *mlx5_ib_get_dma_mr(struct ib_pd *pd, int acc)
 	mr->mmkey.type = MLX5_MKEY_MR;
 	mr->ibmr.lkey = mr->mmkey.key;
 	mr->ibmr.rkey = mr->mmkey.key;
-	mr->umem = NULL;
+	mr->ibmr.device = pd->device;
+	rdma_rt_set_type(&mr->ibmr.res, RDMA_RESTRACK_MR);
+	err = rdma_restrack_add(&mr->ibmr.res);
+	if (err)
+		goto err_rt;
 
 	return &mr->ibmr;
 
+err_rt: mlx5_core_destroy_mkey(mdev, &mr->mmkey);
 err_in:
 	kfree(in);
 
@@ -1673,6 +1679,7 @@ static void dereg_mr(struct mlx5_ib_dev *dev, struct mlx5_ib_mr *mr)
 
 int mlx5_ib_dereg_mr(struct ib_mr *ibmr)
 {
+	rdma_restrack_del(&ibmr->res);
 	dereg_mr(to_mdev(ibmr->device), to_mmr(ibmr));
 	return 0;
 }
@@ -1693,6 +1700,13 @@ struct ib_mr *mlx5_ib_alloc_mr(struct ib_pd *pd,
 	if (!mr)
 		return ERR_PTR(-ENOMEM);
 
+	mr->ibmr.device = pd->device;
+	rdma_rt_set_type(&mr->ibmr.res, RDMA_RESTRACK_MR);
+	err = rdma_restrack_add(&mr->ibmr.res);
+	if (err)
+		goto err_rt;
+
+
 	in = kzalloc(inlen, GFP_KERNEL);
 	if (!in) {
 		err = -ENOMEM;
@@ -1789,6 +1803,8 @@ struct ib_mr *mlx5_ib_alloc_mr(struct ib_pd *pd,
 err_free_in:
 	kfree(in);
 err_free:
+	rdma_restrack_del(&mr->ibmr.res);
+err_rt:
 	kfree(mr);
 	return ERR_PTR(err);
 }
-- 
2.19.1