[RFC PATCH v3 11/14] mm/memcontrol, migrate: Transfer tier charge on migration

[Joshua Hahn <joshua.hahnjy@gmail.com>]

Folio migration is charge-neutral, and existing migration paths take
advantage of this fact to simply force destination folio charges or
just transfer memcg data across folios.

Per-tier memcg limits break this assumption. A migration across tiers
(i.e. promotion or demotion) keeps the memcg-level charge neutral,
but the per-memcg tier charges change. As a result, the destination tier
may go over the limit.

Charge the destination separately instead, from migrate_folio_unmap
where the destination folio has just been allocated but can still be
rolled back. This charge attempts a single pass at reclaim if it goes
over the hard limit, and fails the migration if not enough headroom is
created on the destination memcg tier.

Note that this source of migration failure returns -EBUSY and not
-ENOMEM since -ENOMEM will attempt the migration again by splitting the
folio and aborting the batch, which both do nothing to reduce the
memory usage of the memcg tier. We also don't try too hard to reclaim
here (__GFP_NORETRY) since failing migrations is cheap, and we don't
want to OOM kill because of a promotion attempt.

One side effect is that cross-tier migrations now hold both folios'
charges until the source is freed, the same way mem_cgroup_replace_folio
temporarily holds a duplicate charge.

No-op unless the system has tiered memcg limits enabled.

Suggested-by: Johannes Weiner <hannes@cmpxchg.org>
Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 include/linux/memcontrol.h |  8 ++++++++
 mm/memcontrol.c            | 35 +++++++++++++++++++++++++++++++++++
 mm/migrate.c               | 30 +++++++++++++++++++++++++++---
 3 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index ceba0fd6de184..9c2f11191a499 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -708,6 +708,8 @@ static inline void mem_cgroup_uncharge_folios(struct folio_batch *folios)
 
 void mem_cgroup_replace_folio(struct folio *old, struct folio *new);
 void mem_cgroup_migrate(struct folio *old, struct folio *new);
+int mem_cgroup_migrate_charge(struct folio *src, struct folio *dst,
+			      bool force);
 
 /**
  * mem_cgroup_lruvec - get the lru list vector for a memcg & node
@@ -1204,6 +1206,12 @@ static inline void mem_cgroup_migrate(struct folio *old, struct folio *new)
 {
 }
 
+static inline int mem_cgroup_migrate_charge(struct folio *src,
+					    struct folio *dst, bool force)
+{
+	return 0;
+}
+
 static inline struct lruvec *mem_cgroup_lruvec(struct mem_cgroup *memcg,
 					       struct pglist_data *pgdat)
 {
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index de6762520f475..1161934e81380 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -5716,6 +5716,41 @@ void mem_cgroup_replace_folio(struct folio *old, struct folio *new)
 	rcu_read_unlock();
 }
 
+/**
+ * mem_cgroup_migrate_charge - Charge a migration destination up front.
+ * @src: Folio being migrated away from.
+ * @dst: Folio being migrated to.
+ * @force: Charge even if the destination tier is at its limit.
+ *
+ * Folio migrations result in a net 0 memcg charge, but the node location of the
+ * charge may change during promotions or demotions. When this happens, charge
+ * @dst in its own right instead of inheriting @src's charge.
+ *
+ * Return: 0, or -ENOMEM if @dst could not be charged.
+ */
+int mem_cgroup_migrate_charge(struct folio *src, struct folio *dst, bool force)
+{
+	struct mem_cgroup *memcg;
+	gfp_t gfp = GFP_KERNEL;
+	int ret;
+
+	if (mem_cgroup_disabled() || !folio_memcg_charged(src))
+		return 0;
+
+	if (!mem_cgroup_tiered_limits() ||
+	    nid_tier_slot(folio_nid(src)) == nid_tier_slot(folio_nid(dst)))
+		return 0;
+
+	/* Refuse the migration if the first reclaim round fails */
+	gfp |= force ? __GFP_NOFAIL : __GFP_NORETRY;
+
+	memcg = get_mem_cgroup_from_folio(src);
+	ret = charge_memcg(dst, memcg, gfp);
+	mem_cgroup_put(memcg);
+
+	return ret;
+}
+
 /**
  * mem_cgroup_migrate - Transfer the memcg data from the old to the new folio.
  * @old: Currently circulating folio.
diff --git a/mm/migrate.c b/mm/migrate.c
index ab15a4dddd047..45d6d23d53859 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -862,7 +862,13 @@ void folio_migrate_flags(struct folio *newfolio, struct folio *folio)
 	folio_copy_owner(newfolio, folio);
 	pgalloc_tag_swap(newfolio, folio);
 
-	mem_cgroup_migrate(folio, newfolio);
+	/*
+	 * For failable memcg charge transfers (demotion / promotion) the charge
+	 * has already been transferred at this point. For everyone else simply
+	 * transfer the charge here, where it can no longer fail.
+	 */
+	if (!folio_memcg_charged(newfolio))
+		mem_cgroup_migrate(folio, newfolio);
 }
 EXPORT_SYMBOL(folio_migrate_flags);
 
@@ -1216,7 +1222,7 @@ static void migrate_folio_done(struct folio *src,
 static int migrate_folio_unmap(new_folio_t get_new_folio,
 		free_folio_t put_new_folio, unsigned long private,
 		struct folio *src, struct folio **dstp, enum migrate_mode mode,
-		struct list_head *ret)
+		bool force_charge, struct list_head *ret)
 {
 	struct folio *dst;
 	int rc = -EAGAIN;
@@ -1228,6 +1234,17 @@ static int migrate_folio_unmap(new_folio_t get_new_folio,
 	dst = get_new_folio(src, private);
 	if (!dst)
 		return -ENOMEM;
+
+	if (mem_cgroup_migrate_charge(src, dst, force_charge)) {
+		if (put_new_folio)
+			put_new_folio(dst, private);
+		else
+			folio_put(dst);
+		if (ret)
+			list_move_tail(&src->lru, ret);
+		return -EBUSY;
+	}
+
 	*dstp = dst;
 
 	dst->migrate_info = 0;
@@ -1918,8 +1935,15 @@ static int migrate_pages_batch(struct list_head *from,
 				continue;
 			}
 
+			/*
+			 * Hotplug must not be refused: offline_pages() retries
+			 * indefinitely and ignores migration failures, so a
+			 * refusal would hang it rather than fail it.
+			 */
 			rc = migrate_folio_unmap(get_new_folio, put_new_folio,
-					private, folio, &dst, mode, ret_folios);
+					private, folio, &dst, mode,
+					reason == MR_MEMORY_HOTPLUG,
+					ret_folios);
 			/*
 			 * The rules are:
 			 *	0: folio will be put on unmap_folios list,
-- 
2.53.0-Meta