[PATCH RFC 02/14] mm: memcg: introduce mem_cgroup_ptr

[Roman Gushchin]

This commit introduces mem_cgroup_ptr structure and corresponding API.
It implements a pointer to a memory cgroup with a built-in reference
counter. The main goal of it is to implement reparenting efficiently.

If a number of objects (e.g. slab pages) have to keep a pointer and
a reference to a memory cgroup, they can use mem_cgroup_ptr instead.
On reparenting, only one mem_cgroup_ptr->memcg pointer has to be
changed, instead of walking over all accounted objects.

mem_cgroup_ptr holds a single reference to the corresponding memory
cgroup. Because it's initialized before the css reference counter,
css's refcounter can't be bumped at allocation time. Instead, it's
bumped on reparenting which happens during offlining. A cgroup is
never released online, so it's fine.

mem_cgroup_ptr is released using rcu, so memcg->kmem_memcg_ptr can
be accessed in a rcu read section. On reparenting it's atomically
switched to NULL. If the reader gets NULL, it can just read parent's
kmem_memcg_ptr instead.

Each memory cgroup contains a list of kmem_memcg_ptrs. On reparenting
the list is spliced into the parent's list. The list is protected
using the css set lock.

Signed-off-by: Roman Gushchin <guro@xxxxxx>
---
 include/linux/memcontrol.h | 50 ++++++++++++++++++++++
 mm/memcontrol.c            | 87 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 133 insertions(+), 4 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 120d39066148..dd5ebfe5a86c 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -23,6 +23,7 @@
 #include <linux/page-flags.h>
 
 struct mem_cgroup;
+struct mem_cgroup_ptr;
 struct page;
 struct mm_struct;
 struct kmem_cache;
@@ -197,6 +198,22 @@ struct memcg_cgwb_frn {
 	int memcg_id;			/* memcg->css.id of foreign inode */
 	u64 at;				/* jiffies_64 at the time of dirtying */
 	struct wb_completion done;	/* tracks in-flight foreign writebacks */
+}
+
+/*
+ * A pointer to a memory cgroup with a built-in reference counter.
+ * For a use as an intermediate object to simplify reparenting of
+ * objects charged to the cgroup. The memcg pointer can be switched
+ * to the parent cgroup without a need to modify all objects
+ * which hold the reference to the cgroup.
+ */
+struct mem_cgroup_ptr {
+	struct percpu_ref refcnt;
+	struct mem_cgroup *memcg;
+	union {
+		struct list_head list;
+		struct rcu_head rcu;
+	};
 };
 
 /*
@@ -312,6 +329,8 @@ struct mem_cgroup {
 	int kmemcg_id;
 	enum memcg_kmem_state kmem_state;
 	struct list_head kmem_caches;
+	struct mem_cgroup_ptr __rcu *kmem_memcg_ptr;
+	struct list_head kmem_memcg_ptr_list;
 #endif
 
 	int last_scanned_node;
@@ -440,6 +459,21 @@ struct mem_cgroup *mem_cgroup_from_css(struct cgroup_subsys_state *css){
 	return css ? container_of(css, struct mem_cgroup, css) : NULL;
 }
 
+static inline bool mem_cgroup_ptr_tryget(struct mem_cgroup_ptr *ptr)
+{
+	return percpu_ref_tryget(&ptr->refcnt);
+}
+
+static inline void mem_cgroup_ptr_get(struct mem_cgroup_ptr *ptr)
+{
+	percpu_ref_get(&ptr->refcnt);
+}
+
+static inline void mem_cgroup_ptr_put(struct mem_cgroup_ptr *ptr)
+{
+	percpu_ref_put(&ptr->refcnt);
+}
+
 static inline void mem_cgroup_put(struct mem_cgroup *memcg)
 {
 	if (memcg)
@@ -1433,6 +1467,22 @@ static inline int memcg_cache_id(struct mem_cgroup *memcg)
 	return memcg ? memcg->kmemcg_id : -1;
 }
 
+static inline struct mem_cgroup_ptr *
+mem_cgroup_get_kmem_ptr(struct mem_cgroup *memcg)
+{
+	struct mem_cgroup_ptr *memcg_ptr;
+
+	rcu_read_lock();
+	do {
+		memcg_ptr = rcu_dereference(memcg->kmem_memcg_ptr);
+		if (memcg_ptr && mem_cgroup_ptr_tryget(memcg_ptr))
+			break;
+	} while ((memcg = parent_mem_cgroup(memcg)));
+	rcu_read_unlock();
+
+	return memcg_ptr;
+}
+
 #else
 
 static inline int memcg_kmem_charge(struct page *page, gfp_t gfp, int order)
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index effefcec47b3..cb9adb31360e 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -266,6 +266,77 @@ struct cgroup_subsys_state *vmpressure_to_css(struct vmpressure *vmpr)
 }
 
 #ifdef CONFIG_MEMCG_KMEM
+extern spinlock_t css_set_lock;
+
+static void memcg_ptr_release(struct percpu_ref *ref)
+{
+	struct mem_cgroup_ptr *ptr = container_of(ref, struct mem_cgroup_ptr,
+						  refcnt);
+	unsigned long flags;
+
+	spin_lock_irqsave(&css_set_lock, flags);
+	list_del(&ptr->list);
+	spin_unlock_irqrestore(&css_set_lock, flags);
+
+	mem_cgroup_put(ptr->memcg);
+	percpu_ref_exit(ref);
+	kfree_rcu(ptr, rcu);
+}
+
+static int memcg_init_kmem_memcg_ptr(struct mem_cgroup *memcg)
+{
+	struct mem_cgroup_ptr *kmem_memcg_ptr;
+	int ret;
+
+	kmem_memcg_ptr = kmalloc(sizeof(struct mem_cgroup_ptr), GFP_KERNEL);
+	if (!kmem_memcg_ptr)
+		return -ENOMEM;
+
+	ret = percpu_ref_init(&kmem_memcg_ptr->refcnt, memcg_ptr_release,
+			      0, GFP_KERNEL);
+	if (ret) {
+		kfree(kmem_memcg_ptr);
+		return ret;
+	}
+
+	kmem_memcg_ptr->memcg = memcg;
+	INIT_LIST_HEAD(&kmem_memcg_ptr->list);
+	rcu_assign_pointer(memcg->kmem_memcg_ptr, kmem_memcg_ptr);
+	list_add(&kmem_memcg_ptr->list, &memcg->kmem_memcg_ptr_list);
+	return 0;
+}
+
+static void memcg_reparent_kmem_memcg_ptr(struct mem_cgroup *memcg,
+					  struct mem_cgroup *parent)
+{
+	unsigned int nr_reparented = 0;
+	struct mem_cgroup_ptr *memcg_ptr = NULL;
+
+	rcu_swap_protected(memcg->kmem_memcg_ptr, memcg_ptr, true);
+	percpu_ref_kill(&memcg_ptr->refcnt);
+
+	/*
+	 * kmem_memcg_ptr is initialized before css refcounter, so until now
+	 * it doesn't hold a reference to the memcg. Bump it here.
+	 */
+	css_get(&memcg->css);
+
+	spin_lock_irq(&css_set_lock);
+	list_for_each_entry(memcg_ptr, &memcg->kmem_memcg_ptr_list, list) {
+		xchg(&memcg_ptr->memcg, parent);
+		nr_reparented++;
+	}
+	if (nr_reparented)
+		list_splice(&memcg->kmem_memcg_ptr_list,
+			    &parent->kmem_memcg_ptr_list);
+	spin_unlock_irq(&css_set_lock);
+
+	if (nr_reparented) {
+		css_get_many(&parent->css, nr_reparented);
+		css_put_many(&memcg->css, nr_reparented);
+	}
+}
+
 /*
  * This will be the memcg's index in each cache's ->memcg_params.memcg_caches.
  * The main reason for not using cgroup id for this:
@@ -3554,7 +3625,7 @@ static void memcg_flush_percpu_vmevents(struct mem_cgroup *memcg)
 #ifdef CONFIG_MEMCG_KMEM
 static int memcg_online_kmem(struct mem_cgroup *memcg)
 {
-	int memcg_id;
+	int memcg_id, ret;
 
 	if (cgroup_memory_nokmem)
 		return 0;
@@ -3566,6 +3637,12 @@ static int memcg_online_kmem(struct mem_cgroup *memcg)
 	if (memcg_id < 0)
 		return memcg_id;
 
+	ret = memcg_init_kmem_memcg_ptr(memcg);
+	if (ret) {
+		memcg_free_cache_id(memcg_id);
+		return ret;
+	}
+
 	static_branch_inc(&memcg_kmem_enabled_key);
 	/*
 	 * A memory cgroup is considered kmem-online as soon as it gets
@@ -3601,12 +3678,13 @@ static void memcg_offline_kmem(struct mem_cgroup *memcg)
 		parent = root_mem_cgroup;
 
 	/*
-	 * Deactivate and reparent kmem_caches. Then flush percpu
-	 * slab statistics to have precise values at the parent and
-	 * all ancestor levels. It's required to keep slab stats
+	 * Deactivate and reparent kmem_caches and reparent kmem_memcg_ptr.
+	 * Then flush percpu slab statistics to have precise values at the
+	 * parent and all ancestor levels. It's required to keep slab stats
 	 * accurate after the reparenting of kmem_caches.
 	 */
 	memcg_deactivate_kmem_caches(memcg, parent);
+	memcg_reparent_kmem_memcg_ptr(memcg, parent);
 	memcg_flush_percpu_vmstats(memcg, true);
 
 	kmemcg_id = memcg->kmemcg_id;
@@ -5171,6 +5249,7 @@ static struct mem_cgroup *mem_cgroup_alloc(void)
 	memcg->socket_pressure = jiffies;
 #ifdef CONFIG_MEMCG_KMEM
 	memcg->kmemcg_id = -1;
+	INIT_LIST_HEAD(&memcg->kmem_memcg_ptr_list);
 #endif
 #ifdef CONFIG_CGROUP_WRITEBACK
 	INIT_LIST_HEAD(&memcg->cgwb_list);
-- 
2.21.0