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

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.

Link: https://lkml.kernel.org/r/20240803053306.2685541-1-yosryahmed@google.com
Signed-off-by: Yosry Ahmed <yosryahmed@google.com>
Acked-by: Johannes Weiner <hannes@cmpxchg.org>
Reviewed-by: Chengming Zhou <chengming.zhou@linux.dev>
Reviewed-by: Nhat Pham <nphamcs@gmail.com>
Cc: Takero Funaki <flintglass@gmail.com>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/mm/zswap.c b/mm/zswap.c
index d46caa42ed4f8752cf53ae9c86418b6cfa3c0350..71b75ff1f3fb7c770573a76831c9c7e9b21e1ffb 100644 (file)
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1348,24 +1348,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) {