2 files changed, 36 insertions, 8 deletions

diff --git a/mm/ksm.c b/mm/ksm.c
index c3035fee8080..91b8cb35f7cc 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -2032,12 +2032,19 @@ again:
                         if ((rmap_item->mm == vma->vm_mm) == search_new_forks)
                                 continue;
 
+                        if (rwc->invalid_vma && rwc->invalid_vma(vma, rwc->arg))
+                                continue;
+
                         ret = rwc->rmap_one(page, vma,
                                         rmap_item->address, rwc->arg);
                         if (ret != SWAP_AGAIN) {
                                 anon_vma_unlock_read(anon_vma);
                                 goto out;
                         }
+                        if (rwc->done && rwc->done(page)) {
+                                anon_vma_unlock_read(anon_vma);
+                                goto out;
+                        }
                 }
                 anon_vma_unlock_read(anon_vma);
         }
diff --git a/mm/rmap.c b/mm/rmap.c
index f8f10ad5d359..97bf8f0396f8 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1684,10 +1684,14 @@ void __put_anon_vma(struct anon_vma *anon_vma)
 }
 
 #ifdef CONFIG_MIGRATION
-static struct anon_vma *rmap_walk_anon_lock(struct page *page)
+static struct anon_vma *rmap_walk_anon_lock(struct page *page,
+                                        struct rmap_walk_control *rwc)
 {
         struct anon_vma *anon_vma;
 
+        if (rwc->anon_lock)
+                return rwc->anon_lock(page);
+
         /*
          * Note: remove_migration_ptes() cannot use page_lock_anon_vma_read()
          * because that depends on page_mapped(); but not all its usages
@@ -1713,16 +1717,22 @@ static int rmap_walk_anon(struct page *page, struct rmap_walk_control *rwc)
         struct anon_vma_chain *avc;
         int ret = SWAP_AGAIN;
 
-        anon_vma = rmap_walk_anon_lock(page);
+        anon_vma = rmap_walk_anon_lock(page, rwc);
         if (!anon_vma)
                 return ret;
 
         anon_vma_interval_tree_foreach(avc, &anon_vma->rb_root, pgoff, pgoff) {
                 struct vm_area_struct *vma = avc->vma;
                 unsigned long address = vma_address(page, vma);
+
+                if (rwc->invalid_vma && rwc->invalid_vma(vma, rwc->arg))
+                        continue;
+
                 ret = rwc->rmap_one(page, vma, address, rwc->arg);
                 if (ret != SWAP_AGAIN)
                         break;
+                if (rwc->done && rwc->done(page))
+                        break;
         }
         anon_vma_unlock_read(anon_vma);
         return ret;
@@ -1740,15 +1750,26 @@ static int rmap_walk_file(struct page *page, struct rmap_walk_control *rwc)
         mutex_lock(&mapping->i_mmap_mutex);
         vma_interval_tree_foreach(vma, &mapping->i_mmap, pgoff, pgoff) {
                 unsigned long address = vma_address(page, vma);
+
+                if (rwc->invalid_vma && rwc->invalid_vma(vma, rwc->arg))
+                        continue;
+
                 ret = rwc->rmap_one(page, vma, address, rwc->arg);
                 if (ret != SWAP_AGAIN)
-                        break;
+                        goto done;
+                if (rwc->done && rwc->done(page))
+                        goto done;
         }
-        /*
-         * No nonlinear handling: being always shared, nonlinear vmas
-         * never contain migration ptes.  Decide what to do about this
-         * limitation to linear when we need rmap_walk() on nonlinear.
-         */
+
+        if (!rwc->file_nonlinear)
+                goto done;
+
+        if (list_empty(&mapping->i_mmap_nonlinear))
+                goto done;
+
+        ret = rwc->file_nonlinear(page, mapping, vma);
+
+done:
         mutex_unlock(&mapping->i_mmap_mutex);
         return ret;
 }