1 files changed, 7 insertions, 6 deletions

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 2b552224f5cf..661a2c679f64 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -1100,12 +1100,9 @@ struct mem_cgroup *mem_cgroup_iter(struct mem_cgroup *root,
         if (prev && !reclaim)
                 id = css_id(&prev->css);
 
-        if (prev && prev != root)
-                css_put(&prev->css);
-
         if (!root->use_hierarchy && root != root_mem_cgroup) {
                 if (prev)
-                        return NULL;
+                        goto out_css_put;
                 return root;
         }
 
@@ -1121,7 +1118,7 @@ struct mem_cgroup *mem_cgroup_iter(struct mem_cgroup *root,
                         mz = mem_cgroup_zoneinfo(root, nid, zid);
                         iter = &mz->reclaim_iter[reclaim->priority];
                         if (prev && reclaim->generation != iter->generation)
-                                return NULL;
+                                goto out_css_put;
                         id = iter->position;
                 }
 
@@ -1143,8 +1140,12 @@ struct mem_cgroup *mem_cgroup_iter(struct mem_cgroup *root,
                 }
 
                 if (prev && !css)
-                        return NULL;
+                        goto out_css_put;
         }
+out_css_put:
+        if (prev && prev != root)
+                css_put(&prev->css);
+
         return memcg;
 }