[RFC PATCH bpf-next 2/9] mm: Allow using active vm in all contexts

[Yafang Shao <laoar.shao@gmail.com>]

We can use active vm in task, softirq and irq to account specific memory
usage. Two new helpers active_vm_{add,sub} are introduced.

A dummy item is introduced here, which will be removed when we introduce
real item.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/active_vm.h | 50 +++++++++++++++++++++++++++++++++++++++
 include/linux/sched.h     |  5 ++++
 kernel/fork.c             |  4 ++++
 mm/active_vm.c            | 23 ++++++++++++++++++
 mm/active_vm.h            | 38 +++++++++++++++++++++++++++++
 5 files changed, 120 insertions(+)

diff --git a/include/linux/active_vm.h b/include/linux/active_vm.h
index 899e578e94fa..21f9aaca12c4 100644
--- a/include/linux/active_vm.h
+++ b/include/linux/active_vm.h
@@ -4,6 +4,9 @@
 
 #ifdef CONFIG_ACTIVE_VM
 #include <linux/jump_label.h>
+#include <linux/preempt.h>
+#include <linux/percpu-defs.h>
+#include <linux/sched.h>
 
 extern struct static_key_true active_vm_disabled;
 
@@ -14,10 +17,57 @@ static inline bool active_vm_enabled(void)
 
 	return true;
 }
+
+enum active_vm_item {
+	DUMMY_ITEM = 1,
+	NR_ACTIVE_VM_ITEM = DUMMY_ITEM,
+};
+
+struct active_vm_stat {
+	long stat[NR_ACTIVE_VM_ITEM];
+};
+
+DECLARE_PER_CPU(struct active_vm_stat, active_vm_stats);
+DECLARE_PER_CPU(int, irq_active_vm_item);
+DECLARE_PER_CPU(int, soft_active_vm_item);
+
+static inline int
+active_vm_item_set(int item)
+{
+	int old_item;
+
+	if (in_irq()) {
+		old_item = this_cpu_read(irq_active_vm_item);
+		this_cpu_write(irq_active_vm_item, item);
+	} else if (in_softirq()) {
+		old_item = this_cpu_read(soft_active_vm_item);
+		this_cpu_write(soft_active_vm_item, item);
+	} else {
+		old_item = current->active_vm_item;
+		current->active_vm_item = item;
+	}
+
+	return old_item;
+}
+
+long active_vm_item_sum(int item);
+
 #else
 static inline bool active_vm_enabled(void)
 {
 	return false;
 }
+
+static inline int
+active_vm_item_set(int item)
+{
+	return 0;
+}
+
+static inline long active_vm_item_sum(int item)
+{
+	return 0;
+}
+
 #endif /* CONFIG_ACTIVE_VM */
 #endif /* __INCLUDE_ACTIVE_VM_H */
diff --git a/include/linux/sched.h b/include/linux/sched.h
index ffb6eb55cd13..05acefd383d4 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1441,6 +1441,11 @@ struct task_struct {
 	struct mem_cgroup		*active_memcg;
 #endif
 
+#ifdef CONFIG_ACTIVE_VM
+	/* Used for scope-based memory accounting */
+	int				active_vm_item;
+#endif
+
 #ifdef CONFIG_BLK_CGROUP
 	struct request_queue		*throttle_queue;
 #endif
diff --git a/kernel/fork.c b/kernel/fork.c
index 08969f5aa38d..590d949ff131 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -1043,6 +1043,10 @@ static struct task_struct *dup_task_struct(struct task_struct *orig, int node)
 	tsk->active_memcg = NULL;
 #endif
 
+#ifdef CONFIG_ACTIVE_VM
+	tsk->active_vm_item = 0;
+#endif
+
 #ifdef CONFIG_CPU_SUP_INTEL
 	tsk->reported_split_lock = 0;
 #endif
diff --git a/mm/active_vm.c b/mm/active_vm.c
index 60849930a7d3..541b2ba22da9 100644
--- a/mm/active_vm.c
+++ b/mm/active_vm.c
@@ -1,5 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0
 #include <linux/page_ext.h>
+#include <linux/active_vm.h>
 
 static bool __active_vm_enabled __initdata =
 				IS_ENABLED(CONFIG_ACTIVE_VM);
@@ -31,3 +32,25 @@ struct page_ext_operations active_vm_ops = {
 	.need = need_active_vm,
 	.init = init_active_vm,
 };
+
+DEFINE_PER_CPU(int, irq_active_vm_item);
+DEFINE_PER_CPU(int, soft_active_vm_item);
+EXPORT_PER_CPU_SYMBOL(irq_active_vm_item);
+EXPORT_PER_CPU_SYMBOL(soft_active_vm_item);
+DEFINE_PER_CPU(struct active_vm_stat, active_vm_stats);
+EXPORT_PER_CPU_SYMBOL(active_vm_stats);
+
+long active_vm_item_sum(int item)
+{
+	struct active_vm_stat *this;
+	long sum = 0;
+	int cpu;
+
+	WARN_ON_ONCE(item <= 0);
+	for_each_online_cpu(cpu) {
+		this = &per_cpu(active_vm_stats, cpu);
+		sum += this->stat[item - 1];
+	}
+
+	return sum;
+}
diff --git a/mm/active_vm.h b/mm/active_vm.h
index 72978955833e..1df088d768ef 100644
--- a/mm/active_vm.h
+++ b/mm/active_vm.h
@@ -3,6 +3,44 @@
 #define __MM_ACTIVE_VM_H
 
 #ifdef CONFIG_ACTIVE_VM
+#include <linux/active_vm.h>
+
 extern struct page_ext_operations active_vm_ops;
+
+static inline int active_vm_item(void)
+{
+	if (in_irq())
+		return this_cpu_read(irq_active_vm_item);
+
+	if (in_softirq())
+		return this_cpu_read(soft_active_vm_item);
+
+	return current->active_vm_item;
+}
+
+static inline void active_vm_item_add(int item, long delta)
+{
+	WARN_ON_ONCE(item <= 0);
+	this_cpu_add(active_vm_stats.stat[item - 1], delta);
+}
+
+static inline void active_vm_item_sub(int item, long delta)
+{
+	WARN_ON_ONCE(item <= 0);
+	this_cpu_sub(active_vm_stats.stat[item - 1], delta);
+}
+#else /* CONFIG_ACTIVE_VM */
+static inline int active_vm_item(void)
+{
+	return 0;
+}
+
+static inline void active_vm_item_add(int item, long delta)
+{
+}
+
+static inline void active_vm_item_sub(int item, long delta)
+{
+}
 #endif /* CONFIG_ACTIVE_VM */
 #endif /* __MM_ACTIVE_VM_H */
-- 
2.30.1 (Apple Git-130)