[RFC PATCH 3/7] mm/swap_state: do zswap conversion in __swap_cache_do_del_folio

[Baoquan He <baoquan.he@linux.dev>]

Add zswap_try_convert_to_pointer(), called from __swap_cache_do_del_folio
under cluster lock. When a folio is evicted from swap cache and a Shadow
entry is written, this helper checks if zswap has a compressed copy of
the slot and replaces the Shadow with a Pointer entry.

Signed-off-by: Baoquan He <baoquan.he@linux.dev>
---
 mm/swap.h       |  5 +++++
 mm/swap_state.c |  7 +++++++
 mm/zswap.c      | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/mm/swap.h b/mm/swap.h
index afe2200bd160..1b07a08e5277 100644
--- a/mm/swap.h
+++ b/mm/swap.h
@@ -211,6 +211,8 @@ static inline void swap_cluster_unlock_irq(struct swap_cluster_info *ci)
 extern int swap_retry_table_alloc(swp_entry_t entry, gfp_t gfp);
 
 #ifdef CONFIG_ZSWAP
+void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+				  unsigned int ci_off, int type, pgoff_t offset);
 unsigned char zswap_swp_tb_get_count(unsigned long swp_tb);
 unsigned char zswap_swp_tb_get_flags(unsigned long swp_tb);
 int zswap_swp_tb_dup_count(unsigned long swp_tb,
@@ -218,6 +220,9 @@ int zswap_swp_tb_dup_count(unsigned long swp_tb,
 void zswap_swp_tb_put_count(unsigned long swp_tb,
 			    struct swap_cluster_info *ci, unsigned int ci_off);
 #else
+static inline void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+						unsigned int ci_off, int type,
+						pgoff_t offset) {}
 static inline unsigned char zswap_swp_tb_get_count(unsigned long swp_tb)
 {
 	return 0;
diff --git a/mm/swap_state.c b/mm/swap_state.c
index 6fd6e3415b71..20fb57adbb7f 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -278,6 +278,13 @@ static void __swap_cache_do_del_folio(struct swap_cluster_info *ci,
 		/* If shadow is NULL, we set an empty shadow. */
 		__swap_table_set(ci, ci_off, shadow_to_swp_tb(shadow,
 				 __swp_tb_get_flags(old_tb)));
+		/*
+		 * If zswap has a compressed copy of this slot, convert the
+		 * just-written Shadow to a Pointer entry referencing the
+		 * zswap_entry, so the data can be found via the swap table.
+		 */
+		zswap_try_convert_to_pointer(ci, ci_off, si->type,
+					     swp_offset(entry) + ci_off - ci_start);
 	} while (++ci_off < ci_end);
 
 	folio->swap.val = 0;
diff --git a/mm/zswap.c b/mm/zswap.c
index 81115ca4d10a..f6e8259ab56a 100644
--- a/mm/zswap.c
+++ b/mm/zswap.c
@@ -1652,6 +1652,41 @@ void zswap_invalidate(swp_entry_t swp)
 		zswap_entry_free(entry);
 }
 
+/*
+ * This is called after the Shadow entry has been written. We replace the
+ * just-written Shadow with a Pointer entry pointing to the zswap_entry.
+ *
+ * Context: cluster lock must be held. Cannot sleep.
+ */
+void zswap_try_convert_to_pointer(struct swap_cluster_info *ci,
+				  unsigned int ci_off, int type, pgoff_t offset)
+{
+	struct xarray *tree;
+	struct zswap_entry *entry;
+	unsigned long swp_tb;
+
+	tree = &zswap_trees[type][offset >> ZSWAP_ADDRESS_SPACE_SHIFT];
+	if (xa_empty(tree))
+		return;
+
+	entry = xa_load(tree, offset);
+	if (!entry)
+		return;
+
+	swp_tb = __swap_table_get(ci, ci_off);
+	/* Entry must be a Shadow (just written by __swap_cache_do_del_folio) */
+	if (!swp_tb_is_countable(swp_tb))
+		return;
+
+	/*
+	 * Save the complete original swap table entry value, which includes
+	 * the swap count, zero flag, and working set shadow information.
+	 */
+	entry->swp_tb_val = swp_tb;
+
+	__swap_table_set(ci, ci_off, pointer_to_swp_tb(entry));
+}
+
 /*
  * Read the swap count from a Pointer-type swap table entry. The count
  * and flags are stored in zswap_entry->swp_tb_val which preserves the
-- 
2.54.0