1 files changed, 25 insertions, 27 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 408a5c75d77d..614d0d9c0deb 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -1116,7 +1116,6 @@ struct mem_cgroup *mem_cgroup_iter(struct mem_cgroup *root,
         rcu_read_lock();
         while (!memcg) {
                 struct mem_cgroup_reclaim_iter *uninitialized_var(iter);
-                struct cgroup_subsys_state *css = NULL;
 
                 if (reclaim) {
                         int nid = zone_to_nid(reclaim->zone);
@@ -1159,51 +1158,50 @@ struct mem_cgroup *mem_cgroup_iter(struct mem_cgroup *root,
                  * explicit visit.
                  */
                 if (!last_visited) {
-                        css = &root->css;
+                        memcg = root;
                 } else {
                         struct cgroup *prev_cgroup, *next_cgroup;
 
                         prev_cgroup = (last_visited == root) ? NULL
                                 : last_visited->css.cgroup;
-                        next_cgroup = cgroup_next_descendant_pre(prev_cgroup,
-                                        root->css.cgroup);
-                        if (next_cgroup)
-                                css = cgroup_subsys_state(next_cgroup,
-                                                mem_cgroup_subsys_id);
-                }
+skip_node:
+                        next_cgroup = cgroup_next_descendant_pre(
+                                        prev_cgroup, root->css.cgroup);
 
-                /*
-                 * Even if we found a group we have to make sure it is alive.
-                 * css && !memcg means that the groups should be skipped and
-                 * we should continue the tree walk.
-                 * last_visited css is safe to use because it is protected by
-                 * css_get and the tree walk is rcu safe.
-                 */
-                if (css == &root->css || (css && css_tryget(css)))
-                        memcg = mem_cgroup_from_css(css);
+                        /*
+                         * Even if we found a group we have to make sure it is
+                         * alive. css && !memcg means that the groups should be
+                         * skipped and we should continue the tree walk.
+                         * last_visited css is safe to use because it is
+                         * protected by css_get and the tree walk is rcu safe.
+                         */
+                        if (next_cgroup) {
+                                struct mem_cgroup *mem = mem_cgroup_from_cont(
+                                                next_cgroup);
+                                if (css_tryget(&mem->css))
+                                        memcg = mem;
+                                else {
+                                        prev_cgroup = next_cgroup;
+                                        goto skip_node;
+                                }
+                        }
+                }
 
                 if (reclaim) {
-                        struct mem_cgroup *curr = memcg;
-
                         if (last_visited)
                                 css_put(&last_visited->css);
 
-                        if (css && !memcg)
-                                curr = mem_cgroup_from_css(css);
-
-                        iter->last_visited = curr;
+                        iter->last_visited = memcg;
                         smp_wmb();
                         iter->last_dead_count = dead_count;
 
-                        if (!css)
+                        if (!memcg)
                                 iter->generation++;
                         else if (!prev && memcg)
                                 reclaim->generation = iter->generation;
-                } else if (css && !memcg) {
-                        last_visited = mem_cgroup_from_css(css);
                 }
 
-                if (prev && !css)
+                if (prev && !memcg)
                         goto out_unlock;
         }
 out_unlock: