1 files changed, 22 insertions, 23 deletions

diff --git a/mm/oom_kill.c b/mm/oom_kill.c
index 2c93502cfcb4..e41504aa5da9 100644
--- a/mm/oom_kill.c
+++ b/mm/oom_kill.c
@@ -176,7 +176,7 @@ static inline enum oom_constraint constrained_alloc(struct zonelist *zonelist,
 {
 #ifdef CONFIG_NUMA
         struct zone *zone;
-        struct zone **z;
+        struct zoneref *z;
         enum zone_type high_zoneidx = gfp_zone(gfp_mask);
         nodemask_t nodes = node_states[N_HIGH_MEMORY];
 
@@ -462,29 +462,29 @@ EXPORT_SYMBOL_GPL(unregister_oom_notifier);
  * if a parallel OOM killing is already taking place that includes a zone in
  * the zonelist.  Otherwise, locks all zones in the zonelist and returns 1.
  */
-int try_set_zone_oom(struct zonelist *zonelist)
+int try_set_zone_oom(struct zonelist *zonelist, gfp_t gfp_mask)
 {
-        struct zone **z;
+        struct zoneref *z;
+        struct zone *zone;
         int ret = 1;
 
-        z = zonelist->zones;
-
         spin_lock(&zone_scan_mutex);
-        do {
-                if (zone_is_oom_locked(*z)) {
+        for_each_zone_zonelist(zone, z, zonelist, gfp_zone(gfp_mask)) {
+                if (zone_is_oom_locked(zone)) {
                         ret = 0;
                         goto out;
                 }
-        } while (*(++z) != NULL);
+        }
+
+        for_each_zone_zonelist(zone, z, zonelist, gfp_zone(gfp_mask)) {
+                /*
+                 * Lock each zone in the zonelist under zone_scan_mutex so a
+                 * parallel invocation of try_set_zone_oom() doesn't succeed
+                 * when it shouldn't.
+                 */
+                zone_set_flag(zone, ZONE_OOM_LOCKED);
+        }
 
-        /*
-         * Lock each zone in the zonelist under zone_scan_mutex so a parallel
-         * invocation of try_set_zone_oom() doesn't succeed when it shouldn't.
-         */
-        z = zonelist->zones;
-        do {
-                zone_set_flag(*z, ZONE_OOM_LOCKED);
-        } while (*(++z) != NULL);
 out:
         spin_unlock(&zone_scan_mutex);
         return ret;
@@ -495,16 +495,15 @@ out:
  * allocation attempts with zonelists containing them may now recall the OOM
  * killer, if necessary.
  */
-void clear_zonelist_oom(struct zonelist *zonelist)
+void clear_zonelist_oom(struct zonelist *zonelist, gfp_t gfp_mask)
 {
-        struct zone **z;
-
-        z = zonelist->zones;
+        struct zoneref *z;
+        struct zone *zone;
 
         spin_lock(&zone_scan_mutex);
-        do {
-                zone_clear_flag(*z, ZONE_OOM_LOCKED);
-        } while (*(++z) != NULL);
+        for_each_zone_zonelist(zone, z, zonelist, gfp_zone(gfp_mask)) {
+                zone_clear_flag(zone, ZONE_OOM_LOCKED);
+        }
         spin_unlock(&zone_scan_mutex);
 }