mm/zswap: Implement per-cgroup proactive writeback

From: Hao Jia <[email protected]>

Currently, shrink_memcg() writes back at most one entry per-node during its traversal. This makes shrink_worker() inefficient, as it must repeatedly re-enter shrink_memcg() to make any substantial progress.

To address this, extend shrink_memcg() and rewrite its LRU iteration logic to support batch writeback. Introduce the nr_to_writeback parameter to support a writeback budget based on compressed size. This enables batch writeback in the shrink_worker() path, while maintaining a low writeback budget in the zswap_store() path.

Additionally, to prepare for future proactive writeback, update the return value semantics of shrink_memcg(): a positive value now represents the actual number of compressed bytes written back, 0 indicates that candidates existed but no writeback succeeded, and a negative value represents an error code.

Suggested-by Yosry Ahmed <[email protected]>

Signed-off-by Hao Jia <[email protected]>

---
 mm/zswap.c | 116 ++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 97 insertions(+), 19 deletions(-)

mm/zswap.c +7 -1

index 761cd699e0a3..d7d031dee4cd 100644
@@ -160,6 +160,11 @@ struct zswap_pool {
 	char tfm_name[CRYPTO_MAX_ALG_NAME];
 };
 
+struct zswap_shrink_walk_arg {
+	unsigned long bytes_written;
+	bool encountered_page_in_swapcache;
+};
+
 /* Global LRU lists shared by all zswap pools. */
 static struct list_lru zswap_list_lru;
 

@@ -1089,8 +1094,9 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
 				       void *arg)
 {
 	struct zswap_entry *entry = container_of(item, struct zswap_entry, lru);
-	bool *encountered_page_in_swapcache = (bool *)arg;
+	struct zswap_shrink_walk_arg *walk_arg = arg;
 	swp_entry_t swpentry;
+	unsigned int length;
 	enum lru_status ret = LRU_REMOVED_RETRY;
 	int writeback_result;
 

mm/zswap.c +16 -8

 	 */
 	swpentry = entry->swpentry;
+	length = entry->length;
 
 	/*
 	 * It's safe to drop the lock here because we return either

@@ -1155,12 +1166,13 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
 		 * into the warmer region. We should terminate shrinking (if we're in the dynamic
 		 * shrinker context).
 		 */
-		if (writeback_result == -EEXIST && encountered_page_in_swapcache) {
+		if (writeback_result == -EEXIST) {
 			ret = LRU_STOP;
-			*encountered_page_in_swapcache = true;
+			walk_arg->encountered_page_in_swapcache = true;
 		}
 	} else {
 		zswap_written_back_pages++;
+		walk_arg->bytes_written += length;
 	}
 
 	return ret;

@@ -1169,8 +1181,11 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
 static unsigned long zswap_shrinker_scan(struct shrinker *shrinker,
 		struct shrink_control *sc)
 {
+	struct zswap_shrink_walk_arg walk_arg = {
+		.bytes_written = 0,
+		.encountered_page_in_swapcache = false,
+	};
 	unsigned long shrink_ret;
-	bool encountered_page_in_swapcache = false;
 
 	if (!zswap_shrinker_enabled ||
 			!mem_cgroup_zswap_writeback_enabled(sc->memcg)) {

@@ -1179,9 +1194,9 @@ static unsigned long zswap_shrinker_scan(struct shrinker *shrinker,
 	}
 
 	shrink_ret = list_lru_shrink_walk(&zswap_list_lru, sc, &shrink_memcg_cb,
-		&encountered_page_in_swapcache);
+		&walk_arg);
 
-	if (encountered_page_in_swapcache)
+	if (walk_arg.encountered_page_in_swapcache)
 		return SHRINK_STOP;
 
 	return shrink_ret ? shrink_ret : SHRINK_STOP;

@@ -1275,10 +1290,32 @@ static struct shrinker *zswap_alloc_shrinker(void)
 	return shrinker;
 }
 
-static int shrink_memcg(struct mem_cgroup *memcg)
-{
-	int nid, shrunk = 0, scanned = 0;
+/*
+ * The maximum acceptable scan cost factor for writing back
+ * PAGE_SIZE bytes of compressed data.
+ */
+#define ZSWAP_WB_SCAN_FACTOR	16UL
+#define NR_ZSWAP_WB_BATCH	64UL

mm/zswap.c +0 -1

 	if (!mem_cgroup_zswap_writeback_enabled(memcg))
 		return -ENOENT;
 

@@ -1290,24 +1327,65 @@ static int shrink_memcg(struct mem_cgroup *memcg)
 		return -ENOENT;
 
 	for_each_node_state(nid, N_NORMAL_MEMORY) {
-		unsigned long nr_to_walk = 1;

mm/zswap.c +14 -3

-		shrunk += list_lru_walk_one(&zswap_list_lru, nid, memcg,
-					    &shrink_memcg_cb, NULL, &nr_to_walk);
-		scanned += 1 - nr_to_walk;
+		while (nr_scanned < nr_to_scan) {
+			unsigned long nr_to_walk = min(NR_ZSWAP_WB_BATCH,
+						       nr_to_scan - nr_scanned);
+
+			/*
+			 * Account for the committed budget rather than the walker's
+			 * actual delta. If the list is emptied concurrently, the
+			 * walker visits nothing and nr_scanned would never advance.
+			 */
+			nr_scanned += nr_to_walk;
+
+			list_lru_walk_one(&zswap_list_lru, nid, memcg,
+					  &shrink_memcg_cb,
+					  &walk_arg,

mm/zswap.c +8 -6

 	}
 
-	if (!scanned)
+	/* Return -ENOENT if all zswap LRU lists are empty. */
+	if (memcg_list_is_empty)
 		return -ENOENT;
 
-	return shrunk ? 0 : -EAGAIN;
+	return walk_arg.bytes_written;
 }
 
 static void shrink_worker(struct work_struct *w)
 {
 	struct mem_cgroup *memcg;
-	int ret, failures = 0, attempts = 0;
+	int failures = 0, attempts = 0;
 	unsigned long thr;
+	long ret;
 
 	/* Reclaim down to the accept threshold */
 	thr = zswap_accept_thr_pages();

@@ -1368,7 +1446,7 @@ static void shrink_worker(struct work_struct *w)
 			goto resched;
 		}
 
-		ret = shrink_memcg(memcg);
+		ret = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
 		/* drop the extra reference */
 		mem_cgroup_put(memcg);
 

@@ -1382,7 +1460,7 @@ static void shrink_worker(struct work_struct *w)
 			continue;
 		++attempts;
 
-		if (ret && ++failures == MAX_RECLAIM_RETRIES)
+		if (ret <= 0 && ++failures == MAX_RECLAIM_RETRIES)
 			break;
 resched:
 		cond_resched();

@@ -1492,7 +1570,7 @@ bool zswap_store(struct folio *folio)
 	objcg = get_obj_cgroup_from_folio(folio);
 	if (objcg && !obj_cgroup_may_zswap(objcg)) {
 		memcg = get_mem_cgroup_from_objcg(objcg);
-		if (shrink_memcg(memcg)) {
+		if (shrink_memcg(memcg, 1) <= 0) {
 			mem_cgroup_put(memcg);
 			goto put_objcg;
 		}

Signature · 1 line

2.34.1

Quoted 29 lines · Hao Jia

From: Hao Jia <[email protected]>

mm/zswap.c +7 -1

index 761cd699e0a3..d7d031dee4cd 100644
@@ -160,6 +160,11 @@ struct zswap_pool {
  	char tfm_name[CRYPTO_MAX_ALG_NAME];
  };
 
+struct zswap_shrink_walk_arg {
+	unsigned long bytes_written;
+	bool encountered_page_in_swapcache;
+};
+
  /* Global LRU lists shared by all zswap pools. */
  static struct list_lru zswap_list_lru;
 

@@ -1089,8 +1094,9 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
  				       void *arg)
  {
  	struct zswap_entry *entry = container_of(item, struct zswap_entry, lru);
-	bool *encountered_page_in_swapcache = (bool *)arg;
+	struct zswap_shrink_walk_arg *walk_arg = arg;
  	swp_entry_t swpentry;
+	unsigned int length;
  	enum lru_status ret = LRU_REMOVED_RETRY;
  	int writeback_result;
 

mm/zswap.c static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o · line 1147

Quoted 8 lines

mm/zswap.c +4 -0

@@ -1135,8 +1141,13 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
  	 * Once the lru lock is dropped, the entry might get freed. The
  	 * swpentry is copied to the stack, and entry isn't deref'd again
  	 * until the entry is verified to still be alive in the tree.
+	 *
+	 * entry->length is also copied while the lock is held, because
+	 * zswap_writeback_entry() frees the entry on success and we still
+	 * need its compressed size to account for writeback.

Hmm that's unnecessary, just update "The swpentry is copied to the stack.." above to "Copy neded fields to the stack.." or something.

Quoted 61 lines

mm/zswap.c +16 -8

  	 */
  	swpentry = entry->swpentry;
+	length = entry->length;
 
  	/*
  	 * It's safe to drop the lock here because we return either

@@ -1155,12 +1166,13 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
  		 * into the warmer region. We should terminate shrinking (if we're in the dynamic
  		 * shrinker context).
  		 */
-		if (writeback_result == -EEXIST && encountered_page_in_swapcache) {
+		if (writeback_result == -EEXIST) {
  			ret = LRU_STOP;
-			*encountered_page_in_swapcache = true;
+			walk_arg->encountered_page_in_swapcache = true;
  		}
  	} else {
  		zswap_written_back_pages++;
+		walk_arg->bytes_written += length;
  	}
 
  	return ret;

@@ -1169,8 +1181,11 @@ static enum lru_status shrink_memcg_cb(struct list_head *item, struct list_lru_o
  static unsigned long zswap_shrinker_scan(struct shrinker *shrinker,
  		struct shrink_control *sc)
  {
+	struct zswap_shrink_walk_arg walk_arg = {
+		.bytes_written = 0,
+		.encountered_page_in_swapcache = false,
+	};
  	unsigned long shrink_ret;
-	bool encountered_page_in_swapcache = false;
 
  	if (!zswap_shrinker_enabled ||
  			!mem_cgroup_zswap_writeback_enabled(sc->memcg)) {

@@ -1179,9 +1194,9 @@ static unsigned long zswap_shrinker_scan(struct shrinker *shrinker,
  	}
 
  	shrink_ret = list_lru_shrink_walk(&zswap_list_lru, sc, &shrink_memcg_cb,
-		&encountered_page_in_swapcache);
+		&walk_arg);
 
-	if (encountered_page_in_swapcache)
+	if (walk_arg.encountered_page_in_swapcache)
  		return SHRINK_STOP;
 
  	return shrink_ret ? shrink_ret : SHRINK_STOP;

@@ -1275,10 +1290,32 @@ static struct shrinker *zswap_alloc_shrinker(void)
  	return shrinker;
  }
 
-static int shrink_memcg(struct mem_cgroup *memcg)
-{
-	int nid, shrunk = 0, scanned = 0;
+/*
+ * The maximum acceptable scan cost factor for writing back
+ * PAGE_SIZE bytes of compressed data.
+ */
+#define ZSWAP_WB_SCAN_FACTOR	16UL
+#define NR_ZSWAP_WB_BATCH	64UL

mm/zswap.c static struct shrinker *zswap_alloc_shrinker(void) · line 1308

Quoted 10 lines

mm/zswap.c +9 -0

 
+/*
+ * Iterate over the per-node zswap LRUs of @memcg in batches, writing back
+ * up to @nr_to_writeback * PAGE_SIZE bytes of compressed data.
+ *
+ * Return: The number of bytes written back, or -ENOENT if @memcg has
+ * writeback disabled, is a zombie cgroup, or has empty zswap LRUs.
+ */
+static long shrink_memcg(struct mem_cgroup *memcg,
+			 unsigned long nr_to_writeback)

Is nr_to_writeback supposed to be the number of pages we want to writeback (regardless of their compressed size), or the compressed bytes we want to writeback divided by PAGE_SIZE?

The way it's being used below seems like it's the latter, but the batch size should be in terms of scanned pages (i.e. uncompressed pages). So this is confusing.

The zswap_store() path expects to reclaim one uncompressed page, but this will reclaim PAGE_SIZE worth of compressed memory when passing 1 IIUC (actually maybe more, see below).

mm/zswap.c static struct shrinker *zswap_alloc_shrinker(void) · line 1318

Quoted 10 lines

mm/zswap.c +10 -0

+{
+	struct zswap_shrink_walk_arg walk_arg = {
+		.bytes_written = 0,
+		.encountered_page_in_swapcache = false,
+	};
+	u64 bytes_to_writeback = nr_to_writeback << PAGE_SHIFT;
+	bool memcg_list_is_empty = true;
+	int nid;
+
+	/* Memcg with zswap writeback disabled are not candidates. */

The comment is unnecessary here, it should be obvious.

Quoted 9 lines

mm/zswap.c +0 -1

  	if (!mem_cgroup_zswap_writeback_enabled(memcg))
  		return -ENOENT;
 

@@ -1290,24 +1327,65 @@ static int shrink_memcg(struct mem_cgroup *memcg)
  		return -ENOENT;
 
  	for_each_node_state(nid, N_NORMAL_MEMORY) {
-		unsigned long nr_to_walk = 1;

mm/zswap.c static int shrink_memcg(struct mem_cgroup *memcg) · line 1339

Quoted 10 lines

mm/zswap.c +10 -0

+		unsigned long nr_to_scan, nr_scanned = 0;
+		unsigned long remain;
+		walk_arg.encountered_page_in_swapcache = false;
+		/*
+		 * Cap by LRU length: bounds rewalks when referenced
+		 * entries keep rotating to the tail.
+		 */
+		nr_to_scan = list_lru_count_one(&zswap_list_lru, nid, memcg);
+		if (!nr_to_scan)
+			continue;

Hmm generally if we are running out of pages to scan then we should scan the rotated entries, and reclaim them on the second pass, right? So this should be working as intended. But I guess this doesn't work well when iterating multiple memcgs, as we don't want to drain referenced entries in one memcg before reclaiming already rotated entries on another.

So I think the assumption here is that the caller will retry if needed, handling balancing scanning between multiple memcgs if needed. Maybe we should document this in the function doc above? We should explain that referenced entries will be rotated but not reclaimed as part of the same call.

mm/zswap.c static int shrink_memcg(struct mem_cgroup *memcg) · line 1348

Quoted 9 lines

mm/zswap.c +9 -0

+		memcg_list_is_empty = false;
+
+		/*
+		 * Cap by SCAN_FACTOR * remain budget: bounds scan cost
+		 * to the remaining writeback budget.
+		 */
+		remain = DIV_ROUND_UP(bytes_to_writeback - walk_arg.bytes_written, PAGE_SIZE);
+		nr_to_scan = min(nr_to_scan,
+				 remain * ZSWAP_WB_SCAN_FACTOR);

For the zswap_store() path bytes_to_writeback=PAGE_SIZE, so remain will initially be 1. But then we multiply by this factor and now to scan 16 pages? Also, where did this factor and equation come from?

We'll also loop over nodes, so we may end up scanning 32 or more pages depending on the number of nodes in the system.

If this is just a heuristic, we should really just start simple and add heuristics later as needed. The caller should probably pass in the number of pages to scan (i.e. uncompressed pages), and leave it to the caller to decide when to retry if the actual memory savings are realized.

Quoted 18 lines

mm/zswap.c +14 -3

-		shrunk += list_lru_walk_one(&zswap_list_lru, nid, memcg,
-					    &shrink_memcg_cb, NULL, &nr_to_walk);
-		scanned += 1 - nr_to_walk;
+		while (nr_scanned < nr_to_scan) {
+			unsigned long nr_to_walk = min(NR_ZSWAP_WB_BATCH,
+						       nr_to_scan - nr_scanned);
+
+			/*
+			 * Account for the committed budget rather than the walker's
+			 * actual delta. If the list is emptied concurrently, the
+			 * walker visits nothing and nr_scanned would never advance.
+			 */
+			nr_scanned += nr_to_walk;
+
+			list_lru_walk_one(&zswap_list_lru, nid, memcg,
+					  &shrink_memcg_cb,
+					  &walk_arg,

mm/zswap.c static int shrink_memcg(struct mem_cgroup *memcg) · line 1373

Quoted 10 lines

mm/zswap.c +10 -0

+					  &nr_to_walk);
+
+			if (walk_arg.bytes_written >= bytes_to_writeback)
+				return walk_arg.bytes_written;
+
+			if (walk_arg.encountered_page_in_swapcache)
+				break;
+
+			cond_resched();
+		}

If the caller is expected to have a retry loop anyway, should we simplify this and just scan each per-node LRU once?

We should also probably bail early if the number of scanned pages has already been reached? Currently shrink_memcg() scans one page at a time, so if it scans a bit more to balance between the nodes it's probably fine.

But with batching, we could end up scanning hundres of extra pages just to balance between all nodes. Is node imbalance a real issue?

My initial thought was that if cold memory is evenly distributed across nodes and we are doing a large writeback, it would be better to balance the zswap entry writeback across all nodes rather than just draining node 0 first. However, since we currently lack a proper metric to represent hot/cold memory (such as age-based tracking), doing this probably doesn't make much sense right now.

So, perhaps we want something like this? Please correct me if I'm wrong.

static long shrink_memcg(struct mem_cgroup *memcg,
unsigned long nr_to_scan)

{
   struct zswap_shrink_walk_arg walk_arg = {
     .bytes_written = 0,
     .encountered_page_in_swapcache = false,
   };
   unsigned long nr_remaining = nr_to_scan;
   bool memcg_list_is_empty = true;
   int nid;

   if (!mem_cgroup_zswap_writeback_enabled(memcg))
     return -ENOENT;

   if (memcg && !mem_cgroup_online(memcg))
     return -ENOENT;

   for_each_node_state(nid, N_NORMAL_MEMORY) {

     unsigned long nr_to_walk;

     /*
      * Cap the per-node scan by the current LRU length. A referenced
      * entry is only rotated to the tail (second chance) and may be
      * revisited within a single walk; without this cap those rotated
      * entries could drain the shared scan budget on one node.
      */

     nr_to_walk = min(nr_remaining,
          list_lru_count_one(&zswap_list_lru, nid, memcg));
     if (!nr_to_walk)
       continue;
     memcg_list_is_empty = false;

     nr_remaining -= nr_to_walk;
     list_lru_walk_one(&zswap_list_lru, nid, memcg,
           &shrink_memcg_cb, &walk_arg, &nr_to_walk);
     /* Return the unused share of the budget to the pool. */
     nr_remaining += nr_to_walk;

     /* Bail out once the whole scan budget has been spent. */

     if (!nr_remaining)
       break;

     cond_resched();

   }

if (memcg_list_is_empty)

     return -ENOENT;

   return walk_arg.bytes_written;
}

Thanks,
Hao

Quoted 6 lines

My initial thought was that if cold memory is evenly distributed across nodes and we are doing a large writeback, it would be better to balance the zswap entry writeback across all nodes rather than just draining node 0 first. However, since we currently lack a proper metric to represent hot/cold memory (such as age-based tracking), doing this probably doesn't make much sense right now.

Yeah let's start simple and go from there.

Quoted 21 lines

So, perhaps we want something like this? Please correct me if I'm wrong.

static long shrink_memcg(struct mem_cgroup *memcg,
unsigned long nr_to_scan)

{
    struct zswap_shrink_walk_arg walk_arg = {
      .bytes_written = 0,
      .encountered_page_in_swapcache = false,
    };
    unsigned long nr_remaining = nr_to_scan;
    bool memcg_list_is_empty = true;
    int nid;

    if (!mem_cgroup_zswap_writeback_enabled(memcg))
      return -ENOENT;

    if (memcg && !mem_cgroup_online(memcg))
      return -ENOENT;

    for_each_node_state(nid, N_NORMAL_MEMORY) {

Quoted 8 lines

      unsigned long nr_to_walk;

      /*
       * Cap the per-node scan by the current LRU length. A referenced
       * entry is only rotated to the tail (second chance) and may be
       * revisited within a single walk; without this cap those rotated
       * entries could drain the shared scan budget on one node.
       */

The comment here is a bit misleading. It's not just about draining one node. One call to shrink_memcg() should only scan entries once. The caller can then choose to scan the memcg again, or scan a different one. In this case, the caller should iterate all memcgs first before retrying memcgs again and reclaiming rotated entries.

Quoted 13 lines

      nr_to_walk = min(nr_remaining,
           list_lru_count_one(&zswap_list_lru, nid, memcg));
      if (!nr_to_walk)
        continue;
      memcg_list_is_empty = false;

      nr_remaining -= nr_to_walk;
      list_lru_walk_one(&zswap_list_lru, nid, memcg,
            &shrink_memcg_cb, &walk_arg, &nr_to_walk);
      /* Return the unused share of the budget to the pool. */
      nr_remaining += nr_to_walk;

Quoted 1 line

      /* Bail out once the whole scan budget has been spent. */

The comment is unnecessary.

Quoted 5 lines

      if (!nr_remaining)
        break;

      cond_resched();

Did you observe a problem here or did you just add this due to an abundance of caution?

Quoted 4 lines

    }

    if (memcg_list_is_empty)

Do we need memcg_list_is_empty? Can we just check if nr_remaining matches nr_to_scan?

Quoted 4 lines

      return -ENOENT;

    return walk_arg.bytes_written;
}

/*
* Scan up to @nr_to_scan pages across the per-node zswap LRUs of @memcg
* and write back the reclaimable ones.
*
* Since the second-chance algorithm rotates referenced entries to the
* LRU tail, the per-node scan is capped at the current LRU length so
* each entry is scanned at most once per call. It is up to the caller
* to handle retries, deciding whether to scan the next memcg to complete

* the full iteration, or to rescan the current memcg to drain its zswap
* entries.
*
* Return: The number of compressed bytes written back (>= 0), or -ENOENT
* if @memcg has writeback disabled, is a zombie cgroup, or has empty
* zswap LRUs.
*/
static long shrink_memcg(struct mem_cgroup *memcg, unsigned long nr_to_scan)

{
     struct zswap_shrink_walk_arg walk_arg = {
         .bytes_written = 0,
         .encountered_page_in_swapcache = false,
     };
     unsigned long nr_remaining = nr_to_scan;
     int nid;

     if (!mem_cgroup_zswap_writeback_enabled(memcg))
         return -ENOENT;

     /*
      * Skip zombies because their LRUs are reparented and we would be
      * reclaiming from the parent instead of the dead memcg.
      */
     if (memcg && !mem_cgroup_online(memcg))
         return -ENOENT;

     for_each_node_state(nid, N_NORMAL_MEMORY) {
         unsigned long nr_to_walk;

         /*
          * Cap the walk at the current LRU length to ensure each entry is
          * scanned at most once per call. Referenced entries are rotated
          * to the tail for a second chance, and this bound prevents them
          * from being revisited within a single call. Retries are left to
          * the caller, which can choose to rescan the current memcg or
          * move on to the next one.
          */

         nr_to_walk = min(nr_remaining,
                  list_lru_count_one(&zswap_list_lru, nid, memcg));
         if (!nr_to_walk)
             continue;

         nr_remaining -= nr_to_walk;
         list_lru_walk_one(&zswap_list_lru, nid, memcg, &shrink_memcg_cb,
                   &walk_arg, &nr_to_walk);
         /* Return the unused share of the budget to the pool. */
         nr_remaining += nr_to_walk;

         if (!nr_remaining)
             break;
     }

     /* Nothing was scanned: every LRU under @memcg was empty. */
     if (nr_remaining == nr_to_scan)
         return -ENOENT;

     return walk_arg.bytes_written;
}

Thanks,
Hao

Quoted 9 lines

/*
* Scan up to @nr_to_scan pages across the per-node zswap LRUs of @memcg
* and write back the reclaimable ones.
*
* Since the second-chance algorithm rotates referenced entries to the
* LRU tail, the per-node scan is capped at the current LRU length so
* each entry is scanned at most once per call. It is up to the caller
* to handle retries, deciding whether to scan the next memcg to complete

Nit: "whether to scan another memcg to complete.."

Quoted 30 lines

* the full iteration, or to rescan the current memcg to drain its zswap
* entries.
*
* Return: The number of compressed bytes written back (>= 0), or -ENOENT
* if @memcg has writeback disabled, is a zombie cgroup, or has empty
* zswap LRUs.
*/
static long shrink_memcg(struct mem_cgroup *memcg, unsigned long nr_to_scan)

{
      struct zswap_shrink_walk_arg walk_arg = {
          .bytes_written = 0,
          .encountered_page_in_swapcache = false,
      };
      unsigned long nr_remaining = nr_to_scan;
      int nid;

      if (!mem_cgroup_zswap_writeback_enabled(memcg))
          return -ENOENT;

      /*
       * Skip zombies because their LRUs are reparented and we would be
       * reclaiming from the parent instead of the dead memcg.
       */
      if (memcg && !mem_cgroup_online(memcg))
          return -ENOENT;

      for_each_node_state(nid, N_NORMAL_MEMORY) {
          unsigned long nr_to_walk;

Quoted 8 lines

          /*
           * Cap the walk at the current LRU length to ensure each entry is
           * scanned at most once per call. Referenced entries are rotated
           * to the tail for a second chance, and this bound prevents them
           * from being revisited within a single call. Retries are left to
           * the caller, which can choose to rescan the current memcg or
           * move on to the next one.
           */

Nit: Make this more concise since it's already explained above.

Otherwise this looks good to me, thank you!

Quoted 22 lines

          nr_to_walk = min(nr_remaining,
                   list_lru_count_one(&zswap_list_lru, nid, memcg));
          if (!nr_to_walk)
              continue;

          nr_remaining -= nr_to_walk;
          list_lru_walk_one(&zswap_list_lru, nid, memcg, &shrink_memcg_cb,
                    &walk_arg, &nr_to_walk);
          /* Return the unused share of the budget to the pool. */
          nr_remaining += nr_to_walk;

          if (!nr_remaining)
              break;
      }

      /* Nothing was scanned: every LRU under @memcg was empty. */
      if (nr_remaining == nr_to_scan)
          return -ENOENT;

      return walk_arg.bytes_written;
}

+/*
+ * Walk the memcg tree and write back zswap pages until the
+ * (lower_pages, upper_pages) window closes, or abort encounter
+ * MAX_RECLAIM_RETRIES times of the following conditions:
+ * - No writeback-candidate memcgs found in a memcg tree walk.
+ * - Shrinking a writeback-candidate memcg failed.
+ *
+ * For shrink_worker(), it passes lower=thr and upper=zswap_total_pages().
+ * The @upper limit is refreshed in each iteration by re-evaluating
+ * zswap_total_pages(), and the window closes once the total falls
+ * below the threshold.

+ */

mm/zswap.c +33 -9

+static void zswap_try_to_writeback(unsigned long lower_pages,
+				   unsigned long upper_pages)
+{
+	int failures = 0, attempts = 0;
+	struct mem_cgroup *iter_memcg;
+
+	while (lower_pages < upper_pages) {
+		unsigned long batch_size;
+		long shrunk;
 
-		if (!memcg) {
+		cond_resched();
+
+		iter_memcg = zswap_iter_global();
+		if (!iter_memcg) {
 			/*
 			 * Continue shrinking without incrementing failures if
 			 * we found candidate memcgs in the last tree walk.

@@ -1443,12 +1457,16 @@ static void shrink_worker(struct work_struct *w)
 				break;
 
 			attempts = 0;
-			goto resched;
+			continue;
 		}
 
-		ret = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
+		batch_size = min(upper_pages - lower_pages, NR_ZSWAP_WB_BATCH);
+		shrunk = shrink_memcg(iter_memcg, batch_size);
 		/* drop the extra reference */
-		mem_cgroup_put(memcg);
+		mem_cgroup_put(iter_memcg);
+
+		/* zswap total pages might have changed, refresh it. */
+		upper_pages = zswap_total_pages();
 
 		/*
 		 * There are no writeback-candidate pages in the memcg.

@@ -1456,15 +1474,23 @@ static void shrink_worker(struct work_struct *w)
 		 * with pages in zswap. Skip this without incrementing attempts
 		 * and failures.
 		 */
-		if (ret == -ENOENT)
+		if (shrunk == -ENOENT)
 			continue;
 		++attempts;
 
-		if (ret <= 0 && ++failures == MAX_RECLAIM_RETRIES)
+		if (shrunk <= 0 && ++failures == MAX_RECLAIM_RETRIES)
 			break;
-resched:
-		cond_resched();
-	} while (zswap_total_pages() > thr);
+	}
+}
+
+static void shrink_worker(struct work_struct *w)
+{
+	unsigned long thr;
+
+	/* Reclaim down to the accept threshold */
+	thr = zswap_accept_thr_pages();
+
+	zswap_try_to_writeback(thr, zswap_total_pages());
 }
 
 /*********************************

Signature · 1 line

2.34.1

Quoted 11 lines

+/*
+ * Walk the memcg tree and write back zswap pages until the
+ * (lower_pages, upper_pages) window closes, or abort encounter
+ * MAX_RECLAIM_RETRIES times of the following conditions:
+ * - No writeback-candidate memcgs found in a memcg tree walk.
+ * - Shrinking a writeback-candidate memcg failed.
+ *
+ * For shrink_worker(), it passes lower=thr and upper=zswap_total_pages().
+ * The @upper limit is refreshed in each iteration by re-evaluating
+ * zswap_total_pages(), and the window closes once the total falls
+ * below the threshold.

This is the wrong abstraction level, and it's obvious by the fact that the function calls zswap_total_pages() again to recalcualte 'upper_pages'. It gets much worse in the next patch as well.

The lower_pages and upper_pages thing is also unnecessarily hard to follow.

The core of the reuse here is the retry logic. So maybe keep the memcg iteration in the callers, and define a function that takes in one memcg and reclaims one batch from it? failures and attempts can be passed into the function to maintain the state across scans of different memcgs, like zswap_shrink_walk_arg?

WDYT?

Quoted 19 lines · Yosry Ahmed

Quoted 11 lines

+/*
+ * Walk the memcg tree and write back zswap pages until the
+ * (lower_pages, upper_pages) window closes, or abort encounter
+ * MAX_RECLAIM_RETRIES times of the following conditions:
+ * - No writeback-candidate memcgs found in a memcg tree walk.
+ * - Shrinking a writeback-candidate memcg failed.
+ *
+ * For shrink_worker(), it passes lower=thr and upper=zswap_total_pages().
+ * The @upper limit is refreshed in each iteration by re-evaluating
+ * zswap_total_pages(), and the window closes once the total falls
+ * below the threshold.

This is the wrong abstraction level, and it's obvious by the fact that the function calls zswap_total_pages() again to recalcualte 'upper_pages'. It gets much worse in the next patch as well.

The lower_pages and upper_pages thing is also unnecessarily hard to follow.

Quoted 7 lines

The core of the reuse here is the retry logic. So maybe keep the memcg iteration in the callers, and define a function that takes in one memcg and reclaims one batch from it? failures and attempts can be passed into the function to maintain the state across scans of different memcgs, like zswap_shrink_walk_arg?

WDYT?

Perhaps something like this?

struct zswap_shrink_state {
     int attempts;
     int failures;
     bool stop;
};

static bool zswap_shrink_no_candidate(struct zswap_shrink_state *s)

{
     if (!s->attempts && ++s->failures == MAX_RECLAIM_RETRIES)
         return true;

     s->attempts = 0;
     return false;
}

static long zswap_shrink_one(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
     long shrunk;

     shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
     if (shrunk == -ENOENT)
         return 0;

     s->attempts++;
     if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
         s->stop = true;

     return shrunk;
}

static void shrink_worker(struct work_struct *w)

{
     struct zswap_shrink_state s = {};
     unsigned long thr;

     /* Reclaim down to the accept threshold */
     thr = zswap_accept_thr_pages();

     while (zswap_total_pages() > thr) {
         struct mem_cgroup *memcg;

         cond_resched();

         memcg = zswap_iter_global();
         if (!memcg) {
             if (zswap_shrink_no_candidate(&s))
                 break;
             continue;
         }

         zswap_shrink_one(memcg, &s);
         /* Drop the extra reference taken by the iterator. */
         mem_cgroup_put(memcg);
         if (s.stop)
             break;
     }
}

We could also fold the logic of zswap_shrink_no_candidate() into zswap_shrink_one(), but adding a !memcg check inside zswap_shrink_one() feels a bit awkward.

WDYT?

Quoted 33 lines · Hao Jia

On 2026/6/23 07:36, Yosry Ahmed wrote:

Perhaps something like this?

struct zswap_shrink_state {
      int attempts;
      int failures;
      bool stop;
};

static bool zswap_shrink_no_candidate(struct zswap_shrink_state *s)

{
      if (!s->attempts && ++s->failures == MAX_RECLAIM_RETRIES)
          return true;

      s->attempts = 0;
      return false;
}

static long zswap_shrink_one(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
      long shrunk;

      shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
      if (shrunk == -ENOENT)
          return 0;

Quoted 3 lines

      s->attempts++;
      if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
          s->stop = true;

Do we need 'stop' or can we just return a value here to indicate that we should stop (e.g. -EBUSY)?

Quoted 24 lines

      return shrunk;
}

static void shrink_worker(struct work_struct *w)

{
      struct zswap_shrink_state s = {};
      unsigned long thr;

      /* Reclaim down to the accept threshold */
      thr = zswap_accept_thr_pages();

      while (zswap_total_pages() > thr) {
          struct mem_cgroup *memcg;

          cond_resched();

          memcg = zswap_iter_global();
          if (!memcg) {
              if (zswap_shrink_no_candidate(&s))
                  break;
              continue;
          }

Quoted 7 lines

          zswap_shrink_one(memcg, &s);
          /* Drop the extra reference taken by the iterator. */
          mem_cgroup_put(memcg);
          if (s.stop)
              break;
      }
}

I think splitting the shrink/retry logic over 2 functions makes it more difficult to follow, so yeah I think fold zswap_shrink_no_candidate() into zswap_shrink_one(). Then the callers only need to iterate memcgs (depending on the context) and call zswap_shrink_one() for each of them.

Quoted 13 lines

Quoted 12 lines

static long zswap_shrink_one(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
      long shrunk;

      shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
      if (shrunk == -ENOENT)
          return 0;

      s->attempts++;
      if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
          s->stop = true;

Quoted 3 lines

Do we need 'stop' or can we just return a value here to indicate that we should stop (e.g. -EBUSY)?

Perhaps we could return -EAGAIN instead of -EBUSY? This would align with the semantics of the memory.reclaim interface, which returns -EAGAIN when it reclaims fewer bytes than requested.

Quoted 6 lines

I think splitting the shrink/retry logic over 2 functions makes it more difficult to follow, so yeah I think fold zswap_shrink_no_candidate() into zswap_shrink_one(). Then the callers only need to iterate memcgs (depending on the context) and call zswap_shrink_one() for each of them.

So, something like this?

/* Track progress of a memcg-tree writeback walk. */

struct zswap_shrink_state {
     int attempts;

     int failures;
};

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EAGAIN to signal the caller to abort the walk
* after encountering the following conditions MAX_RECLAIM_RETRIES times:
* - No writeback-candidate memcgs were found in a memcg tree walk.
* - Shrinking a writeback-candidate memcg failed.

*
* Return: The number of compressed bytes written back (>= 0), or -EAGAIN
* once the retry budget is exhausted and the caller should abort the walk. */
static long zswap_shrink_one(struct mem_cgroup *memcg,

struct zswap_shrink_state *s)

{
     long shrunk;

     /*
      * If the iterator has completed a full pass, update the shrink state
      * and check whether we should keep going.
      */
     if (!memcg) {
         /*
          * Continue shrinking without incrementing failures if we found
          * candidate memcgs in the last tree walk.
          */
         if (!s->attempts && ++s->failures == MAX_RECLAIM_RETRIES)
             return -EAGAIN;
         s->attempts = 0;
         return 0;
     }

     shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);

     /*
      * There are no writeback-candidate pages in the memcg. This is not an
      * issue as long as we can find another memcg with pages in zswap. Skip
      * this without incrementing attempts and failures.
      */
     if (shrunk == -ENOENT)
         return 0;
     s->attempts++;

     if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
         return -EAGAIN;

     return shrunk;
}

static void shrink_worker(struct work_struct *w)

{
     struct zswap_shrink_state s = {};
     unsigned long thr;

     /* Reclaim down to the accept threshold */
     thr = zswap_accept_thr_pages();

     while (zswap_total_pages() > thr) {
         struct mem_cgroup *memcg;
         long ret;

         cond_resched();

         memcg = zswap_iter_global();

         ret = zswap_shrink_one(memcg, &s);
         /* drop the extra reference taken by zswap_iter_global() */
         mem_cgroup_put(memcg);
         if (ret == -EAGAIN)
             break;
     }
}

Quoted 17 lines

Quoted 16 lines

Quoted 12 lines

static long zswap_shrink_one(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
       long shrunk;

       shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);
       if (shrunk == -ENOENT)
           return 0;

       s->attempts++;
       if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
           s->stop = true;

Do we need 'stop' or can we just return a value here to indicate that we should stop (e.g. -EBUSY)?

Quoted 3 lines

Perhaps we could return -EAGAIN instead of -EBUSY? This would align with the semantics of the memory.reclaim interface, which returns -EAGAIN when it reclaims fewer bytes than requested.

Hmm but -EAGAIN tells the caller to try again, while here -EAGAIN tells the caller *not* to try again because we exhausted all retries?

Quoted 8 lines

Quoted 6 lines

I think splitting the shrink/retry logic over 2 functions makes it more difficult to follow, so yeah I think fold zswap_shrink_no_candidate() into zswap_shrink_one(). Then the callers only need to iterate memcgs (depending on the context) and call zswap_shrink_one() for each of them.

So, something like this?

Yeah, something like this :)

Quoted 3 lines

/* Track progress of a memcg-tree writeback walk. */

struct zswap_shrink_state {
      int attempts;

While at it, I think "attempts" is really the number of scans, right? Should we rename it? Maybe "scans" or similar?

Quoted 13 lines

      int failures;
};

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EAGAIN to signal the caller to abort the walk
* after encountering the following conditions MAX_RECLAIM_RETRIES times:
* - No writeback-candidate memcgs were found in a memcg tree walk.
* - Shrinking a writeback-candidate memcg failed.

Orthogonal to this patch, but I wonder if this can be simplified. I wonder if these two conditions can be replaced with "shrinking a memcg that has zswap entries failed". The "no writeback-candidate memcgs in the tree" case seems like we should abort right away instead of retrying?

Nhat, WDYT?

Quoted 5 lines

*
* Return: The number of compressed bytes written back (>= 0), or -EAGAIN
* once the retry budget is exhausted and the caller should abort the walk. */
static long zswap_shrink_one(struct mem_cgroup *memcg,

Nit: zswap_shrink_one_memcg()

BTW, the existing writeback logic has been broken for a while now when memcg is disabled. I think we constantly hit the !memcg case and run out of retries. Not sure if your patch changes this in any way, or if you want to fix that while you're at it :)

Quoted 49 lines

struct zswap_shrink_state *s)

{
      long shrunk;

      /*
       * If the iterator has completed a full pass, update the shrink state
       * and check whether we should keep going.
       */
      if (!memcg) {
          /*
           * Continue shrinking without incrementing failures if we found
           * candidate memcgs in the last tree walk.
           */
          if (!s->attempts && ++s->failures == MAX_RECLAIM_RETRIES)
              return -EAGAIN;
          s->attempts = 0;
          return 0;
      }

      shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);

      /*
       * There are no writeback-candidate pages in the memcg. This is not an
       * issue as long as we can find another memcg with pages in zswap. Skip
       * this without incrementing attempts and failures.
       */
      if (shrunk == -ENOENT)
          return 0;
      s->attempts++;

      if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
          return -EAGAIN;

      return shrunk;
}

static void shrink_worker(struct work_struct *w)

{
      struct zswap_shrink_state s = {};
      unsigned long thr;

      /* Reclaim down to the accept threshold */
      thr = zswap_accept_thr_pages();

      while (zswap_total_pages() > thr) {
          struct mem_cgroup *memcg;
          long ret;

Quoted 3 lines

          cond_resched();

          memcg = zswap_iter_global();

Do we still need this helper? Or should we just keep the memcg iteration open-coded?

Quoted 8 lines

          ret = zswap_shrink_one(memcg, &s);
          /* drop the extra reference taken by zswap_iter_global() */
          mem_cgroup_put(memcg);
          if (ret == -EAGAIN)
              break;
      }
}

Quoted 17 lines

Quoted 10 lines

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EAGAIN to signal the caller to abort the walk
* after encountering the following conditions MAX_RECLAIM_RETRIES times:
* - No writeback-candidate memcgs were found in a memcg tree walk.
* - Shrinking a writeback-candidate memcg failed.

Orthogonal to this patch, but I wonder if this can be simplified. I wonder if these two conditions can be replaced with "shrinking a memcg that has zswap entries failed". The "no writeback-candidate memcgs in the tree" case seems like we should abort right away instead of retrying?

Quoted 2 lines

Nhat, WDYT?

Perhaps something like the following is what you had in mind? I've drafted the implementation below to make it easier for Nhat to compare with the previous behavior.

Quoted 10 lines

Quoted 7 lines

*
* Return: The number of compressed bytes written back (>= 0), or -EAGAIN
* once the retry budget is exhausted and the caller should abort the walk. */
static long zswap_shrink_one(struct mem_cgroup *memcg,

Nit: zswap_shrink_one_memcg()

Quoted 4 lines

BTW, the existing writeback logic has been broken for a while now when memcg is disabled. I think we constantly hit the !memcg case and run out of retries. Not sure if your patch changes this in any way, or if you want to fix that while you're at it :)

Yes, I'd be happy to do that. However, would it be better to submit a separate fix patch or combine it with this one?

/* Track progress of a memcg-tree writeback walk. */

struct zswap_shrink_state {
     int scans;
     int failures;
};

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EBUSY to signal the caller to abort the walk when
* either of the following occurs:
* - A full pass over the tree found no writeback-candidate memcg.
* - Shrinking a writeback-candidate memcg failed MAX_RECLAIM_RETRIES
times.
*
* When memory cgroup is disabled, the iterator always yields NULL. All
* zswap entries then live on the root list_lru, so NULL is treated as the
* root memcg and shrunk directly rather than as a completed tree pass.

*
* Return: The number of compressed bytes written back (>= 0), or -EBUSY
* when the caller should abort the walk.
*/
static long zswap_shrink_one_memcg(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
     bool disabled = mem_cgroup_disabled();

     long shrunk;

     /*
      * If the iterator has completed a full pass, update the shrink state
      * and check whether we should keep going.
      * With memcg disabled the iterator always yields NULL, so fall through
      * and shrink the root memcg directly instead.
      */
     if (!memcg && !disabled) {
         /*
          * Abort if no writeback-candidate memcgs in the last tree walk.
          * Otherwise reset the scans count and continue.
          */
         if (!s->scans)
             return -EBUSY;
         s->scans = 0;
         return 0;
     }

     shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);

     /*
      * There are no writeback-candidate pages in the memcg. With memcg
      * enabled this is not an issue as long as we can find another memcg
      * with pages in zswap, so skip without counting it as a candidate.
      * With memcg disabled the root LRU is the only target, so we should
      * abort if it has no writeback-candidate pages.
      */
     if (shrunk == -ENOENT)
         return disabled ? -EBUSY : 0;
     s->scans++;

     if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
         return -EBUSY;

     return shrunk;
}

static void shrink_worker(struct work_struct *w)

{
     struct zswap_shrink_state s = {};
     unsigned long thr;

     /* Reclaim down to the accept threshold */
     thr = zswap_accept_thr_pages();

     /*
      * Global reclaim will select cgroup in a round-robin fashion from all
      * online memcgs, but memcgs that have no pages in zswap and
      * writeback-disabled memcgs (memory.zswap.writeback=0) are not
      * candidates for shrinking.
      *
      * We save iteration cursor memcg into zswap_next_shrink,
      * which can be modified by the offline memcg cleaner
      * zswap_memcg_offline_cleanup().
      *
      * Since the offline cleaner is called only once, we cannot leave an
      * offline memcg reference in zswap_next_shrink.
      * We can rely on the cleaner only if we get online memcg under lock.
      *
      * If we get an offline memcg, we cannot determine if the cleaner has
      * already been called or will be called later. We must put back the
      * reference before returning from this function. Otherwise, the
      * offline memcg left in zswap_next_shrink will hold the reference
      * until the next run of shrink_worker().
      */
     while (zswap_total_pages() > thr) {
         struct mem_cgroup *memcg;
         long ret;

         cond_resched();
         /*
          * 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));
         spin_unlock(&zswap_shrink_lock);

         ret = zswap_shrink_one_memcg(memcg, &s);
         /* drop the extra reference taken above */
         mem_cgroup_put(memcg);
         if (ret == -EBUSY)
             break;
     }
}

Thanks,
Hao

Quoted 20 lines

Quoted 19 lines

Quoted 10 lines

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EAGAIN to signal the caller to abort the walk
* after encountering the following conditions MAX_RECLAIM_RETRIES times:
* - No writeback-candidate memcgs were found in a memcg tree walk.
* - Shrinking a writeback-candidate memcg failed.

Orthogonal to this patch, but I wonder if this can be simplified. I wonder if these two conditions can be replaced with "shrinking a memcg that has zswap entries failed". The "no writeback-candidate memcgs in the tree" case seems like we should abort right away instead of retrying?

Nhat, WDYT?

Quoted 3 lines

Perhaps something like the following is what you had in mind? I've drafted the implementation below to make it easier for Nhat to compare with the previous behavior.

Hmm I think if we pursue this it should be in a separate patch or even outside of this series, ideally with numbers/proof that it's not introducing regressions to the scenario that lead to its introduction.

Quoted 15 lines

Quoted 14 lines

Quoted 7 lines

*
* Return: The number of compressed bytes written back (>= 0), or -EAGAIN
* once the retry budget is exhausted and the caller should abort the walk. */
static long zswap_shrink_one(struct mem_cgroup *memcg,

Nit: zswap_shrink_one_memcg()

BTW, the existing writeback logic has been broken for a while now when memcg is disabled. I think we constantly hit the !memcg case and run out of retries. Not sure if your patch changes this in any way, or if you want to fix that while you're at it :)

Quoted 2 lines

Yes, I'd be happy to do that. However, would it be better to submit a separate fix patch or combine it with this one?

A separate patch. Feel free to send it with this series to avoid conflicts, but probably as patch 1 as we'll want to CC stable on it.

[..]

Quoted 6 lines

/* Track progress of a memcg-tree writeback walk. */

struct zswap_shrink_state {
      int scans;
      int failures;
};

Quoted 15 lines

/*
* Take one step of a memcg-tree writeback walk driven by the caller's
* iterator, and fold the result into @s, the retry bookkeeping shared
* across steps. @memcg is the iterator's current memcg, or NULL once
* it has wrapped around after a full pass over the tree.
*
* The function returns -EBUSY to signal the caller to abort the walk when
* either of the following occurs:
* - A full pass over the tree found no writeback-candidate memcg.
* - Shrinking a writeback-candidate memcg failed MAX_RECLAIM_RETRIES
times.
*
* When memory cgroup is disabled, the iterator always yields NULL. All
* zswap entries then live on the root list_lru, so NULL is treated as the
* root memcg and shrunk directly rather than as a completed tree pass.

I think this chunk should be moved above the code returning -EBUSY when mem_cgroup_disabled() is true, and probably made more succinct as it should be obvious.

Quoted 9 lines

*
* Return: The number of compressed bytes written back (>= 0), or -EBUSY
* when the caller should abort the walk.
*/
static long zswap_shrink_one_memcg(struct mem_cgroup *memcg,
struct zswap_shrink_state *s)

{
      bool disabled = mem_cgroup_disabled();

No need to store this in a variable AFAICT, it's a static branch and it's clearer to just call it directly in both call sites imo.

Quoted 38 lines

      long shrunk;

      /*
       * If the iterator has completed a full pass, update the shrink state
       * and check whether we should keep going.
       * With memcg disabled the iterator always yields NULL, so fall through
       * and shrink the root memcg directly instead.
       */
      if (!memcg && !disabled) {
          /*
           * Abort if no writeback-candidate memcgs in the last tree walk.
           * Otherwise reset the scans count and continue.
           */
          if (!s->scans)
              return -EBUSY;
          s->scans = 0;
          return 0;
      }

      shrunk = shrink_memcg(memcg, NR_ZSWAP_WB_BATCH);

      /*
       * There are no writeback-candidate pages in the memcg. With memcg
       * enabled this is not an issue as long as we can find another memcg
       * with pages in zswap, so skip without counting it as a candidate.
       * With memcg disabled the root LRU is the only target, so we should
       * abort if it has no writeback-candidate pages.
       */
      if (shrunk == -ENOENT)
          return disabled ? -EBUSY : 0;
      s->scans++;

      if (shrunk <= 0 && ++s->failures == MAX_RECLAIM_RETRIES)
          return -EBUSY;

      return shrunk;
}

From: Hao Jia <[email protected]>

Zswap currently writes back pages to backing swap reactively, triggered either by the shrinker or when the pool reaches its size limit. There is no mechanism to control the amount of writeback for a specific memory cgroup. However, users may want to proactively write back zswap pages, e.g., to free up memory for other applications or to prepare for memory-intensive workloads.

Introduce a "zswap_writeback_only" key to the memory.reclaim cgroup interface. When specified, this key bypasses standard memory reclaim and exclusively performs proactive zswap writeback up to the requested budget. If omitted, the default reclaim behavior remains unchanged.

Example usage:
# Write back 10MB of compressed data from zswap to the backing swap echo "10M zswap_writeback_only" > memory.reclaim

Note that the actual amount of compressed data written back may be less than requested due to the zswap second-chance algorithm: referenced entries are rotated on the LRU on the first encounter and only written back on a second pass. If fewer bytes are written back than requested, -EAGAIN is returned, matching the existing memory.reclaim semantics.

Internally, extend user_proactive_reclaim() to parse the new "zswap_writeback_only" token and invoke the dedicated handler zswap_proactive_writeback(). This handler reuses zswap_try_to_writeback() to walk the target memcg subtree, draining per-node zswap LRUs through list_lru_walk_one() with the shrink_memcg_cb() callback.

Suggested-by Yosry Ahmed <[email protected]>

Suggested-by Nhat Pham <[email protected]>

Signed-off-by Hao Jia <[email protected]>

mm/zswap.c +26 -4

index e29f8a61412d..28200552dde3 100644
@@ -1423,6 +1423,27 @@ static struct mem_cgroup *zswap_iter_global(void)
 	return memcg;
 }
 
+/*
+ * Local iteration uses a local cursor to select from online memcgs
+ * under @root in a round-robin fashion.
+ *
+ * Pass the previous return value as @prev to advance the round-robin
+ * iteration, or pass NULL to start a new walk. If exiting early before
+ * the iteration completes, the caller must call mem_cgroup_iter_break()
+ * to release the cursor reference.
+ */
+static struct mem_cgroup *zswap_iter_local(struct mem_cgroup *root,
+					   struct mem_cgroup *prev)
+{
+	struct mem_cgroup *memcg;
+
+	do {
+		memcg = mem_cgroup_iter(root, prev, NULL);
+		prev = memcg;
+	} while (memcg && !mem_cgroup_tryget_online(memcg));
+	return memcg;
+}
+
 /*
  * Walk the memcg tree and write back zswap pages until the
  * (lower_pages, upper_pages) window closes, or abort encounter

@@ -1430,16 +1451,23 @@ static struct mem_cgroup *zswap_iter_global(void)
  * - No writeback-candidate memcgs found in a memcg tree walk.
  * - Shrinking a writeback-candidate memcg failed.
  *
- * For shrink_worker(), it passes lower=thr and upper=zswap_total_pages().
- * The @upper limit is refreshed in each iteration by re-evaluating
- * zswap_total_pages(), and the window closes once the total falls
- * below the threshold.
+ * For shrink_worker() (proactive=false), it passes lower=thr and
+ * upper=zswap_total_pages(). The @upper limit is refreshed in each
+ * iteration by re-evaluating zswap_total_pages(), and the window
+ * closes once the total falls below the threshold.
+ *

mm/zswap.c +2 -1

 static void shrink_worker(struct work_struct *w)

@@ -1490,7 +1536,7 @@ static void shrink_worker(struct work_struct *w)
 	/* Reclaim down to the accept threshold */
 	thr = zswap_accept_thr_pages();
 
-	zswap_try_to_writeback(thr, zswap_total_pages());
+	zswap_try_to_writeback(NULL, thr, zswap_total_pages(), false);
 }
 
 /*********************************

@@ -1736,6 +1782,19 @@ int zswap_load(struct folio *folio)
 	return 0;
 }
 
+int zswap_proactive_writeback(struct mem_cgroup *memcg,

mm/zswap.c +2 -0

+}
+
 void zswap_invalidate(swp_entry_t swp)
 {
 	pgoff_t offset = swp_offset(swp);

Signature · 1 line

2.34.1

{

mm/zswap.c +0 -0

 	struct xarray *tree;
 	pgoff_t offset = swp_offset(swpentry);

@@ -1045,6 +1047,15 @@ static int zswap_writeback_entry(struct zswap_entry *entry,
 	if (entry->objcg)

From: Hao Jia <[email protected]>

Zswap currently writes back pages to backing swap reactively, triggered either by the shrinker or by the pool reaching its size limit. Although proactive memory reclaim can automatically write back a portion of zswap pages via the shrinker, it cannot explicitly control the amount of writeback for a specific memory cgroup. Moreover, proactive memory reclaim may not always be triggered during a steady state.

In certain scenarios, it is desirable to trigger writeback in advance to free up memory. For example, users may want to prepare for an upcoming memory-intensive workload by flushing cold memory to the backing storage when the system is relatively idle.

This patch series introduces a "zswap_writeback_only" key to memory.reclaim cgroup interface, allowing users to proactively write back cold compressed data from zswap to the backing swap device. When specified, this key bypasses standard memory reclaim and exclusively performs proactive zswap writeback up to the requested budget. If omitted, the default reclaim behavior remains unchanged.

Example usage:
# Write back 10MB of compressed data from zswap to the backing swap echo "10M zswap_writeback_only" > memory.reclaim

Patch 1: Extend shrink_memcg() to support batch writeback based on a
compressed-size budget and update its return value semantics, thereby improving the writeback efficiency in the shrink_worker() path.
Patch 2: Extract the memcg iteration and writeback loop into helper
functions to prepare for proactive writeback.
Patch 3: Extend the memory.reclaim cgroup v2 interface with a new
"zswap_writeback_only" key, allowing users to trigger proactive zswap writeback up to a requested budget.
Patch 4: Add the zswpwb_proactive_b stat to track the compressed bytes
of proactive writeback for better monitoring and tuning.
Patch 5:
Add tests for zswap proactive writeback.

v3->v4:
- Drop the per-memcg cursor and keep the root cgroup cursor (zswap_next_shrink) logic intact.
- Stick to using the zswap_writeback_only key, and change the proactive writeback size to use the compressed size.
- Consolidate and reuse the logic between shrink_worker() and shrink_memcg(). Enable batch writeback in the shrink_worker() path, while maintaining a low writeback budget in the zswap_store() path.

v2->v3:
- Align the return value of zswap_proactive_writeback() with memory.reclaim and update the corresponding documentation accordingly.
- Resolve conflicts in test_zswap.c on the mm-unstable branch.
- Enhance the zswap proactive writeback selftests to guard against potential future regressions.

v1->v2:
- As suggested by Yosry and Nhat, extend the memory.reclaim cgroup v2 interface with a "zswap_writeback_only" key instead of adding a new dedicated cgroup interface.
- Update the zswap documentation and add selftests for proactive writeback.

[v3] https://lore.kernel.org/all/[email protected]
[v2] https://lore.kernel.org/all/[email protected]
[v1] https://lore.kernel.org/all/[email protected]

Hao Jia (5):
mm/zswap: Extend shrink_memcg() writeback capability mm/zswap: Factor writeback loop out of shrink_worker() mm/zswap: Implement proactive writeback mm/zswap: Add per-memcg stat for proactive writeback selftests/cgroup: Add tests for zswap proactive writeback

Documentation/admin-guide/cgroup-v2.rst     |  22 +-
Documentation/admin-guide/mm/zswap.rst      |  11 +-
include/linux/memcontrol.h                  |   1 +
include/linux/zswap.h                       |   7 +
mm/memcontrol.c                             |   3 +
mm/vmscan.c                                 |  14 +
mm/zswap.c                                  | 322 +++++++++++++++-----
tools/testing/selftests/cgroup/test_zswap.c | 153 +++++++++-
8 files changed, 456 insertions(+), 77 deletions(-)

Signature · 1 line

2.34.1

Quoted 25 lines · Hao Jia

From: Hao Jia <[email protected]>

Zswap currently writes back pages to backing swap reactively, triggered either by the shrinker or by the pool reaching its size limit. Although proactive memory reclaim can automatically write back a portion of zswap pages via the shrinker, it cannot explicitly control the amount of writeback for a specific memory cgroup. Moreover, proactive memory reclaim may not always be triggered during a steady state.

In certain scenarios, it is desirable to trigger writeback in advance to free up memory. For example, users may want to prepare for an upcoming memory-intensive workload by flushing cold memory to the backing storage when the system is relatively idle.

This patch series introduces a "zswap_writeback_only" key to memory.reclaim cgroup interface, allowing users to proactively write back cold compressed data from zswap to the backing swap device. When specified, this key bypasses standard memory reclaim and exclusively performs proactive zswap writeback up to the requested budget. If omitted, the default reclaim behavior remains unchanged.

Quoted 3 lines

Example usage:
# Write back 10MB of compressed data from zswap to the backing swap echo "10M zswap_writeback_only" > memory.reclaim

I’m not entirely sure if other candidate names were already brought up in previous discussions, so my apologies if I'm repeating something here! I do think expanding memory.reclaim is a great approach. That said, I was wondering if we could make the interface a bit more concise while keeping it flexible for future extensions.

Essentially, what we want is to control the specific targets of the reclaim process—such as file, anon, or zswap. What do you think about using something like "source=zswap"? For instance, if we want to reclaim 10M from zswap, the command would look like this:

echo "10M source=zswap" > memory.reclaim

If we only want to reclaim 10M from file pages, we could easily extend the syntax:

echo "10M source=file" > memory.reclaim

And of course, we could even combine them down the road:

echo "10M source=anon,file" > memory.reclaim

to only reclaim anon and file but bypass zswap.

Just some thoughts of mine.

Muchun,
Thanks

Quoted 35 lines · Muchun Song

Quoted 25 lines · Hao Jia

From: Hao Jia <[email protected]>

Zswap currently writes back pages to backing swap reactively, triggered either by the shrinker or by the pool reaching its size limit. Although proactive memory reclaim can automatically write back a portion of zswap pages via the shrinker, it cannot explicitly control the amount of writeback for a specific memory cgroup. Moreover, proactive memory reclaim may not always be triggered during a steady state.

In certain scenarios, it is desirable to trigger writeback in advance to free up memory. For example, users may want to prepare for an upcoming memory-intensive workload by flushing cold memory to the backing storage when the system is relatively idle.

This patch series introduces a "zswap_writeback_only" key to memory.reclaim cgroup interface, allowing users to proactively write back cold compressed data from zswap to the backing swap device. When specified, this key bypasses standard memory reclaim and exclusively performs proactive zswap writeback up to the requested budget. If omitted, the default reclaim behavior remains unchanged.

Quoted 3 lines

Example usage:
# Write back 10MB of compressed data from zswap to the backing swap echo "10M zswap_writeback_only" > memory.reclaim

I’m not entirely sure if other candidate names were already brought up in previous discussions, so my apologies if I'm repeating something here! I do think expanding memory.reclaim is a great approach. That said, I was wondering if we could make the interface a bit more concise while keeping it flexible for future extensions.

Quoted 7 lines

Essentially, what we want is to control the specific targets of the reclaim process—such as file, anon, or zswap. What do you think about using something like "source=zswap"? For instance, if we want to reclaim 10M from zswap, the command would look like this:

echo "10M source=zswap" > memory.reclaim

Thanks for the suggestion. TBH, I personally think your approach makes more sense than "zswap_writeback_only".

Hi YoungJun and Yosry,

I am not sure if this suggestion from Muchun could decouple zswap proactive writeback from the swap tiers, or make it easier to migrate to swap tiers in the future:

echo "10M source=zswap" > memory.reclaim
For now, we only specify the source. Later on, the swap tiers feature could extend this to control whether to demote to SSD swap, HDD swap, or other tiers.

Thanks,
Hao

If we only want to reclaim 10M from file pages, we could easily extend the syntax:

echo "10M source=file" > memory.reclaim

And of course, we could even combine them down the road:

echo "10M source=anon,file" > memory.reclaim

to only reclaim anon and file but bypass zswap.

Just some thoughts of mine.

Quoted 38 lines · Hao Jia

Quoted 35 lines · Muchun Song

Quoted 25 lines · Hao Jia

From: Hao Jia <[email protected]>

Zswap currently writes back pages to backing swap reactively, triggered either by the shrinker or by the pool reaching its size limit. Although proactive memory reclaim can automatically write back a portion of zswap pages via the shrinker, it cannot explicitly control the amount of writeback for a specific memory cgroup. Moreover, proactive memory reclaim may not always be triggered during a steady state.

In certain scenarios, it is desirable to trigger writeback in advance to free up memory. For example, users may want to prepare for an upcoming memory-intensive workload by flushing cold memory to the backing storage when the system is relatively idle.

This patch series introduces a "zswap_writeback_only" key to memory.reclaim cgroup interface, allowing users to proactively write back cold compressed data from zswap to the backing swap device. When specified, this key bypasses standard memory reclaim and exclusively performs proactive zswap writeback up to the requested budget. If omitted, the default reclaim behavior remains unchanged.

Quoted 3 lines

Example usage:
# Write back 10MB of compressed data from zswap to the backing swap echo "10M zswap_writeback_only" > memory.reclaim

I’m not entirely sure if other candidate names were already brought up in previous discussions, so my apologies if I'm repeating something here! I do think expanding memory.reclaim is a great approach. That said, I was wondering if we could make the interface a bit more concise while keeping it flexible for future extensions.

Quoted 6 lines

Quoted 6 lines

Essentially, what we want is to control the specific targets of the reclaim process—such as file, anon, or zswap. What do you think about using something like "source=zswap"? For instance, if we want to reclaim 10M from zswap, the command would look like this:

echo "10M source=zswap" > memory.reclaim

Quoted 1 line

Quoted 1 line

I also preferred sharing the memory.reclaim interface in the future swap demotion, since it already takes zswap_writeback_only.
https://lore.kernel.org/all/aieUQUBHI+E3uNPW@yjaykim-PowerEdge-T330/

Alternatively, we could use a separate interface as Yosry suggested (e.g. 'swap.tiers.demote'?).

But as Nhat pointed out, allowing user-triggered demotion from the swap tier perspective could lead to issues like LRU inversion. We probably need to discuss whether this kind of user-triggered tier demotion will actually be supported at all.
https://lore.kernel.org/linux-mm/CAKEwX=NfSy0XiD_UMsDOHGCwpE7sYmBmhV4Y9vk_cbnnr6J6PQ@mail.gmail.com/

So, IMHO..

1. If swap tier demotion is NOT exposed.

We can simply choose between "source=" and zswap_writeback_only based on preference. (since there is no need to consider "swap_tier" demotion.)

However, "source=" seems to offer better extensibility if it is expanded to file and anon use cases in the future.

2. If swap tier demotion IS exposed.
We need to consider integration vs decoupling.

(In my view, This is a design consideration. avoiding potentially redundant interfaces vs adding a new one if it is architecturally correct.)

2.1 Integration
- Integrating into 'memory.reclaim':
- "source=": Seems easier to integrate by explicitly specifying the target. (Your suggestion)
- 'zswap_writeback_only': Harder to integrate than "source=".

- Integrating into 'memory.swap.tiers.demote'
- 'memory.swap.tiers.demote' could absorb the memory.reclaim functionality. (But since we only want to allow tiering for vswap+zswap cases like the zswap writeback feature as we discussed, the reclaim interface behavior might still need to stay for zswap only.)

2.2 Decoupling
- 'memory.swap.tiers.demote' handles other swap devices (excluding zswap),
while "source=" or 'zswap_writeback_only' handles only zswap.