[bpf,v4,5/7] bpf: Optimize the free of inner map

Commit Message

Hou Tao Nov. 30, 2023, 2:01 p.m. UTC

From: Hou Tao <[e-mail]>

When removing the inner map from the outer map, the inner map will be
freed after one RCU grace period and one RCU tasks trace grace
period, so it is certain that the bpf program, which may access the
inner map, has exited before the inner map is freed.

However there is no need to wait for one RCU tasks trace grace period if
the outer map is only accessed by non-sleepable program. So adding
sleepable_refcnt in bpf_map and increasing sleepable_refcnt when adding
the outer map into env->used_maps for sleepable program. Considering the
max number of bpf program is INT_MAX - 1, atomic_t instead of atomic64_t
is used for sleepable_refcnt. When removing the inner map from the outer
map, using sleepable_refcnt to decide whether or not a RCU tasks trace
grace period is needed before freeing the inner map.

Signed-off-by: Hou Tao <[e-mail]>
---
 include/linux/bpf.h     |  2 ++
 kernel/bpf/core.c       |  4 ++++
 kernel/bpf/map_in_map.c | 14 +++++++++-----
 kernel/bpf/syscall.c    |  8 ++++++++
 kernel/bpf/verifier.c   |  4 +++-
 5 files changed, 26 insertions(+), 6 deletions(-)

Comments

Alexei Starovoitov Dec. 3, 2023, 6:54 p.m. UTC | #1

On Thu, Nov 30, 2023 at 10:01:18PM +0800, Hou Tao wrote:
> From: Hou Tao <[e-mail]>
> 
> When removing the inner map from the outer map, the inner map will be
> freed after one RCU grace period and one RCU tasks trace grace
> period, so it is certain that the bpf program, which may access the
> inner map, has exited before the inner map is freed.
> 
> However there is no need to wait for one RCU tasks trace grace period if
> the outer map is only accessed by non-sleepable program. So adding
> sleepable_refcnt in bpf_map and increasing sleepable_refcnt when adding
> the outer map into env->used_maps for sleepable program. Considering the
> max number of bpf program is INT_MAX - 1, atomic_t instead of atomic64_t
> is used for sleepable_refcnt. When removing the inner map from the outer
> map, using sleepable_refcnt to decide whether or not a RCU tasks trace
> grace period is needed before freeing the inner map.

Optimizing too soon as usual?
I bet you saw that we use:
 atomic64_t refcnt
for bpf_map, but you probably didn't dig into git history and
missed commit 92117d8443bc ("bpf: fix refcnt overflow") ?

Hou Tao Dec. 4, 2023, 8:57 a.m. UTC | #2

Hi,

On 12/4/2023 2:54 AM, Alexei Starovoitov wrote:
> On Thu, Nov 30, 2023 at 10:01:18PM +0800, Hou Tao wrote:
>> From: Hou Tao <[e-mail]>
>>
>> When removing the inner map from the outer map, the inner map will be
>> freed after one RCU grace period and one RCU tasks trace grace
>> period, so it is certain that the bpf program, which may access the
>> inner map, has exited before the inner map is freed.
>>
>> However there is no need to wait for one RCU tasks trace grace period if
>> the outer map is only accessed by non-sleepable program. So adding
>> sleepable_refcnt in bpf_map and increasing sleepable_refcnt when adding
>> the outer map into env->used_maps for sleepable program. Considering the
>> max number of bpf program is INT_MAX - 1, atomic_t instead of atomic64_t
>> is used for sleepable_refcnt. When removing the inner map from the outer
>> map, using sleepable_refcnt to decide whether or not a RCU tasks trace
>> grace period is needed before freeing the inner map.
> Optimizing too soon as usual?
> I bet you saw that we use:
>  atomic64_t refcnt
> for bpf_map, but you probably didn't dig into git history and
> missed commit 92117d8443bc ("bpf: fix refcnt overflow") ?

Yes. I didn't think it thoroughly. Although the max number of bpf
program INT_MAX - 1, but the allocation of id happens after the increase
of sleepable_refcnt, so the overflow is still possible. Will fix it in v5.

Patch

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 15a6bb951b70..c84095677519 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -297,6 +297,8 @@  struct bpf_map {
 	bool bypass_spec_v1;
 	bool frozen; /* write-once; write-protected by freeze_mutex */
 	bool free_after_mult_rcu_gp;
+	bool free_after_rcu_gp;
+	atomic_t sleepable_refcnt;
 	s64 __percpu *elem_count;
 };
 
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index cd3afe57ece3..f9f86ab83748 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -2664,12 +2664,16 @@  void __bpf_free_used_maps(struct bpf_prog_aux *aux,
 			  struct bpf_map **used_maps, u32 len)
 {
 	struct bpf_map *map;
+	bool sleepable;
 	u32 i;
 
+	sleepable = aux->sleepable;
 	for (i = 0; i < len; i++) {
 		map = used_maps[i];
 		if (map->ops->map_poke_untrack)
 			map->ops->map_poke_untrack(map, aux);
+		if (sleepable)
+			atomic_dec(&map->sleepable_refcnt);
 		bpf_map_put(map);
 	}
 }
diff --git a/kernel/bpf/map_in_map.c b/kernel/bpf/map_in_map.c
index 3248ff5d8161..7a8fe233555f 100644
--- a/kernel/bpf/map_in_map.c
+++ b/kernel/bpf/map_in_map.c
@@ -131,12 +131,16 @@  void bpf_map_fd_put_ptr(struct bpf_map *map, void *ptr, bool need_defer)
 {
 	struct bpf_map *inner_map = ptr;
 
-	/* The inner map may still be used by both non-sleepable and sleepable
-	 * bpf program, so free it after one RCU grace period and one tasks
-	 * trace RCU grace period.
+	/* Defer the freeing of inner map according to the sleepable attribute
+	 * of bpf program which owns the outer map, so unnecessary waiting for
+	 * RCU tasks trace grace period can be avoided.
 	 */
-	if (need_defer)
-		WRITE_ONCE(inner_map->free_after_mult_rcu_gp, true);
+	if (need_defer) {
+		if (atomic_read(&map->sleepable_refcnt))
+			WRITE_ONCE(inner_map->free_after_mult_rcu_gp, true);
+		else
+			WRITE_ONCE(inner_map->free_after_rcu_gp, true);
+	}
 	bpf_map_put(inner_map);
 }
 
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index 88882cb58121..7880bb3daffb 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -750,8 +750,11 @@  void bpf_map_put(struct bpf_map *map)
 		bpf_map_free_id(map);
 		btf_put(map->btf);
 
+		WARN_ON_ONCE(atomic_read(&map->sleepable_refcnt));
 		if (READ_ONCE(map->free_after_mult_rcu_gp))
 			call_rcu_tasks_trace(&map->rcu, bpf_map_free_mult_rcu_gp);
+		else if (READ_ONCE(map->free_after_rcu_gp))
+			call_rcu(&map->rcu, bpf_map_free_rcu_gp);
 		else
 			bpf_map_free_in_work(map);
 	}
@@ -5343,6 +5346,11 @@  static int bpf_prog_bind_map(union bpf_attr *attr)
 		goto out_unlock;
 	}
 
+	/* The bpf program will not access the bpf map, but for the sake of
+	 * simplicity, increase sleepable_refcnt for sleepable program as well.
+	 */
+	if (prog->aux->sleepable)
+		atomic_inc(&map->sleepable_refcnt);
 	memcpy(used_maps_new, used_maps_old,
 	       sizeof(used_maps_old[0]) * prog->aux->used_map_cnt);
 	used_maps_new[prog->aux->used_map_cnt] = map;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 6da370a047fe..8fa28f969609 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -18051,10 +18051,12 @@  static int resolve_pseudo_ldimm64(struct bpf_verifier_env *env)
 				return -E2BIG;
 			}
 
+			if (env->prog->aux->sleepable)
+				atomic_inc(&map->sleepable_refcnt);
 			/* hold the map. If the program is rejected by verifier,
 			 * the map will be released by release_maps() or it
 			 * will be used by the valid program until it's unloaded
-			 * and all maps are released in free_used_maps()
+			 * and all maps are released in bpf_free_used_maps()
 			 */
 			bpf_map_inc(map);