[RFC PATCH bpf-next 08/10] bpf: Recharge memory when reuse bpf map

[Yafang Shao <laoar.shao@xxxxxxxxx>]

When we reuse a pinned bpf map, if it belongs to a memcg which needs to
be recharged, we will uncharge the pages of this bpf map from its
original memcg and then charge its pages to the current memcg.

We have to explicitly tell the kernel if it is a reuse path as the
kernel can't detect it intelligently. That can be done in libbpf, then the
user code don't need to be changed.

Signed-off-by: Yafang Shao <laoar.shao@xxxxxxxxx>
---
 include/linux/bpf.h            |  2 ++
 include/uapi/linux/bpf.h       |  2 +-
 kernel/bpf/syscall.c           | 10 ++++++++++
 tools/include/uapi/linux/bpf.h |  2 +-
 tools/lib/bpf/libbpf.c         |  2 +-
 5 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 0edd7d2c0064..b18a30e70507 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -152,6 +152,8 @@ struct bpf_map_ops {
 				     bpf_callback_t callback_fn,
 				     void *callback_ctx, u64 flags);
 
+	bool (*map_memcg_recharge)(struct bpf_map *map);
+
 	/* BTF id of struct allocated by map_alloc */
 	int *map_btf_id;
 
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index f2f658e224a7..ffbe15c1c8c6 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -6093,7 +6093,7 @@ struct bpf_map_info {
 	__u32 btf_key_type_id;
 	__u32 btf_value_type_id;
 	__s8  memcg_state;
-	__s8  :8;	/* alignment pad */
+	__s8  memcg_recharge;
 	__u16 :16;	/* alignment pad */
 	__u64 map_extra;
 } __attribute__((aligned(8)));
diff --git a/kernel/bpf/syscall.c b/kernel/bpf/syscall.c
index d4659d58d288..8817c40275f3 100644
--- a/kernel/bpf/syscall.c
+++ b/kernel/bpf/syscall.c
@@ -4170,12 +4170,22 @@ static int bpf_map_get_info_by_fd(struct file *file,
 
 #ifdef CONFIG_MEMCG_KMEM
 	if (map->memcg) {
+		size_t offset = offsetof(struct bpf_map_info, memcg_recharge);
 		struct mem_cgroup *memcg = map->memcg;
+		char recharge;
 
 		if (memcg == root_mem_cgroup)
 			info.memcg_state = 0;
 		else
 			info.memcg_state = memcg_need_recharge(memcg) ? -1 : 1;
+
+		if (copy_from_user(&recharge, (char __user *)uinfo + offset, sizeof(char)))
+			return -EFAULT;
+
+		if (recharge && memcg_need_recharge(memcg)) {
+			if (map->ops->map_memcg_recharge)
+				map->ops->map_memcg_recharge(map);
+		}
 	}
 #endif
 
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index f2f658e224a7..ffbe15c1c8c6 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -6093,7 +6093,7 @@ struct bpf_map_info {
 	__u32 btf_key_type_id;
 	__u32 btf_value_type_id;
 	__s8  memcg_state;
-	__s8  :8;	/* alignment pad */
+	__s8  memcg_recharge;
 	__u16 :16;	/* alignment pad */
 	__u64 map_extra;
 } __attribute__((aligned(8)));
diff --git a/tools/lib/bpf/libbpf.c b/tools/lib/bpf/libbpf.c
index 49e359cd34df..f0eb67c983d8 100644
--- a/tools/lib/bpf/libbpf.c
+++ b/tools/lib/bpf/libbpf.c
@@ -4488,7 +4488,7 @@ int bpf_map__set_autocreate(struct bpf_map *map, bool autocreate)
 
 int bpf_map__reuse_fd(struct bpf_map *map, int fd)
 {
-	struct bpf_map_info info = {};
+	struct bpf_map_info info = {.memcg_recharge = 1};
 	__u32 len = sizeof(info);
 	int new_fd, err;
 	char *new_name;
-- 
2.17.1