[PATCH bpf-next 00/16] bpf: Use bpf_mem_cache_alloc/free in bpf_local_storage

[PATCH bpf-next 00/16] bpf: Use bpf_mem_cache_alloc/free in bpf_local_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This set is to use bpf_mem_cache_alloc/free in bpf_local_storage.
The primary motivation is to solve the deadlock/recursion issue
when bpf_task_storage is used in a bpf tracing prog [1]. This set
also comes with a micro-benchmark to test the storage creation.

Patch 1 to 4 are some general cleanup after bpf_local_storage
has been extended to multiple kernel objects (sk, task, inode,
and then cgrp).

Patch 5 to 11 is to refactor the memory free logic into the new
bpf_selem_free() and bpf_local_storage_free() functions. Together
with the existing bpf_selem_alloc() and bpf_local_storage_alloc(),
it should provide an easier way to change the alloc/free path in
the future.

Patch 12 to 13 is to use bpf_mem_cache_alloc/free.

The remaining patches are selftests and benchmark.

[1]: https://lore.kernel.org/bpf/20221118190109.1512674-1-namhyung@kernel.org/

Martin KaFai Lau (16):
  bpf: Move a few bpf_local_storage functions to static scope
  bpf: Refactor codes into bpf_local_storage_destroy
  bpf: Remove __bpf_local_storage_map_alloc
  bpf: Remove the preceding __ from __bpf_selem_unlink_storage
  bpf: Remember smap in bpf_local_storage
  bpf: Repurpose use_trace_rcu to reuse_now in bpf_local_storage
  bpf: Remove bpf_selem_free_fields*_rcu
  bpf: Add bpf_selem_free_rcu callback
  bpf: Add bpf_selem_free()
  bpf: Add bpf_local_storage_rcu callback
  bpf: Add bpf_local_storage_free()
  bpf: Use bpf_mem_cache_alloc/free in bpf_selem_alloc/free
  bpf: Use bpf_mem_cache_alloc/free for bpf_local_storage
  selftests/bpf: Replace CHECK with ASSERT in test_local_storage
  selftests/bpf: Check freeing sk->sk_local_storage with
    sk_local_storage->smap is NULL
  selftests/bpf: Add local-storage-create benchmark

 include/linux/bpf_local_storage.h             |  22 +-
 include/linux/bpf_mem_alloc.h                 |   5 +
 kernel/bpf/bpf_cgrp_storage.c                 |  11 +-
 kernel/bpf/bpf_inode_storage.c                |  10 +-
 kernel/bpf/bpf_local_storage.c                | 292 ++++++++++--------
 kernel/bpf/bpf_task_storage.c                 |  11 +-
 net/core/bpf_sk_storage.c                     |  12 +-
 tools/testing/selftests/bpf/Makefile          |   2 +
 tools/testing/selftests/bpf/bench.c           |   2 +
 .../bpf/benchs/bench_local_storage_create.c   | 141 +++++++++
 .../bpf/prog_tests/test_local_storage.c       |  49 ++-
 .../bpf/progs/bench_local_storage_create.c    |  57 ++++
 .../selftests/bpf/progs/local_storage.c       |  29 +-
 13 files changed, 429 insertions(+), 214 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/benchs/bench_local_storage_create.c
 create mode 100644 tools/testing/selftests/bpf/progs/bench_local_storage_create.c

-- 
2.30.2

[PATCH bpf-next 01/16] bpf: Move a few bpf_local_storage functions to static scope

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch moves the bpf_local_storage_free_rcu() and
bpf_selem_unlink_map() to static because they are
not used outside of bpf_local_storage.c.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h | 4 ----
 kernel/bpf/bpf_local_storage.c    | 4 ++--
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index 6d37a40cd90e..6917c9a408a1 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -147,8 +147,6 @@ void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool use_trace_rcu);
 void bpf_selem_link_map(struct bpf_local_storage_map *smap,
 			struct bpf_local_storage_elem *selem);
 
-void bpf_selem_unlink_map(struct bpf_local_storage_elem *selem);
-
 struct bpf_local_storage_elem *
 bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner, void *value,
 		bool charge_mem, gfp_t gfp_flags);
@@ -163,6 +161,4 @@ struct bpf_local_storage_data *
 bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 			 void *value, u64 map_flags, gfp_t gfp_flags);
 
-void bpf_local_storage_free_rcu(struct rcu_head *rcu);
-
 #endif /* _BPF_LOCAL_STORAGE_H */
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 3d320393a12c..0510f50bd3ea 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -95,7 +95,7 @@ bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner,
 	return NULL;
 }
 
-void bpf_local_storage_free_rcu(struct rcu_head *rcu)
+static void bpf_local_storage_free_rcu(struct rcu_head *rcu)
 {
 	struct bpf_local_storage *local_storage;
 
@@ -251,7 +251,7 @@ void bpf_selem_link_storage_nolock(struct bpf_local_storage *local_storage,
 	hlist_add_head_rcu(&selem->snode, &local_storage->list);
 }
 
-void bpf_selem_unlink_map(struct bpf_local_storage_elem *selem)
+static void bpf_selem_unlink_map(struct bpf_local_storage_elem *selem)
 {
 	struct bpf_local_storage_map *smap;
 	struct bpf_local_storage_map_bucket *b;
-- 
2.30.2

[PATCH bpf-next 02/16] bpf: Refactor codes into bpf_local_storage_destroy

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch first renames bpf_local_storage_unlink_nolock to
bpf_local_storage_destroy(). It better reflects that it is only
used when the storage's owner (sk/task/cgrp/inode) is being kfree().

All bpf_local_storage_destroy's caller is taking the spin lock and
then free the storage. This patch also moves these two steps into
the bpf_local_storage_destroy.

This is a preparation work for a later patch that uses
bpf_mem_cache_alloc/free in the bpf_local_storage.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h | 2 +-
 kernel/bpf/bpf_cgrp_storage.c     | 9 +--------
 kernel/bpf/bpf_inode_storage.c    | 8 +-------
 kernel/bpf/bpf_local_storage.c    | 8 ++++++--
 kernel/bpf/bpf_task_storage.c     | 9 +--------
 net/core/bpf_sk_storage.c         | 8 +-------
 6 files changed, 11 insertions(+), 33 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index 6917c9a408a1..c8dcf6f40497 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -128,7 +128,7 @@ bpf_local_storage_lookup(struct bpf_local_storage *local_storage,
 			 struct bpf_local_storage_map *smap,
 			 bool cacheit_lockit);
 
-bool bpf_local_storage_unlink_nolock(struct bpf_local_storage *local_storage);
+void bpf_local_storage_destroy(struct bpf_local_storage *local_storage);
 
 void bpf_local_storage_map_free(struct bpf_map *map,
 				struct bpf_local_storage_cache *cache,
diff --git a/kernel/bpf/bpf_cgrp_storage.c b/kernel/bpf/bpf_cgrp_storage.c
index 6cdf6d9ed91d..1d00f1d9bdb7 100644
--- a/kernel/bpf/bpf_cgrp_storage.c
+++ b/kernel/bpf/bpf_cgrp_storage.c
@@ -46,8 +46,6 @@ static struct bpf_local_storage __rcu **cgroup_storage_ptr(void *owner)
 void bpf_cgrp_storage_free(struct cgroup *cgroup)
 {
 	struct bpf_local_storage *local_storage;
-	bool free_cgroup_storage = false;
-	unsigned long flags;
 
 	rcu_read_lock();
 	local_storage = rcu_dereference(cgroup->bpf_cgrp_storage);
@@ -57,14 +55,9 @@ void bpf_cgrp_storage_free(struct cgroup *cgroup)
 	}
 
 	bpf_cgrp_storage_lock();
-	raw_spin_lock_irqsave(&local_storage->lock, flags);
-	free_cgroup_storage = bpf_local_storage_unlink_nolock(local_storage);
-	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
+	bpf_local_storage_destroy(local_storage);
 	bpf_cgrp_storage_unlock();
 	rcu_read_unlock();
-
-	if (free_cgroup_storage)
-		kfree_rcu(local_storage, rcu);
 }
 
 static struct bpf_local_storage_data *
diff --git a/kernel/bpf/bpf_inode_storage.c b/kernel/bpf/bpf_inode_storage.c
index 05f4c66c9089..b4a9904df54e 100644
--- a/kernel/bpf/bpf_inode_storage.c
+++ b/kernel/bpf/bpf_inode_storage.c
@@ -57,7 +57,6 @@ static struct bpf_local_storage_data *inode_storage_lookup(struct inode *inode,
 void bpf_inode_storage_free(struct inode *inode)
 {
 	struct bpf_local_storage *local_storage;
-	bool free_inode_storage = false;
 	struct bpf_storage_blob *bsb;
 
 	bsb = bpf_inode(inode);
@@ -72,13 +71,8 @@ void bpf_inode_storage_free(struct inode *inode)
 		return;
 	}
 
-	raw_spin_lock_bh(&local_storage->lock);
-	free_inode_storage = bpf_local_storage_unlink_nolock(local_storage);
-	raw_spin_unlock_bh(&local_storage->lock);
+	bpf_local_storage_destroy(local_storage);
 	rcu_read_unlock();
-
-	if (free_inode_storage)
-		kfree_rcu(local_storage, rcu);
 }
 
 static void *bpf_fd_inode_storage_lookup_elem(struct bpf_map *map, void *key)
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 0510f50bd3ea..4d2bc7c97f7d 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -652,11 +652,12 @@ int bpf_local_storage_map_check_btf(const struct bpf_map *map,
 	return 0;
 }
 
-bool bpf_local_storage_unlink_nolock(struct bpf_local_storage *local_storage)
+void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 {
 	struct bpf_local_storage_elem *selem;
 	bool free_storage = false;
 	struct hlist_node *n;
+	unsigned long flags;
 
 	/* Neither the bpf_prog nor the bpf_map's syscall
 	 * could be modifying the local_storage->list now.
@@ -667,6 +668,7 @@ bool bpf_local_storage_unlink_nolock(struct bpf_local_storage *local_storage)
 	 * when unlinking elem from the local_storage->list and
 	 * the map's bucket->list.
 	 */
+	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	hlist_for_each_entry_safe(selem, n, &local_storage->list, snode) {
 		/* Always unlink from map before unlinking from
 		 * local_storage.
@@ -681,8 +683,10 @@ bool bpf_local_storage_unlink_nolock(struct bpf_local_storage *local_storage)
 		free_storage = bpf_selem_unlink_storage_nolock(
 			local_storage, selem, false, false);
 	}
+	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
-	return free_storage;
+	if (free_storage)
+		kfree_rcu(local_storage, rcu);
 }
 
 struct bpf_map *
diff --git a/kernel/bpf/bpf_task_storage.c b/kernel/bpf/bpf_task_storage.c
index 1e486055a523..b5f404fe146c 100644
--- a/kernel/bpf/bpf_task_storage.c
+++ b/kernel/bpf/bpf_task_storage.c
@@ -72,8 +72,6 @@ task_storage_lookup(struct task_struct *task, struct bpf_map *map,
 void bpf_task_storage_free(struct task_struct *task)
 {
 	struct bpf_local_storage *local_storage;
-	bool free_task_storage = false;
-	unsigned long flags;
 
 	rcu_read_lock();
 
@@ -84,14 +82,9 @@ void bpf_task_storage_free(struct task_struct *task)
 	}
 
 	bpf_task_storage_lock();
-	raw_spin_lock_irqsave(&local_storage->lock, flags);
-	free_task_storage = bpf_local_storage_unlink_nolock(local_storage);
-	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
+	bpf_local_storage_destroy(local_storage);
 	bpf_task_storage_unlock();
 	rcu_read_unlock();
-
-	if (free_task_storage)
-		kfree_rcu(local_storage, rcu);
 }
 
 static void *bpf_pid_task_storage_lookup_elem(struct bpf_map *map, void *key)
diff --git a/net/core/bpf_sk_storage.c b/net/core/bpf_sk_storage.c
index bb378c33f542..42569d0904a5 100644
--- a/net/core/bpf_sk_storage.c
+++ b/net/core/bpf_sk_storage.c
@@ -49,7 +49,6 @@ static int bpf_sk_storage_del(struct sock *sk, struct bpf_map *map)
 void bpf_sk_storage_free(struct sock *sk)
 {
 	struct bpf_local_storage *sk_storage;
-	bool free_sk_storage = false;
 
 	rcu_read_lock();
 	sk_storage = rcu_dereference(sk->sk_bpf_storage);
@@ -58,13 +57,8 @@ void bpf_sk_storage_free(struct sock *sk)
 		return;
 	}
 
-	raw_spin_lock_bh(&sk_storage->lock);
-	free_sk_storage = bpf_local_storage_unlink_nolock(sk_storage);
-	raw_spin_unlock_bh(&sk_storage->lock);
+	bpf_local_storage_destroy(sk_storage);
 	rcu_read_unlock();
-
-	if (free_sk_storage)
-		kfree_rcu(sk_storage, rcu);
 }
 
 static void bpf_sk_storage_map_free(struct bpf_map *map)
-- 
2.30.2

[PATCH bpf-next 03/16] bpf: Remove __bpf_local_storage_map_alloc

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

bpf_local_storage_map_alloc() is the only caller of
__bpf_local_storage_map_alloc().  The remaining logic in
bpf_local_storage_map_alloc() is only a one liner setting
the smap->cache_idx.

Remove __bpf_local_storage_map_alloc() to simplify code.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 63 ++++++++++++++--------------------
 1 file changed, 26 insertions(+), 37 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 4d2bc7c97f7d..acedf6b07c54 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -601,40 +601,6 @@ int bpf_local_storage_map_alloc_check(union bpf_attr *attr)
 	return 0;
 }
 
-static struct bpf_local_storage_map *__bpf_local_storage_map_alloc(union bpf_attr *attr)
-{
-	struct bpf_local_storage_map *smap;
-	unsigned int i;
-	u32 nbuckets;
-
-	smap = bpf_map_area_alloc(sizeof(*smap), NUMA_NO_NODE);
-	if (!smap)
-		return ERR_PTR(-ENOMEM);
-	bpf_map_init_from_attr(&smap->map, attr);
-
-	nbuckets = roundup_pow_of_two(num_possible_cpus());
-	/* Use at least 2 buckets, select_bucket() is undefined behavior with 1 bucket */
-	nbuckets = max_t(u32, 2, nbuckets);
-	smap->bucket_log = ilog2(nbuckets);
-
-	smap->buckets = bpf_map_kvcalloc(&smap->map, sizeof(*smap->buckets),
-					 nbuckets, GFP_USER | __GFP_NOWARN);
-	if (!smap->buckets) {
-		bpf_map_area_free(smap);
-		return ERR_PTR(-ENOMEM);
-	}
-
-	for (i = 0; i < nbuckets; i++) {
-		INIT_HLIST_HEAD(&smap->buckets[i].list);
-		raw_spin_lock_init(&smap->buckets[i].lock);
-	}
-
-	smap->elem_size = offsetof(struct bpf_local_storage_elem,
-				   sdata.data[attr->value_size]);
-
-	return smap;
-}
-
 int bpf_local_storage_map_check_btf(const struct bpf_map *map,
 				    const struct btf *btf,
 				    const struct btf_type *key_type,
@@ -694,10 +660,33 @@ bpf_local_storage_map_alloc(union bpf_attr *attr,
 			    struct bpf_local_storage_cache *cache)
 {
 	struct bpf_local_storage_map *smap;
+	unsigned int i;
+	u32 nbuckets;
+
+	smap = bpf_map_area_alloc(sizeof(*smap), NUMA_NO_NODE);
+	if (!smap)
+		return ERR_PTR(-ENOMEM);
+	bpf_map_init_from_attr(&smap->map, attr);
+
+	nbuckets = roundup_pow_of_two(num_possible_cpus());
+	/* Use at least 2 buckets, select_bucket() is undefined behavior with 1 bucket */
+	nbuckets = max_t(u32, 2, nbuckets);
+	smap->bucket_log = ilog2(nbuckets);
 
-	smap = __bpf_local_storage_map_alloc(attr);
-	if (IS_ERR(smap))
-		return ERR_CAST(smap);
+	smap->buckets = bpf_map_kvcalloc(&smap->map, sizeof(*smap->buckets),
+					 nbuckets, GFP_USER | __GFP_NOWARN);
+	if (!smap->buckets) {
+		bpf_map_area_free(smap);
+		return ERR_PTR(-ENOMEM);
+	}
+
+	for (i = 0; i < nbuckets; i++) {
+		INIT_HLIST_HEAD(&smap->buckets[i].list);
+		raw_spin_lock_init(&smap->buckets[i].lock);
+	}
+
+	smap->elem_size = offsetof(struct bpf_local_storage_elem,
+				   sdata.data[attr->value_size]);
 
 	smap->cache_idx = bpf_local_storage_cache_idx_get(cache);
 	return &smap->map;
-- 
2.30.2

[PATCH bpf-next 04/16] bpf: Remove the preceding __ from __bpf_selem_unlink_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

__bpf_selem_unlink_storage is taking the spin lock and there is
no name collision also. Having the preceding '__' is confusing
when reviewing the later patch.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index acedf6b07c54..fef75beaf66d 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -216,8 +216,8 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 	return free_local_storage;
 }
 
-static void __bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
-				       bool use_trace_rcu)
+static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
+				     bool use_trace_rcu)
 {
 	struct bpf_local_storage *local_storage;
 	bool free_local_storage = false;
@@ -288,7 +288,7 @@ void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool use_trace_rcu)
 	 * the local_storage.
 	 */
 	bpf_selem_unlink_map(selem);
-	__bpf_selem_unlink_storage(selem, use_trace_rcu);
+	bpf_selem_unlink_storage(selem, use_trace_rcu);
 }
 
 /* If cacheit_lockit is false, this lookup function is lockless */
-- 
2.30.2

[PATCH bpf-next 05/16] bpf: Remember smap in bpf_local_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch remembers which smap triggers the allocation
of a 'struct bpf_local_storage' object. The local_storage is
allocated during the very first selem added to the owner.
The smap pointer is needed when using the bpf_mem_cache_free
in a later patch because it needs to free to the correct
smap's bpf_mem_alloc object.

When a selem is being removed, it needs to check if it is
the selem that triggers the creation of the local_storage.
If it is, the local_storage->smap pointer will be reset to NULL.
This NULL reset is done under the local_storage->lock in
bpf_selem_unlink_storage_nolock() when a selem is being removed.
Also note that the local_storage may not go away even
local_storage->smap is NULL because there may be other
selem still stored in the local_storage.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h | 1 +
 kernel/bpf/bpf_local_storage.c    | 4 ++++
 2 files changed, 5 insertions(+)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index c8dcf6f40497..31ee681b4c65 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -83,6 +83,7 @@ struct bpf_local_storage_elem {
 
 struct bpf_local_storage {
 	struct bpf_local_storage_data __rcu *cache[BPF_LOCAL_STORAGE_CACHE_SIZE];
+	struct bpf_local_storage_map __rcu *smap;
 	struct hlist_head list; /* List of bpf_local_storage_elem */
 	void *owner;		/* The object that owns the above "list" of
 				 * bpf_local_storage_elem.
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index fef75beaf66d..c8ca1886374e 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -213,6 +213,9 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 			kfree_rcu(selem, rcu);
 	}
 
+	if (rcu_access_pointer(local_storage->smap) == smap)
+		RCU_INIT_POINTER(local_storage->smap, NULL);
+
 	return free_local_storage;
 }
 
@@ -368,6 +371,7 @@ int bpf_local_storage_alloc(void *owner,
 		goto uncharge;
 	}
 
+	RCU_INIT_POINTER(storage->smap, smap);
 	INIT_HLIST_HEAD(&storage->list);
 	raw_spin_lock_init(&storage->lock);
 	storage->owner = owner;
-- 
2.30.2

[PATCH bpf-next 06/16] bpf: Repurpose use_trace_rcu to reuse_now in bpf_local_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch re-purpose the use_trace_rcu to mean
if the freed memory can be reused immediately or not.
The use_trace_rcu is renamed to reuse_now. Other than
the boolean test is reversed, it should be a no-op.

The following explains the reason for the rename and how it will
be used in a later patch.

In a later patch, bpf_mem_cache_alloc/free will be used
in the bpf_local_storage. The bpf mem allocator will reuse
the freed memory immediately. Some of the free paths in
bpf_local_storage does not support memory to be reused immediately.
These paths are the "delete" elem cases from the bpf_*_storage_delete()
helper and the map_delete_elem() syscall. Note that "delete" elem
before the owner's (sk/task/cgrp/inode) lifetime ended is not
the common usage for the local storage.

The common free path, bpf_local_storage_destroy(), can reuse the
memory immediately. This common path means the storage stays with
its owner until the owner is destroyed.

The above mentioned "delete" elem paths that cannot
reuse immediately always has the 'use_trace_rcu ==  true'.
The cases that is safe for immediate reuse always have
'use_trace_rcu == false'. Instead of adding another arg
in a later patch, this patch re-purpose this arg
to reuse_now and have the test logic reversed.

In a later patch, 'reuse_now == true' will free to the
bpf_mem_cache_free() where the memory can be reused
immediately. 'reuse_now == false' will go through the
call_rcu_tasks_trace().

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h |  2 +-
 kernel/bpf/bpf_cgrp_storage.c     |  2 +-
 kernel/bpf/bpf_inode_storage.c    |  2 +-
 kernel/bpf/bpf_local_storage.c    | 24 ++++++++++++------------
 kernel/bpf/bpf_task_storage.c     |  2 +-
 net/core/bpf_sk_storage.c         |  2 +-
 6 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index 31ee681b4c65..fad09f42a2f4 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -143,7 +143,7 @@ int bpf_local_storage_map_check_btf(const struct bpf_map *map,
 void bpf_selem_link_storage_nolock(struct bpf_local_storage *local_storage,
 				   struct bpf_local_storage_elem *selem);
 
-void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool use_trace_rcu);
+void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool reuse_now);
 
 void bpf_selem_link_map(struct bpf_local_storage_map *smap,
 			struct bpf_local_storage_elem *selem);
diff --git a/kernel/bpf/bpf_cgrp_storage.c b/kernel/bpf/bpf_cgrp_storage.c
index 1d00f1d9bdb7..617e6111a51f 100644
--- a/kernel/bpf/bpf_cgrp_storage.c
+++ b/kernel/bpf/bpf_cgrp_storage.c
@@ -121,7 +121,7 @@ static int cgroup_storage_delete(struct cgroup *cgroup, struct bpf_map *map)
 	if (!sdata)
 		return -ENOENT;
 
-	bpf_selem_unlink(SELEM(sdata), true);
+	bpf_selem_unlink(SELEM(sdata), false);
 	return 0;
 }
 
diff --git a/kernel/bpf/bpf_inode_storage.c b/kernel/bpf/bpf_inode_storage.c
index b4a9904df54e..e37d9d2e28f4 100644
--- a/kernel/bpf/bpf_inode_storage.c
+++ b/kernel/bpf/bpf_inode_storage.c
@@ -122,7 +122,7 @@ static int inode_storage_delete(struct inode *inode, struct bpf_map *map)
 	if (!sdata)
 		return -ENOENT;
 
-	bpf_selem_unlink(SELEM(sdata), true);
+	bpf_selem_unlink(SELEM(sdata), false);
 
 	return 0;
 }
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index c8ca1886374e..8c1401a24c7d 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -147,7 +147,7 @@ static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
  */
 static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_storage,
 					    struct bpf_local_storage_elem *selem,
-					    bool uncharge_mem, bool use_trace_rcu)
+					    bool uncharge_mem, bool reuse_now)
 {
 	struct bpf_local_storage_map *smap;
 	bool free_local_storage;
@@ -201,7 +201,7 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 	 * any special fields.
 	 */
 	rec = smap->map.record;
-	if (use_trace_rcu) {
+	if (!reuse_now) {
 		if (!IS_ERR_OR_NULL(rec))
 			call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_fields_trace_rcu);
 		else
@@ -220,7 +220,7 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 }
 
 static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
-				     bool use_trace_rcu)
+				     bool reuse_now)
 {
 	struct bpf_local_storage *local_storage;
 	bool free_local_storage = false;
@@ -235,11 +235,11 @@ static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	if (likely(selem_linked_to_storage(selem)))
 		free_local_storage = bpf_selem_unlink_storage_nolock(
-			local_storage, selem, true, use_trace_rcu);
+			local_storage, selem, true, reuse_now);
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
 	if (free_local_storage) {
-		if (use_trace_rcu)
+		if (!reuse_now)
 			call_rcu_tasks_trace(&local_storage->rcu,
 				     bpf_local_storage_free_rcu);
 		else
@@ -284,14 +284,14 @@ void bpf_selem_link_map(struct bpf_local_storage_map *smap,
 	raw_spin_unlock_irqrestore(&b->lock, flags);
 }
 
-void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool use_trace_rcu)
+void bpf_selem_unlink(struct bpf_local_storage_elem *selem, bool reuse_now)
 {
 	/* Always unlink from map before unlinking from local_storage
 	 * because selem will be freed after successfully unlinked from
 	 * the local_storage.
 	 */
 	bpf_selem_unlink_map(selem);
-	bpf_selem_unlink_storage(selem, use_trace_rcu);
+	bpf_selem_unlink_storage(selem, reuse_now);
 }
 
 /* If cacheit_lockit is false, this lookup function is lockless */
@@ -538,7 +538,7 @@ bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 	if (old_sdata) {
 		bpf_selem_unlink_map(SELEM(old_sdata));
 		bpf_selem_unlink_storage_nolock(local_storage, SELEM(old_sdata),
-						false, true);
+						false, false);
 	}
 
 unlock:
@@ -651,7 +651,7 @@ void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 		 * of the loop will set the free_cgroup_storage to true.
 		 */
 		free_storage = bpf_selem_unlink_storage_nolock(
-			local_storage, selem, false, false);
+			local_storage, selem, false, true);
 	}
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
@@ -735,7 +735,7 @@ void bpf_local_storage_map_free(struct bpf_map *map,
 				migrate_disable();
 				this_cpu_inc(*busy_counter);
 			}
-			bpf_selem_unlink(selem, false);
+			bpf_selem_unlink(selem, true);
 			if (busy_counter) {
 				this_cpu_dec(*busy_counter);
 				migrate_enable();
@@ -773,8 +773,8 @@ void bpf_local_storage_map_free(struct bpf_map *map,
 		/* We cannot skip rcu_barrier() when rcu_trace_implies_rcu_gp()
 		 * is true, because while call_rcu invocation is skipped in that
 		 * case in bpf_selem_free_fields_trace_rcu (and all local
-		 * storage maps pass use_trace_rcu = true), there can be
-		 * call_rcu callbacks based on use_trace_rcu = false in the
+		 * storage maps pass reuse_now = false), there can be
+		 * call_rcu callbacks based on reuse_now = true in the
 		 * while ((selem = ...)) loop above or when owner's free path
 		 * calls bpf_local_storage_unlink_nolock.
 		 */
diff --git a/kernel/bpf/bpf_task_storage.c b/kernel/bpf/bpf_task_storage.c
index b5f404fe146c..65aeb042c263 100644
--- a/kernel/bpf/bpf_task_storage.c
+++ b/kernel/bpf/bpf_task_storage.c
@@ -168,7 +168,7 @@ static int task_storage_delete(struct task_struct *task, struct bpf_map *map,
 	if (!nobusy)
 		return -EBUSY;
 
-	bpf_selem_unlink(SELEM(sdata), true);
+	bpf_selem_unlink(SELEM(sdata), false);
 
 	return 0;
 }
diff --git a/net/core/bpf_sk_storage.c b/net/core/bpf_sk_storage.c
index 42569d0904a5..8b0c9e4341eb 100644
--- a/net/core/bpf_sk_storage.c
+++ b/net/core/bpf_sk_storage.c
@@ -40,7 +40,7 @@ static int bpf_sk_storage_del(struct sock *sk, struct bpf_map *map)
 	if (!sdata)
 		return -ENOENT;
 
-	bpf_selem_unlink(SELEM(sdata), true);
+	bpf_selem_unlink(SELEM(sdata), false);
 
 	return 0;
 }
-- 
2.30.2

[PATCH bpf-next 07/16] bpf: Remove bpf_selem_free_fields*_rcu

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch removes the bpf_selem_free_fields*_rcu. The
bpf_obj_free_fields() can be done before the call_rcu_trasks_trace()
and kfree_rcu(). It is needed when a later patch uses
bpf_mem_cache_alloc/free. In bpf hashtab, bpf_obj_free_fields()
is also called before calling bpf_mem_cache_free. The discussion
can be found in
https://lore.kernel.org/bpf/f67021ee-21d9-bfae-6134-4ca542fab843@linux.dev/

Cc: Kumar Kartikeya Dwivedi <memxor@gmail.com>
Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 67 +++-------------------------------
 1 file changed, 5 insertions(+), 62 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 8c1401a24c7d..9a1febcc9565 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -109,27 +109,6 @@ static void bpf_local_storage_free_rcu(struct rcu_head *rcu)
 		kfree_rcu(local_storage, rcu);
 }
 
-static void bpf_selem_free_fields_rcu(struct rcu_head *rcu)
-{
-	struct bpf_local_storage_elem *selem;
-	struct bpf_local_storage_map *smap;
-
-	selem = container_of(rcu, struct bpf_local_storage_elem, rcu);
-	/* protected by the rcu_barrier*() */
-	smap = rcu_dereference_protected(SDATA(selem)->smap, true);
-	bpf_obj_free_fields(smap->map.record, SDATA(selem)->data);
-	kfree(selem);
-}
-
-static void bpf_selem_free_fields_trace_rcu(struct rcu_head *rcu)
-{
-	/* Free directly if Tasks Trace RCU GP also implies RCU GP */
-	if (rcu_trace_implies_rcu_gp())
-		bpf_selem_free_fields_rcu(rcu);
-	else
-		call_rcu(rcu, bpf_selem_free_fields_rcu);
-}
-
 static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
 {
 	struct bpf_local_storage_elem *selem;
@@ -151,7 +130,6 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 {
 	struct bpf_local_storage_map *smap;
 	bool free_local_storage;
-	struct btf_record *rec;
 	void *owner;
 
 	smap = rcu_dereference_check(SDATA(selem)->smap, bpf_rcu_lock_held());
@@ -192,26 +170,11 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 	    SDATA(selem))
 		RCU_INIT_POINTER(local_storage->cache[smap->cache_idx], NULL);
 
-	/* A different RCU callback is chosen whenever we need to free
-	 * additional fields in selem data before freeing selem.
-	 * bpf_local_storage_map_free only executes rcu_barrier to wait for RCU
-	 * callbacks when it has special fields, hence we can only conditionally
-	 * dereference smap, as by this time the map might have already been
-	 * freed without waiting for our call_rcu callback if it did not have
-	 * any special fields.
-	 */
-	rec = smap->map.record;
-	if (!reuse_now) {
-		if (!IS_ERR_OR_NULL(rec))
-			call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_fields_trace_rcu);
-		else
-			call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
-	} else {
-		if (!IS_ERR_OR_NULL(rec))
-			call_rcu(&selem->rcu, bpf_selem_free_fields_rcu);
-		else
-			kfree_rcu(selem, rcu);
-	}
+	bpf_obj_free_fields(smap->map.record, SDATA(selem)->data);
+	if (!reuse_now)
+		call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
+	else
+		kfree_rcu(selem, rcu);
 
 	if (rcu_access_pointer(local_storage->smap) == smap)
 		RCU_INIT_POINTER(local_storage->smap, NULL);
@@ -759,26 +722,6 @@ void bpf_local_storage_map_free(struct bpf_map *map,
 	 */
 	synchronize_rcu();
 
-	/* Only delay freeing of smap, buckets are not needed anymore */
 	kvfree(smap->buckets);
-
-	/* When local storage has special fields, callbacks for
-	 * bpf_selem_free_fields_rcu and bpf_selem_free_fields_trace_rcu will
-	 * keep using the map BTF record, we need to execute an RCU barrier to
-	 * wait for them as the record will be freed right after our map_free
-	 * callback.
-	 */
-	if (!IS_ERR_OR_NULL(smap->map.record)) {
-		rcu_barrier_tasks_trace();
-		/* We cannot skip rcu_barrier() when rcu_trace_implies_rcu_gp()
-		 * is true, because while call_rcu invocation is skipped in that
-		 * case in bpf_selem_free_fields_trace_rcu (and all local
-		 * storage maps pass reuse_now = false), there can be
-		 * call_rcu callbacks based on reuse_now = true in the
-		 * while ((selem = ...)) loop above or when owner's free path
-		 * calls bpf_local_storage_unlink_nolock.
-		 */
-		rcu_barrier();
-	}
 	bpf_map_area_free(smap);
 }
-- 
2.30.2

Re: [PATCH bpf-next 07/16] bpf: Remove bpf_selem_free_fields*_rcu

[Kumar Kartikeya Dwivedi]

On Mon, Mar 06, 2023 at 09:42:07AM CET, Martin KaFai Lau wrote:
> From: Martin KaFai Lau <martin.lau@kernel.org>
>
> This patch removes the bpf_selem_free_fields*_rcu. The
> bpf_obj_free_fields() can be done before the call_rcu_trasks_trace()
> and kfree_rcu(). It is needed when a later patch uses
> bpf_mem_cache_alloc/free. In bpf hashtab, bpf_obj_free_fields()
> is also called before calling bpf_mem_cache_free. The discussion
> can be found in
> https://lore.kernel.org/bpf/f67021ee-21d9-bfae-6134-4ca542fab843@linux.dev/
>
> Cc: Kumar Kartikeya Dwivedi <memxor@gmail.com>
> Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
> ---

Acked-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

[PATCH bpf-next 08/16] bpf: Add bpf_selem_free_rcu callback

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

Add bpf_selem_free_rcu() callback to do the kfree() instead
of using kfree_rcu. It is a preparation work for using
bpf_mem_cache_alloc/free in a later patch.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 9a1febcc9565..528579c9f60f 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -109,15 +109,20 @@ static void bpf_local_storage_free_rcu(struct rcu_head *rcu)
 		kfree_rcu(local_storage, rcu);
 }
 
-static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
+static void bpf_selem_free_rcu(struct rcu_head *rcu)
 {
 	struct bpf_local_storage_elem *selem;
 
 	selem = container_of(rcu, struct bpf_local_storage_elem, rcu);
+	kfree(selem);
+}
+
+static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
+{
 	if (rcu_trace_implies_rcu_gp())
-		kfree(selem);
+		bpf_selem_free_rcu(rcu);
 	else
-		kfree_rcu(selem, rcu);
+		call_rcu(rcu, bpf_selem_free_rcu);
 }
 
 /* local_storage->lock must be held and selem->local_storage == local_storage.
@@ -174,7 +179,7 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 	if (!reuse_now)
 		call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
 	else
-		kfree_rcu(selem, rcu);
+		call_rcu(&selem->rcu, bpf_selem_free_rcu);
 
 	if (rcu_access_pointer(local_storage->smap) == smap)
 		RCU_INIT_POINTER(local_storage->smap, NULL);
-- 
2.30.2

[PATCH bpf-next 09/16] bpf: Add bpf_selem_free()

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch refactors the selem freeing logic into bpf_selem_free().
It is a preparation work for a later patch using
bpf_mem_cache_alloc/free. The other kfree(selem) cases
are also changed to bpf_selem_free(..., reuse_now = true).

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h |  4 ++++
 kernel/bpf/bpf_local_storage.c    | 21 ++++++++++++++-------
 net/core/bpf_sk_storage.c         |  2 +-
 3 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index fad09f42a2f4..adb5023a1af5 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -152,6 +152,10 @@ struct bpf_local_storage_elem *
 bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner, void *value,
 		bool charge_mem, gfp_t gfp_flags);
 
+void bpf_selem_free(struct bpf_local_storage_elem *selem,
+		    struct bpf_local_storage_map *smap,
+		    bool reuse_now);
+
 int
 bpf_local_storage_alloc(void *owner,
 			struct bpf_local_storage_map *smap,
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 528579c9f60f..f611668f8a0b 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -125,6 +125,17 @@ static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
 		call_rcu(rcu, bpf_selem_free_rcu);
 }
 
+void bpf_selem_free(struct bpf_local_storage_elem *selem,
+		    struct bpf_local_storage_map *smap,
+		    bool reuse_now)
+{
+	bpf_obj_free_fields(smap->map.record, SDATA(selem)->data);
+	if (!reuse_now)
+		call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
+	else
+		call_rcu(&selem->rcu, bpf_selem_free_rcu);
+}
+
 /* local_storage->lock must be held and selem->local_storage == local_storage.
  * The caller must ensure selem->smap is still valid to be
  * dereferenced for its smap->elem_size and smap->cache_idx.
@@ -175,11 +186,7 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 	    SDATA(selem))
 		RCU_INIT_POINTER(local_storage->cache[smap->cache_idx], NULL);
 
-	bpf_obj_free_fields(smap->map.record, SDATA(selem)->data);
-	if (!reuse_now)
-		call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
-	else
-		call_rcu(&selem->rcu, bpf_selem_free_rcu);
+	bpf_selem_free(selem, smap, reuse_now);
 
 	if (rcu_access_pointer(local_storage->smap) == smap)
 		RCU_INIT_POINTER(local_storage->smap, NULL);
@@ -423,7 +430,7 @@ bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 
 		err = bpf_local_storage_alloc(owner, smap, selem, gfp_flags);
 		if (err) {
-			kfree(selem);
+			bpf_selem_free(selem, smap, true);
 			mem_uncharge(smap, owner, smap->elem_size);
 			return ERR_PTR(err);
 		}
@@ -517,7 +524,7 @@ bpf_local_storage_update(void *owner, struct bpf_local_storage_map *smap,
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 	if (selem) {
 		mem_uncharge(smap, owner, smap->elem_size);
-		kfree(selem);
+		bpf_selem_free(selem, smap, true);
 	}
 	return ERR_PTR(err);
 }
diff --git a/net/core/bpf_sk_storage.c b/net/core/bpf_sk_storage.c
index 8b0c9e4341eb..4fc078e8e9ca 100644
--- a/net/core/bpf_sk_storage.c
+++ b/net/core/bpf_sk_storage.c
@@ -197,7 +197,7 @@ int bpf_sk_storage_clone(const struct sock *sk, struct sock *newsk)
 		} else {
 			ret = bpf_local_storage_alloc(newsk, smap, copy_selem, GFP_ATOMIC);
 			if (ret) {
-				kfree(copy_selem);
+				bpf_selem_free(selem, smap, true);
 				atomic_sub(smap->elem_size,
 					   &newsk->sk_omem_alloc);
 				bpf_map_put(map);
-- 
2.30.2

Re: [PATCH bpf-next 09/16] bpf: Add bpf_selem_free()

[Martin KaFai Lau]

On 3/6/23 12:42 AM, Martin KaFai Lau wrote:
> diff --git a/net/core/bpf_sk_storage.c b/net/core/bpf_sk_storage.c
> index 8b0c9e4341eb..4fc078e8e9ca 100644
> --- a/net/core/bpf_sk_storage.c
> +++ b/net/core/bpf_sk_storage.c
> @@ -197,7 +197,7 @@ int bpf_sk_storage_clone(const struct sock *sk, struct sock *newsk)
>   		} else {
>   			ret = bpf_local_storage_alloc(newsk, smap, copy_selem, GFP_ATOMIC);
>   			if (ret) {
> -				kfree(copy_selem);
> +				bpf_selem_free(selem, smap, true);
noticed there is a bug here, should be copy_selem. will fix.

>   				atomic_sub(smap->elem_size,
>   					   &newsk->sk_omem_alloc);
>   				bpf_map_put(map);

[PATCH bpf-next 10/16] bpf: Add bpf_local_storage_rcu callback

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

The existing bpf_local_storage_free_rcu is renamed to
bpf_local_storage_free_trace_rcu. A new bpf_local_storage_rcu
callback is added to do the kfree instead of using kfree_rcu.
It is a preparation work for a later patch using
bpf_mem_cache_alloc/free.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index f611668f8a0b..0568b479bdb0 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -99,14 +99,19 @@ static void bpf_local_storage_free_rcu(struct rcu_head *rcu)
 {
 	struct bpf_local_storage *local_storage;
 
+	local_storage = container_of(rcu, struct bpf_local_storage, rcu);
+	kfree(local_storage);
+}
+
+static void bpf_local_storage_free_trace_rcu(struct rcu_head *rcu)
+{
 	/* If RCU Tasks Trace grace period implies RCU grace period, do
 	 * kfree(), else do kfree_rcu().
 	 */
-	local_storage = container_of(rcu, struct bpf_local_storage, rcu);
 	if (rcu_trace_implies_rcu_gp())
-		kfree(local_storage);
+		bpf_local_storage_free_rcu(rcu);
 	else
-		kfree_rcu(local_storage, rcu);
+		call_rcu(rcu, bpf_local_storage_free_rcu);
 }
 
 static void bpf_selem_free_rcu(struct rcu_head *rcu)
@@ -216,9 +221,9 @@ static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 	if (free_local_storage) {
 		if (!reuse_now)
 			call_rcu_tasks_trace(&local_storage->rcu,
-				     bpf_local_storage_free_rcu);
+				     bpf_local_storage_free_trace_rcu);
 		else
-			kfree_rcu(local_storage, rcu);
+			call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
 	}
 }
 
@@ -631,7 +636,7 @@ void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
 	if (free_storage)
-		kfree_rcu(local_storage, rcu);
+		call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
 }
 
 struct bpf_map *
-- 
2.30.2

[PATCH bpf-next 11/16] bpf: Add bpf_local_storage_free()

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch refactors local_storage freeing logic into
bpf_local_storage_free(). It is a preparation work for a later
patch that uses bpf_mem_cache_alloc/free. The other kfree(local_storage)
cases are also changed to bpf_local_storage_free(..., reuse_now = true).

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 kernel/bpf/bpf_local_storage.c | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 0568b479bdb0..532b82084ba7 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -114,6 +114,16 @@ static void bpf_local_storage_free_trace_rcu(struct rcu_head *rcu)
 		call_rcu(rcu, bpf_local_storage_free_rcu);
 }
 
+static void bpf_local_storage_free(struct bpf_local_storage *local_storage,
+				   bool reuse_now)
+{
+	if (!reuse_now)
+		call_rcu_tasks_trace(&local_storage->rcu,
+				     bpf_local_storage_free_trace_rcu);
+	else
+		call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
+}
+
 static void bpf_selem_free_rcu(struct rcu_head *rcu)
 {
 	struct bpf_local_storage_elem *selem;
@@ -218,13 +228,8 @@ static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 			local_storage, selem, true, reuse_now);
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
-	if (free_local_storage) {
-		if (!reuse_now)
-			call_rcu_tasks_trace(&local_storage->rcu,
-				     bpf_local_storage_free_trace_rcu);
-		else
-			call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
-	}
+	if (free_local_storage)
+		bpf_local_storage_free(local_storage, reuse_now);
 }
 
 void bpf_selem_link_storage_nolock(struct bpf_local_storage *local_storage,
@@ -391,7 +396,7 @@ int bpf_local_storage_alloc(void *owner,
 	return 0;
 
 uncharge:
-	kfree(storage);
+	bpf_local_storage_free(storage, true);
 	mem_uncharge(smap, owner, sizeof(*storage));
 	return err;
 }
@@ -636,7 +641,7 @@ void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
 	if (free_storage)
-		call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
+		bpf_local_storage_free(local_storage, true);
 }
 
 struct bpf_map *
-- 
2.30.2

[PATCH bpf-next 12/16] bpf: Use bpf_mem_cache_alloc/free in bpf_selem_alloc/free

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch uses bpf_mem_cache_alloc/free in bpf_selem_alloc/free.

The ____cacheline_aligned attribute is no longer needed
in 'struct bpf_local_storage_elem'. bpf_mem_cache_alloc will
have 'struct llist_node' in front of the 'struct bpf_local_storage_elem'.
It will use the 8bytes hole in the bpf_local_storage_elem.

After bpf_mem_cache_alloc(), the SDATA(selem)->data is zero-ed because
bpf_mem_cache_alloc() could return a reused selem. It is to keep
the existing bpf_map_kzalloc() behavior. Only SDATA(selem)->data
is zero-ed. SDATA(selem)->data is the visible part to the bpf prog.
No need to use zero_map_value() to do the zeroing because
bpf_selem_free() ensures no bpf prog is using the selem before
returning the selem through bpf_mem_cache_free(). For the internal
fields of selem, they will be initialized when linking to the
new smap and the new local_storage.

When bpf_mem_cache_alloc() fails, bpf_selem_alloc() will try to
fallback to kzalloc only if the caller has GFP_KERNEL flag set (ie. from
sleepable bpf prog that should not cause deadlock). BPF_MA_SIZE
and BPF_MA_PTR macro are added for that.

For the common selem free path where the selem is freed when its owner
is also being freed, reuse_now == true and selem can be reused
immediately. bpf_selem_free() uses bpf_mem_cache_free() where
selem will be considered for immediate reuse.

For the uncommon path that the bpf prog explicitly deletes the selem (by
using the helper bpf_*_storage_delete), the selem cannot be reused
immediately. reuse_now == false and bpf_selem_free() will stay with
the current call_rcu_tasks_trace. BPF_MA_NODE macro is added to get
the correct address for the kfree.

mem_charge and mem_uncharge are changed to use the BPF_MA_SIZE
macro. It will have a temporarily over-charge for the
bpf_local_storage_alloc() because bpf_local_storage is not
moved to bpf_mem_cache_alloc in this patch but it will be done
in the next patch.

Cc: Namhyung Kim <namhyung@kernel.org>
Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h |  8 ++---
 include/linux/bpf_mem_alloc.h     |  5 +++
 kernel/bpf/bpf_local_storage.c    | 56 +++++++++++++++++++++++++------
 3 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index adb5023a1af5..a236c9b964cf 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -13,6 +13,7 @@
 #include <linux/list.h>
 #include <linux/hash.h>
 #include <linux/types.h>
+#include <linux/bpf_mem_alloc.h>
 #include <uapi/linux/btf.h>
 
 #define BPF_LOCAL_STORAGE_CACHE_SIZE	16
@@ -55,6 +56,7 @@ struct bpf_local_storage_map {
 	u32 bucket_log;
 	u16 elem_size;
 	u16 cache_idx;
+	struct bpf_mem_alloc selem_ma;
 };
 
 struct bpf_local_storage_data {
@@ -74,11 +76,7 @@ struct bpf_local_storage_elem {
 	struct hlist_node snode;	/* Linked to bpf_local_storage */
 	struct bpf_local_storage __rcu *local_storage;
 	struct rcu_head rcu;
-	/* 8 bytes hole */
-	/* The data is stored in another cacheline to minimize
-	 * the number of cachelines access during a cache hit.
-	 */
-	struct bpf_local_storage_data sdata ____cacheline_aligned;
+	struct bpf_local_storage_data sdata;
 };
 
 struct bpf_local_storage {
diff --git a/include/linux/bpf_mem_alloc.h b/include/linux/bpf_mem_alloc.h
index a7104af61ab4..0ab16fb0ab50 100644
--- a/include/linux/bpf_mem_alloc.h
+++ b/include/linux/bpf_mem_alloc.h
@@ -5,6 +5,11 @@
 #include <linux/compiler_types.h>
 #include <linux/workqueue.h>
 
+#define BPF_MA_NODE_SZ sizeof(struct llist_node)
+#define BPF_MA_SIZE(_size) ((_size) + BPF_MA_NODE_SZ)
+#define BPF_MA_PTR(_node) ((void *)(_node) + BPF_MA_NODE_SZ)
+#define BPF_MA_NODE(_ptr) ((void *)(_ptr) - BPF_MA_NODE_SZ)
+
 struct bpf_mem_cache;
 struct bpf_mem_caches;
 
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index 532b82084ba7..d3c0dd5737d6 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -31,7 +31,7 @@ static int mem_charge(struct bpf_local_storage_map *smap, void *owner, u32 size)
 	if (!map->ops->map_local_storage_charge)
 		return 0;
 
-	return map->ops->map_local_storage_charge(smap, owner, size);
+	return map->ops->map_local_storage_charge(smap, owner, BPF_MA_SIZE(size));
 }
 
 static void mem_uncharge(struct bpf_local_storage_map *smap, void *owner,
@@ -40,7 +40,7 @@ static void mem_uncharge(struct bpf_local_storage_map *smap, void *owner,
 	struct bpf_map *map = &smap->map;
 
 	if (map->ops->map_local_storage_uncharge)
-		map->ops->map_local_storage_uncharge(smap, owner, size);
+		map->ops->map_local_storage_uncharge(smap, owner, BPF_MA_SIZE(size));
 }
 
 static struct bpf_local_storage __rcu **
@@ -80,12 +80,32 @@ bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner,
 	if (charge_mem && mem_charge(smap, owner, smap->elem_size))
 		return NULL;
 
-	selem = bpf_map_kzalloc(&smap->map, smap->elem_size,
-				gfp_flags | __GFP_NOWARN);
+	migrate_disable();
+	selem = bpf_mem_cache_alloc(&smap->selem_ma);
+	migrate_enable();
+	if (!selem && (gfp_flags & GFP_KERNEL)) {
+		void *ma_node;
+
+		ma_node = bpf_map_kzalloc(&smap->map,
+					  BPF_MA_SIZE(smap->elem_size),
+					  gfp_flags | __GFP_NOWARN);
+		if (ma_node)
+			selem = BPF_MA_PTR(ma_node);
+	}
+
 	if (selem) {
 		if (value)
 			copy_map_value(&smap->map, SDATA(selem)->data, value);
-		/* No need to call check_and_init_map_value as memory is zero init */
+		else
+			/* Keep the original bpf_map_kzalloc behavior
+			 * before started using the bpf_mem_cache_alloc.
+			 *
+			 * No need to use zero_map_value. The bpf_selem_free()
+			 * only does bpf_mem_cache_free when there is
+			 * no other bpf prog is using the selem.
+			 */
+			memset(SDATA(selem)->data, 0, smap->map.value_size);
+
 		return selem;
 	}
 
@@ -129,7 +149,7 @@ static void bpf_selem_free_rcu(struct rcu_head *rcu)
 	struct bpf_local_storage_elem *selem;
 
 	selem = container_of(rcu, struct bpf_local_storage_elem, rcu);
-	kfree(selem);
+	kfree(BPF_MA_NODE(selem));
 }
 
 static void bpf_selem_free_trace_rcu(struct rcu_head *rcu)
@@ -145,10 +165,13 @@ void bpf_selem_free(struct bpf_local_storage_elem *selem,
 		    bool reuse_now)
 {
 	bpf_obj_free_fields(smap->map.record, SDATA(selem)->data);
-	if (!reuse_now)
+	if (!reuse_now) {
 		call_rcu_tasks_trace(&selem->rcu, bpf_selem_free_trace_rcu);
-	else
-		call_rcu(&selem->rcu, bpf_selem_free_rcu);
+	} else {
+		migrate_disable();
+		bpf_mem_cache_free(&smap->selem_ma, selem);
+		migrate_enable();
+	}
 }
 
 /* local_storage->lock must be held and selem->local_storage == local_storage.
@@ -651,6 +674,7 @@ bpf_local_storage_map_alloc(union bpf_attr *attr,
 	struct bpf_local_storage_map *smap;
 	unsigned int i;
 	u32 nbuckets;
+	int err;
 
 	smap = bpf_map_area_alloc(sizeof(*smap), NUMA_NO_NODE);
 	if (!smap)
@@ -665,8 +689,8 @@ bpf_local_storage_map_alloc(union bpf_attr *attr,
 	smap->buckets = bpf_map_kvcalloc(&smap->map, sizeof(*smap->buckets),
 					 nbuckets, GFP_USER | __GFP_NOWARN);
 	if (!smap->buckets) {
-		bpf_map_area_free(smap);
-		return ERR_PTR(-ENOMEM);
+		err = -ENOMEM;
+		goto free_smap;
 	}
 
 	for (i = 0; i < nbuckets; i++) {
@@ -677,8 +701,17 @@ bpf_local_storage_map_alloc(union bpf_attr *attr,
 	smap->elem_size = offsetof(struct bpf_local_storage_elem,
 				   sdata.data[attr->value_size]);
 
+	err = bpf_mem_alloc_init(&smap->selem_ma, smap->elem_size, false);
+	if (err)
+		goto free_smap;
+
 	smap->cache_idx = bpf_local_storage_cache_idx_get(cache);
 	return &smap->map;
+
+free_smap:
+	kvfree(smap->buckets);
+	bpf_map_area_free(smap);
+	return ERR_PTR(err);
 }
 
 void bpf_local_storage_map_free(struct bpf_map *map,
@@ -744,6 +777,7 @@ void bpf_local_storage_map_free(struct bpf_map *map,
 	 */
 	synchronize_rcu();
 
+	bpf_mem_alloc_destroy(&smap->selem_ma);
 	kvfree(smap->buckets);
 	bpf_map_area_free(smap);
 }
-- 
2.30.2

Re: [PATCH bpf-next 12/16] bpf: Use bpf_mem_cache_alloc/free in bpf_selem_alloc/free

[Alexei Starovoitov]

On Mon, Mar 6, 2023 at 12:43 AM Martin KaFai Lau <martin.lau@linux.dev> wrote:
>
> From: Martin KaFai Lau <martin.lau@kernel.org>
>
> This patch uses bpf_mem_cache_alloc/free in bpf_selem_alloc/free.
>
> The ____cacheline_aligned attribute is no longer needed
> in 'struct bpf_local_storage_elem'. bpf_mem_cache_alloc will
> have 'struct llist_node' in front of the 'struct bpf_local_storage_elem'.
> It will use the 8bytes hole in the bpf_local_storage_elem.
>
> After bpf_mem_cache_alloc(), the SDATA(selem)->data is zero-ed because
> bpf_mem_cache_alloc() could return a reused selem. It is to keep
> the existing bpf_map_kzalloc() behavior. Only SDATA(selem)->data
> is zero-ed. SDATA(selem)->data is the visible part to the bpf prog.
> No need to use zero_map_value() to do the zeroing because
> bpf_selem_free() ensures no bpf prog is using the selem before
> returning the selem through bpf_mem_cache_free(). For the internal
> fields of selem, they will be initialized when linking to the
> new smap and the new local_storage.
>
> When bpf_mem_cache_alloc() fails, bpf_selem_alloc() will try to
> fallback to kzalloc only if the caller has GFP_KERNEL flag set (ie. from
> sleepable bpf prog that should not cause deadlock). BPF_MA_SIZE
> and BPF_MA_PTR macro are added for that.
>
> For the common selem free path where the selem is freed when its owner
> is also being freed, reuse_now == true and selem can be reused
> immediately. bpf_selem_free() uses bpf_mem_cache_free() where
> selem will be considered for immediate reuse.
>
> For the uncommon path that the bpf prog explicitly deletes the selem (by
> using the helper bpf_*_storage_delete), the selem cannot be reused
> immediately. reuse_now == false and bpf_selem_free() will stay with
> the current call_rcu_tasks_trace. BPF_MA_NODE macro is added to get
> the correct address for the kfree.
>
> mem_charge and mem_uncharge are changed to use the BPF_MA_SIZE
> macro. It will have a temporarily over-charge for the
> bpf_local_storage_alloc() because bpf_local_storage is not
> moved to bpf_mem_cache_alloc in this patch but it will be done
> in the next patch.
>
> Cc: Namhyung Kim <namhyung@kernel.org>
> Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
> ---
>  include/linux/bpf_local_storage.h |  8 ++---
>  include/linux/bpf_mem_alloc.h     |  5 +++
>  kernel/bpf/bpf_local_storage.c    | 56 +++++++++++++++++++++++++------
>  3 files changed, 53 insertions(+), 16 deletions(-)
>
> diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
> index adb5023a1af5..a236c9b964cf 100644
> --- a/include/linux/bpf_local_storage.h
> +++ b/include/linux/bpf_local_storage.h
> @@ -13,6 +13,7 @@
>  #include <linux/list.h>
>  #include <linux/hash.h>
>  #include <linux/types.h>
> +#include <linux/bpf_mem_alloc.h>
>  #include <uapi/linux/btf.h>
>
>  #define BPF_LOCAL_STORAGE_CACHE_SIZE   16
> @@ -55,6 +56,7 @@ struct bpf_local_storage_map {
>         u32 bucket_log;
>         u16 elem_size;
>         u16 cache_idx;
> +       struct bpf_mem_alloc selem_ma;
>  };
>
>  struct bpf_local_storage_data {
> @@ -74,11 +76,7 @@ struct bpf_local_storage_elem {
>         struct hlist_node snode;        /* Linked to bpf_local_storage */
>         struct bpf_local_storage __rcu *local_storage;
>         struct rcu_head rcu;
> -       /* 8 bytes hole */
> -       /* The data is stored in another cacheline to minimize
> -        * the number of cachelines access during a cache hit.
> -        */
> -       struct bpf_local_storage_data sdata ____cacheline_aligned;
> +       struct bpf_local_storage_data sdata;
>  };
>
>  struct bpf_local_storage {
> diff --git a/include/linux/bpf_mem_alloc.h b/include/linux/bpf_mem_alloc.h
> index a7104af61ab4..0ab16fb0ab50 100644
> --- a/include/linux/bpf_mem_alloc.h
> +++ b/include/linux/bpf_mem_alloc.h
> @@ -5,6 +5,11 @@
>  #include <linux/compiler_types.h>
>  #include <linux/workqueue.h>
>
> +#define BPF_MA_NODE_SZ sizeof(struct llist_node)
> +#define BPF_MA_SIZE(_size) ((_size) + BPF_MA_NODE_SZ)
> +#define BPF_MA_PTR(_node) ((void *)(_node) + BPF_MA_NODE_SZ)
> +#define BPF_MA_NODE(_ptr) ((void *)(_ptr) - BPF_MA_NODE_SZ)
> +
>  struct bpf_mem_cache;
>  struct bpf_mem_caches;
>
> diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
> index 532b82084ba7..d3c0dd5737d6 100644
> --- a/kernel/bpf/bpf_local_storage.c
> +++ b/kernel/bpf/bpf_local_storage.c
> @@ -31,7 +31,7 @@ static int mem_charge(struct bpf_local_storage_map *smap, void *owner, u32 size)
>         if (!map->ops->map_local_storage_charge)
>                 return 0;
>
> -       return map->ops->map_local_storage_charge(smap, owner, size);
> +       return map->ops->map_local_storage_charge(smap, owner, BPF_MA_SIZE(size));
>  }
>
>  static void mem_uncharge(struct bpf_local_storage_map *smap, void *owner,
> @@ -40,7 +40,7 @@ static void mem_uncharge(struct bpf_local_storage_map *smap, void *owner,
>         struct bpf_map *map = &smap->map;
>
>         if (map->ops->map_local_storage_uncharge)
> -               map->ops->map_local_storage_uncharge(smap, owner, size);
> +               map->ops->map_local_storage_uncharge(smap, owner, BPF_MA_SIZE(size));
>  }
>
>  static struct bpf_local_storage __rcu **
> @@ -80,12 +80,32 @@ bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner,
>         if (charge_mem && mem_charge(smap, owner, smap->elem_size))
>                 return NULL;
>
> -       selem = bpf_map_kzalloc(&smap->map, smap->elem_size,
> -                               gfp_flags | __GFP_NOWARN);
> +       migrate_disable();
> +       selem = bpf_mem_cache_alloc(&smap->selem_ma);
> +       migrate_enable();
> +       if (!selem && (gfp_flags & GFP_KERNEL)) {
> +               void *ma_node;
> +
> +               ma_node = bpf_map_kzalloc(&smap->map,
> +                                         BPF_MA_SIZE(smap->elem_size),
> +                                         gfp_flags | __GFP_NOWARN);
> +               if (ma_node)
> +                       selem = BPF_MA_PTR(ma_node);
> +       }

If I understand it correctly the code is not trying
to free selem the same way it allocated it.
So we can have kzalloc-ed selems freed into bpf_mem_cache_alloc free-list.
That feels dangerous.
I don't think we can do such things in local storage,
but if we add this api to bpf_mem_alloc it might be acceptable.
I mean mem alloc will try to take from the free list and if empty
and GFP_KERNEL it will kzalloc it.
The knowledge of hidden llist_node shouldn't leave the bpf/memalloc.c file.
reuse_now should probably be a memalloc api flag too.
The implementation detail that it's scary but ok-ish to kfree or
bpf_mem_cache_free depending on circumstances should stay in memalloc.c

Re: [PATCH bpf-next 12/16] bpf: Use bpf_mem_cache_alloc/free in bpf_selem_alloc/free

[Martin KaFai Lau]

On 3/6/23 7:47 PM, Alexei Starovoitov wrote:
>> @@ -80,12 +80,32 @@ bpf_selem_alloc(struct bpf_local_storage_map *smap, void *owner,
>>          if (charge_mem && mem_charge(smap, owner, smap->elem_size))
>>                  return NULL;
>>
>> -       selem = bpf_map_kzalloc(&smap->map, smap->elem_size,
>> -                               gfp_flags | __GFP_NOWARN);
>> +       migrate_disable();
>> +       selem = bpf_mem_cache_alloc(&smap->selem_ma);
>> +       migrate_enable();
>> +       if (!selem && (gfp_flags & GFP_KERNEL)) {
>> +               void *ma_node;
>> +
>> +               ma_node = bpf_map_kzalloc(&smap->map,
>> +                                         BPF_MA_SIZE(smap->elem_size),
>> +                                         gfp_flags | __GFP_NOWARN);
>> +               if (ma_node)
>> +                       selem = BPF_MA_PTR(ma_node);
>> +       }
> 
> If I understand it correctly the code is not trying
> to free selem the same way it allocated it.
> So we can have kzalloc-ed selems freed into bpf_mem_cache_alloc free-list.
> That feels dangerous.
> I don't think we can do such things in local storage,
> but if we add this api to bpf_mem_alloc it might be acceptable.
> I mean mem alloc will try to take from the free list and if empty
> and GFP_KERNEL it will kzalloc it.
> The knowledge of hidden llist_node shouldn't leave the bpf/memalloc.c file.
> reuse_now should probably be a memalloc api flag too.
> The implementation detail that it's scary but ok-ish to kfree or
> bpf_mem_cache_free depending on circumstances should stay in memalloc.c

All make sense. I will create a bpf_mem_cache_alloc_flags(..., gfp_t flags) to 
hide the llist_node and kzalloc details. For free, local storage still needs to 
use the selem->rcu head in its call_rcu_tasks_trace(), so I will create a 
bpf_mem_cache_raw_free(void *ptr) to hide the llist_node details, like:

/* 'struct bpf_mem_alloc *ma' is not available at this
  * point but the caller knows it is percpu or not and
  * call different raw_free function.
  */
void bpf_mem_cache_raw_free(void *ptr)
{
         kfree(ptr - LLIST_NODE_SZ);
}

[PATCH bpf-next 13/16] bpf: Use bpf_mem_cache_alloc/free for bpf_local_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch uses bpf_mem_cache_alloc/free for allocating and freeing
bpf_local_storage.

The changes are similar to the previous patch when switching
bpf_local_storage_elem to bpf_mem_cache_alloc/free.

A few things that worth to mention for bpf_local_storage:

The local_storage is deleted when the last selem is deleted.
Before deleting a selem from local_storage, it needs to retrieve the
local_storage->smap because the bpf_selem_unlink_storage_nolock()
may have set it to NULL. Note that local_storage->smap may have
already been NULL when the selem created this local_storage has
been removed. In this case, call_rcu will be used to free the
local_storage.

When bpf_local_storage_alloc getting a reused memory, all
fields are either in the correct values or will be initialized.
'cache[]' must already be all NULLs. 'list' must be empty.
Others will be initialized.

Cc: Namhyung Kim <namhyung@kernel.org>
Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 include/linux/bpf_local_storage.h |  1 +
 kernel/bpf/bpf_local_storage.c    | 54 ++++++++++++++++++++++++++-----
 2 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/include/linux/bpf_local_storage.h b/include/linux/bpf_local_storage.h
index a236c9b964cf..d6d4ec248c00 100644
--- a/include/linux/bpf_local_storage.h
+++ b/include/linux/bpf_local_storage.h
@@ -57,6 +57,7 @@ struct bpf_local_storage_map {
 	u16 elem_size;
 	u16 cache_idx;
 	struct bpf_mem_alloc selem_ma;
+	struct bpf_mem_alloc storage_ma;
 };
 
 struct bpf_local_storage_data {
diff --git a/kernel/bpf/bpf_local_storage.c b/kernel/bpf/bpf_local_storage.c
index d3c0dd5737d6..79f84485069d 100644
--- a/kernel/bpf/bpf_local_storage.c
+++ b/kernel/bpf/bpf_local_storage.c
@@ -120,7 +120,7 @@ static void bpf_local_storage_free_rcu(struct rcu_head *rcu)
 	struct bpf_local_storage *local_storage;
 
 	local_storage = container_of(rcu, struct bpf_local_storage, rcu);
-	kfree(local_storage);
+	kfree(BPF_MA_NODE(local_storage));
 }
 
 static void bpf_local_storage_free_trace_rcu(struct rcu_head *rcu)
@@ -135,13 +135,26 @@ static void bpf_local_storage_free_trace_rcu(struct rcu_head *rcu)
 }
 
 static void bpf_local_storage_free(struct bpf_local_storage *local_storage,
+				   struct bpf_local_storage_map *smap,
 				   bool reuse_now)
 {
-	if (!reuse_now)
+	if (!reuse_now) {
 		call_rcu_tasks_trace(&local_storage->rcu,
 				     bpf_local_storage_free_trace_rcu);
-	else
+		return;
+	}
+
+	/* smap could be NULL if the selem that triggered this 'local_storage'
+	 * creation had been long gone. In this case, directly do
+	 * call_rcu().
+	 */
+	if (smap) {
+		migrate_disable();
+		bpf_mem_cache_free(&smap->storage_ma, local_storage);
+		migrate_enable();
+	} else {
 		call_rcu(&local_storage->rcu, bpf_local_storage_free_rcu);
+	}
 }
 
 static void bpf_selem_free_rcu(struct rcu_head *rcu)
@@ -235,6 +248,7 @@ static bool bpf_selem_unlink_storage_nolock(struct bpf_local_storage *local_stor
 static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 				     bool reuse_now)
 {
+	struct bpf_local_storage_map *storage_smap;
 	struct bpf_local_storage *local_storage;
 	bool free_local_storage = false;
 	unsigned long flags;
@@ -245,6 +259,8 @@ static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 
 	local_storage = rcu_dereference_check(selem->local_storage,
 					      bpf_rcu_lock_held());
+	storage_smap = rcu_dereference_check(local_storage->smap,
+					     bpf_rcu_lock_held());
 	raw_spin_lock_irqsave(&local_storage->lock, flags);
 	if (likely(selem_linked_to_storage(selem)))
 		free_local_storage = bpf_selem_unlink_storage_nolock(
@@ -252,7 +268,8 @@ static void bpf_selem_unlink_storage(struct bpf_local_storage_elem *selem,
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
 	if (free_local_storage)
-		bpf_local_storage_free(local_storage, reuse_now);
+		bpf_local_storage_free(local_storage, storage_smap,
+				       reuse_now);
 }
 
 void bpf_selem_link_storage_nolock(struct bpf_local_storage *local_storage,
@@ -372,8 +389,18 @@ int bpf_local_storage_alloc(void *owner,
 	if (err)
 		return err;
 
-	storage = bpf_map_kzalloc(&smap->map, sizeof(*storage),
-				  gfp_flags | __GFP_NOWARN);
+	migrate_disable();
+	storage = bpf_mem_cache_alloc(&smap->storage_ma);
+	migrate_enable();
+	if (!storage && (gfp_flags & GFP_KERNEL)) {
+		void *ma_obj;
+
+		ma_obj = bpf_map_kzalloc(&smap->map, BPF_MA_SIZE(sizeof(*storage)),
+					 gfp_flags | __GFP_NOWARN);
+		if (ma_obj)
+			storage = BPF_MA_PTR(ma_obj);
+	}
+
 	if (!storage) {
 		err = -ENOMEM;
 		goto uncharge;
@@ -419,7 +446,7 @@ int bpf_local_storage_alloc(void *owner,
 	return 0;
 
 uncharge:
-	bpf_local_storage_free(storage, true);
+	bpf_local_storage_free(storage, smap, true);
 	mem_uncharge(smap, owner, sizeof(*storage));
 	return err;
 }
@@ -632,11 +659,15 @@ int bpf_local_storage_map_check_btf(const struct bpf_map *map,
 
 void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 {
+	struct bpf_local_storage_map *storage_smap;
 	struct bpf_local_storage_elem *selem;
 	bool free_storage = false;
 	struct hlist_node *n;
 	unsigned long flags;
 
+	storage_smap = rcu_dereference_check(local_storage->smap,
+					     bpf_rcu_lock_held());
+
 	/* Neither the bpf_prog nor the bpf_map's syscall
 	 * could be modifying the local_storage->list now.
 	 * Thus, no elem can be added to or deleted from the
@@ -664,7 +695,7 @@ void bpf_local_storage_destroy(struct bpf_local_storage *local_storage)
 	raw_spin_unlock_irqrestore(&local_storage->lock, flags);
 
 	if (free_storage)
-		bpf_local_storage_free(local_storage, true);
+		bpf_local_storage_free(local_storage, storage_smap, true);
 }
 
 struct bpf_map *
@@ -705,6 +736,12 @@ bpf_local_storage_map_alloc(union bpf_attr *attr,
 	if (err)
 		goto free_smap;
 
+	err = bpf_mem_alloc_init(&smap->storage_ma, sizeof(struct bpf_local_storage), false);
+	if (err) {
+		bpf_mem_alloc_destroy(&smap->selem_ma);
+		goto free_smap;
+	}
+
 	smap->cache_idx = bpf_local_storage_cache_idx_get(cache);
 	return &smap->map;
 
@@ -778,6 +815,7 @@ void bpf_local_storage_map_free(struct bpf_map *map,
 	synchronize_rcu();
 
 	bpf_mem_alloc_destroy(&smap->selem_ma);
+	bpf_mem_alloc_destroy(&smap->storage_ma);
 	kvfree(smap->buckets);
 	bpf_map_area_free(smap);
 }
-- 
2.30.2

[PATCH bpf-next 14/16] selftests/bpf: Replace CHECK with ASSERT in test_local_storage

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch migrates the CHECK macro to ASSERT macro.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 .../bpf/prog_tests/test_local_storage.c       | 49 +++++++------------
 1 file changed, 19 insertions(+), 30 deletions(-)

diff --git a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
index 9c77cd6b1eaf..c33f840f4880 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
@@ -13,8 +13,6 @@
 #include "network_helpers.h"
 #include "task_local_storage_helpers.h"
 
-static unsigned int duration;
-
 #define TEST_STORAGE_VALUE 0xbeefdead
 
 struct storage {
@@ -60,36 +58,32 @@ static bool check_syscall_operations(int map_fd, int obj_fd)
 
 	/* Looking up an existing element should fail initially */
 	err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
-	if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
-		  "err:%d errno:%d\n", err, errno))
+	if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
+	    !ASSERT_EQ(errno, ENOENT, "errno"))
 		return false;
 
 	/* Create a new element */
 	err = bpf_map_update_elem(map_fd, &obj_fd, &val, BPF_NOEXIST);
-	if (CHECK(err < 0, "bpf_map_update_elem", "err:%d errno:%d\n", err,
-		  errno))
+	if (!ASSERT_OK(err, "bpf_map_update_elem"))
 		return false;
 
 	/* Lookup the newly created element */
 	err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
-	if (CHECK(err < 0, "bpf_map_lookup_elem", "err:%d errno:%d", err,
-		  errno))
+	if (!ASSERT_OK(err, "bpf_map_lookup_elem"))
 		return false;
 
 	/* Check the value of the newly created element */
-	if (CHECK(lookup_val.value != val.value, "bpf_map_lookup_elem",
-		  "value got = %x errno:%d", lookup_val.value, val.value))
+	if (!ASSERT_EQ(lookup_val.value, val.value, "bpf_map_lookup_elem"))
 		return false;
 
 	err = bpf_map_delete_elem(map_fd, &obj_fd);
-	if (CHECK(err, "bpf_map_delete_elem()", "err:%d errno:%d\n", err,
-		  errno))
+	if (!ASSERT_OK(err, "bpf_map_delete_elem()"))
 		return false;
 
 	/* The lookup should fail, now that the element has been deleted */
 	err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
-	if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
-		  "err:%d errno:%d\n", err, errno))
+	if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
+	    !ASSERT_EQ(errno, ENOENT, "errno"))
 		return false;
 
 	return true;
@@ -104,35 +98,32 @@ void test_test_local_storage(void)
 	char cmd[256];
 
 	skel = local_storage__open_and_load();
-	if (CHECK(!skel, "skel_load", "lsm skeleton failed\n"))
+	if (!ASSERT_OK_PTR(skel, "skel_load"))
 		goto close_prog;
 
 	err = local_storage__attach(skel);
-	if (CHECK(err, "attach", "lsm attach failed: %d\n", err))
+	if (!ASSERT_OK(err, "attach"))
 		goto close_prog;
 
 	task_fd = sys_pidfd_open(getpid(), 0);
-	if (CHECK(task_fd < 0, "pidfd_open",
-		  "failed to get pidfd err:%d, errno:%d", task_fd, errno))
+	if (!ASSERT_GE(task_fd, 0, "pidfd_open"))
 		goto close_prog;
 
 	if (!check_syscall_operations(bpf_map__fd(skel->maps.task_storage_map),
 				      task_fd))
 		goto close_prog;
 
-	if (CHECK(!mkdtemp(tmp_dir_path), "mkdtemp",
-		  "unable to create tmpdir: %d\n", errno))
+	if (!ASSERT_OK_PTR(mkdtemp(tmp_dir_path), "mkdtemp"))
 		goto close_prog;
 
 	snprintf(tmp_exec_path, sizeof(tmp_exec_path), "%s/copy_of_rm",
 		 tmp_dir_path);
 	snprintf(cmd, sizeof(cmd), "cp /bin/rm %s", tmp_exec_path);
-	if (CHECK_FAIL(system(cmd)))
+	if (!ASSERT_OK(system(cmd), "system(cp)"))
 		goto close_prog_rmdir;
 
 	rm_fd = open(tmp_exec_path, O_RDONLY);
-	if (CHECK(rm_fd < 0, "open", "failed to open %s err:%d, errno:%d",
-		  tmp_exec_path, rm_fd, errno))
+	if (!ASSERT_GE(rm_fd, 0, "open(tmp_exec_path)"))
 		goto close_prog_rmdir;
 
 	if (!check_syscall_operations(bpf_map__fd(skel->maps.inode_storage_map),
@@ -145,7 +136,7 @@ void test_test_local_storage(void)
 	 * LSM program.
 	 */
 	err = run_self_unlink(&skel->bss->monitored_pid, tmp_exec_path);
-	if (CHECK(err != EPERM, "run_self_unlink", "err %d want EPERM\n", err))
+	if (!ASSERT_EQ(err, EPERM, "run_self_unlink"))
 		goto close_prog_rmdir;
 
 	/* Set the process being monitored to be the current process */
@@ -156,18 +147,16 @@ void test_test_local_storage(void)
 	 */
 	snprintf(cmd, sizeof(cmd), "mv %s/copy_of_rm %s/check_null_ptr",
 		 tmp_dir_path, tmp_dir_path);
-	if (CHECK_FAIL(system(cmd)))
+	if (!ASSERT_OK(system(cmd), "system(mv)"))
 		goto close_prog_rmdir;
 
-	CHECK(skel->data->inode_storage_result != 0, "inode_storage_result",
-	      "inode_local_storage not set\n");
+	ASSERT_EQ(skel->data->inode_storage_result, 0, "inode_storage_result");
 
 	serv_sk = start_server(AF_INET6, SOCK_STREAM, NULL, 0, 0);
-	if (CHECK(serv_sk < 0, "start_server", "failed to start server\n"))
+	if (!ASSERT_GE(serv_sk, 0, "start_server"))
 		goto close_prog_rmdir;
 
-	CHECK(skel->data->sk_storage_result != 0, "sk_storage_result",
-	      "sk_local_storage not set\n");
+	ASSERT_EQ(skel->data->sk_storage_result, 0, "sk_storage_result");
 
 	if (!check_syscall_operations(bpf_map__fd(skel->maps.sk_storage_map),
 				      serv_sk))
-- 
2.30.2

Re: [PATCH bpf-next 14/16] selftests/bpf: Replace CHECK with ASSERT in test_local_storage

[Andrii Nakryiko]

On Mon, Mar 6, 2023 at 12:43 AM Martin KaFai Lau <martin.lau@linux.dev> wrote:
>
> From: Martin KaFai Lau <martin.lau@kernel.org>
>
> This patch migrates the CHECK macro to ASSERT macro.
>
> Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
> ---

Thanks for the cleanup!

>  .../bpf/prog_tests/test_local_storage.c       | 49 +++++++------------
>  1 file changed, 19 insertions(+), 30 deletions(-)
>
> diff --git a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
> index 9c77cd6b1eaf..c33f840f4880 100644
> --- a/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
> +++ b/tools/testing/selftests/bpf/prog_tests/test_local_storage.c
> @@ -13,8 +13,6 @@
>  #include "network_helpers.h"
>  #include "task_local_storage_helpers.h"
>
> -static unsigned int duration;
> -
>  #define TEST_STORAGE_VALUE 0xbeefdead
>
>  struct storage {
> @@ -60,36 +58,32 @@ static bool check_syscall_operations(int map_fd, int obj_fd)
>
>         /* Looking up an existing element should fail initially */
>         err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
> -       if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
> -                 "err:%d errno:%d\n", err, errno))
> +       if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
> +           !ASSERT_EQ(errno, ENOENT, "errno"))

all libbpf APIs since v1.0 always return actual error number directly,
so no need to check errno anymore, you can simplify this further

>                 return false;
>
>         /* Create a new element */
>         err = bpf_map_update_elem(map_fd, &obj_fd, &val, BPF_NOEXIST);
> -       if (CHECK(err < 0, "bpf_map_update_elem", "err:%d errno:%d\n", err,
> -                 errno))
> +       if (!ASSERT_OK(err, "bpf_map_update_elem"))
>                 return false;
>
>         /* Lookup the newly created element */
>         err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
> -       if (CHECK(err < 0, "bpf_map_lookup_elem", "err:%d errno:%d", err,
> -                 errno))
> +       if (!ASSERT_OK(err, "bpf_map_lookup_elem"))
>                 return false;
>
>         /* Check the value of the newly created element */
> -       if (CHECK(lookup_val.value != val.value, "bpf_map_lookup_elem",
> -                 "value got = %x errno:%d", lookup_val.value, val.value))
> +       if (!ASSERT_EQ(lookup_val.value, val.value, "bpf_map_lookup_elem"))
>                 return false;
>
>         err = bpf_map_delete_elem(map_fd, &obj_fd);
> -       if (CHECK(err, "bpf_map_delete_elem()", "err:%d errno:%d\n", err,
> -                 errno))
> +       if (!ASSERT_OK(err, "bpf_map_delete_elem()"))
>                 return false;
>
>         /* The lookup should fail, now that the element has been deleted */
>         err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
> -       if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
> -                 "err:%d errno:%d\n", err, errno))
> +       if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
> +           !ASSERT_EQ(errno, ENOENT, "errno"))

same here and probably in other places (I haven't checked everything)

>                 return false;
>
>         return true;
> @@ -104,35 +98,32 @@ void test_test_local_storage(void)
>         char cmd[256];
>
>         skel = local_storage__open_and_load();
> -       if (CHECK(!skel, "skel_load", "lsm skeleton failed\n"))
> +       if (!ASSERT_OK_PTR(skel, "skel_load"))
>                 goto close_prog;
>
>         err = local_storage__attach(skel);
> -       if (CHECK(err, "attach", "lsm attach failed: %d\n", err))
> +       if (!ASSERT_OK(err, "attach"))
>                 goto close_prog;
>
>         task_fd = sys_pidfd_open(getpid(), 0);
> -       if (CHECK(task_fd < 0, "pidfd_open",
> -                 "failed to get pidfd err:%d, errno:%d", task_fd, errno))
> +       if (!ASSERT_GE(task_fd, 0, "pidfd_open"))
>                 goto close_prog;
>
>         if (!check_syscall_operations(bpf_map__fd(skel->maps.task_storage_map),
>                                       task_fd))
>                 goto close_prog;
>
> -       if (CHECK(!mkdtemp(tmp_dir_path), "mkdtemp",
> -                 "unable to create tmpdir: %d\n", errno))
> +       if (!ASSERT_OK_PTR(mkdtemp(tmp_dir_path), "mkdtemp"))
>                 goto close_prog;
>
>         snprintf(tmp_exec_path, sizeof(tmp_exec_path), "%s/copy_of_rm",
>                  tmp_dir_path);
>         snprintf(cmd, sizeof(cmd), "cp /bin/rm %s", tmp_exec_path);
> -       if (CHECK_FAIL(system(cmd)))
> +       if (!ASSERT_OK(system(cmd), "system(cp)"))
>                 goto close_prog_rmdir;
>
>         rm_fd = open(tmp_exec_path, O_RDONLY);
> -       if (CHECK(rm_fd < 0, "open", "failed to open %s err:%d, errno:%d",
> -                 tmp_exec_path, rm_fd, errno))
> +       if (!ASSERT_GE(rm_fd, 0, "open(tmp_exec_path)"))
>                 goto close_prog_rmdir;
>
>         if (!check_syscall_operations(bpf_map__fd(skel->maps.inode_storage_map),
> @@ -145,7 +136,7 @@ void test_test_local_storage(void)
>          * LSM program.
>          */
>         err = run_self_unlink(&skel->bss->monitored_pid, tmp_exec_path);
> -       if (CHECK(err != EPERM, "run_self_unlink", "err %d want EPERM\n", err))
> +       if (!ASSERT_EQ(err, EPERM, "run_self_unlink"))
>                 goto close_prog_rmdir;
>
>         /* Set the process being monitored to be the current process */
> @@ -156,18 +147,16 @@ void test_test_local_storage(void)
>          */
>         snprintf(cmd, sizeof(cmd), "mv %s/copy_of_rm %s/check_null_ptr",
>                  tmp_dir_path, tmp_dir_path);
> -       if (CHECK_FAIL(system(cmd)))
> +       if (!ASSERT_OK(system(cmd), "system(mv)"))
>                 goto close_prog_rmdir;
>
> -       CHECK(skel->data->inode_storage_result != 0, "inode_storage_result",
> -             "inode_local_storage not set\n");
> +       ASSERT_EQ(skel->data->inode_storage_result, 0, "inode_storage_result");
>
>         serv_sk = start_server(AF_INET6, SOCK_STREAM, NULL, 0, 0);
> -       if (CHECK(serv_sk < 0, "start_server", "failed to start server\n"))
> +       if (!ASSERT_GE(serv_sk, 0, "start_server"))
>                 goto close_prog_rmdir;
>
> -       CHECK(skel->data->sk_storage_result != 0, "sk_storage_result",
> -             "sk_local_storage not set\n");
> +       ASSERT_EQ(skel->data->sk_storage_result, 0, "sk_storage_result");
>
>         if (!check_syscall_operations(bpf_map__fd(skel->maps.sk_storage_map),
>                                       serv_sk))
> --
> 2.30.2
>

Re: [PATCH bpf-next 14/16] selftests/bpf: Replace CHECK with ASSERT in test_local_storage

[Martin KaFai Lau]

On 3/7/23 5:15 PM, Andrii Nakryiko wrote:
>> @@ -60,36 +58,32 @@ static bool check_syscall_operations(int map_fd, int obj_fd)
>>
>>          /* Looking up an existing element should fail initially */
>>          err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
>> -       if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
>> -                 "err:%d errno:%d\n", err, errno))
>> +       if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
>> +           !ASSERT_EQ(errno, ENOENT, "errno"))
> 
> all libbpf APIs since v1.0 always return actual error number directly,
> so no need to check errno anymore, you can simplify this further
> 
>>                  return false;
>>
>>          /* Create a new element */
>>          err = bpf_map_update_elem(map_fd, &obj_fd, &val, BPF_NOEXIST);
>> -       if (CHECK(err < 0, "bpf_map_update_elem", "err:%d errno:%d\n", err,
>> -                 errno))
>> +       if (!ASSERT_OK(err, "bpf_map_update_elem"))
>>                  return false;
>>
>>          /* Lookup the newly created element */
>>          err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
>> -       if (CHECK(err < 0, "bpf_map_lookup_elem", "err:%d errno:%d", err,
>> -                 errno))
>> +       if (!ASSERT_OK(err, "bpf_map_lookup_elem"))
>>                  return false;
>>
>>          /* Check the value of the newly created element */
>> -       if (CHECK(lookup_val.value != val.value, "bpf_map_lookup_elem",
>> -                 "value got = %x errno:%d", lookup_val.value, val.value))
>> +       if (!ASSERT_EQ(lookup_val.value, val.value, "bpf_map_lookup_elem"))
>>                  return false;
>>
>>          err = bpf_map_delete_elem(map_fd, &obj_fd);
>> -       if (CHECK(err, "bpf_map_delete_elem()", "err:%d errno:%d\n", err,
>> -                 errno))
>> +       if (!ASSERT_OK(err, "bpf_map_delete_elem()"))
>>                  return false;
>>
>>          /* The lookup should fail, now that the element has been deleted */
>>          err = bpf_map_lookup_elem_flags(map_fd, &obj_fd, &lookup_val, 0);
>> -       if (CHECK(!err || errno != ENOENT, "bpf_map_lookup_elem",
>> -                 "err:%d errno:%d\n", err, errno))
>> +       if (!ASSERT_ERR(err, "bpf_map_lookup_elem") ||
>> +           !ASSERT_EQ(errno, ENOENT, "errno"))
> 
> same here and probably in other places (I haven't checked everything)

Ack. will simplify.

[PATCH bpf-next 15/16] selftests/bpf: Check freeing sk->sk_local_storage with sk_local_storage->smap is NULL

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch tweats the socket_bind bpf prog to test the
local_storage->smap == NULL case in the bpf_local_storage_free()
code path. The idea is to create the local_storage with
the sk_storage_map's selem first. Then add the sk_storage_map2's selem
and then delete the earlier sk_storeage_map's selem.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 .../selftests/bpf/progs/local_storage.c       | 29 +++++++++++++------
 1 file changed, 20 insertions(+), 9 deletions(-)

diff --git a/tools/testing/selftests/bpf/progs/local_storage.c b/tools/testing/selftests/bpf/progs/local_storage.c
index 19423ed862e3..797c81655a47 100644
--- a/tools/testing/selftests/bpf/progs/local_storage.c
+++ b/tools/testing/selftests/bpf/progs/local_storage.c
@@ -109,18 +109,17 @@ int BPF_PROG(socket_bind, struct socket *sock, struct sockaddr *address,
 {
 	__u32 pid = bpf_get_current_pid_tgid() >> 32;
 	struct local_storage *storage;
-	int err;
 
 	if (pid != monitored_pid)
 		return 0;
 
-	storage = bpf_sk_storage_get(&sk_storage_map, sock->sk, 0,
-				     BPF_LOCAL_STORAGE_GET_F_CREATE);
+	storage = bpf_sk_storage_get(&sk_storage_map, sock->sk, 0, 0);
 	if (!storage)
 		return 0;
 
+	sk_storage_result = -1;
 	if (storage->value != DUMMY_STORAGE_VALUE)
-		sk_storage_result = -1;
+		return 0;
 
 	/* This tests that we can associate multiple elements
 	 * with the local storage.
@@ -130,14 +129,26 @@ int BPF_PROG(socket_bind, struct socket *sock, struct sockaddr *address,
 	if (!storage)
 		return 0;
 
-	err = bpf_sk_storage_delete(&sk_storage_map, sock->sk);
-	if (err)
+	if (bpf_sk_storage_delete(&sk_storage_map2, sock->sk))
 		return 0;
 
-	err = bpf_sk_storage_delete(&sk_storage_map2, sock->sk);
-	if (!err)
-		sk_storage_result = err;
+	storage = bpf_sk_storage_get(&sk_storage_map2, sock->sk, 0,
+				     BPF_LOCAL_STORAGE_GET_F_CREATE);
+	if (!storage)
+		return 0;
+
+	if (bpf_sk_storage_delete(&sk_storage_map, sock->sk))
+		return 0;
+
+	/* Ensure that the sk_storage_map is disconnected from the storage.
+	 * The storage memory should not be freed back to the
+	 * bpf_mem_alloc of the sk_bpf_storage_map because
+	 * sk_bpf_storage_map may have been gone.
+	 */
+	if (!sock->sk->sk_bpf_storage || sock->sk->sk_bpf_storage->smap)
+		return 0;
 
+	sk_storage_result = 0;
 	return 0;
 }
 
-- 
2.30.2

[PATCH bpf-next 16/16] selftests/bpf: Add local-storage-create benchmark

[Martin KaFai Lau]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch tests how many kmallocs is needed to create and free
a batch of UDP sockets and each socket has a 64bytes bpf storage.
It also measures how fast the UDP sockets can be created.

The result is from my qemu setup.

Before bpf_mem_alloc/free:
./bench -p 1 local-storage-create
Setting up benchmark 'local-storage-create'...
Benchmark 'local-storage-create' started.
Iter   0 ( 73.193us): creates  213.552k/s (213.552k/prod), 3.09 kmallocs/create
Iter   1 (-20.724us): creates  211.908k/s (211.908k/prod), 3.09 kmallocs/create
Iter   2 (  9.280us): creates  212.574k/s (212.574k/prod), 3.12 kmallocs/create
Iter   3 ( 11.039us): creates  213.209k/s (213.209k/prod), 3.12 kmallocs/create
Iter   4 (-11.411us): creates  213.351k/s (213.351k/prod), 3.12 kmallocs/create
Iter   5 ( -7.915us): creates  214.754k/s (214.754k/prod), 3.12 kmallocs/create
Iter   6 ( 11.317us): creates  210.942k/s (210.942k/prod), 3.12 kmallocs/create
Summary: creates  212.789 ± 1.310k/s (212.789k/prod), 3.12 kmallocs/create

After bpf_mem_alloc/free:
./bench -p 1 local-storage-create
Setting up benchmark 'local-storage-create'...
Benchmark 'local-storage-create' started.
Iter   0 ( 68.265us): creates  243.984k/s (243.984k/prod), 1.04 kmallocs/create
Iter   1 ( 30.357us): creates  238.424k/s (238.424k/prod), 1.04 kmallocs/create
Iter   2 (-18.712us): creates  232.963k/s (232.963k/prod), 1.04 kmallocs/create
Iter   3 (-15.885us): creates  238.879k/s (238.879k/prod), 1.04 kmallocs/create
Iter   4 (  5.590us): creates  237.490k/s (237.490k/prod), 1.04 kmallocs/create
Iter   5 (  8.577us): creates  237.521k/s (237.521k/prod), 1.04 kmallocs/create
Iter   6 ( -6.263us): creates  238.508k/s (238.508k/prod), 1.04 kmallocs/create
Summary: creates  237.298 ± 2.198k/s (237.298k/prod), 1.04 kmallocs/create

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 tools/testing/selftests/bpf/Makefile          |   2 +
 tools/testing/selftests/bpf/bench.c           |   2 +
 .../bpf/benchs/bench_local_storage_create.c   | 141 ++++++++++++++++++
 .../bpf/progs/bench_local_storage_create.c    |  57 +++++++
 4 files changed, 202 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/benchs/bench_local_storage_create.c
 create mode 100644 tools/testing/selftests/bpf/progs/bench_local_storage_create.c

diff --git a/tools/testing/selftests/bpf/Makefile b/tools/testing/selftests/bpf/Makefile
index 16f404aa1b23..545fff0a6dfc 100644
--- a/tools/testing/selftests/bpf/Makefile
+++ b/tools/testing/selftests/bpf/Makefile
@@ -639,6 +639,7 @@ $(OUTPUT)/bench_strncmp.o: $(OUTPUT)/strncmp_bench.skel.h
 $(OUTPUT)/bench_bpf_hashmap_full_update.o: $(OUTPUT)/bpf_hashmap_full_update_bench.skel.h
 $(OUTPUT)/bench_local_storage.o: $(OUTPUT)/local_storage_bench.skel.h
 $(OUTPUT)/bench_local_storage_rcu_tasks_trace.o: $(OUTPUT)/local_storage_rcu_tasks_trace_bench.skel.h
+$(OUTPUT)/bench_local_storage_create.o: $(OUTPUT)/bench_local_storage_create.skel.h
 $(OUTPUT)/bench_bpf_hashmap_lookup.o: $(OUTPUT)/bpf_hashmap_lookup.skel.h
 $(OUTPUT)/bench.o: bench.h testing_helpers.h $(BPFOBJ)
 $(OUTPUT)/bench: LDLIBS += -lm
@@ -656,6 +657,7 @@ $(OUTPUT)/bench: $(OUTPUT)/bench.o \
 		 $(OUTPUT)/bench_local_storage.o \
 		 $(OUTPUT)/bench_local_storage_rcu_tasks_trace.o \
 		 $(OUTPUT)/bench_bpf_hashmap_lookup.o \
+		 $(OUTPUT)/bench_local_storage_create.o \
 		 #
 	$(call msg,BINARY,,$@)
 	$(Q)$(CC) $(CFLAGS) $(LDFLAGS) $(filter %.a %.o,$^) $(LDLIBS) -o $@
diff --git a/tools/testing/selftests/bpf/bench.c b/tools/testing/selftests/bpf/bench.c
index 0b2a53bb8460..dc3827c1f139 100644
--- a/tools/testing/selftests/bpf/bench.c
+++ b/tools/testing/selftests/bpf/bench.c
@@ -515,6 +515,7 @@ extern const struct bench bench_local_storage_cache_interleaved_get;
 extern const struct bench bench_local_storage_cache_hashmap_control;
 extern const struct bench bench_local_storage_tasks_trace;
 extern const struct bench bench_bpf_hashmap_lookup;
+extern const struct bench bench_local_storage_create;
 
 static const struct bench *benchs[] = {
 	&bench_count_global,
@@ -555,6 +556,7 @@ static const struct bench *benchs[] = {
 	&bench_local_storage_cache_hashmap_control,
 	&bench_local_storage_tasks_trace,
 	&bench_bpf_hashmap_lookup,
+	&bench_local_storage_create,
 };
 
 static void find_benchmark(void)
diff --git a/tools/testing/selftests/bpf/benchs/bench_local_storage_create.c b/tools/testing/selftests/bpf/benchs/bench_local_storage_create.c
new file mode 100644
index 000000000000..f8b2a640ccbe
--- /dev/null
+++ b/tools/testing/selftests/bpf/benchs/bench_local_storage_create.c
@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
+
+#include <sys/types.h>
+#include <sys/socket.h>
+
+#include "bench.h"
+#include "bench_local_storage_create.skel.h"
+
+#define BATCH_SZ 32
+
+struct thread {
+	int fds[BATCH_SZ];
+};
+
+static struct bench_local_storage_create *skel;
+static struct thread *threads;
+static long socket_errs;
+
+static void validate(void)
+{
+	if (env.consumer_cnt > 1) {
+		fprintf(stderr,
+			"local-storage-create benchmark does not need consumer\n");
+		exit(1);
+	}
+}
+
+static void setup(void)
+{
+	skel = bench_local_storage_create__open_and_load();
+	if (!skel) {
+		fprintf(stderr, "error loading skel\n");
+		exit(1);
+	}
+
+	skel->bss->bench_pid = getpid();
+
+	if (!bpf_program__attach(skel->progs.socket_post_create)) {
+		fprintf(stderr, "Error attaching bpf program\n");
+		exit(1);
+	}
+
+	if (!bpf_program__attach(skel->progs.kmalloc)) {
+		fprintf(stderr, "Error attaching bpf program\n");
+		exit(1);
+	}
+
+	threads = calloc(env.producer_cnt, sizeof(*threads));
+
+	if (!threads) {
+		fprintf(stderr, "cannot alloc thread_res\n");
+		exit(1);
+	}
+}
+
+static void measure(struct bench_res *res)
+{
+	res->hits = atomic_swap(&skel->bss->create_cnts, 0);
+	res->drops = atomic_swap(&skel->bss->kmalloc_cnts, 0);
+}
+
+static void *consumer(void *input)
+{
+	return NULL;
+}
+
+static void *producer(void *input)
+{
+	struct thread *t = &threads[(long)(input)];
+	int *fds = t->fds;
+	int i;
+
+	while (true) {
+		for (i = 0; i < BATCH_SZ; i++) {
+			fds[i] = socket(AF_INET6, SOCK_DGRAM, 0);
+			if (fds[i] == -1)
+				atomic_inc(&socket_errs);
+		}
+
+		for (i = 0; i < BATCH_SZ; i++) {
+			if (fds[i] != -1)
+				close(fds[i]);
+		}
+	}
+
+	return NULL;
+}
+
+static void report_progress(int iter, struct bench_res *res, long delta_ns)
+{
+	double creates_per_sec, kmallocs_per_create;
+
+	creates_per_sec = res->hits / 1000.0 / (delta_ns / 1000000000.0);
+	kmallocs_per_create = (double)res->drops / res->hits;
+
+	printf("Iter %3d (%7.3lfus): ",
+	       iter, (delta_ns - 1000000000) / 1000.0);
+	printf("creates %8.3lfk/s (%7.3lfk/prod), ",
+	       creates_per_sec, creates_per_sec / env.producer_cnt);
+	printf("%3.2lf kmallocs/create\n", kmallocs_per_create);
+}
+
+static void report_final(struct bench_res res[], int res_cnt)
+{
+	double creates_mean = 0.0, creates_stddev = 0.0;
+	long total_creates = 0, total_kmallocs = 0;
+	int i;
+
+	for (i = 0; i < res_cnt; i++) {
+		creates_mean += res[i].hits / 1000.0 / (0.0 + res_cnt);
+		total_creates += res[i].hits;
+		total_kmallocs += res[i].drops;
+	}
+
+	if (res_cnt > 1)  {
+		for (i = 0; i < res_cnt; i++)
+			creates_stddev += (creates_mean - res[i].hits / 1000.0) *
+				       (creates_mean - res[i].hits / 1000.0) /
+				       (res_cnt - 1.0);
+		creates_stddev = sqrt(creates_stddev);
+	}
+	printf("Summary: creates %8.3lf \u00B1 %5.3lfk/s (%7.3lfk/prod), ",
+	       creates_mean, creates_stddev, creates_mean / env.producer_cnt);
+	printf("%4.2lf kmallocs/create\n", (double)total_kmallocs / total_creates);
+	if (socket_errs || skel->bss->create_errs)
+		printf("socket() errors %ld create_errs %ld\n", socket_errs,
+		       skel->bss->create_errs);
+}
+
+/* Benchmark performance of creating bpf local storage  */
+const struct bench bench_local_storage_create = {
+	.name = "local-storage-create",
+	.validate = validate,
+	.setup = setup,
+	.producer_thread = producer,
+	.consumer_thread = consumer,
+	.measure = measure,
+	.report_progress = report_progress,
+	.report_final = report_final,
+};
diff --git a/tools/testing/selftests/bpf/progs/bench_local_storage_create.c b/tools/testing/selftests/bpf/progs/bench_local_storage_create.c
new file mode 100644
index 000000000000..2814bab54d28
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/bench_local_storage_create.c
@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */
+
+#include "vmlinux.h"
+#include "bpf_tracing_net.h"
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_helpers.h>
+
+long create_errs = 0;
+long create_cnts = 0;
+long kmalloc_cnts = 0;
+__u32 bench_pid = 0;
+
+struct storage {
+	__u8 data[64];
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__type(key, int);
+	__type(value, struct storage);
+} sk_storage_map SEC(".maps");
+
+SEC("raw_tp/kmalloc")
+int BPF_PROG(kmalloc, unsigned long call_site, const void *ptr,
+	     size_t bytes_req, size_t bytes_alloc, gfp_t gfp_flags,
+	     int node)
+{
+	__sync_fetch_and_add(&kmalloc_cnts, 1);
+
+	return 0;
+}
+
+SEC("lsm.s/socket_post_create")
+int BPF_PROG(socket_post_create, struct socket *sock, int family, int type,
+	     int protocol, int kern)
+{
+	struct storage *stg;
+	__u32 pid;
+
+	pid = bpf_get_current_pid_tgid() >> 32;
+	if (pid != bench_pid)
+		return 0;
+
+	stg = bpf_sk_storage_get(&sk_storage_map, sock->sk, NULL,
+				 BPF_LOCAL_STORAGE_GET_F_CREATE);
+
+	if (stg)
+		__sync_fetch_and_add(&create_cnts, 1);
+	else
+		__sync_fetch_and_add(&create_errs, 1);
+
+	return 0;
+}
+
+char __license[] SEC("license") = "GPL";
-- 
2.30.2