[PATCH] mm: zswap: make the lock critical section obvious in shrink_worker()

[Yosry Ahmed <yosryahmed@google.com>]

Move the comments and spin_{lock/unlock}() calls around in
shrink_worker() to make it obvious the lock is protecting the loop
updating zswap_next_shrink.

Signed-off-by: Yosry Ahmed <yosryahmed@google.com>
---

This is intended to be squashed into "mm: zswap: fix global shrinker
memcg iteration".

---
 mm/zswap.c | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/mm/zswap.c b/mm/zswap.c
index babf0abbcc765..df620eacd1d11 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1364,24 +1364,22 @@ static void shrink_worker(struct work_struct *w)
 	 * until the next run of shrink_worker().
 	 */
 	do {
-		spin_lock(&zswap_shrink_lock);
-
 		/*
 		 * Start shrinking from the next memcg after zswap_next_shrink.
 		 * When the offline cleaner has already advanced the cursor,
 		 * advancing the cursor here overlooks one memcg, but this
 		 * should be negligibly rare.
+		 *
+		 * If we get an online memcg, keep the extra reference in case
+		 * the original one obtained by mem_cgroup_iter() is dropped by
+		 * zswap_memcg_offline_cleanup() while we are shrinking the
+		 * memcg.
 		 */
+		spin_lock(&zswap_shrink_lock);
 		do {
 			memcg = mem_cgroup_iter(NULL, zswap_next_shrink, NULL);
 			zswap_next_shrink = memcg;
 		} while (memcg && !mem_cgroup_tryget_online(memcg));
-		/*
-		 * Note that if we got an online memcg, we will keep the extra
-		 * reference in case the original reference obtained by mem_cgroup_iter
-		 * is dropped by the zswap memcg offlining callback, ensuring that the
-		 * memcg is not killed when we are reclaiming.
-		 */
 		spin_unlock(&zswap_shrink_lock);
 
 		if (!memcg) {
-- 
2.46.0.rc2.264.g509ed76dc8-goog