We will utilize this new kfunc bpf_mm_get_task() to retrieve the
associated task_struct from the given @mm. The obtained task_struct must
be released by calling bpf_task_release() as a paired operation.
Signed-off-by: Yafang Shao <laoar.shao@xxxxxxxxx>
---
mm/bpf_thp.c | 34 ++++++++++++++++++++++++++++++++++
1 file changed, 34 insertions(+)
diff --git a/mm/bpf_thp.c b/mm/bpf_thp.c
index b757e8f425fd..46b3bc96359e 100644
--- a/mm/bpf_thp.c
+++ b/mm/bpf_thp.c
@@ -205,11 +205,45 @@ __bpf_kfunc void bpf_put_mem_cgroup(struct mem_cgroup *memcg)
#endif
}
+/**
+ * bpf_mm_get_task - Get the task struct associated with a mm_struct.
+ * @mm: The mm_struct to query
+ *
+ * The obtained task_struct must be released by calling bpf_task_release().
+ *
+ * Return: The associated task_struct on success, or NULL on failure. Note that
+ * this function depends on CONFIG_MEMCG being enabled - it will always return
+ * NULL if CONFIG_MEMCG is not configured.
+ */
+__bpf_kfunc struct task_struct *bpf_mm_get_task(struct mm_struct *mm)
+{
+#ifdef CONFIG_MEMCG
+ struct task_struct *task;
+
+ if (!mm)
+ return NULL;
+ rcu_read_lock();
+ task = rcu_dereference(mm->owner);
+ if (!task)
+ goto out;
+ if (!refcount_inc_not_zero(&task->rcu_users))
+ goto out;
+
+ rcu_read_unlock();
+ return task;
+
+out:
+ rcu_read_unlock();
+#endif
+ return NULL;
+}
+
__bpf_kfunc_end_defs();
BTF_KFUNCS_START(bpf_thp_ids)
BTF_ID_FLAGS(func, bpf_mm_get_mem_cgroup, KF_TRUSTED_ARGS | KF_ACQUIRE | KF_RET_NULL)
BTF_ID_FLAGS(func, bpf_put_mem_cgroup, KF_RELEASE)
+BTF_ID_FLAGS(func, bpf_mm_get_task, KF_TRUSTED_ARGS | KF_ACQUIRE | KF_RET_NULL)
BTF_KFUNCS_END(bpf_thp_ids)
static const struct btf_kfunc_id_set bpf_thp_set = {
--
2.47.3
