lib/test_hmm: check alloc_page_vma() return value and handle OOM

Check alloc_page_vma() return status for page allocation failures, free
allocated pages and return VM_FAULT_OOM on error.

Handle return codes of dmirror_devmem_fault_alloc_and_copy(), call
migrate_vma_finalize() to remove migration entries from
migrate_vma_setup().

Link: https://lore.kernel.org/20260521021858.21511-1-liuqiangneo@163.com
Signed-off-by: Qiang Liu <liuqiang@kylinos.cn>
Cc: Alistair Popple <apopple@nvidia.com>
Cc: Jason Gunthorpe <jgg@ziepe.ca>
Cc: Leon Romanovsky <leon@kernel.org>
[akpm@linux-foundation.org: fix dmirror_devmem_fault_alloc_and_copy() retval handling]
Link: https://lore.kernel.org/oe-kbuild-all/202606011329.zWs2BKy4-lkp@intel.com/
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/lib/test_hmm.c b/lib/test_hmm.c
index 35f774ed2d999994ce3da5c23614ef1efbcdb6aa..9c59d1ceb5b57d6257f5fe43a3b60586333ad4e3 100644 (file)
--- a/lib/test_hmm.c
+++ b/lib/test_hmm.c
@@ -1063,6 +1063,25 @@ static vm_fault_t dmirror_devmem_fault_alloc_and_copy(struct migrate_vma *args,
                        /* Try with smaller pages if large allocation fails */
                        if (!dpage && order) {
                                dpage = alloc_page_vma(GFP_HIGHUSER_MOVABLE, args->vma, addr);
+                               if (!dpage) {
+                                       /* Unlock and free pages already allocated. */
+                                       while (i > 0) {
+                                               struct page *fpage;
+
+                                               fpage = migrate_pfn_to_page(dst[--i]);
+                                               unlock_page(fpage);
+                                               __free_page(fpage);
+                                       }
+                                       /* Clear remaining dst entries to avoid
+                                        * migrate_vma_pages/finalize() using
+                                        * uninitialized values.
+                                        */
+                                       while (i < (1 << order)) {
+                                               dst[i] = 0;
+                                               i++;
+                                       }
+                                       return VM_FAULT_OOM;
+                               }
                                lock_page(dpage);
                                dst[i] = migrate_pfn(page_to_pfn(dpage));
                                dst_page = pfn_to_page(page_to_pfn(dpage));
@@ -1148,7 +1167,11 @@ static int dmirror_migrate_to_system(struct dmirror *dmirror,
                        goto out;
 
                pr_debug("Migrating from device mem to sys mem\n");
-               dmirror_devmem_fault_alloc_and_copy(&args, dmirror);
+               if (dmirror_devmem_fault_alloc_and_copy(&args, dmirror)) {
+                       migrate_vma_finalize(&args);
+                       ret = -ENOMEM;
+                       goto out;
+               }
 
                migrate_vma_pages(&args);
                cmd->cpages += dmirror_successful_migrated_pages(&args);
@@ -1689,8 +1712,10 @@ static vm_fault_t dmirror_devmem_fault(struct vm_fault *vmf)
        }
 
        ret = dmirror_devmem_fault_alloc_and_copy(&args, dmirror);
-       if (ret)
+       if (ret) {
+               migrate_vma_finalize(&args);
                goto err;
+       }
        migrate_vma_pages(&args);
        /*
         * No device finalize step is needed since