[PATCH] mm/memcontrol: avoid stuck FLUSHING_CACHED_CHARGE bit on isolated cpus

[Rik van Riel <riel@surriel.com>]

drain_all_stock() can leave FLUSHING_CACHED_CHARGE set after the
work is dropped.  It sets the bit before checking isolation and
schedule_drain_work() checks isolation and queues in a separate RCU
critical section, so housekeeping_update()'s synchronize_rcu() can
race the second check.

drain_local_stock() only clears the bit for work that ran, so the bit
remains set and the stock is never drained again.

Reorganize the drain_all_stock() loop, reducing nesting, splitting
out local vs remote cpu handling, and skipping everything on isolated
cpus, which solves the stuck FLUSHING_CACHED_CHARGE flag.

Fixes: 6a792697a53a ("memcg: do not drain charge pcp caches on remote isolated cpus")
Cc: stable@vger.kernel.org
Assisted-by: Hermes:muse-spark-1.2
Signed-off-by: Rik van Riel <riel@surriel.com>

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index a660ea0f820b..b0d8ec042d48 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2330,19 +2330,6 @@ static bool obj_stock_should_drain(struct obj_stock_pcp *stock,
 	return true;
 }
 
-static void schedule_drain_work(int cpu, struct work_struct *work)
-{
-	/*
-	 * Protect housekeeping cpumask read and work enqueue together
-	 * in the same RCU critical section so that later cpuset isolated
-	 * partition update only need to wait for an RCU GP and flush the
-	 * pending work on newly isolated CPUs.
-	 */
-	guard(rcu)();
-	if (!cpu_is_isolated(cpu))
-		queue_work_on(cpu, memcg_wq, work);
-}
-
 /*
  * Drains all per-CPU charge caches for given root_memcg resp. subtree
  * of the hierarchy under it.
@@ -2366,19 +2353,23 @@ void drain_all_stock(struct mem_cgroup *root_memcg)
 		struct memcg_stock_pcp *memcg_st = &per_cpu(memcg_stock, cpu);
 		struct obj_stock_pcp *obj_st = &per_cpu(obj_stock, cpu);
 
-		if (memcg_stock_should_drain(memcg_st, root_memcg)) {
-			if (cpu == curcpu)
+		if (cpu == curcpu) {
+			if (memcg_stock_should_drain(memcg_st, root_memcg))
 				drain_local_memcg_stock(&memcg_st->work);
-			else
-				schedule_drain_work(cpu, &memcg_st->work);
-		}
-
-		if (obj_stock_should_drain(obj_st, root_memcg)) {
-			if (cpu == curcpu)
+			if (obj_stock_should_drain(obj_st, root_memcg))
 				drain_local_obj_stock(&obj_st->work);
-			else
-				schedule_drain_work(cpu, &obj_st->work);
+			continue;
 		}
+
+		/* Pairs with RCU barrier in housekeeping_update(). */
+		guard(rcu)();
+		if (cpu_is_isolated(cpu))
+			continue;
+
+		if (memcg_stock_should_drain(memcg_st, root_memcg))
+			queue_work_on(cpu, memcg_wq, &memcg_st->work);
+		if (obj_stock_should_drain(obj_st, root_memcg))
+			queue_work_on(cpu, memcg_wq, &obj_st->work);
 	}
 	migrate_enable();
 	mutex_unlock(&percpu_charge_mutex);