From: Nhat Pham <nphamcs@gmail.com>

An xswap entry is recorded but not charged at allocation.  Charge
memcg->swap when the entry takes a physical backend slot, in
xswap_backend_alloc(), and uncharge it when the backend is released.

memory.swap.current therefore still counts the pages that do reach the
disk, and memory.swap.max is still enforced on them.

Suggested-by: Johannes Weiner <hannes@cmpxchg.org>
Signed-off-by: Nhat Pham <nphamcs@gmail.com>
Signed-off-by: Baoquan He <hebaoquan@kylinos.cn>
---
 mm/swapfile.c | 60 +++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/mm/swapfile.c b/mm/swapfile.c
index 627f2cab8077..f848c6edc6e7 100644
--- a/mm/swapfile.c
+++ b/mm/swapfile.c
@@ -79,6 +79,11 @@ static void xswap_free_phys_slot(struct swap_info_struct *si,
 static void xswap_release_slot_backend(struct swap_info_struct *si,
 				       struct swap_cluster_info *ci,
 				       unsigned int slot);
+static struct mem_cgroup *xswap_entry_memcg(struct swap_cluster_info *ci,
+					    unsigned long off);
+static void xswap_backend_uncharge(swp_entry_t entry, unsigned int nr);
+static void __xswap_backend_unlink(swp_entry_t entry, swp_entry_t phys,
+				   unsigned int order);
 
 static int xswap_create(int prio);
 static int xswap_destroy(int type);
@@ -3169,6 +3174,8 @@ static void xswap_release_slot_backend(struct swap_info_struct *si,
 	psi = swap_type_to_info(swp_type(phys));
 	if (psi)
 		xswap_free_phys_slot(psi, entry, phys);
+
+	xswap_backend_uncharge(entry, 1);
 }
 
 /*
@@ -4359,6 +4366,7 @@ swp_entry_t xswap_backend_alloc(swp_entry_t entry, unsigned int order)
 {
 	struct swap_info_struct *si = __swap_entry_to_info(entry);
 	struct swap_cluster_info *ci;
+	struct mem_cgroup *memcg = NULL;
 	unsigned long nr = 1UL << order;
 	unsigned long off = swp_offset(entry);
 	swp_entry_t phys;
@@ -4369,6 +4377,21 @@ swp_entry_t xswap_backend_alloc(swp_entry_t entry, unsigned int order)
 	if (!phys.val)
 		return phys;
 
+	/*
+	 * The xswap entry was only recorded at allocation, so charge the
+	 * physical swap here.  On failure, drop the physical run without
+	 * uncharging it.
+	 */
+	ci = swap_cluster_lock(si, off);
+	if (ci) {
+		memcg = xswap_entry_memcg(ci, off);
+		swap_cluster_unlock(ci);
+	}
+	if (mem_cgroup_swap_charge(memcg, nr)) {
+		__xswap_backend_unlink(entry, phys, order);
+		return (swp_entry_t){};
+	}
+
 	/* One IO reads a large folio, so its range has to be contiguous. */
 	VM_WARN_ON_ONCE(swp_offset(phys) % nr);
 	VM_WARN_ON_ONCE(off % SWAPFILE_CLUSTER + nr > SWAPFILE_CLUSTER);
@@ -4400,8 +4423,34 @@ swp_entry_t xswap_backend_alloc(swp_entry_t entry, unsigned int order)
 	return phys;
 }
 
-/* Undo xswap_backend_alloc(): the zswap copy is still in place. */
-void xswap_backend_free(swp_entry_t entry, swp_entry_t phys, unsigned int order)
+/* The memcg that owns an xswap entry, pinned by its recorded ID. */
+static struct mem_cgroup *xswap_entry_memcg(struct swap_cluster_info *ci,
+					    unsigned long off)
+{
+	unsigned short id = __swap_cgroup_get(ci, off % SWAPFILE_CLUSTER);
+
+	return id ? mem_cgroup_from_private_id(id) : NULL;
+}
+
+static void xswap_backend_uncharge(swp_entry_t entry, unsigned int nr)
+{
+	struct swap_info_struct *si = __swap_entry_to_info(entry);
+	struct swap_cluster_info *ci;
+	struct mem_cgroup *memcg = NULL;
+	unsigned long off = swp_offset(entry);
+
+	ci = swap_cluster_lock(si, off);
+	if (ci) {
+		memcg = xswap_entry_memcg(ci, off);
+		swap_cluster_unlock(ci);
+	}
+	if (memcg)
+		mem_cgroup_swap_uncharge(memcg, nr);
+}
+
+/* Drop the physical run and the reverse mapping, without any charging. */
+static void __xswap_backend_unlink(swp_entry_t entry, swp_entry_t phys,
+				   unsigned int order)
 {
 	struct swap_info_struct *si = __swap_entry_to_info(entry);
 	struct swap_info_struct *psi = swap_type_to_info(swp_type(phys));
@@ -4430,6 +4479,13 @@ void xswap_backend_free(swp_entry_t entry, swp_entry_t phys, unsigned int order)
 	}
 }
 
+/* Undo xswap_backend_alloc(): the zswap copy is still in place. */
+void xswap_backend_free(swp_entry_t entry, swp_entry_t phys, unsigned int order)
+{
+	__xswap_backend_unlink(entry, phys, order);
+	xswap_backend_uncharge(entry, 1UL << order);
+}
+
 static int xswap_map_clusters(struct swap_info_struct *si,
 			      unsigned long start_idx, unsigned long nr)
 {
-- 
2.54.0