diff mbox series

[v3,05/16] mm/mmap: Introduce vma_munmap_struct for use in munmap operations

Message ID 20240704182718.2653918-6-Liam.Howlett@oracle.com (mailing list archive)
State New
Headers show
Series Avoid MAP_FIXED gap exposure | expand

Commit Message

Liam R. Howlett July 4, 2024, 6:27 p.m. UTC
Use a structure to pass along all the necessary information and counters
involved in removing vmas from the mm_struct.

Update vmi_ function names to vms_ to indicate the first argument
type change.

Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
Reviewed-by: Suren Baghdasaryan <surenb@google.com>
---
 mm/internal.h |  16 ++++++
 mm/mmap.c     | 137 ++++++++++++++++++++++++++------------------------
 2 files changed, 88 insertions(+), 65 deletions(-)

Comments

Lorenzo Stoakes July 5, 2024, 6:39 p.m. UTC | #1
On Thu, Jul 04, 2024 at 02:27:07PM GMT, Liam R. Howlett wrote:
> Use a structure to pass along all the necessary information and counters
> involved in removing vmas from the mm_struct.
>
> Update vmi_ function names to vms_ to indicate the first argument
> type change.
>
> Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
> Reviewed-by: Suren Baghdasaryan <surenb@google.com>
> ---
>  mm/internal.h |  16 ++++++
>  mm/mmap.c     | 137 ++++++++++++++++++++++++++------------------------
>  2 files changed, 88 insertions(+), 65 deletions(-)
>
> diff --git a/mm/internal.h b/mm/internal.h
> index 2ea9a88dcb95..f1e6dea2efcf 100644
> --- a/mm/internal.h
> +++ b/mm/internal.h
> @@ -1481,6 +1481,22 @@ struct vma_prepare {
>  	struct vm_area_struct *remove2;
>  };
>
> +/*
> + * vma munmap operation
> + */
> +struct vma_munmap_struct {
> +	struct vma_iterator *vmi;
> +	struct mm_struct *mm;
> +	struct vm_area_struct *vma;	/* The first vma to munmap */
> +	struct list_head *uf;		/* Userfaultfd list_head */
> +	unsigned long start;		/* Aligned start addr */
> +	unsigned long end;		/* Aligned end addr */
> +	int vma_count;			/* Number of vmas that will be removed */
> +	unsigned long nr_pages;		/* Number of pages being removed */
> +	unsigned long locked_vm;	/* Number of locked pages */
> +	bool unlock;			/* Unlock after the munmap */
> +};


I'm a big fan of breaking out and threading state like this through some of
these more... verbose VMA functions.

I have a similar idea as part of my long dreamed of 'delete vma_merge()'
patch set. Coming soon :)

> +
>  void __meminit __init_single_page(struct page *page, unsigned long pfn,
>  				unsigned long zone, int nid);
>
> diff --git a/mm/mmap.c b/mm/mmap.c
> index 8dc8ffbf9d8d..76e93146ee9d 100644
> --- a/mm/mmap.c
> +++ b/mm/mmap.c
> @@ -500,6 +500,31 @@ static inline void init_vma_prep(struct vma_prepare *vp,
>  	init_multi_vma_prep(vp, vma, NULL, NULL, NULL);
>  }
>
> +/*
> + * init_vma_munmap() - Initializer wrapper for vma_munmap_struct
> + * @vms: The vma munmap struct
> + * @vmi: The vma iterator
> + * @vma: The first vm_area_struct to munmap
> + * @start: The aligned start address to munmap
> + * @end: The aligned end address to munmap

Maybe worth mentioning if inclusive/exclusive.

> + * @uf: The userfaultfd list_head
> + * @unlock: Unlock after the operation.  Only unlocked on success
> + */
> +static inline void init_vma_munmap(struct vma_munmap_struct *vms,
> +		struct vma_iterator *vmi, struct vm_area_struct *vma,
> +		unsigned long start, unsigned long end, struct list_head *uf,
> +		bool unlock)
> +{
> +	vms->vmi = vmi;
> +	vms->vma = vma;
> +	vms->mm = vma->vm_mm;
> +	vms->start = start;
> +	vms->end = end;
> +	vms->unlock = unlock;
> +	vms->uf = uf;
> +	vms->vma_count = 0;
> +	vms->nr_pages = vms->locked_vm = 0;
> +}
>
>  /*
>   * vma_prepare() - Helper function for handling locking VMAs prior to altering
> @@ -2606,78 +2631,59 @@ static inline void abort_munmap_vmas(struct ma_state *mas_detach)
>  }
>
>  /*
> - * vmi_complete_munmap_vmas() - Finish the munmap() operation
> - * @vmi: The vma iterator
> - * @start: The start address
> - * @end: The end address
> - * @unlock: Unlock the mm or not
> - * @mas_detach: them maple state of the detached vma maple tree
> - * @locked_vm: The locked_vm count in the detached vmas
> - *
> - * This function updates the mm_struct, unmaps the region, frees the resources
> - * used for the munmap() and may downgrade the lock - if requested.  Everything
> - * needed to be done once the vma maple tree is updated.
> + * vms_complete_munmap_vmas() - Finish the munmap() operation
> + * @vms: The vma munmap struct
> + * @mas_detach: The maple state of the detached vmas
>   */
> -static void
> -vmi_complete_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
> -		struct mm_struct *mm, unsigned long start,
> -		unsigned long end, bool unlock, struct ma_state *mas_detach,
> -		unsigned long locked_vm)
> +
> +static void vms_complete_munmap_vmas(struct vma_munmap_struct *vms,
> +		struct ma_state *mas_detach)
>  {
>  	struct vm_area_struct *prev, *next;
> -	int count;
> +	struct mm_struct *mm;
>
> -	count = mas_detach->index + 1;
> -	mm->map_count -= count;
> -	mm->locked_vm -= locked_vm;
> -	if (unlock)
> +	mm = vms->mm;
> +	mm->map_count -= vms->vma_count;
> +	mm->locked_vm -= vms->locked_vm;
> +	if (vms->unlock)
>  		mmap_write_downgrade(mm);
>
> -	prev = vma_iter_prev_range(vmi);
> -	next = vma_next(vmi);
> +	prev = vma_iter_prev_range(vms->vmi);
> +	next = vma_next(vms->vmi);
>  	if (next)
> -		vma_iter_prev_range(vmi);
> +		vma_iter_prev_range(vms->vmi);
>
>  	/*
>  	 * We can free page tables without write-locking mmap_lock because VMAs
>  	 * were isolated before we downgraded mmap_lock.
>  	 */
>  	mas_set(mas_detach, 1);
> -	unmap_region(mm, mas_detach, vma, prev, next, start, end, count,
> -		     !unlock);
> +	unmap_region(mm, mas_detach, vms->vma, prev, next, vms->start, vms->end,
> +		     vms->vma_count, !vms->unlock);
>  	/* Statistics and freeing VMAs */
>  	mas_set(mas_detach, 0);
>  	remove_mt(mm, mas_detach);
>  	validate_mm(mm);
> -	if (unlock)
> +	if (vms->unlock)
>  		mmap_read_unlock(mm);
>
>  	__mt_destroy(mas_detach->tree);
>  }
>
>  /*
> - * vmi_gather_munmap_vmas() - Put all VMAs within a range into a maple tree
> + * vms_gather_munmap_vmas() - Put all VMAs within a range into a maple tree
>   * for removal at a later date.  Handles splitting first and last if necessary
>   * and marking the vmas as isolated.
>   *
> - * @vmi: The vma iterator
> - * @vma: The starting vm_area_struct
> - * @mm: The mm_struct
> - * @start: The aligned start address to munmap.
> - * @end: The aligned end address to munmap.
> - * @uf: The userfaultfd list_head
> + * @vms: The vma munmap struct
>   * @mas_detach: The maple state tracking the detached tree
>   *
>   * Return: 0 on success
>   */
> -static int
> -vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
> -		    struct mm_struct *mm, unsigned long start,
> -		    unsigned long end, struct list_head *uf,
> -		    struct ma_state *mas_detach, unsigned long *locked_vm)
> +static int vms_gather_munmap_vmas(struct vma_munmap_struct *vms,
> +		struct ma_state *mas_detach)
>  {

This is actually glorious.

>  	struct vm_area_struct *next = NULL;
> -	int count = 0;
>  	int error = -ENOMEM;
>
>  	/*
> @@ -2689,17 +2695,18 @@ vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  	 */
>
>  	/* Does it split the first one? */
> -	if (start > vma->vm_start) {
> +	if (vms->start > vms->vma->vm_start) {
>
>  		/*
>  		 * Make sure that map_count on return from munmap() will
>  		 * not exceed its limit; but let map_count go just above
>  		 * its limit temporarily, to help free resources as expected.
>  		 */
> -		if (end < vma->vm_end && mm->map_count >= sysctl_max_map_count)
> +		if (vms->end < vms->vma->vm_end &&
> +		    vms->mm->map_count >= sysctl_max_map_count)
>  			goto map_count_exceeded;
>
> -		error = __split_vma(vmi, vma, start, 1);
> +		error = __split_vma(vms->vmi, vms->vma, vms->start, 1);
>  		if (error)
>  			goto start_split_failed;
>  	}
> @@ -2708,24 +2715,24 @@ vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  	 * Detach a range of VMAs from the mm. Using next as a temp variable as
>  	 * it is always overwritten.
>  	 */
> -	next = vma;
> +	next = vms->vma;
>  	do {
>  		/* Does it split the end? */
> -		if (next->vm_end > end) {
> -			error = __split_vma(vmi, next, end, 0);
> +		if (next->vm_end > vms->end) {
> +			error = __split_vma(vms->vmi, next, vms->end, 0);
>  			if (error)
>  				goto end_split_failed;
>  		}
>  		vma_start_write(next);
> -		mas_set(mas_detach, count++);
> +		mas_set(mas_detach, vms->vma_count++);
>  		if (next->vm_flags & VM_LOCKED)
> -			*locked_vm += vma_pages(next);
> +			vms->locked_vm += vma_pages(next);
>
>  		error = mas_store_gfp(mas_detach, next, GFP_KERNEL);
>  		if (error)
>  			goto munmap_gather_failed;
>  		vma_mark_detached(next, true);
> -		if (unlikely(uf)) {
> +		if (unlikely(vms->uf)) {
>  			/*
>  			 * If userfaultfd_unmap_prep returns an error the vmas
>  			 * will remain split, but userland will get a
> @@ -2735,16 +2742,17 @@ vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  			 * split, despite we could. This is unlikely enough
>  			 * failure that it's not worth optimizing it for.
>  			 */
> -			error = userfaultfd_unmap_prep(next, start, end, uf);
> +			error = userfaultfd_unmap_prep(next, vms->start,
> +						       vms->end, vms->uf);
>
>  			if (error)
>  				goto userfaultfd_error;
>  		}
>  #ifdef CONFIG_DEBUG_VM_MAPLE_TREE
> -		BUG_ON(next->vm_start < start);
> -		BUG_ON(next->vm_start > end);
> +		BUG_ON(next->vm_start < vms->start);
> +		BUG_ON(next->vm_start > vms->end);
>  #endif
> -	} for_each_vma_range(*vmi, next, end);
> +	} for_each_vma_range(*(vms->vmi), next, vms->end);
>
>  #if defined(CONFIG_DEBUG_VM_MAPLE_TREE)
>  	/* Make sure no VMAs are about to be lost. */
> @@ -2753,21 +2761,21 @@ vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  		struct vm_area_struct *vma_mas, *vma_test;
>  		int test_count = 0;
>
> -		vma_iter_set(vmi, start);
> +		vma_iter_set(vms->vmi, vms->start);
>  		rcu_read_lock();
> -		vma_test = mas_find(&test, count - 1);
> -		for_each_vma_range(*vmi, vma_mas, end) {
> +		vma_test = mas_find(&test, vms->vma_count - 1);
> +		for_each_vma_range(*(vms->vmi), vma_mas, vms->end) {
>  			BUG_ON(vma_mas != vma_test);
>  			test_count++;
> -			vma_test = mas_next(&test, count - 1);
> +			vma_test = mas_next(&test, vms->vma_count - 1);
>  		}
>  		rcu_read_unlock();
> -		BUG_ON(count != test_count);
> +		BUG_ON(vms->vma_count != test_count);
>  	}
>  #endif
>
> -	while (vma_iter_addr(vmi) > start)
> -		vma_iter_prev_range(vmi);
> +	while (vma_iter_addr(vms->vmi) > vms->start)
> +		vma_iter_prev_range(vms->vmi);
>
>  	return 0;
>
> @@ -2803,11 +2811,11 @@ do_vmi_align_munmap(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  	MA_STATE(mas_detach, &mt_detach, 0, 0);
>  	mt_init_flags(&mt_detach, vmi->mas.tree->ma_flags & MT_FLAGS_LOCK_MASK);
>  	mt_on_stack(mt_detach);
> +	struct vma_munmap_struct vms;
>  	int error;
> -	unsigned long locked_vm = 0;
>
> -	error = vmi_gather_munmap_vmas(vmi, vma, mm, start, end, uf,
> -				       &mas_detach, &locked_vm);
> +	init_vma_munmap(&vms, vmi, vma, start, end, uf, unlock);
> +	error = vms_gather_munmap_vmas(&vms, &mas_detach);
>  	if (error)
>  		goto gather_failed;
>
> @@ -2816,8 +2824,7 @@ do_vmi_align_munmap(struct vma_iterator *vmi, struct vm_area_struct *vma,
>  		goto clear_tree_failed;
>
>  	/* Point of no return */
> -	vmi_complete_munmap_vmas(vmi, vma, mm, start, end, unlock, &mas_detach,
> -				 locked_vm);
> +	vms_complete_munmap_vmas(&vms, &mas_detach);
>  	return 0;
>
>  clear_tree_failed:
> --
> 2.43.0
>

Very nice refactoring.

Reviewed-by: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>
Liam R. Howlett July 5, 2024, 7:09 p.m. UTC | #2
* Lorenzo Stoakes <lorenzo.stoakes@oracle.com> [240705 14:39]:
> On Thu, Jul 04, 2024 at 02:27:07PM GMT, Liam R. Howlett wrote:
> > Use a structure to pass along all the necessary information and counters
> > involved in removing vmas from the mm_struct.
> >
> > Update vmi_ function names to vms_ to indicate the first argument
> > type change.
> >
> > Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
> > Reviewed-by: Suren Baghdasaryan <surenb@google.com>
> > ---
> >  mm/internal.h |  16 ++++++
> >  mm/mmap.c     | 137 ++++++++++++++++++++++++++------------------------
> >  2 files changed, 88 insertions(+), 65 deletions(-)
> >
> > diff --git a/mm/internal.h b/mm/internal.h
> > index 2ea9a88dcb95..f1e6dea2efcf 100644
> > --- a/mm/internal.h
> > +++ b/mm/internal.h
> > @@ -1481,6 +1481,22 @@ struct vma_prepare {
> >  	struct vm_area_struct *remove2;
> >  };
> >
> > +/*
> > + * vma munmap operation
> > + */
> > +struct vma_munmap_struct {
> > +	struct vma_iterator *vmi;
> > +	struct mm_struct *mm;
> > +	struct vm_area_struct *vma;	/* The first vma to munmap */
> > +	struct list_head *uf;		/* Userfaultfd list_head */
> > +	unsigned long start;		/* Aligned start addr */
> > +	unsigned long end;		/* Aligned end addr */
> > +	int vma_count;			/* Number of vmas that will be removed */
> > +	unsigned long nr_pages;		/* Number of pages being removed */
> > +	unsigned long locked_vm;	/* Number of locked pages */
> > +	bool unlock;			/* Unlock after the munmap */
> > +};
> 
> 
> I'm a big fan of breaking out and threading state like this through some of
> these more... verbose VMA functions.
> 
> I have a similar idea as part of my long dreamed of 'delete vma_merge()'
> patch set. Coming soon :)
> 
> > +
> >  void __meminit __init_single_page(struct page *page, unsigned long pfn,
> >  				unsigned long zone, int nid);
> >
> > diff --git a/mm/mmap.c b/mm/mmap.c
> > index 8dc8ffbf9d8d..76e93146ee9d 100644
> > --- a/mm/mmap.c
> > +++ b/mm/mmap.c
> > @@ -500,6 +500,31 @@ static inline void init_vma_prep(struct vma_prepare *vp,
> >  	init_multi_vma_prep(vp, vma, NULL, NULL, NULL);
> >  }
> >
> > +/*
> > + * init_vma_munmap() - Initializer wrapper for vma_munmap_struct
> > + * @vms: The vma munmap struct
> > + * @vmi: The vma iterator
> > + * @vma: The first vm_area_struct to munmap
> > + * @start: The aligned start address to munmap
> > + * @end: The aligned end address to munmap
> 
> Maybe worth mentioning if inclusive/exclusive.

The "address to munmap" isn't specific enough that we are using the same
logic as the munmap call?  That is, the vma inclusive and exclusive for
start and end, respectively.

Not a big change, either way.

...

> 
> Reviewed-by: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>

Thanks
diff mbox series

Patch

diff --git a/mm/internal.h b/mm/internal.h
index 2ea9a88dcb95..f1e6dea2efcf 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1481,6 +1481,22 @@  struct vma_prepare {
 	struct vm_area_struct *remove2;
 };
 
+/*
+ * vma munmap operation
+ */
+struct vma_munmap_struct {
+	struct vma_iterator *vmi;
+	struct mm_struct *mm;
+	struct vm_area_struct *vma;	/* The first vma to munmap */
+	struct list_head *uf;		/* Userfaultfd list_head */
+	unsigned long start;		/* Aligned start addr */
+	unsigned long end;		/* Aligned end addr */
+	int vma_count;			/* Number of vmas that will be removed */
+	unsigned long nr_pages;		/* Number of pages being removed */
+	unsigned long locked_vm;	/* Number of locked pages */
+	bool unlock;			/* Unlock after the munmap */
+};
+
 void __meminit __init_single_page(struct page *page, unsigned long pfn,
 				unsigned long zone, int nid);
 
diff --git a/mm/mmap.c b/mm/mmap.c
index 8dc8ffbf9d8d..76e93146ee9d 100644
--- a/mm/mmap.c
+++ b/mm/mmap.c
@@ -500,6 +500,31 @@  static inline void init_vma_prep(struct vma_prepare *vp,
 	init_multi_vma_prep(vp, vma, NULL, NULL, NULL);
 }
 
+/*
+ * init_vma_munmap() - Initializer wrapper for vma_munmap_struct
+ * @vms: The vma munmap struct
+ * @vmi: The vma iterator
+ * @vma: The first vm_area_struct to munmap
+ * @start: The aligned start address to munmap
+ * @end: The aligned end address to munmap
+ * @uf: The userfaultfd list_head
+ * @unlock: Unlock after the operation.  Only unlocked on success
+ */
+static inline void init_vma_munmap(struct vma_munmap_struct *vms,
+		struct vma_iterator *vmi, struct vm_area_struct *vma,
+		unsigned long start, unsigned long end, struct list_head *uf,
+		bool unlock)
+{
+	vms->vmi = vmi;
+	vms->vma = vma;
+	vms->mm = vma->vm_mm;
+	vms->start = start;
+	vms->end = end;
+	vms->unlock = unlock;
+	vms->uf = uf;
+	vms->vma_count = 0;
+	vms->nr_pages = vms->locked_vm = 0;
+}
 
 /*
  * vma_prepare() - Helper function for handling locking VMAs prior to altering
@@ -2606,78 +2631,59 @@  static inline void abort_munmap_vmas(struct ma_state *mas_detach)
 }
 
 /*
- * vmi_complete_munmap_vmas() - Finish the munmap() operation
- * @vmi: The vma iterator
- * @start: The start address
- * @end: The end address
- * @unlock: Unlock the mm or not
- * @mas_detach: them maple state of the detached vma maple tree
- * @locked_vm: The locked_vm count in the detached vmas
- *
- * This function updates the mm_struct, unmaps the region, frees the resources
- * used for the munmap() and may downgrade the lock - if requested.  Everything
- * needed to be done once the vma maple tree is updated.
+ * vms_complete_munmap_vmas() - Finish the munmap() operation
+ * @vms: The vma munmap struct
+ * @mas_detach: The maple state of the detached vmas
  */
-static void
-vmi_complete_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
-		struct mm_struct *mm, unsigned long start,
-		unsigned long end, bool unlock, struct ma_state *mas_detach,
-		unsigned long locked_vm)
+
+static void vms_complete_munmap_vmas(struct vma_munmap_struct *vms,
+		struct ma_state *mas_detach)
 {
 	struct vm_area_struct *prev, *next;
-	int count;
+	struct mm_struct *mm;
 
-	count = mas_detach->index + 1;
-	mm->map_count -= count;
-	mm->locked_vm -= locked_vm;
-	if (unlock)
+	mm = vms->mm;
+	mm->map_count -= vms->vma_count;
+	mm->locked_vm -= vms->locked_vm;
+	if (vms->unlock)
 		mmap_write_downgrade(mm);
 
-	prev = vma_iter_prev_range(vmi);
-	next = vma_next(vmi);
+	prev = vma_iter_prev_range(vms->vmi);
+	next = vma_next(vms->vmi);
 	if (next)
-		vma_iter_prev_range(vmi);
+		vma_iter_prev_range(vms->vmi);
 
 	/*
 	 * We can free page tables without write-locking mmap_lock because VMAs
 	 * were isolated before we downgraded mmap_lock.
 	 */
 	mas_set(mas_detach, 1);
-	unmap_region(mm, mas_detach, vma, prev, next, start, end, count,
-		     !unlock);
+	unmap_region(mm, mas_detach, vms->vma, prev, next, vms->start, vms->end,
+		     vms->vma_count, !vms->unlock);
 	/* Statistics and freeing VMAs */
 	mas_set(mas_detach, 0);
 	remove_mt(mm, mas_detach);
 	validate_mm(mm);
-	if (unlock)
+	if (vms->unlock)
 		mmap_read_unlock(mm);
 
 	__mt_destroy(mas_detach->tree);
 }
 
 /*
- * vmi_gather_munmap_vmas() - Put all VMAs within a range into a maple tree
+ * vms_gather_munmap_vmas() - Put all VMAs within a range into a maple tree
  * for removal at a later date.  Handles splitting first and last if necessary
  * and marking the vmas as isolated.
  *
- * @vmi: The vma iterator
- * @vma: The starting vm_area_struct
- * @mm: The mm_struct
- * @start: The aligned start address to munmap.
- * @end: The aligned end address to munmap.
- * @uf: The userfaultfd list_head
+ * @vms: The vma munmap struct
  * @mas_detach: The maple state tracking the detached tree
  *
  * Return: 0 on success
  */
-static int
-vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
-		    struct mm_struct *mm, unsigned long start,
-		    unsigned long end, struct list_head *uf,
-		    struct ma_state *mas_detach, unsigned long *locked_vm)
+static int vms_gather_munmap_vmas(struct vma_munmap_struct *vms,
+		struct ma_state *mas_detach)
 {
 	struct vm_area_struct *next = NULL;
-	int count = 0;
 	int error = -ENOMEM;
 
 	/*
@@ -2689,17 +2695,18 @@  vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
 	 */
 
 	/* Does it split the first one? */
-	if (start > vma->vm_start) {
+	if (vms->start > vms->vma->vm_start) {
 
 		/*
 		 * Make sure that map_count on return from munmap() will
 		 * not exceed its limit; but let map_count go just above
 		 * its limit temporarily, to help free resources as expected.
 		 */
-		if (end < vma->vm_end && mm->map_count >= sysctl_max_map_count)
+		if (vms->end < vms->vma->vm_end &&
+		    vms->mm->map_count >= sysctl_max_map_count)
 			goto map_count_exceeded;
 
-		error = __split_vma(vmi, vma, start, 1);
+		error = __split_vma(vms->vmi, vms->vma, vms->start, 1);
 		if (error)
 			goto start_split_failed;
 	}
@@ -2708,24 +2715,24 @@  vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
 	 * Detach a range of VMAs from the mm. Using next as a temp variable as
 	 * it is always overwritten.
 	 */
-	next = vma;
+	next = vms->vma;
 	do {
 		/* Does it split the end? */
-		if (next->vm_end > end) {
-			error = __split_vma(vmi, next, end, 0);
+		if (next->vm_end > vms->end) {
+			error = __split_vma(vms->vmi, next, vms->end, 0);
 			if (error)
 				goto end_split_failed;
 		}
 		vma_start_write(next);
-		mas_set(mas_detach, count++);
+		mas_set(mas_detach, vms->vma_count++);
 		if (next->vm_flags & VM_LOCKED)
-			*locked_vm += vma_pages(next);
+			vms->locked_vm += vma_pages(next);
 
 		error = mas_store_gfp(mas_detach, next, GFP_KERNEL);
 		if (error)
 			goto munmap_gather_failed;
 		vma_mark_detached(next, true);
-		if (unlikely(uf)) {
+		if (unlikely(vms->uf)) {
 			/*
 			 * If userfaultfd_unmap_prep returns an error the vmas
 			 * will remain split, but userland will get a
@@ -2735,16 +2742,17 @@  vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
 			 * split, despite we could. This is unlikely enough
 			 * failure that it's not worth optimizing it for.
 			 */
-			error = userfaultfd_unmap_prep(next, start, end, uf);
+			error = userfaultfd_unmap_prep(next, vms->start,
+						       vms->end, vms->uf);
 
 			if (error)
 				goto userfaultfd_error;
 		}
 #ifdef CONFIG_DEBUG_VM_MAPLE_TREE
-		BUG_ON(next->vm_start < start);
-		BUG_ON(next->vm_start > end);
+		BUG_ON(next->vm_start < vms->start);
+		BUG_ON(next->vm_start > vms->end);
 #endif
-	} for_each_vma_range(*vmi, next, end);
+	} for_each_vma_range(*(vms->vmi), next, vms->end);
 
 #if defined(CONFIG_DEBUG_VM_MAPLE_TREE)
 	/* Make sure no VMAs are about to be lost. */
@@ -2753,21 +2761,21 @@  vmi_gather_munmap_vmas(struct vma_iterator *vmi, struct vm_area_struct *vma,
 		struct vm_area_struct *vma_mas, *vma_test;
 		int test_count = 0;
 
-		vma_iter_set(vmi, start);
+		vma_iter_set(vms->vmi, vms->start);
 		rcu_read_lock();
-		vma_test = mas_find(&test, count - 1);
-		for_each_vma_range(*vmi, vma_mas, end) {
+		vma_test = mas_find(&test, vms->vma_count - 1);
+		for_each_vma_range(*(vms->vmi), vma_mas, vms->end) {
 			BUG_ON(vma_mas != vma_test);
 			test_count++;
-			vma_test = mas_next(&test, count - 1);
+			vma_test = mas_next(&test, vms->vma_count - 1);
 		}
 		rcu_read_unlock();
-		BUG_ON(count != test_count);
+		BUG_ON(vms->vma_count != test_count);
 	}
 #endif
 
-	while (vma_iter_addr(vmi) > start)
-		vma_iter_prev_range(vmi);
+	while (vma_iter_addr(vms->vmi) > vms->start)
+		vma_iter_prev_range(vms->vmi);
 
 	return 0;
 
@@ -2803,11 +2811,11 @@  do_vmi_align_munmap(struct vma_iterator *vmi, struct vm_area_struct *vma,
 	MA_STATE(mas_detach, &mt_detach, 0, 0);
 	mt_init_flags(&mt_detach, vmi->mas.tree->ma_flags & MT_FLAGS_LOCK_MASK);
 	mt_on_stack(mt_detach);
+	struct vma_munmap_struct vms;
 	int error;
-	unsigned long locked_vm = 0;
 
-	error = vmi_gather_munmap_vmas(vmi, vma, mm, start, end, uf,
-				       &mas_detach, &locked_vm);
+	init_vma_munmap(&vms, vmi, vma, start, end, uf, unlock);
+	error = vms_gather_munmap_vmas(&vms, &mas_detach);
 	if (error)
 		goto gather_failed;
 
@@ -2816,8 +2824,7 @@  do_vmi_align_munmap(struct vma_iterator *vmi, struct vm_area_struct *vma,
 		goto clear_tree_failed;
 
 	/* Point of no return */
-	vmi_complete_munmap_vmas(vmi, vma, mm, start, end, unlock, &mas_detach,
-				 locked_vm);
+	vms_complete_munmap_vmas(&vms, &mas_detach);
 	return 0;
 
 clear_tree_failed: