From: Bingfang Guo <bingfangguo@tencent.com>

__mem_cgroup_uncharge_swap() needs both the memcg and the id refcount
drop.  Right now it looks the memcg up by id, uncharges it, then looks
it up again inside mem_cgroup_private_id_put() to drop the reference.

Make mem_cgroup_private_id_put() resolve the id once, drop the
reference, and return the nearest online memcg with a reference held for
the caller.  __mem_cgroup_uncharge_swap() then uses that memcg directly
and drops the reference after uncharging, avoiding the second xarray
lookup.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol.c | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 048c9bb0fad79..f0503a1e5492d 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4048,14 +4048,28 @@ static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n
 	}
 }
 
-static void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
+/**
+ * mem_cgroup_private_id_put - put memcgid and get the nearest online memcg
+ * @id: the memcg private id got from mem_cgroup_id_get_online
+ * @n: count of references to put
+ */
+static struct mem_cgroup *mem_cgroup_private_id_put(unsigned short id, unsigned int n)
 {
 	struct mem_cgroup *memcg;
 
 	rcu_read_lock();
 	memcg = mem_cgroup_from_private_id(id);
+	if (!memcg)
+		goto out;
+
 	__mem_cgroup_private_id_put(memcg, n);
+
+	while (memcg_is_dying(memcg) || !mem_cgroup_tryget(memcg))
+		memcg = parent_mem_cgroup(memcg);
+
+out:
 	rcu_read_unlock();
+	return memcg;
 }
 
 static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
@@ -5816,7 +5830,7 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 	struct mem_cgroup *memcg;
 
 	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
+	memcg = mem_cgroup_private_id_put(id, nr_pages);
 	if (memcg) {
 		if (!mem_cgroup_is_root(memcg)) {
 			if (do_memsw_account())
@@ -5825,10 +5839,10 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 				page_counter_uncharge(&memcg->swap, nr_pages);
 		}
 		mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
-		mem_cgroup_private_id_put(id, nr_pages);
 	}
 	rcu_read_unlock();
 
+	mem_cgroup_put(memcg);
 }
 
 long mem_cgroup_get_nr_swap_pages(struct mem_cgroup *memcg)

-- 
2.43.7