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

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

[Jose Fernandez]

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 <josef@netflix.com>
Reviewed-by: Tycho Andersen <tycho@tycho.pizza>
---
 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)
 #endif
 BTF_ID_FLAGS(func, bpf_task_from_pid, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_throw)
-- 
2.40.1

[PATCH bpf-next 2/2] selftests/bpf: add selftest for btf_task_get_cgroup_id

[Jose Fernandez]

This patch adds a selftest for the `bpf_task_get_cgroup_id` kfunc. The
test focuses on the use case of obtaining the cgroup ID of the previous
task in a `sched_switch` tracepoint.

The selftest involves creating a test cgroup, attaching a BPF program
that utilizes the `bpf_task_get_cgroup_id` during a `sched_switch`
tracepoint, and validating that the obtained cgroup ID for the previous
task matches the expected cgroup ID.

task_get_cgroup_id:OK
Summary: 1/0 PASSED, 0 SKIPPED, 0 FAILED

Signed-off-by: Jose Fernandez <josef@netflix.com>
Reviewed-by: Tycho Andersen <tycho@tycho.pizza>
---
 .../bpf/prog_tests/task_get_cgroup_id.c       | 58 +++++++++++++++++++
 .../bpf/progs/test_task_get_cgroup_id.c       | 30 ++++++++++
 2 files changed, 88 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/task_get_cgroup_id.c
 create mode 100644 tools/testing/selftests/bpf/progs/test_task_get_cgroup_id.c

diff --git a/tools/testing/selftests/bpf/prog_tests/task_get_cgroup_id.c b/tools/testing/selftests/bpf/prog_tests/task_get_cgroup_id.c
new file mode 100644
index 000000000000..b8c4551195d3
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/task_get_cgroup_id.c
@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: GPL-2.0
+// Copyright 2024 Netflix, Inc.
+
+#include <test_progs.h>
+#include <cgroup_helpers.h>
+#include "test_task_get_cgroup_id.skel.h"
+#include <unistd.h>
+
+#define TEST_CGROUP "/test-task-get-cgroup-id/"
+
+void test_task_get_cgroup_id(void)
+{
+	struct test_task_get_cgroup_id *skel;
+	int err, fd;
+	pid_t pid;
+	__u64 cgroup_id, expected_cgroup_id;
+	const struct timespec req = {
+		.tv_sec = 1,
+		.tv_nsec = 0,
+	};
+
+	fd = test__join_cgroup(TEST_CGROUP);
+	if (!ASSERT_OK(fd < 0, "test_join_cgroup_TEST_CGROUP"))
+		return;
+
+	skel = test_task_get_cgroup_id__open();
+	if (!ASSERT_OK_PTR(skel, "test_task_get_cgroup_id__open"))
+		goto cleanup;
+
+	err = test_task_get_cgroup_id__load(skel);
+	if (!ASSERT_OK(err, "test_task_get_cgroup_id__load"))
+		goto cleanup;
+
+	err = test_task_get_cgroup_id__attach(skel);
+	if (!ASSERT_OK(err, "test_task_get_cgroup_id__attach"))
+		goto cleanup;
+
+	pid = getpid();
+	expected_cgroup_id = get_cgroup_id(TEST_CGROUP);
+	if (!ASSERT_GT(expected_cgroup_id, 0, "get_cgroup_id"))
+		goto cleanup;
+
+	/* Trigger nanosleep to enter the sched_switch tracepoint */
+	/* The previous task should be this process */
+	syscall(__NR_nanosleep, &req, NULL);
+
+	err = bpf_map_lookup_elem(bpf_map__fd(skel->maps.pid_to_cgid_map), &pid,
+				  &cgroup_id);
+
+	if (!ASSERT_OK(err, "bpf_map_lookup_elem"))
+		goto cleanup;
+
+	ASSERT_EQ(cgroup_id, expected_cgroup_id, "cgroup_id");
+
+cleanup:
+	test_task_get_cgroup_id__destroy(skel);
+	close(fd);
+}
diff --git a/tools/testing/selftests/bpf/progs/test_task_get_cgroup_id.c b/tools/testing/selftests/bpf/progs/test_task_get_cgroup_id.c
new file mode 100644
index 000000000000..7e6bc008970f
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/test_task_get_cgroup_id.c
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: GPL-2.0
+// Copyright 2024 Netflix, Inc.
+
+#include "vmlinux.h"
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+u64 bpf_task_get_cgroup_id(struct task_struct *task) __ksym;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__uint(max_entries, 4096);
+	__type(key, __u32);
+	__type(value, __u64);
+} pid_to_cgid_map SEC(".maps");
+
+SEC("tp_btf/sched_switch")
+int BPF_PROG(sched_switch, bool preempt, struct task_struct *prev,
+	     struct task_struct *next)
+{
+	u32 pid = prev->pid;
+	u64 cgroup_id;
+
+	cgroup_id = bpf_task_get_cgroup_id(prev);
+	bpf_map_update_elem(&pid_to_cgid_map, &pid, &cgroup_id, BPF_ANY);
+
+	return 0;
+}
+
+char _license[] SEC("license") = "GPL";
-- 
2.40.1

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

[Stanislav Fomichev]

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 <josef@netflix.com>
> Reviewed-by: Tycho Andersen <tycho@tycho.pizza>
> ---
>  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.

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

[Jose Fernandez]

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 <josef@netflix.com>
> > Reviewed-by: Tycho Andersen <tycho@tycho.pizza>
> > ---
> >  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.