#include <linux/ctype.h>
 #include <linux/blk-cgroup.h>
 #include <linux/tracehook.h>
+#include <linux/psi.h>
 #include "blk.h"
 
 #define MAX_KEY_LEN 100
  */
 static void blkcg_maybe_throttle_blkg(struct blkcg_gq *blkg, bool use_memdelay)
 {
+       unsigned long pflags;
        u64 now = ktime_to_ns(ktime_get());
        u64 exp;
        u64 delay_nsec = 0;
         */
        delay_nsec = min_t(u64, delay_nsec, 250 * NSEC_PER_MSEC);
 
-       /*
-        * TODO: the use_memdelay flag is going to be for the upcoming psi stuff
-        * that hasn't landed upstream yet.  Once that stuff is in place we need
-        * to do a psi_memstall_enter/leave if memdelay is set.
-        */
+       if (use_memdelay)
+               psi_memstall_enter(&pflags);
 
        exp = ktime_add_ns(now, delay_nsec);
        tok = io_schedule_prepare();
                        break;
        } while (!fatal_signal_pending(current));
        io_schedule_finish(tok);
+
+       if (use_memdelay)
+               psi_memstall_leave(&pflags);
 }
 
 /**