[PATCH v2 2/5] memcg: force charges from reclaim before trying the limits

[Shakeel Butt <shakeel.butt@linux.dev>]

Reclaim can charge memory, for example when zswap stores a page.
try_charge_memcg() forces these PF_MEMALLOC charges, but only after
they fail at a limit. A charge that fits still has memory.high
enforced. That can make a task in direct reclaim do more reclaim and
sleep longer later, and kswapd records overage that it never handles.

It is also a problem for the next patches, which make kernel threads
queue high_work. The worker charges memory from its own reclaim, so it
can queue itself again. With those patches but not this one, the worker
queued itself in 118 of 284 runs in a VM test with incompressible zswap
data.

Check PF_MEMALLOC before trying the limits. Reclaim charges then skip
memory.high, as they can already go over memory.max. Drop the
PF_MEMALLOC check in memcg_enforce_high(), which no longer sees them.

Other effects:
- As before, the check reads current, so it also catches interrupts on
  a PF_MEMALLOC task such as kswapd. Softirqs set PF_MEMALLOC
  themselves for swap over network, so the check can't be limited to
  task context.
- These charges no longer charge a batch or refill the per-cpu stock.
- MEMCG_MAX still goes to the first cgroup that the charge puts over
  its limit. On cgroup v1, memory.failcnt no longer counts them.
- high_work can still run again and again when there is nothing left
  to reclaim. This only stops reclaim from queuing more reclaim.

Suggested-by: Michal Hocko <mhocko@suse.com>
Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 mm/memcontrol.c | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 0e484f7f823a..af95a4574a0f 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2762,7 +2762,6 @@ static void memcg_enforce_high(struct mem_cgroup *memcg, unsigned int nr_pages,
 	 * when it rechecks the overage and simply bail out.
 	 */
 	if (current->memcg_nr_pages_over_high > MEMCG_CHARGE_BATCH &&
-	    !(current->flags & PF_MEMALLOC) &&
 	    gfpflags_allow_blocking(gfp_mask))
 		__mem_cgroup_handle_over_high(gfp_mask);
 }
@@ -2787,6 +2786,21 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	if (consume_stock(memcg, nr_pages))
 		return ret;
 
+	/*
+	 * Reclaim may need memory to free memory. Let it go over the
+	 * limits for a while rather than fail or OOM kill. Doing this
+	 * before trying the limits also skips memory.high, so the
+	 * high_work worker can't queue itself from its own reclaim.
+	 */
+	if (unlikely(current->flags & PF_MEMALLOC)) {
+		/* Find the cgroup this charge puts over its limit, if any. */
+		for (mem_over_limit = memcg; mem_over_limit;
+		     mem_over_limit = parent_mem_cgroup(mem_over_limit))
+			if (mem_cgroup_margin(mem_over_limit) < nr_pages)
+				break;
+		goto force;
+	}
+
 	if (!allow_spinning)
 		/* Avoid the refill and flush of the older stock */
 		batch = nr_pages;
@@ -2809,15 +2823,6 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 		goto retry;
 	}
 
-	/*
-	 * Prevent unbounded recursion when reclaim operations need to
-	 * allocate memory. This might exceed the limits temporarily,
-	 * but we prefer facilitating memory reclaim and getting back
-	 * under the limit over triggering OOM kills in these cases.
-	 */
-	if (unlikely(current->flags & PF_MEMALLOC))
-		goto force;
-
 	if (unlikely(task_in_memcg_oom(current)))
 		goto nomem;
 
@@ -2913,9 +2918,9 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 out:
 	/*
 	 * Don't forget to raise a MEMCG_MAX event for forced or rejected
-	 * requests.
+	 * requests. mem_over_limit is NULL if a PF_MEMALLOC charge fits.
 	 */
-	if (!raised_max_event)
+	if (!raised_max_event && mem_over_limit)
 		__memcg_memory_event(mem_over_limit, MEMCG_MAX, allow_spinning);
 
 	return ret;
-- 
2.53.0-Meta