enum compact_priority prio, enum compact_result *compact_result)
 {
        struct page *page;
-       unsigned int noreclaim_flag = current->flags & PF_MEMALLOC;
+       unsigned int noreclaim_flag;
 
        if (!order)
                return NULL;
 
-       current->flags |= PF_MEMALLOC;
+       noreclaim_flag = memalloc_noreclaim_save();
        *compact_result = try_to_compact_pages(gfp_mask, order, alloc_flags, ac,
                                                                        prio);
-       current->flags = (current->flags & ~PF_MEMALLOC) | noreclaim_flag;
+       memalloc_noreclaim_restore(noreclaim_flag);
 
        if (*compact_result <= COMPACT_INACTIVE)
                return NULL;
 {
        struct reclaim_state reclaim_state;
        int progress;
+       unsigned int noreclaim_flag;
 
        cond_resched();
 
        /* We now go into synchronous reclaim */
        cpuset_memory_pressure_bump();
-       current->flags |= PF_MEMALLOC;
+       noreclaim_flag = memalloc_noreclaim_save();
        lockdep_set_current_reclaim_state(gfp_mask);
        reclaim_state.reclaimed_slab = 0;
        current->reclaim_state = &reclaim_state;
 
        current->reclaim_state = NULL;
        lockdep_clear_current_reclaim_state();
-       current->flags &= ~PF_MEMALLOC;
+       memalloc_noreclaim_restore(noreclaim_flag);
 
        cond_resched();
 
 
        struct zonelist *zonelist;
        unsigned long nr_reclaimed;
        int nid;
+       unsigned int noreclaim_flag;
        struct scan_control sc = {
                .nr_to_reclaim = max(nr_pages, SWAP_CLUSTER_MAX),
                .gfp_mask = (current_gfp_context(gfp_mask) & GFP_RECLAIM_MASK) |
                                            sc.gfp_mask,
                                            sc.reclaim_idx);
 
-       current->flags |= PF_MEMALLOC;
+       noreclaim_flag = memalloc_noreclaim_save();
        nr_reclaimed = do_try_to_free_pages(zonelist, &sc);
-       current->flags &= ~PF_MEMALLOC;
+       memalloc_noreclaim_restore(noreclaim_flag);
 
        trace_mm_vmscan_memcg_reclaim_end(nr_reclaimed);
 
        struct zonelist *zonelist = node_zonelist(numa_node_id(), sc.gfp_mask);
        struct task_struct *p = current;
        unsigned long nr_reclaimed;
+       unsigned int noreclaim_flag;
 
-       p->flags |= PF_MEMALLOC;
+       noreclaim_flag = memalloc_noreclaim_save();
        lockdep_set_current_reclaim_state(sc.gfp_mask);
        reclaim_state.reclaimed_slab = 0;
        p->reclaim_state = &reclaim_state;
 
        p->reclaim_state = NULL;
        lockdep_clear_current_reclaim_state();
-       p->flags &= ~PF_MEMALLOC;
+       memalloc_noreclaim_restore(noreclaim_flag);
 
        return nr_reclaimed;
 }
        struct task_struct *p = current;
        struct reclaim_state reclaim_state;
        int classzone_idx = gfp_zone(gfp_mask);
+       unsigned int noreclaim_flag;
        struct scan_control sc = {
                .nr_to_reclaim = max(nr_pages, SWAP_CLUSTER_MAX),
                .gfp_mask = (gfp_mask = current_gfp_context(gfp_mask)),
         * and we also need to be able to write out pages for RECLAIM_WRITE
         * and RECLAIM_UNMAP.
         */
-       p->flags |= PF_MEMALLOC | PF_SWAPWRITE;
+       noreclaim_flag = memalloc_noreclaim_save();
+       p->flags |= PF_SWAPWRITE;
        lockdep_set_current_reclaim_state(gfp_mask);
        reclaim_state.reclaimed_slab = 0;
        p->reclaim_state = &reclaim_state;
        }
 
        p->reclaim_state = NULL;
-       current->flags &= ~(PF_MEMALLOC | PF_SWAPWRITE);
+       current->flags &= ~PF_SWAPWRITE;
+       memalloc_noreclaim_restore(noreclaim_flag);
        lockdep_clear_current_reclaim_state();
        return sc.nr_reclaimed >= nr_pages;
 }