mm, memcg: add missing memory stall section in mem_cgroup_handle_over_high

When backporting commit 0e4b01df8659 ("mm, memcg: throttle allocators
when failing reclaim over memory.high"), memory stall section was
inadvertently missing. Fix this issue by adding it back.

Fixes: eda29cc0 ("mm, memcg: throttle allocators when failing reclaim over memory.high")
Signed-off-by: NCaspar Zhang <caspar@linux.alibaba.com>
Acked-by: NJoseph Qi <joseph.qi@linux.alibaba.com>

mm/memcontrol.c

@@ -65,6 +65,7 @@
 #include <linux/lockdep.h>
 #include <linux/file.h>
 #include <linux/tracehook.h>
+#include <linux/psi.h>
 #include "internal.h"
 #include <net/sock.h>
 #include <net/ip.h>
@@ -2235,6 +2236,7 @@ static void high_work_func(struct work_struct *work)
 void mem_cgroup_handle_over_high(void)
 {
 	unsigned long usage, high, clamped_high;
+	unsigned long pflags;
 	unsigned long penalty_jiffies, overage;
 	unsigned int nr_pages = current->memcg_nr_pages_over_high;
 	struct mem_cgroup *memcg;
@@ -2307,7 +2309,9 @@ void mem_cgroup_handle_over_high(void)
 	 * schedule_timeout_killable sets TASK_KILLABLE). This means we don't
 	 * need to account for any ill-begotten jiffies to pay them off later.
 	 */
+	psi_memstall_enter(&pflags);
 	schedule_timeout_killable(penalty_jiffies);
+	psi_memstall_leave(&pflags);
 
 out:
 	css_put(&memcg->css);