Re: [PATCH 1/3] mm/memfd_luo: don't use memfd_pin_folios()

[Pasha Tatashin]

> memfd_luo takes a snapshot of the folios in the file's mapping. To make
> sure the snapshot does not go out of sync, it pins the folios in the
> file via memfd_pin_folios(). The pinning also migrates the folios out of
> MIGRATE_MOVABLE to make sure doesn't land in KHO boot memory or in a CMA
> region.
> 
> memfd_pin_folios() does not lock the folios. So it can race with
> shmem_writeout(). A folio returned by memfd_pin_folios() might already
> be in swap cache. This happens to work but is fragile. memfd_luo looks
> at folio flags and index on a folio that isn't in the memfd's page
> cache.
> 
> In addition, memfd_pin_folios() needs the caller to allocate an array of
> folios. memfd_luo does not really need that since it will vmalloc
> folios_ser and store the information there anyway. The folios array is
> unnecessary overhead.
> 
> Finally, memfd_pin_folios() fills in holes in the memfd. This makes
> preserving sparse memfds very expensive in both compute and memory.
> While this patch doesn't yet stop filling in the holes, it sets up the
> infrastructure to make it simpler.

+1, this is problematic and should be resolved. In some setups, VMMs 
leave holes for PCI/IO space; with the current approach, single memfd 
preservation does not work for such use cases.

> Fix all this by walking the page cache directly using shmem_get_folio()
> with SGP_CACHE. This returns a locked folio with an elevated refcount.
> memfd_luo holds that reference until the memfd is unpreserved. This
> blocks compaction, migration, etc.
> 
> There is still a need to migrate folios out of MIGRATE_MOVABLE page
> blocks. This is because those might be CMA or KHO scratch. memfd_luo
> will hold a longterm reference to the folios so it will block migration
> in that pageblock. Introduce memfd_luo_migrate_folio() to do this.
> 
> Also factor out folio unpreservation logic into memfd_luo_put_folios()
> and use it from memfd_luo_unpreserve_folios() and from error paths in
> memfd_luo_preserve_folios().
> 
> Since shmem_get_folio() always returns uptodate folios, drop the folio
> zeroing logic.
> 
> This change also improves preservation performance, shown in the table
> below.
> 
> |------------------+----------+------------+-------|
> | memfd            | baseline | this patch | delta |
> |------------------+----------+------------+-------|
> | 4K pages, 32 GiB | 748 ms   | 600 ms     |  -20% |
> | THP, 4 GiB       | 0.59 ms  | 0.59 ms    |    0% |
> | THP, 16 GiB      | 2.99 ms  | 2.61 ms    |  -13% |
> | THP, 32 GiB      | 4.83 ms  | 5.21 ms    |   +8% |
> | sparse, 32 GiB   | 8695 ms  | 7269 ms    |  -16% |
> |------------------+----------+------------+-------|

For THP 16 GiB / 32 GiB, the -13% / +8% swings look like high run-to-run 
variance. However, maybe not: is it because vmalloc_array() is 
overallocating memory for folios_ser?

Also, could you clarify what "sparse, 32 GiB" means in this benchmark? 
Was the memfd completely unfaulted before preserve (so preserve faults 
in all 4K pages), or partially populated?

> 
> Signed-off-by: Pratyush Yadav (Google) <pratyush@xxxxxxxxxx>
> Link: https://patch.msgid.link/20261027173031.763213-2-pratyush@xxxxxxxxxx
>
> diff --git a/mm/memfd_luo.c b/mm/memfd_luo.c
> index 59de210bee5f..fe182d505789 100644
> --- a/mm/memfd_luo.c
> +++ b/mm/memfd_luo.c
> @@ -77,26 +77,126 @@
>  #include <linux/kexec_handover.h>
>  #include <linux/kho/abi/memfd.h>
>  #include <linux/liveupdate.h>
> +#include <linux/migrate.h>
> +#include <linux/sched/mm.h>
> +#include <linux/sched/signal.h>
>  #include <linux/shmem_fs.h>
> +#include <linux/swap.h>
>  #include <linux/vmalloc.h>
>  #include <linux/memfd.h>
>  #include <uapi/linux/memfd.h>
>  
>  #include "internal.h"
>  
> +static void memfd_luo_put_folios(const struct memfd_luo_folio_ser *folios_ser,
> +				 u64 nr_folios)
> +{
> +	u64 i;
> +
> +	for (i = 0; i < nr_folios; i++) {
> +		const struct memfd_luo_folio_ser *pfolio = &folios_ser[i];
> +		struct folio *folio;
> +
> +		if (!pfolio->pfn)
> +			continue;
> +
> +		folio = pfn_folio(pfolio->pfn);
> +
> +		kho_unpreserve_folio(folio);
> +		folio_put(folio);
> +	}
> +}
> +
> +/* How many times to retry isolating a folio someone else has isolated. */
> +#define MEMFD_LUO_MIGRATE_RETRIES	3
> +
> +/*
> + * Move a folio out of ZONE_MOVABLE or CMA, so it can be held until the next
> + * kernel. Called with a reference on the unlocked folio, which is always
> + * dropped. The caller must look up the folio again afterwards.
> + *
> + * Returns 0 on success, -EBUSY if the folio could not be isolated (e.g.
> + * because someone else has it isolated right now), or -ENOMEM if migration
> + * failed.
> + */
> +static int memfd_luo_migrate_folio(struct folio *folio)
> +{
> +	struct migration_target_control mtc = {
> +		.nid = NUMA_NO_NODE,
> +		/* Not movable, so the target is neither in ZONE_MOVABLE nor CMA. */
> +		.gfp_mask = GFP_USER | __GFP_NOWARN,
> +		.reason = MR_LONGTERM_PIN,
> +	};
> +	LIST_HEAD(source);
> +	bool isolated;
> +
> +	/* Flush per-CPU LRU batches holding the folio so isolation works. */
> +	lru_cache_drain_for_folio(folio, 1, NULL);
> +	isolated = isolate_folio_to_list(folio, &source);
> +	/* Isolation holds its own reference. */
> +	folio_put(folio);
> +	if (!isolated)
> +		return -EBUSY;
> +
> +	if (migrate_pages(&source, alloc_migration_target, NULL,
> +			  (unsigned long)&mtc, MIGRATE_SYNC, MR_LONGTERM_PIN,
> +			  NULL)) {
> +		putback_movable_pages(&source);
> +		return -ENOMEM;
> +	}
> +
> +	return 0;
> +}
> +
> +static int memfd_luo_preserve_folio(struct folio *folio,
> +				    struct memfd_luo_folio_ser *pfolio)
> +{
> +	int err;
> +
> +	err = kho_preserve_folio(folio);
> +	if (err)
> +		return err;
> +
> +	/*
> +	 * A dirty folio is one which has been written to. A clean folio
> +	 * is its opposite. Since a clean folio does not carry user
> +	 * data, it can be freed by page reclaim under memory pressure.
> +	 *
> +	 * Saving the dirty flag at prepare() time doesn't work since it
> +	 * can change later. Saving it at freeze() also won't work
> +	 * because the dirty bit is normally synced at unmap and there
> +	 * might still be a mapping of the file at freeze().
> +	 *
> +	 * To see why this is a problem, say a folio is clean at
> +	 * preserve, but gets dirtied later. The pfolio flags will mark
> +	 * it as clean. After retrieve, the next kernel might try to
> +	 * reclaim this folio under memory pressure, losing user data.
> +	 *
> +	 * Unconditionally mark it dirty to avoid this problem. This
> +	 * comes at the cost of making clean folios un-reclaimable after
> +	 * live update.
> +	 */
> +	folio_mark_dirty(folio);
> +
> +	pfolio->pfn = folio_pfn(folio);
> +	pfolio->flags = MEMFD_LUO_FOLIO_DIRTY | MEMFD_LUO_FOLIO_UPTODATE;
> +	pfolio->index = folio->index;
> +
> +	return 0;
> +}
> +
>  static int memfd_luo_preserve_folios(struct file *file,
>  				     struct kho_vmalloc *kho_vmalloc,
>  				     struct memfd_luo_folio_ser **out_folios_ser,
>  				     u64 *nr_foliosp)
>  {
> +	struct memfd_luo_folio_ser *folios_ser, *shrunk;
>  	struct inode *inode = file_inode(file);
> -	struct memfd_luo_folio_ser *folios_ser;
> -	unsigned int max_folios;
> -	long i, size, nr_pinned;
> -	struct folio **folios;
> -	int err = -EINVAL;
> -	pgoff_t offset;
> -	u64 nr_folios;
> +	u64 nr_folios = 0, max_folios;
> +	int err, busy_retries = 0;
> +	unsigned int pin_flags;
> +	pgoff_t index = 0, end;
> +	loff_t size;
>  
>  	size = i_size_read(inode);
>  	/*
> @@ -113,99 +213,77 @@ static int memfd_luo_preserve_folios(struct file *file,
>  	 * Guess the number of folios based on inode size. Real number might end
>  	 * up being smaller if there are higher order folios.
>  	 */
> -	max_folios = PAGE_ALIGN(size) / PAGE_SIZE;
> -	folios = kvmalloc_objs(*folios, max_folios);
> -	if (!folios)
> +	end = DIV_ROUND_UP_ULL(size, PAGE_SIZE);
> +	max_folios = end;
> +	folios_ser = vmalloc_array(max_folios, sizeof(*folios_ser));
> +	if (!folios_ser)
>  		return -ENOMEM;

I bet in practice memfds almost always contain folios of the same size, 
so why not estimate the folios_ser size based on the first folio's 
order, and only if that ends up being too small, vrealloc() to match the 
maximum possible remaining size?

It would be a different change though.

>  
>  	/*
> -	 * Pin the folios so they don't move around behind our back. This also
> -	 * ensures none of the folios are in CMA -- which ensures they don't
> -	 * fall in KHO scratch memory. It also moves swapped out folios back to
> -	 * memory.
> +	 * The folios are held until the next kernel, so keep newly allocated
> +	 * ones out of ZONE_MOVABLE and CMA -- which also ensures they don't
> +	 * fall in KHO bootmem. Existing ones are migrated below.
>  	 *
> -	 * A side effect of doing this is that it allocates a folio for all
> -	 * indices in the file. This might waste memory on sparse memfds. If
> -	 * that is really a problem in the future, we can have a
> -	 * memfd_pin_folios() variant that does not allocate a page on empty
> -	 * slots.
> +	 * A side effect of this is that it allocates a folio for all indices in
> +	 * the file. This might waste memory on sparse memfds. If that is really
> +	 * a problem in the future, this loop can be taught to skip empty slots.

Yeah, as I said earlier, it really should skip empty slots, again a 
different change, not related to this patch, but since you are making 
this change, I am commenting on it.

>  	 */
> -	nr_pinned = memfd_pin_folios(file, 0, size - 1, folios, max_folios,
> -				     &offset);
> -	if (nr_pinned < 0) {
> -		err = nr_pinned;
> -		pr_err("failed to pin folios: %d\n", err);
> -		goto err_free_folios;
> -	}
> -	nr_folios = nr_pinned;
> +	pin_flags = memalloc_pin_save();
> +	while (index < end) {
> +		struct folio *folio;
>  
> -	folios_ser = vcalloc(nr_folios, sizeof(*folios_ser));
> -	if (!folios_ser) {
> -		err = -ENOMEM;
> -		goto err_unpin;
> -	}
> +		cond_resched();
>  
> -	for (i = 0; i < nr_folios; i++) {
> -		struct memfd_luo_folio_ser *pfolio = &folios_ser[i];
> -		struct folio *folio = folios[i];
> +		if (fatal_signal_pending(current)) {
> +			err = -EINTR;
> +			goto err_restore;
> +		}
>  
> -		err = kho_preserve_folio(folio);
> +		if (WARN_ON_ONCE(nr_folios >= max_folios)) {
> +			err = -EINVAL;
> +			goto err_restore;
> +		}
> +
> +		err = shmem_get_folio(inode, index, 0, &folio, SGP_CACHE);
>  		if (err)
> -			goto err_unpreserve;
> -
> -		folio_lock(folio);
> -
> -		/*
> -		 * A dirty folio is one which has been written to. A clean folio
> -		 * is its opposite. Since a clean folio does not carry user
> -		 * data, it can be freed by page reclaim under memory pressure.
> -		 *
> -		 * Saving the dirty flag at prepare() time doesn't work since it
> -		 * can change later. Saving it at freeze() also won't work
> -		 * because the dirty bit is normally synced at unmap and there
> -		 * might still be a mapping of the file at freeze().
> -		 *
> -		 * To see why this is a problem, say a folio is clean at
> -		 * preserve, but gets dirtied later. The pfolio flags will mark
> -		 * it as clean. After retrieve, the next kernel might try to
> -		 * reclaim this folio under memory pressure, losing user data.
> -		 *
> -		 * Unconditionally mark it dirty to avoid this problem. This
> -		 * comes at the cost of making clean folios un-reclaimable after
> -		 * live update.
> -		 */
> -		folio_mark_dirty(folio);
> -
> -		/*
> -		 * If the folio is not uptodate, it was fallocated but never
> -		 * used. Saving this flag at prepare() doesn't work since it
> -		 * might change later when someone uses the folio.
> -		 *
> -		 * Since we have taken the performance penalty of allocating,
> -		 * zeroing, and pinning all the folios in the holes, take a bit
> -		 * more and zero all non-uptodate folios too.
> -		 *
> -		 * NOTE: For someone looking to improve preserve performance,
> -		 * this is a good place to look.
> -		 */
> -		if (!folio_test_uptodate(folio)) {
> -			folio_zero_range(folio, 0, folio_size(folio));
> -			flush_dcache_folio(folio);
> -			folio_mark_uptodate(folio);
> +			goto err_restore;
> +
> +		if (unlikely(!folio_is_longterm_pinnable(folio))) {
> +			folio_unlock(folio);
> +			err = memfd_luo_migrate_folio(folio);
> +			/* Someone else might have it isolated, retry a bit. */
> +			if (err == -EBUSY && ++busy_retries < MEMFD_LUO_MIGRATE_RETRIES)
> +				continue;
> +			if (err)
> +				goto err_restore;
> +			continue;
>  		}

Nit:

I would suggest adding another helper, memfd_luo_get_folio(inode, index, 
&folio), that wraps:

- shmem_get_folio()
- The !folio_is_longterm_pinnable() migration retry loop logic

I think this would make the final code a bit more manageable.

Other than that the patch looks alright, feel free to add
Acked-by: Pasha Tatashin <pasha.tatashin@xxxxxxxxxx>
Once the nit is addressed.

-- 
Pasha Tatashin <pasha.tatashin@xxxxxxxxxx>