> diff --git a/mm/Makefile b/mm/Makefile
> index bf46fe31dc14..c2b887491f09 100644

[ ... ]

> @@ -107,8 +107,13 @@ ifdef CONFIG_SWAP
>  obj-$(CONFIG_MEMCG) += swap_cgroup.o
>  endif
>  ifdef CONFIG_BPF_SYSCALL
> +ifdef CONFIG_NUMA_BALANCING
>  obj-$(CONFIG_MEMCG) += bpf_memcontrol.o
>  endif
> +endif
> +ifdef CONFIG_BPF_SYSCALL
> +obj-$(CONFIG_MEMCG) += bpf_numa_balancing.o
> +endif

Should bpf_numa_balancing.o also require CONFIG_SHRINKER_DEBUG? The code
calls mem_cgroup_get_from_ino() which is only defined when
CONFIG_SHRINKER_DEBUG is enabled. Without it, compilation will fail due to
an undefined reference.

Also, should it require CONFIG_NUMA_BALANCING? Without that config option,
task_numab_enabled() returns false unconditionally, so the BPF hook would
never be invoked even if a program is attached.

> diff --git a/mm/bpf_numa_balancing.c b/mm/bpf_numa_balancing.c
> new file mode 100644
> index 000000000000..aac4eec7c6ba
> --- /dev/null
> +++ b/mm/bpf_numa_balancing.c

[ ... ]

> +static int bpf_numab_reg(void *kdata, struct bpf_link *link)
> +{
> +	struct bpf_numab_ops *ops = kdata;
> +	struct mem_cgroup *memcg;
> +	int err = 0;
> +
> +	/* Only the link mode is supported. */
> +	if (!link)
> +		return -EOPNOTSUPP;
> +
> +	/* Depends on CONFIG_SHRINKER_DEBUG */
> +	memcg = mem_cgroup_get_from_ino(ops->cgroup_id);
> +	if (!memcg || IS_ERR(memcg))
> +		return -ENOENT;
> +
> +	spin_lock(&numab_ops_lock);
> +	/* Each memory cgroup can have at most one attached BPF program to ensure
> +	 * exclusive control and avoid interference between different BPF policies.
> +	 */
> +	if (rcu_access_pointer(memcg->bpf_numab)) {
> +		err = -EBUSY;
> +		goto out;
> +	}
> +	rcu_assign_pointer(memcg->bpf_numab, ops);
> +	spin_unlock(&numab_ops_lock);
> +	static_branch_inc(&bpf_numab_enabled_key);
> +
> +out:
> +	mem_cgroup_put(memcg);
> +	return err;
> +}

When rcu_access_pointer(memcg->bpf_numab) is non-NULL, the code sets err
to -EBUSY and jumps to out:, but this bypasses the spin_unlock(). Is the
spinlock left held when returning -EBUSY?

> +static void bpf_numab_unreg(void *kdata, struct bpf_link *link)
> +{
> +	struct bpf_numab_ops *ops = kdata;
> +	struct mem_cgroup *memcg;
> +
> +	memcg = mem_cgroup_get_from_ino(ops->cgroup_id);
> +	if (!memcg)
> +		return;

mem_cgroup_get_from_ino() returns ERR_PTR() on failure, not NULL. If the
cgroup was deleted between reg and unreg, this function would get an
ERR_PTR value, the if (!memcg) check would pass, and then
rcu_access_pointer(memcg->bpf_numab) would dereference an invalid pointer.
Should this check be if (!memcg || IS_ERR(memcg)) like in bpf_numab_reg()?

Also, mem_cgroup_get_from_ino() takes a reference on the memcg via
cgroup_get_e_css(). Where is the corresponding mem_cgroup_put() call?
The function appears to leak the memcg reference on all paths.

> +	spin_lock(&numab_ops_lock);
> +	if (!rcu_access_pointer(memcg->bpf_numab)) {
> +		spin_unlock(&numab_ops_lock);
> +		return;
> +	}
> +	rcu_replace_pointer(memcg->bpf_numab, NULL, lockdep_is_held(&numab_ops_lock));
> +	spin_unlock(&numab_ops_lock);
> +	static_branch_dec(&bpf_numab_enabled_key);
> +	synchronize_rcu();
> +}
> +
> +static int bpf_numab_update(void *kdata, void *old_kdata, struct bpf_link *link)
> +{
> +	struct bpf_numab_ops *ops = kdata;
> +	struct mem_cgroup *memcg;
> +
> +	memcg = mem_cgroup_get_from_ino(ops->cgroup_id);
> +	if (!memcg)
> +		return -EINVAL;

Same two issues as bpf_numab_unreg(): the check should handle ERR_PTR()
values, and there appears to be no mem_cgroup_put() to release the
reference taken by mem_cgroup_get_from_ino().

> +	spin_lock(&numab_ops_lock);
> +	/* The update can proceed regardless of whether memcg->bpf_numab has been previously set. */
> +	rcu_replace_pointer(memcg->bpf_numab, ops, lockdep_is_held(&numab_ops_lock));
> +	spin_unlock(&numab_ops_lock);
> +	synchronize_rcu();
> +	return 0;
> +}


---
AI reviewed your patch. Please fix the bug or email reply why it's not a bug.
See: https://github.com/kernel-patches/vmtest/blob/master/ci/claude/README.md

CI run summary: https://github.com/kernel-patches/bpf/actions/runs/20956455529