[bpf-next,1/2] bpf: add btf_task_get_cgroup_id kfunc

Commit Message

Jose Fernandez March 15, 2024, 6:28 p.m. UTC

This patch enhances the BPF helpers by adding a kfunc to retrieve the
cgroup v2 ID of a specific task, addressing a previous limitation where
only bpf_task_get_cgroup1 was available for cgroup v1. The new kfunc is
particularly useful for scenarios where obtaining the cgroup ID of a
task other than the "current" one is necessary, which the existing
bpf_get_current_cgroup_id helper cannot accommodate. A specific use case
at Netflix involved the sched_switch tracepoint, where we had to get
the cgroup IDs of both the previous and next tasks.

The bpf_task_get_cgroup_id kfunc returns a task's cgroup ID, correctly
implementing RCU read locking and unlocking for safe data access, and
leverages existing cgroup.h helpers to fetch the cgroup and its ID.

Signed-off-by: Jose Fernandez <[e-mail]>
Reviewed-by: Tycho Andersen <[e-mail]>
---
 kernel/bpf/helpers.c | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

Comments

Stanislav Fomichev March 15, 2024, 6:50 p.m. UTC | #1

On 03/15, Jose Fernandez wrote:
> This patch enhances the BPF helpers by adding a kfunc to retrieve the
> cgroup v2 ID of a specific task, addressing a previous limitation where
> only bpf_task_get_cgroup1 was available for cgroup v1. The new kfunc is
> particularly useful for scenarios where obtaining the cgroup ID of a
> task other than the "current" one is necessary, which the existing
> bpf_get_current_cgroup_id helper cannot accommodate. A specific use case
> at Netflix involved the sched_switch tracepoint, where we had to get
> the cgroup IDs of both the previous and next tasks.
> 
> The bpf_task_get_cgroup_id kfunc returns a task's cgroup ID, correctly
> implementing RCU read locking and unlocking for safe data access, and
> leverages existing cgroup.h helpers to fetch the cgroup and its ID.
> 
> Signed-off-by: Jose Fernandez <[e-mail]>
> Reviewed-by: Tycho Andersen <[e-mail]>
> ---
>  kernel/bpf/helpers.c | 22 ++++++++++++++++++++++
>  1 file changed, 22 insertions(+)
> 
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index a89587859571..8038b2bd3488 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -2266,6 +2266,27 @@ bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id)
>  		return NULL;
>  	return cgrp;
>  }
> +
> +/**
> + * bpf_task_get_cgroup_id - Get the cgroup ID of a task.
> + * @task: The target task
> + *
> + * This function returns the ID of the task's default cgroup, primarily
> + * designed for use with cgroup v2. In cgroup v1, the concept of default
> + * cgroup varies by subsystem, and while this function will work with
> + * cgroup v1, it's recommended to use bpf_task_get_cgroup1 instead.
> + */
> +__bpf_kfunc u64 bpf_task_get_cgroup_id(struct task_struct *task)
> +{
> +	struct cgroup *cgrp;
> +	u64 cgrp_id;
> +
> +	rcu_read_lock();
> +	cgrp = task_dfl_cgroup(task);
> +	cgrp_id = cgroup_id(cgrp);
> +	rcu_read_unlock();
> +	return cgrp_id;
> +}
>  #endif /* CONFIG_CGROUPS */
>  
>  /**
> @@ -2573,6 +2594,7 @@ BTF_ID_FLAGS(func, bpf_cgroup_ancestor, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
>  BTF_ID_FLAGS(func, bpf_cgroup_from_id, KF_ACQUIRE | KF_RET_NULL)
>  BTF_ID_FLAGS(func, bpf_task_under_cgroup, KF_RCU)
>  BTF_ID_FLAGS(func, bpf_task_get_cgroup1, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
> +BTF_ID_FLAGS(func, bpf_task_get_cgroup_id, KF_RCU)

Any reason we are not returning 'struct cgroup' pointer? That should
allow using other kfuncs that are all 'struct cgrop' based as well.

Jose Fernandez March 15, 2024, 10:42 p.m. UTC | #2

On 24/03/15 11:50AM, Stanislav Fomichev wrote:
> On 03/15, Jose Fernandez wrote:
> > This patch enhances the BPF helpers by adding a kfunc to retrieve the
> > cgroup v2 ID of a specific task, addressing a previous limitation where
> > only bpf_task_get_cgroup1 was available for cgroup v1. The new kfunc is
> > particularly useful for scenarios where obtaining the cgroup ID of a
> > task other than the "current" one is necessary, which the existing
> > bpf_get_current_cgroup_id helper cannot accommodate. A specific use case
> > at Netflix involved the sched_switch tracepoint, where we had to get
> > the cgroup IDs of both the previous and next tasks.
> > 
> > The bpf_task_get_cgroup_id kfunc returns a task's cgroup ID, correctly
> > implementing RCU read locking and unlocking for safe data access, and
> > leverages existing cgroup.h helpers to fetch the cgroup and its ID.
> > 
> > Signed-off-by: Jose Fernandez <[e-mail]>
> > Reviewed-by: Tycho Andersen <[e-mail]>
> > ---
> >  kernel/bpf/helpers.c | 22 ++++++++++++++++++++++
> >  1 file changed, 22 insertions(+)
> > 
> > diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> > index a89587859571..8038b2bd3488 100644
> > --- a/kernel/bpf/helpers.c
> > +++ b/kernel/bpf/helpers.c
> > @@ -2266,6 +2266,27 @@ bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id)
> >  		return NULL;
> >  	return cgrp;
> >  }
> > +
> > +/**
> > + * bpf_task_get_cgroup_id - Get the cgroup ID of a task.
> > + * @task: The target task
> > + *
> > + * This function returns the ID of the task's default cgroup, primarily
> > + * designed for use with cgroup v2. In cgroup v1, the concept of default
> > + * cgroup varies by subsystem, and while this function will work with
> > + * cgroup v1, it's recommended to use bpf_task_get_cgroup1 instead.
> > + */
> > +__bpf_kfunc u64 bpf_task_get_cgroup_id(struct task_struct *task)
> > +{
> > +	struct cgroup *cgrp;
> > +	u64 cgrp_id;
> > +
> > +	rcu_read_lock();
> > +	cgrp = task_dfl_cgroup(task);
> > +	cgrp_id = cgroup_id(cgrp);
> > +	rcu_read_unlock();
> > +	return cgrp_id;
> > +}
> >  #endif /* CONFIG_CGROUPS */
> >  
> >  /**
> > @@ -2573,6 +2594,7 @@ BTF_ID_FLAGS(func, bpf_cgroup_ancestor, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
> >  BTF_ID_FLAGS(func, bpf_cgroup_from_id, KF_ACQUIRE | KF_RET_NULL)
> >  BTF_ID_FLAGS(func, bpf_task_under_cgroup, KF_RCU)
> >  BTF_ID_FLAGS(func, bpf_task_get_cgroup1, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
> > +BTF_ID_FLAGS(func, bpf_task_get_cgroup_id, KF_RCU)
> 
> Any reason we are not returning 'struct cgroup' pointer? That should
> allow using other kfuncs that are all 'struct cgrop' based as well.

Returning the cgroup pointer would make this kfunc more flexible, agreed. 
My intention was to make the kfunc more user friendly by returning the cgroup ID, 
but I can see how it would be beneficial to have the cgroup pointer as well. 
I'll update the patch to return the cgroup pointer.

Patch

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index a89587859571..8038b2bd3488 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -2266,6 +2266,27 @@  bpf_task_get_cgroup1(struct task_struct *task, int hierarchy_id)
 		return NULL;
 	return cgrp;
 }
+
+/**
+ * bpf_task_get_cgroup_id - Get the cgroup ID of a task.
+ * @task: The target task
+ *
+ * This function returns the ID of the task's default cgroup, primarily
+ * designed for use with cgroup v2. In cgroup v1, the concept of default
+ * cgroup varies by subsystem, and while this function will work with
+ * cgroup v1, it's recommended to use bpf_task_get_cgroup1 instead.
+ */
+__bpf_kfunc u64 bpf_task_get_cgroup_id(struct task_struct *task)
+{
+	struct cgroup *cgrp;
+	u64 cgrp_id;
+
+	rcu_read_lock();
+	cgrp = task_dfl_cgroup(task);
+	cgrp_id = cgroup_id(cgrp);
+	rcu_read_unlock();
+	return cgrp_id;
+}
 #endif /* CONFIG_CGROUPS */
 
 /**
@@ -2573,6 +2594,7 @@  BTF_ID_FLAGS(func, bpf_cgroup_ancestor, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_cgroup_from_id, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_task_under_cgroup, KF_RCU)
 BTF_ID_FLAGS(func, bpf_task_get_cgroup1, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_task_get_cgroup_id, KF_RCU)
 #endif
 BTF_ID_FLAGS(func, bpf_task_from_pid, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_throw)