[PATCH 3/4] memcg: nmi-safe slab stats updates

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

The objcg based kmem [un]charging can be called in nmi context and it
may need to update NR_SLAB_[UN]RECLAIMABLE_B stats. So, let's correctly
handle the updates of these stats in the nmi context.

Signed-off-by: Shakeel Butt <shakeel.butt@linux.dev>
---
 mm/memcontrol.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index e91e4368650f..bba549c1f18c 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2591,8 +2591,18 @@ static inline void __mod_objcg_mlstate(struct obj_cgroup *objcg,
 
 	rcu_read_lock();
 	memcg = obj_cgroup_memcg(objcg);
-	lruvec = mem_cgroup_lruvec(memcg, pgdat);
-	__mod_memcg_lruvec_state(lruvec, idx, nr);
+	if (likely(!in_nmi())) {
+		lruvec = mem_cgroup_lruvec(memcg, pgdat);
+		__mod_memcg_lruvec_state(lruvec, idx, nr);
+	} else {
+		struct mem_cgroup_per_node *pn = memcg->nodeinfo[pgdat->node_id];
+
+		/* TODO: add to cgroup update tree once it is nmi-safe. */
+		if (idx == NR_SLAB_RECLAIMABLE_B)
+			atomic64_add(nr, &pn->slab_reclaimable);
+		else
+			atomic64_add(nr, &pn->slab_unreclaimable);
+	}
 	rcu_read_unlock();
 }
 
-- 
2.47.1