diff mbox series

[3/5] mm: Rework return value for copy_one_pte()

Message ID 20200921211744.24758-4-peterx@redhat.com (mailing list archive)
State New, archived
Headers show
Series mm: Break COW for pinned pages during fork() | expand

Commit Message

Peter Xu Sept. 21, 2020, 9:17 p.m. UTC
There's one special path for copy_one_pte() with swap entries, in which
add_swap_count_continuation(GFP_ATOMIC) might fail.  In that case we'll return
the swp_entry_t so that the caller will release the locks and redo the same
thing with GFP_KERNEL.

It's confusing when copy_one_pte() must return a swp_entry_t (even if all the
ptes are non-swap entries).  More importantly, we face other requirement to
extend this "we need to do something else, but without the locks held" case.

Rework the return value into something easier to understand, as defined in enum
copy_mm_ret.  We'll pass the swp_entry_t back using the newly introduced union
copy_mm_data parameter.

Another trivial change is to move the reset of the "progress" counter into the
retry path, so that we'll reset it for other reasons too.

This should prepare us with adding new return codes, very soon.

Signed-off-by: Peter Xu <peterx@redhat.com>
---
 mm/memory.c | 42 +++++++++++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 13 deletions(-)

Comments

John Hubbard Sept. 22, 2020, 7:11 a.m. UTC | #1
On 9/21/20 2:17 PM, Peter Xu wrote:
> There's one special path for copy_one_pte() with swap entries, in which
> add_swap_count_continuation(GFP_ATOMIC) might fail.  In that case we'll return

I might be looking at the wrong place, but the existing code seems to call
add_swap_count_continuation(GFP_KERNEL), not with GFP_ATOMIC?

> the swp_entry_t so that the caller will release the locks and redo the same
> thing with GFP_KERNEL.
> 
> It's confusing when copy_one_pte() must return a swp_entry_t (even if all the
> ptes are non-swap entries).  More importantly, we face other requirement to
> extend this "we need to do something else, but without the locks held" case.
> 
> Rework the return value into something easier to understand, as defined in enum
> copy_mm_ret.  We'll pass the swp_entry_t back using the newly introduced union

I like the documentation here, but it doesn't match what you did in the patch.
Actually, the documentation had the right idea (enum, rather than #define, for
COPY_MM_* items). Below...

> copy_mm_data parameter.
> 
> Another trivial change is to move the reset of the "progress" counter into the
> retry path, so that we'll reset it for other reasons too.
> 
> This should prepare us with adding new return codes, very soon.
> 
> Signed-off-by: Peter Xu <peterx@redhat.com>
> ---
>   mm/memory.c | 42 +++++++++++++++++++++++++++++-------------
>   1 file changed, 29 insertions(+), 13 deletions(-)
> 
> diff --git a/mm/memory.c b/mm/memory.c
> index 7525147908c4..1530bb1070f4 100644
> --- a/mm/memory.c
> +++ b/mm/memory.c
> @@ -689,16 +689,24 @@ struct page *vm_normal_page_pmd(struct vm_area_struct *vma, unsigned long addr,
>   }
>   #endif
>   
> +#define  COPY_MM_DONE               0
> +#define  COPY_MM_SWAP_CONT          1

Those should be enums, so as to get a little type safety and other goodness from
using non-macro items.

...
> @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
>   	pte_unmap_unlock(orig_dst_pte, dst_ptl);
>   	cond_resched();
>   
> -	if (entry.val) {
> -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> +	switch (copy_ret) {
> +	case COPY_MM_SWAP_CONT:
> +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
>   			return -ENOMEM;
> -		progress = 0;

Yes. Definitely a little cleaner to reset this above, instead of here.

> +		break;
> +	default:
> +		break;

I assume this no-op noise is to placate the compiler and/or static checkers. :)

I'm unable to find any actual problems with the diffs, aside from the nit about
using an enum.

thanks,
Oleg Nesterov Sept. 22, 2020, 10:08 a.m. UTC | #2
On 09/21, Peter Xu wrote:
>
> @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
>  	pte_unmap_unlock(orig_dst_pte, dst_ptl);
>  	cond_resched();
>  
> -	if (entry.val) {
> -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> +	switch (copy_ret) {
> +	case COPY_MM_SWAP_CONT:
> +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
>  			return -ENOMEM;
> -		progress = 0;
> +		break;

Note that you didn't clear copy_ret, it is still COPY_MM_SWAP_CONT,

> +	default:
> +		break;
>  	}
> +
>  	if (addr != end)
>  		goto again;

After that the main loop can stop again because of need_resched(), and
in this case add_swap_count_continuation(data.entry) will be called again?

Oleg.
Oleg Nesterov Sept. 22, 2020, 10:18 a.m. UTC | #3
On 09/22, Oleg Nesterov wrote:
>
> On 09/21, Peter Xu wrote:
> >
> > @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
> >  	pte_unmap_unlock(orig_dst_pte, dst_ptl);
> >  	cond_resched();
> >
> > -	if (entry.val) {
> > -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> > +	switch (copy_ret) {
> > +	case COPY_MM_SWAP_CONT:
> > +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
> >  			return -ENOMEM;
> > -		progress = 0;
> > +		break;
>
> Note that you didn't clear copy_ret, it is still COPY_MM_SWAP_CONT,
>
> > +	default:
> > +		break;
> >  	}
> > +
> >  	if (addr != end)
> >  		goto again;
>
> After that the main loop can stop again because of need_resched(), and
> in this case add_swap_count_continuation(data.entry) will be called again?

No, this is not possible, copy_one_pte() should be called at least once,
progress = 0 before restart. Sorry for noise.

Oleg.
Peter Xu Sept. 22, 2020, 3:29 p.m. UTC | #4
On Tue, Sep 22, 2020 at 12:11:29AM -0700, John Hubbard wrote:
> On 9/21/20 2:17 PM, Peter Xu wrote:
> > There's one special path for copy_one_pte() with swap entries, in which
> > add_swap_count_continuation(GFP_ATOMIC) might fail.  In that case we'll return
> 
> I might be looking at the wrong place, but the existing code seems to call
> add_swap_count_continuation(GFP_KERNEL), not with GFP_ATOMIC?

Ah, I wanted to reference the one in swap_duplicate().

> 
> > the swp_entry_t so that the caller will release the locks and redo the same
> > thing with GFP_KERNEL.
> > 
> > It's confusing when copy_one_pte() must return a swp_entry_t (even if all the
> > ptes are non-swap entries).  More importantly, we face other requirement to
> > extend this "we need to do something else, but without the locks held" case.
> > 
> > Rework the return value into something easier to understand, as defined in enum
> > copy_mm_ret.  We'll pass the swp_entry_t back using the newly introduced union
> 
> I like the documentation here, but it doesn't match what you did in the patch.
> Actually, the documentation had the right idea (enum, rather than #define, for
> COPY_MM_* items). Below...

Yeah actually my very initial version has it as an enum, then I changed it to
macros because I started to want it return negative as errors.  However funnily
in the current version copy_one_pte() won't return an error anymore... So
probably, yes, it should be a good idea to get the enum back.

Also we should be able to drop the negative handling too with copy_ret, though
it should be in the next patch.

> 
> > copy_mm_data parameter.
> > 
> > Another trivial change is to move the reset of the "progress" counter into the
> > retry path, so that we'll reset it for other reasons too.
> > 
> > This should prepare us with adding new return codes, very soon.
> > 
> > Signed-off-by: Peter Xu <peterx@redhat.com>
> > ---
> >   mm/memory.c | 42 +++++++++++++++++++++++++++++-------------
> >   1 file changed, 29 insertions(+), 13 deletions(-)
> > 
> > diff --git a/mm/memory.c b/mm/memory.c
> > index 7525147908c4..1530bb1070f4 100644
> > --- a/mm/memory.c
> > +++ b/mm/memory.c
> > @@ -689,16 +689,24 @@ struct page *vm_normal_page_pmd(struct vm_area_struct *vma, unsigned long addr,
> >   }
> >   #endif
> > +#define  COPY_MM_DONE               0
> > +#define  COPY_MM_SWAP_CONT          1
> 
> Those should be enums, so as to get a little type safety and other goodness from
> using non-macro items.
> 
> ...
> > @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
> >   	pte_unmap_unlock(orig_dst_pte, dst_ptl);
> >   	cond_resched();
> > -	if (entry.val) {
> > -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> > +	switch (copy_ret) {
> > +	case COPY_MM_SWAP_CONT:
> > +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
> >   			return -ENOMEM;
> > -		progress = 0;
> 
> Yes. Definitely a little cleaner to reset this above, instead of here.
> 
> > +		break;
> > +	default:
> > +		break;
> 
> I assume this no-op noise is to placate the compiler and/or static checkers. :)

This is (so far) for COPY_MM_DONE.  I normally will cover all cases in a
"switch()" and here "default" is for it.  Even if I covered all the
possibilities, I may still tend to keep one "default" and a WARN_ON_ONCE(1) to
make sure nothing I've missed.  Not sure whether that's the ideal way, though.

Thanks,
Peter Xu Sept. 22, 2020, 3:36 p.m. UTC | #5
On Tue, Sep 22, 2020 at 12:18:16PM +0200, Oleg Nesterov wrote:
> On 09/22, Oleg Nesterov wrote:
> >
> > On 09/21, Peter Xu wrote:
> > >
> > > @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
> > >  	pte_unmap_unlock(orig_dst_pte, dst_ptl);
> > >  	cond_resched();
> > >
> > > -	if (entry.val) {
> > > -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> > > +	switch (copy_ret) {
> > > +	case COPY_MM_SWAP_CONT:
> > > +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
> > >  			return -ENOMEM;
> > > -		progress = 0;
> > > +		break;
> >
> > Note that you didn't clear copy_ret, it is still COPY_MM_SWAP_CONT,
> >
> > > +	default:
> > > +		break;
> > >  	}
> > > +
> > >  	if (addr != end)
> > >  		goto again;
> >
> > After that the main loop can stop again because of need_resched(), and
> > in this case add_swap_count_continuation(data.entry) will be called again?
> 
> No, this is not possible, copy_one_pte() should be called at least once,
> progress = 0 before restart. Sorry for noise.

Oh wait, I think you're right... when we get a COPY_MM_SWAP_CONT, goto "again",
then if there're 32 pte_none() ptes _plus_ an need_resched(), then we might
reach again at the same add_swap_count_continuation() with the same swp entry.

However since I didn't change this logic in this patch, it probably means this
bug is also in the original code before this series...  I'm thinking maybe I
should prepare a standalone patch to clear the swp_entry_t and cc stable.

Thanks,
Oleg Nesterov Sept. 22, 2020, 3:48 p.m. UTC | #6
On 09/22, Peter Xu wrote:
>
> On Tue, Sep 22, 2020 at 12:18:16PM +0200, Oleg Nesterov wrote:
> > On 09/22, Oleg Nesterov wrote:
> > >
> > > On 09/21, Peter Xu wrote:
> > > >
> > > > @@ -866,13 +877,18 @@ static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
> > > >  	pte_unmap_unlock(orig_dst_pte, dst_ptl);
> > > >  	cond_resched();
> > > >
> > > > -	if (entry.val) {
> > > > -		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
> > > > +	switch (copy_ret) {
> > > > +	case COPY_MM_SWAP_CONT:
> > > > +		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
> > > >  			return -ENOMEM;
> > > > -		progress = 0;
> > > > +		break;
> > >
> > > Note that you didn't clear copy_ret, it is still COPY_MM_SWAP_CONT,
> > >
> > > > +	default:
> > > > +		break;
> > > >  	}
> > > > +
> > > >  	if (addr != end)
> > > >  		goto again;
> > >
> > > After that the main loop can stop again because of need_resched(), and
> > > in this case add_swap_count_continuation(data.entry) will be called again?
> >
> > No, this is not possible, copy_one_pte() should be called at least once,
> > progress = 0 before restart. Sorry for noise.
>
> Oh wait, I think you're right... when we get a COPY_MM_SWAP_CONT, goto "again",
> then if there're 32 pte_none() ptes _plus_ an need_resched(), then we might
> reach again at the same add_swap_count_continuation() with the same swp entry.

Yes, please see my reply to 4/5 ;)

> However since I didn't change this logic in this patch, it probably means this
> bug is also in the original code before this series...  I'm thinking maybe I
> should prepare a standalone patch to clear the swp_entry_t and cc stable.

Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
will be called at least once.

So _think_ that the current code is fine, but I can be easily wrong and I agree
this doesn't look clean.

Oleg.
Peter Xu Sept. 22, 2020, 4:03 p.m. UTC | #7
On Tue, Sep 22, 2020 at 05:48:46PM +0200, Oleg Nesterov wrote:
> > However since I didn't change this logic in this patch, it probably means this
> > bug is also in the original code before this series...  I'm thinking maybe I
> > should prepare a standalone patch to clear the swp_entry_t and cc stable.
> 
> Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
> pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
> will be called at least once.

Note that we've released the page table locks, so afaict the old swp entry can
be gone under us when we go back to the "do" loop... :) Extremely corner case,
but maybe still good to fix, extra clearness as a (good) side effect.
Oleg Nesterov Sept. 22, 2020, 4:53 p.m. UTC | #8
On 09/22, Peter Xu wrote:
>
> On Tue, Sep 22, 2020 at 05:48:46PM +0200, Oleg Nesterov wrote:
> > > However since I didn't change this logic in this patch, it probably means this
> > > bug is also in the original code before this series...  I'm thinking maybe I
> > > should prepare a standalone patch to clear the swp_entry_t and cc stable.
> >
> > Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
> > pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
> > will be called at least once.
>
> Note that we've released the page table locks, so afaict the old swp entry can
> be gone under us when we go back to the "do" loop... :)

But how?

I am just curious, I don't understand this code enough.

Oleg.
Peter Xu Sept. 22, 2020, 6:13 p.m. UTC | #9
On Tue, Sep 22, 2020 at 06:53:55PM +0200, Oleg Nesterov wrote:
> On 09/22, Peter Xu wrote:
> >
> > On Tue, Sep 22, 2020 at 05:48:46PM +0200, Oleg Nesterov wrote:
> > > > However since I didn't change this logic in this patch, it probably means this
> > > > bug is also in the original code before this series...  I'm thinking maybe I
> > > > should prepare a standalone patch to clear the swp_entry_t and cc stable.
> > >
> > > Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
> > > pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
> > > will be called at least once.
> >
> > Note that we've released the page table locks, so afaict the old swp entry can
> > be gone under us when we go back to the "do" loop... :)
> 
> But how?
> 
> I am just curious, I don't understand this code enough.

Me neither.

The point is I think we can't assume *src_pte will read the same if we have
released the src_ptl in copy_pte_range(), because imho the src_ptl is the only
thing to protect it.  Or to be more explicit, we need pte_alloc_map_lock() to
read a stable pmd/pte or before update (since src_ptl itself could change too).

Thanks,
Oleg Nesterov Sept. 22, 2020, 6:23 p.m. UTC | #10
On 09/22, Peter Xu wrote:
>
> On Tue, Sep 22, 2020 at 06:53:55PM +0200, Oleg Nesterov wrote:
> > On 09/22, Peter Xu wrote:
> > >
> > > On Tue, Sep 22, 2020 at 05:48:46PM +0200, Oleg Nesterov wrote:
> > > > > However since I didn't change this logic in this patch, it probably means this
> > > > > bug is also in the original code before this series...  I'm thinking maybe I
> > > > > should prepare a standalone patch to clear the swp_entry_t and cc stable.
> > > >
> > > > Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
> > > > pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
> > > > will be called at least once.
> > >
> > > Note that we've released the page table locks, so afaict the old swp entry can
> > > be gone under us when we go back to the "do" loop... :)
> >
> > But how?
> >
> > I am just curious, I don't understand this code enough.
>
> Me neither.
>
> The point is I think we can't assume *src_pte will read the same if we have
> released the src_ptl in copy_pte_range(),

This is clear.

But I still think that !pte_none() -> pte_none() transition is not possible
under mmap_write_lock()...

OK, let me repeat I don't understans these code paths enough, let me reword:
I don't see how this transition is possible.

Oleg.
Peter Xu Sept. 22, 2020, 6:49 p.m. UTC | #11
On Tue, Sep 22, 2020 at 08:23:18PM +0200, Oleg Nesterov wrote:
> On 09/22, Peter Xu wrote:
> >
> > On Tue, Sep 22, 2020 at 06:53:55PM +0200, Oleg Nesterov wrote:
> > > On 09/22, Peter Xu wrote:
> > > >
> > > > On Tue, Sep 22, 2020 at 05:48:46PM +0200, Oleg Nesterov wrote:
> > > > > > However since I didn't change this logic in this patch, it probably means this
> > > > > > bug is also in the original code before this series...  I'm thinking maybe I
> > > > > > should prepare a standalone patch to clear the swp_entry_t and cc stable.
> > > > >
> > > > > Well, if copy_one_pte(src_pte) hits a swap entry and returns entry.val != 0, then
> > > > > pte_none(*src_pte) is not possible after restart? This means that copy_one_pte()
> > > > > will be called at least once.
> > > >
> > > > Note that we've released the page table locks, so afaict the old swp entry can
> > > > be gone under us when we go back to the "do" loop... :)
> > >
> > > But how?
> > >
> > > I am just curious, I don't understand this code enough.
> >
> > Me neither.
> >
> > The point is I think we can't assume *src_pte will read the same if we have
> > released the src_ptl in copy_pte_range(),
> 
> This is clear.
> 
> But I still think that !pte_none() -> pte_none() transition is not possible
> under mmap_write_lock()...
> 
> OK, let me repeat I don't understans these code paths enough, let me reword:
> I don't see how this transition is possible.

Though I guess I'll keep my wording, because I still think it's accurate to
me. :)

Can we e.g. punch a page hole without changing vmas?
Oleg Nesterov Sept. 23, 2020, 6:52 a.m. UTC | #12
On 09/22, Peter Xu wrote:
>
> On Tue, Sep 22, 2020 at 08:23:18PM +0200, Oleg Nesterov wrote:
> >
> > But I still think that !pte_none() -> pte_none() transition is not possible
> > under mmap_write_lock()...
> >
> > OK, let me repeat I don't understans these code paths enough, let me reword:
> > I don't see how this transition is possible.
>
> Though I guess I'll keep my wording, because I still think it's accurate to
> me. :)
>
> Can we e.g. punch a page hole without changing vmas?

punch a hole? I don't know what does it mean...

However, I think you are right anyway. I forgot that (at least) truncate can
clear this pte without mmap_sem after pte_unmap_unlock().

So I think you are right, the current code is wrong too.

Thanks!

Oleg.
Linus Torvalds Sept. 23, 2020, 5:16 p.m. UTC | #13
On Mon, Sep 21, 2020 at 2:18 PM Peter Xu <peterx@redhat.com> wrote:
>
> There's one special path for copy_one_pte() with swap entries, in which
> add_swap_count_continuation(GFP_ATOMIC) might fail.  In that case we'll return
> the swp_entry_t so that the caller will release the locks and redo the same
> thing with GFP_KERNEL.
>
> It's confusing when copy_one_pte() must return a swp_entry_t (even if all the
> ptes are non-swap entries).  More importantly, we face other requirement to
> extend this "we need to do something else, but without the locks held" case.
>
> Rework the return value into something easier to understand, as defined in enum
> copy_mm_ret.  We'll pass the swp_entry_t back using the newly introduced union
> copy_mm_data parameter.

Ok, I'm reading this series, and I do hate this.

And I think it's unnecessary.

There's a very simple way to avoid this all: split out the
"!pte_present(pte)" case from the function entirely.

That actually makes the code much more legible: that non-present case
is very different, and it's also unlikely() and causes deeper
indentation etc.

Because it's unlikely, it probably also shouldn't be inline.

That unlikely case is also why when then have that special
"out_set_pte" label, which should just go away and be copied into the
(now uninlined) function.

Once that re-organization has been done, the second step is to then
just move the "pte_present()" check into the caller, and suddenly all
the ugly return value games will go entirely away.

I'm attaching the two patches that do this here, but I do want to note
how that first patch is much more legible with "--ignore-all-space",
and then you really see that the diff is a _pure_ code movement thing.
Otherwise it looks like it's doing a big change.

Comments?

NOTE! The intent here is that now we can easily add new argument (a
pre-allocated page or NULL) and a return value to
"copy_present_page()": it can return "I needed a temporary page but
you hadn't allocated one yet" or "I used up the temporary page you
gave me" or "all good, keep the temporary page around for the future".

But these two patches are very intentionally meant to be just "this
clearly changes NO semantics at all".

                   Linus
Linus Torvalds Sept. 23, 2020, 9:24 p.m. UTC | #14
On Wed, Sep 23, 2020 at 10:16 AM Linus Torvalds
<torvalds@linux-foundation.org> wrote:
>
> But these two patches are very intentionally meant to be just "this
> clearly changes NO semantics at all".

The more I look at these, the more I go "this is a cleanup
regardless", so I'll just keep thes in my tree as-is.

                 Linus
diff mbox series

Patch

diff --git a/mm/memory.c b/mm/memory.c
index 7525147908c4..1530bb1070f4 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -689,16 +689,24 @@  struct page *vm_normal_page_pmd(struct vm_area_struct *vma, unsigned long addr,
 }
 #endif
 
+#define  COPY_MM_DONE               0
+#define  COPY_MM_SWAP_CONT          1
+
+struct copy_mm_data {
+	/* COPY_MM_SWAP_CONT */
+	swp_entry_t entry;
+};
+
 /*
  * copy one vm_area from one task to the other. Assumes the page tables
  * already present in the new task to be cleared in the whole range
  * covered by this vma.
  */
 
-static inline unsigned long
+static inline int
 copy_one_pte(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 		pte_t *dst_pte, pte_t *src_pte, struct vm_area_struct *vma,
-		unsigned long addr, int *rss)
+		unsigned long addr, int *rss, struct copy_mm_data *data)
 {
 	unsigned long vm_flags = vma->vm_flags;
 	pte_t pte = *src_pte;
@@ -709,8 +717,10 @@  copy_one_pte(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 		swp_entry_t entry = pte_to_swp_entry(pte);
 
 		if (likely(!non_swap_entry(entry))) {
-			if (swap_duplicate(entry) < 0)
-				return entry.val;
+			if (swap_duplicate(entry) < 0) {
+				data->entry = entry;
+				return COPY_MM_SWAP_CONT;
+			}
 
 			/* make sure dst_mm is on swapoff's mmlist. */
 			if (unlikely(list_empty(&dst_mm->mmlist))) {
@@ -809,7 +819,7 @@  copy_one_pte(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 
 out_set_pte:
 	set_pte_at(dst_mm, addr, dst_pte, pte);
-	return 0;
+	return COPY_MM_DONE;
 }
 
 static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
@@ -820,9 +830,9 @@  static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 	pte_t *orig_src_pte, *orig_dst_pte;
 	pte_t *src_pte, *dst_pte;
 	spinlock_t *src_ptl, *dst_ptl;
-	int progress = 0;
+	int progress, copy_ret = COPY_MM_DONE;
 	int rss[NR_MM_COUNTERS];
-	swp_entry_t entry = (swp_entry_t){0};
+	struct copy_mm_data data;
 
 again:
 	init_rss_vec(rss);
@@ -837,6 +847,7 @@  static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 	orig_dst_pte = dst_pte;
 	arch_enter_lazy_mmu_mode();
 
+	progress = 0;
 	do {
 		/*
 		 * We are holding two locks at this point - either of them
@@ -852,9 +863,9 @@  static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 			progress++;
 			continue;
 		}
-		entry.val = copy_one_pte(dst_mm, src_mm, dst_pte, src_pte,
-							vma, addr, rss);
-		if (entry.val)
+		copy_ret = copy_one_pte(dst_mm, src_mm, dst_pte, src_pte,
+					vma, addr, rss, &data);
+		if (copy_ret != COPY_MM_DONE)
 			break;
 		progress += 8;
 	} while (dst_pte++, src_pte++, addr += PAGE_SIZE, addr != end);
@@ -866,13 +877,18 @@  static int copy_pte_range(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 	pte_unmap_unlock(orig_dst_pte, dst_ptl);
 	cond_resched();
 
-	if (entry.val) {
-		if (add_swap_count_continuation(entry, GFP_KERNEL) < 0)
+	switch (copy_ret) {
+	case COPY_MM_SWAP_CONT:
+		if (add_swap_count_continuation(data.entry, GFP_KERNEL) < 0)
 			return -ENOMEM;
-		progress = 0;
+		break;
+	default:
+		break;
 	}
+
 	if (addr != end)
 		goto again;
+
 	return 0;
 }