[PATCH v3 0/5] mm: improve performance of accounted kernel memory allocations

[PATCH v3 0/5] mm: improve performance of accounted kernel memory allocations

[Roman Gushchin]

This patchset improves the performance of accounted kernel memory allocations
by ~30% as measured by a micro-benchmark [1]. The benchmark is very
straightforward: 1M of 64 bytes-large kmalloc() allocations.

Below are results with the disabled kernel memory accounting, the original state
and with this patchset applied.

|             | Kmem disabled | Original | Patched |  Delta |
|-------------+---------------+----------+---------+--------|
| User cgroup |         29764 |    84548 |   59078 | -30.0% |
| Root cgroup |         29742 |    48342 |   31501 | -34.8% |

As we can see, the patchset removes the majority of the overhead when there is
no actual accounting (a task belongs to the root memory cgroup) and almost
halves the accounting overhead otherwise.

The main idea is to get rid of unnecessary memcg to objcg conversions and switch
to a scope-based protection of objcgs, which eliminates extra operations with
objcg reference counters under a rcu read lock. More details are provided in
individual commit descriptions.

v3:
	- fixed a bug spotted by Shakeel
	- added some comments, per Shakeel
v2:
	- fixed a bug discovered by Naresh Kamboju
	- code changes asked by Johannes (added comments, open-coded bit ops)
	- merged in a couple of small fixes
v1:
	- made the objcg update fully lockless
	- fixed !CONFIG_MMU build issues
rfc:
	https://lwn.net/Articles/945722/

--
[1]:

static int memory_alloc_test(struct seq_file *m, void *v)
{
       unsigned long i, j;
       void **ptrs;
       ktime_t start, end;
       s64 delta, min_delta = LLONG_MAX;

       ptrs = kvmalloc(sizeof(void *) * 1000000, GFP_KERNEL);
       if (!ptrs)
               return -ENOMEM;

       for (j = 0; j < 100; j++) {
               start = ktime_get();
               for (i = 0; i < 1000000; i++)
                       ptrs[i] = kmalloc(64, GFP_KERNEL_ACCOUNT);
               end = ktime_get();

               delta = ktime_us_delta(end, start);
               if (delta < min_delta)
                       min_delta = delta;

               for (i = 0; i < 1000000; i++)
                       kfree(ptrs[i]);
       }

       kvfree(ptrs);
       seq_printf(m, "%lld us\n", min_delta);

       return 0;
}

--

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>


Roman Gushchin (5):
  mm: kmem: optimize get_obj_cgroup_from_current()
  mm: kmem: add direct objcg pointer to task_struct
  mm: kmem: make memcg keep a reference to the original objcg
  mm: kmem: scoped objcg protection
  percpu: scoped objcg protection

 include/linux/memcontrol.h |  17 +++-
 include/linux/sched.h      |   4 +
 include/linux/sched/mm.h   |   4 +
 mm/memcontrol.c            | 204 ++++++++++++++++++++++++++++++++-----
 mm/percpu.c                |   8 +-
 mm/slab.h                  |  15 +--
 6 files changed, 214 insertions(+), 38 deletions(-)

-- 
2.42.0

[PATCH v3 1/5] mm: kmem: optimize get_obj_cgroup_from_current()

[Roman Gushchin]

Manually inline memcg_kmem_bypass() and active_memcg() to speed up
get_obj_cgroup_from_current() by avoiding duplicate in_task() checks
and active_memcg() readings.

Also add a likely() macro to __get_obj_cgroup_from_memcg():
obj_cgroup_tryget() should succeed at almost all times except a very
unlikely race with the memcg deletion path.

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
Acked-by: Shakeel Butt <shakeelb@google.com>
Acked-by: Johannes Weiner <hannes@cmpxchg.org>
---
 mm/memcontrol.c | 34 ++++++++++++++--------------------
 1 file changed, 14 insertions(+), 20 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 9741d62d0424..16ac2a5838fb 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -1068,19 +1068,6 @@ struct mem_cgroup *get_mem_cgroup_from_mm(struct mm_struct *mm)
 }
 EXPORT_SYMBOL(get_mem_cgroup_from_mm);
 
-static __always_inline bool memcg_kmem_bypass(void)
-{
-	/* Allow remote memcg charging from any context. */
-	if (unlikely(active_memcg()))
-		return false;
-
-	/* Memcg to charge can't be determined. */
-	if (!in_task() || !current->mm || (current->flags & PF_KTHREAD))
-		return true;
-
-	return false;
-}
-
 /**
  * mem_cgroup_iter - iterate over memory cgroup hierarchy
  * @root: hierarchy root
@@ -3007,7 +2994,7 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
 
 	for (; !mem_cgroup_is_root(memcg); memcg = parent_mem_cgroup(memcg)) {
 		objcg = rcu_dereference(memcg->objcg);
-		if (objcg && obj_cgroup_tryget(objcg))
+		if (likely(objcg && obj_cgroup_tryget(objcg)))
 			break;
 		objcg = NULL;
 	}
@@ -3016,16 +3003,23 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
 
 __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
 {
-	struct obj_cgroup *objcg = NULL;
 	struct mem_cgroup *memcg;
+	struct obj_cgroup *objcg;
 
-	if (memcg_kmem_bypass())
-		return NULL;
+	if (in_task()) {
+		memcg = current->active_memcg;
+
+		/* Memcg to charge can't be determined. */
+		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
+			return NULL;
+	} else {
+		memcg = this_cpu_read(int_active_memcg);
+		if (likely(!memcg))
+			return NULL;
+	}
 
 	rcu_read_lock();
-	if (unlikely(active_memcg()))
-		memcg = active_memcg();
-	else
+	if (!memcg)
 		memcg = mem_cgroup_from_task(current);
 	objcg = __get_obj_cgroup_from_memcg(memcg);
 	rcu_read_unlock();
-- 
2.42.0

[PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Roman Gushchin]

To charge a freshly allocated kernel object to a memory cgroup, the
kernel needs to obtain an objcg pointer. Currently it does it
indirectly by obtaining the memcg pointer first and then calling to
__get_obj_cgroup_from_memcg().

Usually tasks spend their entire life belonging to the same object
cgroup. So it makes sense to save the objcg pointer on task_struct
directly, so it can be obtained faster. It requires some work on fork,
exit and cgroup migrate paths, but these paths are way colder.

To avoid any costly synchronization the following rules are applied:
1) A task sets it's objcg pointer itself.

2) If a task is being migrated to another cgroup, the least
   significant bit of the objcg pointer is set atomically.

3) On the allocation path the objcg pointer is obtained locklessly
   using the READ_ONCE() macro and the least significant bit is
   checked. If it's set, the following procedure is used to update
   it locklessly:
       - task->objcg is zeroed using cmpxcg
       - new objcg pointer is obtained
       - task->objcg is updated using try_cmpxchg
       - operation is repeated if try_cmpxcg fails
   It guarantees that no updates will be lost if task migration
   is racing against objcg pointer update. It also allows to keep
   both read and write paths fully lockless.

Because the task is keeping a reference to the objcg, it can't go away
while the task is alive.

This commit doesn't change the way the remote memcg charging works.

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
Acked-by: Johannes Weiner <hannes@cmpxchg.org>
---
 include/linux/sched.h |   4 ++
 mm/memcontrol.c       | 130 +++++++++++++++++++++++++++++++++++++++---
 2 files changed, 125 insertions(+), 9 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 77f01ac385f7..60de42715b56 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1443,6 +1443,10 @@ struct task_struct {
 	struct mem_cgroup		*active_memcg;
 #endif
 
+#ifdef CONFIG_MEMCG_KMEM
+	struct obj_cgroup		*objcg;
+#endif
+
 #ifdef CONFIG_BLK_CGROUP
 	struct gendisk			*throttle_disk;
 #endif
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 16ac2a5838fb..0605e45bd4a2 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -249,6 +249,8 @@ struct mem_cgroup *vmpressure_to_memcg(struct vmpressure *vmpr)
 	return container_of(vmpr, struct mem_cgroup, vmpressure);
 }
 
+#define CURRENT_OBJCG_UPDATE_FLAG 0x1UL
+
 #ifdef CONFIG_MEMCG_KMEM
 static DEFINE_SPINLOCK(objcg_lock);
 
@@ -3001,6 +3003,50 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
 	return objcg;
 }
 
+static struct obj_cgroup *current_objcg_update(void)
+{
+	struct mem_cgroup *memcg;
+	struct obj_cgroup *old, *objcg = NULL;
+
+	do {
+		/* Atomically drop the update bit. */
+		old = xchg(&current->objcg, NULL);
+		if (old) {
+			old = (struct obj_cgroup *)
+				((unsigned long)old & ~CURRENT_OBJCG_UPDATE_FLAG);
+			if (old)
+				obj_cgroup_put(old);
+
+			old = NULL;
+		}
+
+		/* Obtain the new objcg pointer. */
+		rcu_read_lock();
+		memcg = mem_cgroup_from_task(current);
+		/*
+		 * The current task can be asynchronously moved to another
+		 * memcg and the previous memcg can be offlined. So let's
+		 * get the memcg pointer and try get a reference to objcg
+		 * under a rcu read lock.
+		 */
+		for (; memcg != root_mem_cgroup; memcg = parent_mem_cgroup(memcg)) {
+			objcg = rcu_dereference(memcg->objcg);
+			if (likely(objcg && obj_cgroup_tryget(objcg)))
+				break;
+			objcg = NULL;
+		}
+		rcu_read_unlock();
+
+		/*
+		 * Try set up a new objcg pointer atomically. If it
+		 * fails, it means the update flag was set concurrently, so
+		 * the whole procedure should be repeated.
+		 */
+	} while (!try_cmpxchg(&current->objcg, &old, objcg));
+
+	return objcg;
+}
+
 __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
 {
 	struct mem_cgroup *memcg;
@@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
 
 	if (in_task()) {
 		memcg = current->active_memcg;
+		if (unlikely(memcg))
+			goto from_memcg;
 
-		/* Memcg to charge can't be determined. */
-		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
-			return NULL;
+		objcg = READ_ONCE(current->objcg);
+		if (unlikely((unsigned long)objcg & CURRENT_OBJCG_UPDATE_FLAG))
+			objcg = current_objcg_update();
+
+		if (objcg) {
+			obj_cgroup_get(objcg);
+			return objcg;
+		}
 	} else {
 		memcg = this_cpu_read(int_active_memcg);
-		if (likely(!memcg))
-			return NULL;
+		if (unlikely(memcg))
+			goto from_memcg;
 	}
+	return NULL;
 
+from_memcg:
 	rcu_read_lock();
-	if (!memcg)
-		memcg = mem_cgroup_from_task(current);
 	objcg = __get_obj_cgroup_from_memcg(memcg);
 	rcu_read_unlock();
 	return objcg;
@@ -6345,6 +6398,7 @@ static void mem_cgroup_move_task(void)
 		mem_cgroup_clear_mc();
 	}
 }
+
 #else	/* !CONFIG_MMU */
 static int mem_cgroup_can_attach(struct cgroup_taskset *tset)
 {
@@ -6358,8 +6412,39 @@ static void mem_cgroup_move_task(void)
 }
 #endif
 
+#ifdef CONFIG_MEMCG_KMEM
+static void mem_cgroup_fork(struct task_struct *task)
+{
+	/*
+	 * Set the update flag to cause task->objcg to be initialized lazily
+	 * on the first allocation. It can be done without any synchronization
+	 * because it's always performed on the current task, so does
+	 * current_objcg_update().
+	 */
+	task->objcg = (struct obj_cgroup *)CURRENT_OBJCG_UPDATE_FLAG;
+}
+
+static void mem_cgroup_exit(struct task_struct *task)
+{
+	struct obj_cgroup *objcg = task->objcg;
+
+	objcg = (struct obj_cgroup *)
+		((unsigned long)objcg & ~CURRENT_OBJCG_UPDATE_FLAG);
+	if (objcg)
+		obj_cgroup_put(objcg);
+
+	/*
+	 * Some kernel allocations can happen after this point,
+	 * but let's ignore them. It can be done without any synchronization
+	 * because it's always performed on the current task, so does
+	 * current_objcg_update().
+	 */
+	task->objcg = NULL;
+}
+#endif
+
 #ifdef CONFIG_LRU_GEN
-static void mem_cgroup_attach(struct cgroup_taskset *tset)
+static void mem_cgroup_lru_gen_attach(struct cgroup_taskset *tset)
 {
 	struct task_struct *task;
 	struct cgroup_subsys_state *css;
@@ -6377,10 +6462,31 @@ static void mem_cgroup_attach(struct cgroup_taskset *tset)
 	task_unlock(task);
 }
 #else
+static void mem_cgroup_lru_gen_attach(struct cgroup_taskset *tset) {}
+#endif /* CONFIG_LRU_GEN */
+
+#ifdef CONFIG_MEMCG_KMEM
+static void mem_cgroup_kmem_attach(struct cgroup_taskset *tset)
+{
+	struct task_struct *task;
+	struct cgroup_subsys_state *css;
+
+	cgroup_taskset_for_each(task, css, tset) {
+		/* atomically set the update bit */
+		set_bit(0, (unsigned long *)&task->objcg);
+	}
+}
+#else
+static void mem_cgroup_kmem_attach(struct cgroup_taskset *tset) {}
+#endif /* CONFIG_MEMCG_KMEM */
+
+#if defined(CONFIG_LRU_GEN) || defined(CONFIG_MEMCG_KMEM)
 static void mem_cgroup_attach(struct cgroup_taskset *tset)
 {
+	mem_cgroup_lru_gen_attach(tset);
+	mem_cgroup_kmem_attach(tset);
 }
-#endif /* CONFIG_LRU_GEN */
+#endif
 
 static int seq_puts_memcg_tunable(struct seq_file *m, unsigned long value)
 {
@@ -6824,9 +6930,15 @@ struct cgroup_subsys memory_cgrp_subsys = {
 	.css_reset = mem_cgroup_css_reset,
 	.css_rstat_flush = mem_cgroup_css_rstat_flush,
 	.can_attach = mem_cgroup_can_attach,
+#if defined(CONFIG_LRU_GEN) || defined(CONFIG_MEMCG_KMEM)
 	.attach = mem_cgroup_attach,
+#endif
 	.cancel_attach = mem_cgroup_cancel_attach,
 	.post_attach = mem_cgroup_move_task,
+#ifdef CONFIG_MEMCG_KMEM
+	.fork = mem_cgroup_fork,
+	.exit = mem_cgroup_exit,
+#endif
 	.dfl_cftypes = memory_files,
 	.legacy_cftypes = mem_cgroup_legacy_files,
 	.early_init = 0,
-- 
2.42.0

[PATCH v3 3/5] mm: kmem: make memcg keep a reference to the original objcg

[Roman Gushchin]

Keep a reference to the original objcg object for the entire life
of a memcg structure.

This allows to simplify the synchronization on the kernel memory
allocation paths: pinning a (live) memcg will also pin the
corresponding objcg.

The memory overhead of this change is minimal because object cgroups
usually outlive their corresponding memory cgroups even without this
change, so it's only an additional pointer per memcg.

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
Acked-by: Shakeel Butt <shakeelb@google.com>
---
 include/linux/memcontrol.h | 8 +++++++-
 mm/memcontrol.c            | 5 +++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index ab94ad4597d0..277690af383d 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -299,7 +299,13 @@ struct mem_cgroup {
 
 #ifdef CONFIG_MEMCG_KMEM
 	int kmemcg_id;
-	struct obj_cgroup __rcu *objcg;
+	/*
+	 * memcg->objcg is wiped out as a part of the objcg repaprenting
+	 * process. memcg->orig_objcg preserves a pointer (and a reference)
+	 * to the original objcg until the end of live of memcg.
+	 */
+	struct obj_cgroup __rcu	*objcg;
+	struct obj_cgroup	*orig_objcg;
 	/* list of inherited objcgs, protected by objcg_lock */
 	struct list_head objcg_list;
 #endif
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 0605e45bd4a2..d90cc19e4113 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3808,6 +3808,8 @@ static int memcg_online_kmem(struct mem_cgroup *memcg)
 
 	objcg->memcg = memcg;
 	rcu_assign_pointer(memcg->objcg, objcg);
+	obj_cgroup_get(objcg);
+	memcg->orig_objcg = objcg;
 
 	static_branch_enable(&memcg_kmem_online_key);
 
@@ -5302,6 +5304,9 @@ static void __mem_cgroup_free(struct mem_cgroup *memcg)
 {
 	int node;
 
+	if (memcg->orig_objcg)
+		obj_cgroup_put(memcg->orig_objcg);
+
 	for_each_node(node)
 		free_mem_cgroup_per_node_info(memcg, node);
 	kfree(memcg->vmstats);
-- 
2.42.0

[PATCH v3 4/5] mm: kmem: scoped objcg protection

[Roman Gushchin]

Switch to a scope-based protection of the objcg pointer on slab/kmem
allocation paths. Instead of using the get_() semantics in the
pre-allocation hook and put the reference afterwards, let's rely
on the fact that objcg is pinned by the scope.

It's possible because:
1) if the objcg is received from the current task struct, the task is
   keeping a reference to the objcg.
2) if the objcg is received from an active memcg (remote charging),
   the memcg is pinned by the scope and has a reference to the
   corresponding objcg.

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
Acked-by: Shakeel Butt <shakeelb@google.com>
---
 include/linux/memcontrol.h |  9 ++++++++
 include/linux/sched/mm.h   |  4 ++++
 mm/memcontrol.c            | 47 ++++++++++++++++++++++++++++++++++++--
 mm/slab.h                  | 15 ++++++------
 4 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 277690af383d..a89df289144d 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -1769,6 +1769,15 @@ bool mem_cgroup_kmem_disabled(void);
 int __memcg_kmem_charge_page(struct page *page, gfp_t gfp, int order);
 void __memcg_kmem_uncharge_page(struct page *page, int order);
 
+/*
+ * The returned objcg pointer is safe to use without additional
+ * protection within a scope. The scope is defined either by
+ * the current task (similar to the "current" global variable)
+ * or by set_active_memcg() pair.
+ * Please, use obj_cgroup_get() to get a reference if the pointer
+ * needs to be used outside of the local scope.
+ */
+struct obj_cgroup *current_obj_cgroup(void);
 struct obj_cgroup *get_obj_cgroup_from_current(void);
 struct obj_cgroup *get_obj_cgroup_from_folio(struct folio *folio);
 
diff --git a/include/linux/sched/mm.h b/include/linux/sched/mm.h
index 8d89c8c4fac1..9a19f1b42f64 100644
--- a/include/linux/sched/mm.h
+++ b/include/linux/sched/mm.h
@@ -403,6 +403,10 @@ DECLARE_PER_CPU(struct mem_cgroup *, int_active_memcg);
  * __GFP_ACCOUNT allocations till the end of the scope will be charged to the
  * given memcg.
  *
+ * Please, make sure that caller has a reference to the passed memcg structure,
+ * so its lifetime is guaranteed to exceed the scope between two
+ * set_active_memcg() calls.
+ *
  * NOTE: This function can nest. Users must save the return value and
  * reset the previous value after their own charging scope is over.
  */
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index d90cc19e4113..852fe6918f82 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3079,6 +3079,49 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
 	return objcg;
 }
 
+__always_inline struct obj_cgroup *current_obj_cgroup(void)
+{
+	struct mem_cgroup *memcg;
+	struct obj_cgroup *objcg;
+
+	if (in_task()) {
+		memcg = current->active_memcg;
+		if (unlikely(memcg))
+			goto from_memcg;
+
+		objcg = READ_ONCE(current->objcg);
+		if (unlikely((unsigned long)objcg & CURRENT_OBJCG_UPDATE_FLAG))
+			objcg = current_objcg_update();
+		/*
+		 * Objcg reference is kept by the task, so it's safe
+		 * to use the objcg by the current task.
+		 */
+		return objcg;
+	}
+
+	memcg = this_cpu_read(int_active_memcg);
+	if (unlikely(memcg))
+		goto from_memcg;
+
+	return NULL;
+
+from_memcg:
+	for (; !mem_cgroup_is_root(memcg); memcg = parent_mem_cgroup(memcg)) {
+		/*
+		 * Memcg pointer is protected by scope (see set_active_memcg())
+		 * and is pinning the corresponding objcg, so objcg can't go
+		 * away and can be used within the scope without any additional
+		 * protection.
+		 */
+		objcg = rcu_dereference_check(memcg->objcg, 1);
+		if (likely(objcg))
+			break;
+		objcg = NULL;
+	}
+
+	return objcg;
+}
+
 struct obj_cgroup *get_obj_cgroup_from_folio(struct folio *folio)
 {
 	struct obj_cgroup *objcg;
@@ -3173,15 +3216,15 @@ int __memcg_kmem_charge_page(struct page *page, gfp_t gfp, int order)
 	struct obj_cgroup *objcg;
 	int ret = 0;
 
-	objcg = get_obj_cgroup_from_current();
+	objcg = current_obj_cgroup();
 	if (objcg) {
 		ret = obj_cgroup_charge_pages(objcg, gfp, 1 << order);
 		if (!ret) {
+			obj_cgroup_get(objcg);
 			page->memcg_data = (unsigned long)objcg |
 				MEMCG_DATA_KMEM;
 			return 0;
 		}
-		obj_cgroup_put(objcg);
 	}
 	return ret;
 }
diff --git a/mm/slab.h b/mm/slab.h
index 799a315695c6..3d07fb428393 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -484,7 +484,12 @@ static inline bool memcg_slab_pre_alloc_hook(struct kmem_cache *s,
 	if (!(flags & __GFP_ACCOUNT) && !(s->flags & SLAB_ACCOUNT))
 		return true;
 
-	objcg = get_obj_cgroup_from_current();
+	/*
+	 * The obtained objcg pointer is safe to use within the current scope,
+	 * defined by current task or set_active_memcg() pair.
+	 * obj_cgroup_get() is used to get a permanent reference.
+	 */
+	objcg = current_obj_cgroup();
 	if (!objcg)
 		return true;
 
@@ -497,17 +502,14 @@ static inline bool memcg_slab_pre_alloc_hook(struct kmem_cache *s,
 		css_put(&memcg->css);
 
 		if (ret)
-			goto out;
+			return false;
 	}
 
 	if (obj_cgroup_charge(objcg, flags, objects * obj_full_size(s)))
-		goto out;
+		return false;
 
 	*objcgp = objcg;
 	return true;
-out:
-	obj_cgroup_put(objcg);
-	return false;
 }
 
 static inline void memcg_slab_post_alloc_hook(struct kmem_cache *s,
@@ -542,7 +544,6 @@ static inline void memcg_slab_post_alloc_hook(struct kmem_cache *s,
 			obj_cgroup_uncharge(objcg, obj_full_size(s));
 		}
 	}
-	obj_cgroup_put(objcg);
 }
 
 static inline void memcg_slab_free_hook(struct kmem_cache *s, struct slab *slab,
-- 
2.42.0

[PATCH v3 5/5] percpu: scoped objcg protection

[Roman Gushchin]

Similar to slab and kmem, switch to a scope-based protection of the
objcg pointer to avoid.

Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
Acked-by: Shakeel Butt <shakeelb@google.com>
---
 mm/percpu.c | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/mm/percpu.c b/mm/percpu.c
index a7665de8485f..f53ba692d67a 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -1628,14 +1628,12 @@ static bool pcpu_memcg_pre_alloc_hook(size_t size, gfp_t gfp,
 	if (!memcg_kmem_online() || !(gfp & __GFP_ACCOUNT))
 		return true;
 
-	objcg = get_obj_cgroup_from_current();
+	objcg = current_obj_cgroup();
 	if (!objcg)
 		return true;
 
-	if (obj_cgroup_charge(objcg, gfp, pcpu_obj_full_size(size))) {
-		obj_cgroup_put(objcg);
+	if (obj_cgroup_charge(objcg, gfp, pcpu_obj_full_size(size)))
 		return false;
-	}
 
 	*objcgp = objcg;
 	return true;
@@ -1649,6 +1647,7 @@ static void pcpu_memcg_post_alloc_hook(struct obj_cgroup *objcg,
 		return;
 
 	if (likely(chunk && chunk->obj_cgroups)) {
+		obj_cgroup_get(objcg);
 		chunk->obj_cgroups[off >> PCPU_MIN_ALLOC_SHIFT] = objcg;
 
 		rcu_read_lock();
@@ -1657,7 +1656,6 @@ static void pcpu_memcg_post_alloc_hook(struct obj_cgroup *objcg,
 		rcu_read_unlock();
 	} else {
 		obj_cgroup_uncharge(objcg, pcpu_obj_full_size(size));
-		obj_cgroup_put(objcg);
 	}
 }
 
-- 
2.42.0

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Shakeel Butt]

On Mon, Oct 16, 2023 at 3:19 PM Roman Gushchin <roman.gushchin@linux.dev> wrote:
>
> To charge a freshly allocated kernel object to a memory cgroup, the
> kernel needs to obtain an objcg pointer. Currently it does it
> indirectly by obtaining the memcg pointer first and then calling to
> __get_obj_cgroup_from_memcg().
>
> Usually tasks spend their entire life belonging to the same object
> cgroup. So it makes sense to save the objcg pointer on task_struct
> directly, so it can be obtained faster. It requires some work on fork,
> exit and cgroup migrate paths, but these paths are way colder.
>
> To avoid any costly synchronization the following rules are applied:
> 1) A task sets it's objcg pointer itself.
>
> 2) If a task is being migrated to another cgroup, the least
>    significant bit of the objcg pointer is set atomically.
>
> 3) On the allocation path the objcg pointer is obtained locklessly
>    using the READ_ONCE() macro and the least significant bit is
>    checked. If it's set, the following procedure is used to update
>    it locklessly:
>        - task->objcg is zeroed using cmpxcg
>        - new objcg pointer is obtained
>        - task->objcg is updated using try_cmpxchg
>        - operation is repeated if try_cmpxcg fails
>    It guarantees that no updates will be lost if task migration
>    is racing against objcg pointer update. It also allows to keep
>    both read and write paths fully lockless.
>
> Because the task is keeping a reference to the objcg, it can't go away
> while the task is alive.
>
> This commit doesn't change the way the remote memcg charging works.
>
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Johannes Weiner <hannes@cmpxchg.org>

Acked-by: Shakeel Butt <shakeelb@google.com>

Re: [PATCH v3 1/5] mm: kmem: optimize get_obj_cgroup_from_current()

[Vlastimil Babka]

On 10/17/23 00:18, Roman Gushchin wrote:
> Manually inline memcg_kmem_bypass() and active_memcg() to speed up
> get_obj_cgroup_from_current() by avoiding duplicate in_task() checks
> and active_memcg() readings.

I guess the compiler could figure out most of it, except maybe the
active_memcg() after rcu_read_lock(), as I note below.

But anyway it's a single caller of memcg_kmem_bypass() so this makes sense.

> Also add a likely() macro to __get_obj_cgroup_from_memcg():
> obj_cgroup_tryget() should succeed at almost all times except a very
> unlikely race with the memcg deletion path.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Shakeel Butt <shakeelb@google.com>
> Acked-by: Johannes Weiner <hannes@cmpxchg.org>

Reviewed-by: Vlastimil Babka <vbabka@suse.cz>

> ---
>  mm/memcontrol.c | 34 ++++++++++++++--------------------
>  1 file changed, 14 insertions(+), 20 deletions(-)
> 
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index 9741d62d0424..16ac2a5838fb 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -1068,19 +1068,6 @@ struct mem_cgroup *get_mem_cgroup_from_mm(struct mm_struct *mm)
>  }
>  EXPORT_SYMBOL(get_mem_cgroup_from_mm);
>  
> -static __always_inline bool memcg_kmem_bypass(void)
> -{
> -	/* Allow remote memcg charging from any context. */
> -	if (unlikely(active_memcg()))
> -		return false;
> -
> -	/* Memcg to charge can't be determined. */
> -	if (!in_task() || !current->mm || (current->flags & PF_KTHREAD))
> -		return true;
> -
> -	return false;
> -}
> -
>  /**
>   * mem_cgroup_iter - iterate over memory cgroup hierarchy
>   * @root: hierarchy root
> @@ -3007,7 +2994,7 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
>  
>  	for (; !mem_cgroup_is_root(memcg); memcg = parent_mem_cgroup(memcg)) {
>  		objcg = rcu_dereference(memcg->objcg);
> -		if (objcg && obj_cgroup_tryget(objcg))
> +		if (likely(objcg && obj_cgroup_tryget(objcg)))
>  			break;
>  		objcg = NULL;
>  	}
> @@ -3016,16 +3003,23 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
>  
>  __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
>  {
> -	struct obj_cgroup *objcg = NULL;
>  	struct mem_cgroup *memcg;
> +	struct obj_cgroup *objcg;
>  
> -	if (memcg_kmem_bypass())
> -		return NULL;
> +	if (in_task()) {
> +		memcg = current->active_memcg;
> +
> +		/* Memcg to charge can't be determined. */
> +		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> +			return NULL;
> +	} else {
> +		memcg = this_cpu_read(int_active_memcg);
> +		if (likely(!memcg))
> +			return NULL;
> +	}
>  
>  	rcu_read_lock();
> -	if (unlikely(active_memcg()))
> -		memcg = active_memcg();

I guess the rcu_read_lock() has in theory prevented us from being migrated
to a different cpu after doing the !in_task()
this_cpu_read(int_active_memcg); and now we trust the reading outside of the
rcu lock, but it's ok because a) nothing prevented us from migrating cpu
after getting to the objcg anyway so we could end up with a wrong one, and
b) if we're not in_task() and thus read int_active_memcg, it means our
context is some interrupt handler that can't be migrated to another cpu
anyway, right?

> -	else
> +	if (!memcg)
>  		memcg = mem_cgroup_from_task(current);
>  	objcg = __get_obj_cgroup_from_memcg(memcg);
>  	rcu_read_unlock();

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Vlastimil Babka]

On 10/17/23 00:18, Roman Gushchin wrote:
> To charge a freshly allocated kernel object to a memory cgroup, the
> kernel needs to obtain an objcg pointer. Currently it does it
> indirectly by obtaining the memcg pointer first and then calling to
> __get_obj_cgroup_from_memcg().
> 
> Usually tasks spend their entire life belonging to the same object
> cgroup. So it makes sense to save the objcg pointer on task_struct
> directly, so it can be obtained faster. It requires some work on fork,
> exit and cgroup migrate paths, but these paths are way colder.
> 
> To avoid any costly synchronization the following rules are applied:
> 1) A task sets it's objcg pointer itself.
> 
> 2) If a task is being migrated to another cgroup, the least
>    significant bit of the objcg pointer is set atomically.
> 
> 3) On the allocation path the objcg pointer is obtained locklessly
>    using the READ_ONCE() macro and the least significant bit is
>    checked. If it's set, the following procedure is used to update
>    it locklessly:
>        - task->objcg is zeroed using cmpxcg
>        - new objcg pointer is obtained
>        - task->objcg is updated using try_cmpxchg
>        - operation is repeated if try_cmpxcg fails
>    It guarantees that no updates will be lost if task migration
>    is racing against objcg pointer update. It also allows to keep
>    both read and write paths fully lockless.
> 
> Because the task is keeping a reference to the objcg, it can't go away
> while the task is alive.
> 
> This commit doesn't change the way the remote memcg charging works.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Johannes Weiner <hannes@cmpxchg.org>
> ---
>  include/linux/sched.h |   4 ++
>  mm/memcontrol.c       | 130 +++++++++++++++++++++++++++++++++++++++---
>  2 files changed, 125 insertions(+), 9 deletions(-)
> 
> diff --git a/include/linux/sched.h b/include/linux/sched.h
> index 77f01ac385f7..60de42715b56 100644
> --- a/include/linux/sched.h
> +++ b/include/linux/sched.h
> @@ -1443,6 +1443,10 @@ struct task_struct {
>  	struct mem_cgroup		*active_memcg;
>  #endif
>  
> +#ifdef CONFIG_MEMCG_KMEM
> +	struct obj_cgroup		*objcg;
> +#endif
> +
>  #ifdef CONFIG_BLK_CGROUP
>  	struct gendisk			*throttle_disk;
>  #endif
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index 16ac2a5838fb..0605e45bd4a2 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -249,6 +249,8 @@ struct mem_cgroup *vmpressure_to_memcg(struct vmpressure *vmpr)
>  	return container_of(vmpr, struct mem_cgroup, vmpressure);
>  }
>  
> +#define CURRENT_OBJCG_UPDATE_FLAG 0x1UL

There's a silent relation between this and set_bit(0, ...) in
mem_cgroup_kmem_attach(), maybe worth a comment at least, or defining the
bit number first and from that the flag?

> +
>  #ifdef CONFIG_MEMCG_KMEM
>  static DEFINE_SPINLOCK(objcg_lock);
>  
> @@ -3001,6 +3003,50 @@ static struct obj_cgroup *__get_obj_cgroup_from_memcg(struct mem_cgroup *memcg)
>  	return objcg;
>  }
>  
> +static struct obj_cgroup *current_objcg_update(void)
> +{
> +	struct mem_cgroup *memcg;
> +	struct obj_cgroup *old, *objcg = NULL;
> +
> +	do {
> +		/* Atomically drop the update bit. */
> +		old = xchg(&current->objcg, NULL);
> +		if (old) {
> +			old = (struct obj_cgroup *)
> +				((unsigned long)old & ~CURRENT_OBJCG_UPDATE_FLAG);
> +			if (old)
> +				obj_cgroup_put(old);
> +
> +			old = NULL;
> +		}
> +
> +		/* Obtain the new objcg pointer. */
> +		rcu_read_lock();
> +		memcg = mem_cgroup_from_task(current);
> +		/*
> +		 * The current task can be asynchronously moved to another
> +		 * memcg and the previous memcg can be offlined. So let's
> +		 * get the memcg pointer and try get a reference to objcg
> +		 * under a rcu read lock.
> +		 */
> +		for (; memcg != root_mem_cgroup; memcg = parent_mem_cgroup(memcg)) {
> +			objcg = rcu_dereference(memcg->objcg);
> +			if (likely(objcg && obj_cgroup_tryget(objcg)))

So IIUC here we increase objcg refcount.

> +				break;
> +			objcg = NULL;
> +		}
> +		rcu_read_unlock();
> +
> +		/*
> +		 * Try set up a new objcg pointer atomically. If it
> +		 * fails, it means the update flag was set concurrently, so
> +		 * the whole procedure should be repeated.
> +		 */
> +	} while (!try_cmpxchg(&current->objcg, &old, objcg));

And if this fails we throw objcg away and try again, but we should do
obj_cgroup_put(objcg) first, as otherwise it would cause a leak?

> +
> +	return objcg;
> +}
> +
>  __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
>  {
>  	struct mem_cgroup *memcg;
> @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
>  
>  	if (in_task()) {
>  		memcg = current->active_memcg;
> +		if (unlikely(memcg))
> +			goto from_memcg;
>  
> -		/* Memcg to charge can't be determined. */
> -		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))

The checks for current->mm and PF_KTHREAD seem to be gone completely after
the patch, was that intended and why?

> -			return NULL;
> +		objcg = READ_ONCE(current->objcg);
> +		if (unlikely((unsigned long)objcg & CURRENT_OBJCG_UPDATE_FLAG))
> +			objcg = current_objcg_update();
> +
> +		if (objcg) {
> +			obj_cgroup_get(objcg);
> +			return objcg;
> +		}
>  	} else {
>  		memcg = this_cpu_read(int_active_memcg);
> -		if (likely(!memcg))
> -			return NULL;
> +		if (unlikely(memcg))
> +			goto from_memcg;
>  	}
> +	return NULL;
>  
> +from_memcg:
>  	rcu_read_lock();
> -	if (!memcg)
> -		memcg = mem_cgroup_from_task(current);
>  	objcg = __get_obj_cgroup_from_memcg(memcg);
>  	rcu_read_unlock();
>  	return objcg;
> @@ -6345,6 +6398,7 @@ static void mem_cgroup_move_task(void)
>  		mem_cgroup_clear_mc();
>  	}
>  }
> +

Re: [PATCH v3 3/5] mm: kmem: make memcg keep a reference to the original objcg

[Vlastimil Babka]

On 10/17/23 00:18, Roman Gushchin wrote:
> Keep a reference to the original objcg object for the entire life
> of a memcg structure.
> 
> This allows to simplify the synchronization on the kernel memory
> allocation paths: pinning a (live) memcg will also pin the
> corresponding objcg.
> 
> The memory overhead of this change is minimal because object cgroups
> usually outlive their corresponding memory cgroups even without this
> change, so it's only an additional pointer per memcg.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Shakeel Butt <shakeelb@google.com>

Reviewed-by: Vlastimil Babka <vbabka@suse.cz>

Re: [PATCH v3 4/5] mm: kmem: scoped objcg protection

[Vlastimil Babka]

On 10/17/23 00:18, Roman Gushchin wrote:
> Switch to a scope-based protection of the objcg pointer on slab/kmem
> allocation paths. Instead of using the get_() semantics in the
> pre-allocation hook and put the reference afterwards, let's rely
> on the fact that objcg is pinned by the scope.
> 
> It's possible because:
> 1) if the objcg is received from the current task struct, the task is
>    keeping a reference to the objcg.
> 2) if the objcg is received from an active memcg (remote charging),
>    the memcg is pinned by the scope and has a reference to the
>    corresponding objcg.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Shakeel Butt <shakeelb@google.com>

Reviewed-by: Vlastimil Babka <vbabka@suse.cz>

Re: [PATCH v3 3/5] mm: kmem: make memcg keep a reference to the original objcg

[Vlastimil Babka]

On 10/17/23 00:18, Roman Gushchin wrote:
> Keep a reference to the original objcg object for the entire life
> of a memcg structure.
> 
> This allows to simplify the synchronization on the kernel memory
> allocation paths: pinning a (live) memcg will also pin the
> corresponding objcg.
> 
> The memory overhead of this change is minimal because object cgroups
> usually outlive their corresponding memory cgroups even without this
> change, so it's only an additional pointer per memcg.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Shakeel Butt <shakeelb@google.com>

Forgot to point out typo:

> ---
>  include/linux/memcontrol.h | 8 +++++++-
>  mm/memcontrol.c            | 5 +++++
>  2 files changed, 12 insertions(+), 1 deletion(-)
> 
> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
> index ab94ad4597d0..277690af383d 100644
> --- a/include/linux/memcontrol.h
> +++ b/include/linux/memcontrol.h
> @@ -299,7 +299,13 @@ struct mem_cgroup {
>  
>  #ifdef CONFIG_MEMCG_KMEM
>  	int kmemcg_id;
> -	struct obj_cgroup __rcu *objcg;
> +	/*
> +	 * memcg->objcg is wiped out as a part of the objcg repaprenting

						reparenting ^

> +	 * process. memcg->orig_objcg preserves a pointer (and a reference)
> +	 * to the original objcg until the end of live of memcg.
> +	 */
> +	struct obj_cgroup __rcu	*objcg;
> +	struct obj_cgroup	*orig_objcg;
>  	/* list of inherited objcgs, protected by objcg_lock */
>  	struct list_head objcg_list;
>  #endif
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index 0605e45bd4a2..d90cc19e4113 100644

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Vlastimil Babka]

On 10/18/23 11:52, Vlastimil Babka wrote:
>> +
>> +	return objcg;
>> +}
>> +
>>  __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
>>  {
>>  	struct mem_cgroup *memcg;
>> @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
>>  
>>  	if (in_task()) {
>>  		memcg = current->active_memcg;
>> +		if (unlikely(memcg))
>> +			goto from_memcg;
>>  
>> -		/* Memcg to charge can't be determined. */
>> -		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> The checks for current->mm and PF_KTHREAD seem to be gone completely after
> the patch, was that intended and why?

And also they are not present in the new flavor that's current_obj_cgroup().

Re: [PATCH v3 5/5] percpu: scoped objcg protection

[Vlastimil Babka]

On 10/17/23 00:19, Roman Gushchin wrote:
> Similar to slab and kmem, switch to a scope-based protection of the
> objcg pointer to avoid.
> 
> Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> Acked-by: Shakeel Butt <shakeelb@google.com>

Reviewed-by: Vlastimil Babka <vbabka@suse.cz>

Do you plan to convert also the bpf users of get_obj_cgroup_from_current()
so it could be removed?

Thanks!

> ---
>  mm/percpu.c | 8 +++-----
>  1 file changed, 3 insertions(+), 5 deletions(-)
> 
> diff --git a/mm/percpu.c b/mm/percpu.c
> index a7665de8485f..f53ba692d67a 100644
> --- a/mm/percpu.c
> +++ b/mm/percpu.c
> @@ -1628,14 +1628,12 @@ static bool pcpu_memcg_pre_alloc_hook(size_t size, gfp_t gfp,
>  	if (!memcg_kmem_online() || !(gfp & __GFP_ACCOUNT))
>  		return true;
>  
> -	objcg = get_obj_cgroup_from_current();
> +	objcg = current_obj_cgroup();
>  	if (!objcg)
>  		return true;
>  
> -	if (obj_cgroup_charge(objcg, gfp, pcpu_obj_full_size(size))) {
> -		obj_cgroup_put(objcg);
> +	if (obj_cgroup_charge(objcg, gfp, pcpu_obj_full_size(size)))
>  		return false;
> -	}
>  
>  	*objcgp = objcg;
>  	return true;
> @@ -1649,6 +1647,7 @@ static void pcpu_memcg_post_alloc_hook(struct obj_cgroup *objcg,
>  		return;
>  
>  	if (likely(chunk && chunk->obj_cgroups)) {
> +		obj_cgroup_get(objcg);
>  		chunk->obj_cgroups[off >> PCPU_MIN_ALLOC_SHIFT] = objcg;
>  
>  		rcu_read_lock();
> @@ -1657,7 +1656,6 @@ static void pcpu_memcg_post_alloc_hook(struct obj_cgroup *objcg,
>  		rcu_read_unlock();
>  	} else {
>  		obj_cgroup_uncharge(objcg, pcpu_obj_full_size(size));
> -		obj_cgroup_put(objcg);
>  	}
>  }
>  

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Shakeel Butt]

On Wed, Oct 18, 2023 at 2:52 AM Vlastimil Babka <vbabka@suse.cz> wrote:
[...]
> >
> > +static struct obj_cgroup *current_objcg_update(void)
> > +{
> > +     struct mem_cgroup *memcg;
> > +     struct obj_cgroup *old, *objcg = NULL;
> > +
> > +     do {
> > +             /* Atomically drop the update bit. */
> > +             old = xchg(&current->objcg, NULL);
> > +             if (old) {
> > +                     old = (struct obj_cgroup *)
> > +                             ((unsigned long)old & ~CURRENT_OBJCG_UPDATE_FLAG);
> > +                     if (old)
> > +                             obj_cgroup_put(old);
> > +
> > +                     old = NULL;
> > +             }
> > +
> > +             /* Obtain the new objcg pointer. */
> > +             rcu_read_lock();
> > +             memcg = mem_cgroup_from_task(current);
> > +             /*
> > +              * The current task can be asynchronously moved to another
> > +              * memcg and the previous memcg can be offlined. So let's
> > +              * get the memcg pointer and try get a reference to objcg
> > +              * under a rcu read lock.
> > +              */
> > +             for (; memcg != root_mem_cgroup; memcg = parent_mem_cgroup(memcg)) {
> > +                     objcg = rcu_dereference(memcg->objcg);
> > +                     if (likely(objcg && obj_cgroup_tryget(objcg)))
>
> So IIUC here we increase objcg refcount.
>
> > +                             break;
> > +                     objcg = NULL;
> > +             }
> > +             rcu_read_unlock();
> > +
> > +             /*
> > +              * Try set up a new objcg pointer atomically. If it
> > +              * fails, it means the update flag was set concurrently, so
> > +              * the whole procedure should be repeated.
> > +              */
> > +     } while (!try_cmpxchg(&current->objcg, &old, objcg));
>
> And if this fails we throw objcg away and try again, but we should do
> obj_cgroup_put(objcg) first, as otherwise it would cause a leak?
>

Indeed there is a reference leak here.

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Roman Gushchin]

On Wed, Oct 18, 2023 at 11:52:27AM +0200, Vlastimil Babka wrote:
> On 10/17/23 00:18, Roman Gushchin wrote:
> > To charge a freshly allocated kernel object to a memory cgroup, the
> > kernel needs to obtain an objcg pointer. Currently it does it
> > indirectly by obtaining the memcg pointer first and then calling to
> > __get_obj_cgroup_from_memcg().
> > 
> > Usually tasks spend their entire life belonging to the same object
> > cgroup. So it makes sense to save the objcg pointer on task_struct
> > directly, so it can be obtained faster. It requires some work on fork,
> > exit and cgroup migrate paths, but these paths are way colder.
> > 
> > To avoid any costly synchronization the following rules are applied:
> > 1) A task sets it's objcg pointer itself.
> > 
> > 2) If a task is being migrated to another cgroup, the least
> >    significant bit of the objcg pointer is set atomically.
> > 
> > 3) On the allocation path the objcg pointer is obtained locklessly
> >    using the READ_ONCE() macro and the least significant bit is
> >    checked. If it's set, the following procedure is used to update
> >    it locklessly:
> >        - task->objcg is zeroed using cmpxcg
> >        - new objcg pointer is obtained
> >        - task->objcg is updated using try_cmpxchg
> >        - operation is repeated if try_cmpxcg fails
> >    It guarantees that no updates will be lost if task migration
> >    is racing against objcg pointer update. It also allows to keep
> >    both read and write paths fully lockless.
> > 
> > Because the task is keeping a reference to the objcg, it can't go away
> > while the task is alive.
> > 
> > This commit doesn't change the way the remote memcg charging works.
> > 
> > Signed-off-by: Roman Gushchin (Cruise) <roman.gushchin@linux.dev>
> > Tested-by: Naresh Kamboju <naresh.kamboju@linaro.org>
> > Acked-by: Johannes Weiner <hannes@cmpxchg.org>
> > ---
> >  include/linux/sched.h |   4 ++
> >  mm/memcontrol.c       | 130 +++++++++++++++++++++++++++++++++++++++---
> >  2 files changed, 125 insertions(+), 9 deletions(-)
> > 
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index 16ac2a5838fb..0605e45bd4a2 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> 
> So IIUC here we increase objcg refcount.
> 
> > +				break;
> > +			objcg = NULL;
> > +		}
> > +		rcu_read_unlock();
> > +
> > +		/*
> > +		 * Try set up a new objcg pointer atomically. If it
> > +		 * fails, it means the update flag was set concurrently, so
> > +		 * the whole procedure should be repeated.
> > +		 */
> > +	} while (!try_cmpxchg(&current->objcg, &old, objcg));
> 
> And if this fails we throw objcg away and try again, but we should do
> obj_cgroup_put(objcg) first, as otherwise it would cause a leak?

Great catch! Thanks!

> 
> > +
> > +	return objcg;
> > +}
> > +
> >  __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> >  {
> >  	struct mem_cgroup *memcg;
> > @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> >  
> >  	if (in_task()) {
> >  		memcg = current->active_memcg;
> > +		if (unlikely(memcg))
> > +			goto from_memcg;
> >  
> > -		/* Memcg to charge can't be determined. */
> > -		if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> 
> The checks for current->mm and PF_KTHREAD seem to be gone completely after
> the patch, was that intended and why?

There is no need for those anymore because it's as cheap or cheaper
to check task->objcg for being NULL. Those were primarily used to rule out
kernel threads allocations early.

I gonna fix the objcg ref leak, add the comment you asked above and post v4
of this particular patch.

Thank you for reviewing the series!

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Shakeel Butt]

On Wed, Oct 18, 2023 at 10:22 AM Roman Gushchin
<roman.gushchin@linux.dev> wrote:
>
[...]
> > >     struct mem_cgroup *memcg;
> > > @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> > >
> > >     if (in_task()) {
> > >             memcg = current->active_memcg;
> > > +           if (unlikely(memcg))
> > > +                   goto from_memcg;
> > >
> > > -           /* Memcg to charge can't be determined. */
> > > -           if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> >
> > The checks for current->mm and PF_KTHREAD seem to be gone completely after
> > the patch, was that intended and why?
>
> There is no need for those anymore because it's as cheap or cheaper
> to check task->objcg for being NULL. Those were primarily used to rule out
> kernel threads allocations early.
>

I have the same understanding but please correct my suspicions here.
We can echo the kernel thread's pid to cgroup.procs which have
PF_NO_SETAFFINITY and thus this will cause the lower bit of the kernel
thread's task->objcg to be set. Please correct me if I am missing
something.

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Roman Gushchin]

On Wed, Oct 18, 2023 at 11:26:59AM -0700, Shakeel Butt wrote:
> On Wed, Oct 18, 2023 at 10:22 AM Roman Gushchin
> <roman.gushchin@linux.dev> wrote:
> >
> [...]
> > > >     struct mem_cgroup *memcg;
> > > > @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> > > >
> > > >     if (in_task()) {
> > > >             memcg = current->active_memcg;
> > > > +           if (unlikely(memcg))
> > > > +                   goto from_memcg;
> > > >
> > > > -           /* Memcg to charge can't be determined. */
> > > > -           if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> > >
> > > The checks for current->mm and PF_KTHREAD seem to be gone completely after
> > > the patch, was that intended and why?
> >
> > There is no need for those anymore because it's as cheap or cheaper
> > to check task->objcg for being NULL. Those were primarily used to rule out
> > kernel threads allocations early.
> >
> 
> I have the same understanding but please correct my suspicions here.
> We can echo the kernel thread's pid to cgroup.procs which have
> PF_NO_SETAFFINITY and thus this will cause the lower bit of the kernel
> thread's task->objcg to be set. Please correct me if I am missing
> something.

Yes, you seem to be right. It's a gray zone because moving kernel threads out of
the root cgroup doesn't sound like a good idea, but I agree it's better to keep
the old behavior in place.

Does this fixlet look good to you?

Thanks!

--

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 1a2835448028..0b0d2dc7a7d4 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3021,6 +3021,10 @@ static struct obj_cgroup *current_objcg_update(void)
                        old = NULL;
                }

+               /* If new objcg is NULL, no reason for the second atomic update. */
+               if (!current->mm || (current->flags & PF_KTHREAD))
+                       return NULL;
+
                /*
                 * Release the objcg pointer from the previous iteration,
                 * if try_cmpxcg() below fails.

Re: [PATCH v3 2/5] mm: kmem: add direct objcg pointer to task_struct

[Shakeel Butt]

On Wed, Oct 18, 2023 at 3:38 PM Roman Gushchin <roman.gushchin@linux.dev> wrote:
>
> On Wed, Oct 18, 2023 at 11:26:59AM -0700, Shakeel Butt wrote:
> > On Wed, Oct 18, 2023 at 10:22 AM Roman Gushchin
> > <roman.gushchin@linux.dev> wrote:
> > >
> > [...]
> > > > >     struct mem_cgroup *memcg;
> > > > > @@ -3008,19 +3054,26 @@ __always_inline struct obj_cgroup *get_obj_cgroup_from_current(void)
> > > > >
> > > > >     if (in_task()) {
> > > > >             memcg = current->active_memcg;
> > > > > +           if (unlikely(memcg))
> > > > > +                   goto from_memcg;
> > > > >
> > > > > -           /* Memcg to charge can't be determined. */
> > > > > -           if (likely(!memcg) && (!current->mm || (current->flags & PF_KTHREAD)))
> > > >
> > > > The checks for current->mm and PF_KTHREAD seem to be gone completely after
> > > > the patch, was that intended and why?
> > >
> > > There is no need for those anymore because it's as cheap or cheaper
> > > to check task->objcg for being NULL. Those were primarily used to rule out
> > > kernel threads allocations early.
> > >
> >
> > I have the same understanding but please correct my suspicions here.
> > We can echo the kernel thread's pid to cgroup.procs which have
> > PF_NO_SETAFFINITY and thus this will cause the lower bit of the kernel
> > thread's task->objcg to be set. Please correct me if I am missing
> > something.
>
> Yes, you seem to be right. It's a gray zone because moving kernel threads out of
> the root cgroup doesn't sound like a good idea, but I agree it's better to keep
> the old behavior in place.
>
> Does this fixlet look good to you?
>

This looks fine. Another option is not to set the bit for such
task_structs in fork/attach.

> Thanks!
>
> --
>
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index 1a2835448028..0b0d2dc7a7d4 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -3021,6 +3021,10 @@ static struct obj_cgroup *current_objcg_update(void)
>                         old = NULL;
>                 }
>
> +               /* If new objcg is NULL, no reason for the second atomic update. */
> +               if (!current->mm || (current->flags & PF_KTHREAD))
> +                       return NULL;
> +
>                 /*
>                  * Release the objcg pointer from the previous iteration,
>                  * if try_cmpxcg() below fails.