[RFC PATCH 9/9 v2] mm/memcontrol: Make memory.max tier-aware

[Joshua Hahn <joshua.hahnjy@gmail.com>]

On machines serving multiple workloads whose memory is isolated via the
memory cgroup controller, it is currently impossible to enforce a fair
distribution of toptier memory among the workloads, as the limits only
enforce total memory footprint, but not where that memory resides.

This makes ensuring consistent baseline performance difficult, as each
workload's performance is heavily impacted by workload-external factors
such as which other workloads are co-located in the same host, and the
order in which the workloads are started.

Extend the existing memory.max protection to be tier-aware.

Depending on the combination of limit breaches, selectively reclaim on
toptier nodes: when memory.max is breached, perform reclaim on all
nodes.  When memory.max is safe but toptier.max is breached, perform
targeted reclaim on toptier nodes only.

Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 mm/memcontrol.c | 56 ++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 44 insertions(+), 12 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index e5f39830d250d..d8d67ada993ff 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -1518,6 +1518,15 @@ static unsigned long mem_cgroup_margin(struct mem_cgroup *memcg)
 	if (count < limit)
 		margin = limit - count;
 
+	if (mem_cgroup_tiered_limits()) {
+		count = page_counter_read(&memcg->toptier);
+		limit = READ_ONCE(memcg->toptier.max);
+		if (count < limit)
+			margin = min(margin, limit - count);
+		else
+			margin = 0;
+	}
+
 	if (do_memsw_account()) {
 		count = page_counter_read(&memcg->memsw);
 		limit = READ_ONCE(memcg->memsw.max);
@@ -2424,11 +2433,12 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	bool raised_max_event = false;
 	unsigned long pflags;
 	bool allow_spinning = gfpflags_allow_spinning(gfp_mask);
-	bool toptier_charged;
+	nodemask_t toptier_nodes;
+	nodemask_t *reclaim_nodes;
 
 retry:
 	reclaim_options = MEMCG_RECLAIM_MAY_SWAP;
-	toptier_charged = false;
+	reclaim_nodes = NULL;
 
 	if (do_memsw_account() &&
 	    !page_counter_try_charge(&memcg->memsw, nr_pages, &counter)) {
@@ -2438,13 +2448,20 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	}
 
 	if (toptier &&
-	    page_counter_try_charge(&memcg->toptier, nr_pages, &counter))
-		toptier_charged = true;
+	    !page_counter_try_charge(&memcg->toptier, nr_pages, &counter)) {
+		get_toptier_nodemask(&toptier_nodes);
+		reclaim_nodes = &toptier_nodes;
+		mem_over_limit = mem_cgroup_from_counter(counter, toptier);
+
+		if (do_memsw_account())
+			page_counter_uncharge(&memcg->memsw, nr_pages);
+		goto reclaim;
+	}
 
 	if (page_counter_try_charge(&memcg->memory, nr_pages, &counter))
 		goto done_restock;
 
-	if (toptier_charged)
+	if (toptier)
 		page_counter_uncharge(&memcg->toptier, nr_pages);
 	if (do_memsw_account())
 		page_counter_uncharge(&memcg->memsw, nr_pages);
@@ -2473,7 +2490,7 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	psi_memstall_enter(&pflags);
 	nr_reclaimed = try_to_free_mem_cgroup_pages(mem_over_limit, nr_pages,
 						    gfp_mask, reclaim_options,
-						    NULL, NULL);
+						    NULL, reclaim_nodes);
 	psi_memstall_leave(&pflags);
 
 	if (mem_cgroup_margin(mem_over_limit) >= nr_pages)
@@ -4683,7 +4700,8 @@ static ssize_t memory_max_write(struct kernfs_open_file *of,
 	struct mem_cgroup *memcg = mem_cgroup_from_css(of_css(of));
 	unsigned int nr_reclaims = MAX_RECLAIM_RETRIES;
 	bool drained = false;
-	unsigned long max;
+	unsigned long max, toptier_max = PAGE_COUNTER_MAX;
+	nodemask_t toptier_nodes;
 	int err;
 
 	buf = strstrip(buf);
@@ -4692,16 +4710,30 @@ static ssize_t memory_max_write(struct kernfs_open_file *of,
 		return err;
 
 	xchg(&memcg->memory.max, max);
-	if (mem_cgroup_tiered_limits())
-		xchg(&memcg->toptier.max, page_counter_max_or_scale(max));
+	if (mem_cgroup_tiered_limits()) {
+		toptier_max = page_counter_max_or_scale(max);
+		xchg(&memcg->toptier.max, toptier_max);
+		get_toptier_nodemask(&toptier_nodes);
+	}
 
 	if (of->file->f_flags & O_NONBLOCK)
 		goto out;
 
 	for (;;) {
 		unsigned long nr_pages = page_counter_read(&memcg->memory);
+		unsigned long nr_toptier = page_counter_read(&memcg->toptier);
+		unsigned long to_reclaim = 0;
+		nodemask_t *reclaim_nodes = NULL;
+
+		if (nr_pages > max) {
+			to_reclaim = nr_pages - max;
+		} else if (mem_cgroup_tiered_limits() &&
+				nr_toptier > toptier_max) {
+			to_reclaim = nr_toptier - toptier_max;
+			reclaim_nodes = &toptier_nodes;
+		}
 
-		if (nr_pages <= max)
+		if (!to_reclaim)
 			break;
 
 		if (signal_pending(current))
@@ -4714,9 +4746,9 @@ static ssize_t memory_max_write(struct kernfs_open_file *of,
 		}
 
 		if (nr_reclaims) {
-			if (!try_to_free_mem_cgroup_pages(memcg, nr_pages - max,
+			if (!try_to_free_mem_cgroup_pages(memcg, to_reclaim,
 					GFP_KERNEL, MEMCG_RECLAIM_MAY_SWAP,
-					NULL, NULL))
+					NULL, reclaim_nodes))
 				nr_reclaims--;
 			continue;
 		}
-- 
2.52.0