Re: [PATCH v7 1/4] userfaultfd: Add UFFD WP Async support

[Peter Xu]

Hi, Muhammad,

On Mon, Jan 09, 2023 at 11:45:16AM +0500, Muhammad Usama Anjum wrote:
> Add new WP Async mode (UFFDIO_WRITEPROTECT_MODE_ASYNC_WP) which resolves
> the page faults on its own. It can be used to track that which pages have
> been written to from the time the pages were write protected. It is very
> efficient way to track the changes as uffd is by nature pte/pmd based.
> 
> UFFD WP (UFFDIO_WRITEPROTECT_MODE_WP) sends the page faults to the
> userspace where the pages which have been written-to can be tracked. But
> it is not efficient. This is why this async version is being added.
> After setting the WP Async, the pages which have been written to can be
> found in the pagemap file or information can be obtained from the
> PAGEMAP_IOCTL (see next patches).
> 
> Signed-off-by: Muhammad Usama Anjum <usama.anjum@xxxxxxxxxxxxx>
> ---
>  fs/userfaultfd.c                 | 150 +++++++++++++++++--------------
>  include/uapi/linux/userfaultfd.h |   6 ++
>  2 files changed, 90 insertions(+), 66 deletions(-)
> 
> diff --git a/fs/userfaultfd.c b/fs/userfaultfd.c
> index 15a5bf765d43..be5e10d15058 100644
> --- a/fs/userfaultfd.c
> +++ b/fs/userfaultfd.c
> @@ -69,6 +69,7 @@ struct userfaultfd_ctx {
>  	unsigned int features;
>  	/* released */
>  	bool released;
> +	bool async;

Let's just make it a feature flag,

  UFFD_FEATURE_WP_ASYNC

>  	/* memory mappings are changing because of non-cooperative event */
>  	atomic_t mmap_changing;
>  	/* mm with one ore more vmas attached to this userfaultfd_ctx */
> @@ -497,80 +498,93 @@ vm_fault_t handle_userfault(struct vm_fault *vmf, unsigned long reason)
>  
>  	/* take the reference before dropping the mmap_lock */
>  	userfaultfd_ctx_get(ctx);
> +	if (ctx->async) {

Firstly, please consider not touching the existing code/indent as much as
what this patch did.  Hopefully we can keep the major part of sync uffd be
there with its git log, it also helps reviewing your code.  You can add the
async block before that, handle the fault and return just earlier.

And, I think this is a bit too late because we're going to return with
VM_FAULT_RETRY here, while maybe we don't need to retry at all here because
we're going to resolve the page fault immediately.

I assume you added this because you wanted userfaultfd_ctx_get() to make
sure the uffd context will not go away from under us, but it's not needed
if we're still holding the mmap read lock.  I'd expect for async mode we
don't really need to release it at all.

> +		// Resolve page fault of this page

Please use "/* ... */" as that's the common pattern of commenting in the
Linux kernel, at least what I see in mm/.

> +		unsigned long addr = (ctx->features & UFFD_FEATURE_EXACT_ADDRESS) ?
> +				      vmf->real_address : vmf->address;
> +		struct vm_area_struct *dst_vma = find_vma(ctx->mm, addr);
> +		size_t s = PAGE_SIZE;

This is weird - if we want async uffd-wp, let's consider huge page from the
1st day.

> +
> +		if (dst_vma->vm_flags & VM_HUGEPAGE) {

VM_HUGEPAGE is only a hint.  It doesn't mean this page is always a huge
page.  For anon, we can have thp wr-protected as a whole, not happening for
!anon because we'll split already.

For anon, if a write happens to a thp being uffd-wp-ed, we'll keep that pmd
wr-protected and report the uffd message.  The pmd split happens when the
user invokes UFFDIO_WRITEPROTECT on the small page.  I think it'll stop
working for async uffd-wp because we're going to resolve the page faults
right away.

So for async uffd-wp (note: this will be different from hugetlb), you may
want to consider having a pre-requisite patch to change wp_huge_pmd()
behavior: rather than calling handle_userfault(), IIUC you can also just
fallback to the split path right below (__split_huge_pmd) so the thp will
split now even before the uffd message is generated.

I think it should be transparent to the user and it'll start working for
you with async uffd-wp here, because it means when reaching
handle_userfault, it should not be possible to have thp at all since they
should have all split up.

> +			s = HPAGE_SIZE;
> +			addr &= HPAGE_MASK;
> +		}
>  
> -	init_waitqueue_func_entry(&uwq.wq, userfaultfd_wake_function);
> -	uwq.wq.private = current;
> -	uwq.msg = userfault_msg(vmf->address, vmf->real_address, vmf->flags,
> -				reason, ctx->features);
> -	uwq.ctx = ctx;
> -	uwq.waken = false;
> -
> -	blocking_state = userfaultfd_get_blocking_state(vmf->flags);
> +		ret = mwriteprotect_range(ctx->mm, addr, s, false, &ctx->mmap_changing);

This is an overkill - we're pretty sure it's a single page, no need to call
a range function here.

> +	} else {
> +		init_waitqueue_func_entry(&uwq.wq, userfaultfd_wake_function);
> +		uwq.wq.private = current;
> +		uwq.msg = userfault_msg(vmf->address, vmf->real_address, vmf->flags,
> +					reason, ctx->features);
> +		uwq.ctx = ctx;
> +		uwq.waken = false;
>  
> -        /*
> -         * Take the vma lock now, in order to safely call
> -         * userfaultfd_huge_must_wait() later. Since acquiring the
> -         * (sleepable) vma lock can modify the current task state, that
> -         * must be before explicitly calling set_current_state().
> -         */
> -	if (is_vm_hugetlb_page(vma))
> -		hugetlb_vma_lock_read(vma);
> +		blocking_state = userfaultfd_get_blocking_state(vmf->flags);
>  
> -	spin_lock_irq(&ctx->fault_pending_wqh.lock);
> -	/*
> -	 * After the __add_wait_queue the uwq is visible to userland
> -	 * through poll/read().
> -	 */
> -	__add_wait_queue(&ctx->fault_pending_wqh, &uwq.wq);
> -	/*
> -	 * The smp_mb() after __set_current_state prevents the reads
> -	 * following the spin_unlock to happen before the list_add in
> -	 * __add_wait_queue.
> -	 */
> -	set_current_state(blocking_state);
> -	spin_unlock_irq(&ctx->fault_pending_wqh.lock);
> +		/*
> +		 * Take the vma lock now, in order to safely call
> +		 * userfaultfd_huge_must_wait() later. Since acquiring the
> +		 * (sleepable) vma lock can modify the current task state, that
> +		 * must be before explicitly calling set_current_state().
> +		 */
> +		if (is_vm_hugetlb_page(vma))
> +			hugetlb_vma_lock_read(vma);
>  
> -	if (!is_vm_hugetlb_page(vma))
> -		must_wait = userfaultfd_must_wait(ctx, vmf->address, vmf->flags,
> -						  reason);
> -	else
> -		must_wait = userfaultfd_huge_must_wait(ctx, vma,
> -						       vmf->address,
> -						       vmf->flags, reason);
> -	if (is_vm_hugetlb_page(vma))
> -		hugetlb_vma_unlock_read(vma);
> -	mmap_read_unlock(mm);
> +		spin_lock_irq(&ctx->fault_pending_wqh.lock);
> +		/*
> +		 * After the __add_wait_queue the uwq is visible to userland
> +		 * through poll/read().
> +		 */
> +		__add_wait_queue(&ctx->fault_pending_wqh, &uwq.wq);
> +		/*
> +		 * The smp_mb() after __set_current_state prevents the reads
> +		 * following the spin_unlock to happen before the list_add in
> +		 * __add_wait_queue.
> +		 */
> +		set_current_state(blocking_state);
> +		spin_unlock_irq(&ctx->fault_pending_wqh.lock);
>  
> -	if (likely(must_wait && !READ_ONCE(ctx->released))) {
> -		wake_up_poll(&ctx->fd_wqh, EPOLLIN);
> -		schedule();
> -	}
> +		if (!is_vm_hugetlb_page(vma))
> +			must_wait = userfaultfd_must_wait(ctx, vmf->address, vmf->flags,
> +							  reason);
> +		else
> +			must_wait = userfaultfd_huge_must_wait(ctx, vma,
> +							       vmf->address,
> +							       vmf->flags, reason);
> +		if (is_vm_hugetlb_page(vma))
> +			hugetlb_vma_unlock_read(vma);
> +		mmap_read_unlock(mm);
> +
> +		if (likely(must_wait && !READ_ONCE(ctx->released))) {
> +			wake_up_poll(&ctx->fd_wqh, EPOLLIN);
> +			schedule();
> +		}
>  
> -	__set_current_state(TASK_RUNNING);
> +		__set_current_state(TASK_RUNNING);
>  
> -	/*
> -	 * Here we race with the list_del; list_add in
> -	 * userfaultfd_ctx_read(), however because we don't ever run
> -	 * list_del_init() to refile across the two lists, the prev
> -	 * and next pointers will never point to self. list_add also
> -	 * would never let any of the two pointers to point to
> -	 * self. So list_empty_careful won't risk to see both pointers
> -	 * pointing to self at any time during the list refile. The
> -	 * only case where list_del_init() is called is the full
> -	 * removal in the wake function and there we don't re-list_add
> -	 * and it's fine not to block on the spinlock. The uwq on this
> -	 * kernel stack can be released after the list_del_init.
> -	 */
> -	if (!list_empty_careful(&uwq.wq.entry)) {
> -		spin_lock_irq(&ctx->fault_pending_wqh.lock);
>  		/*
> -		 * No need of list_del_init(), the uwq on the stack
> -		 * will be freed shortly anyway.
> +		 * Here we race with the list_del; list_add in
> +		 * userfaultfd_ctx_read(), however because we don't ever run
> +		 * list_del_init() to refile across the two lists, the prev
> +		 * and next pointers will never point to self. list_add also
> +		 * would never let any of the two pointers to point to
> +		 * self. So list_empty_careful won't risk to see both pointers
> +		 * pointing to self at any time during the list refile. The
> +		 * only case where list_del_init() is called is the full
> +		 * removal in the wake function and there we don't re-list_add
> +		 * and it's fine not to block on the spinlock. The uwq on this
> +		 * kernel stack can be released after the list_del_init.
>  		 */
> -		list_del(&uwq.wq.entry);
> -		spin_unlock_irq(&ctx->fault_pending_wqh.lock);
> +		if (!list_empty_careful(&uwq.wq.entry)) {
> +			spin_lock_irq(&ctx->fault_pending_wqh.lock);
> +			/*
> +			 * No need of list_del_init(), the uwq on the stack
> +			 * will be freed shortly anyway.
> +			 */
> +			list_del(&uwq.wq.entry);
> +			spin_unlock_irq(&ctx->fault_pending_wqh.lock);
> +		}
>  	}
> -
>  	/*
>  	 * ctx may go away after this if the userfault pseudo fd is
>  	 * already released.
> @@ -1861,11 +1875,14 @@ static int userfaultfd_writeprotect(struct userfaultfd_ctx *ctx,
>  		return ret;
>  
>  	if (uffdio_wp.mode & ~(UFFDIO_WRITEPROTECT_MODE_DONTWAKE |
> -			       UFFDIO_WRITEPROTECT_MODE_WP))
> +			       UFFDIO_WRITEPROTECT_MODE_WP |
> +			       UFFDIO_WRITEPROTECT_MODE_ASYNC_WP))
>  		return -EINVAL;
>  
> -	mode_wp = uffdio_wp.mode & UFFDIO_WRITEPROTECT_MODE_WP;
> +	mode_wp = uffdio_wp.mode & (UFFDIO_WRITEPROTECT_MODE_WP |
> +				    UFFDIO_WRITEPROTECT_MODE_ASYNC_WP);
>  	mode_dontwake = uffdio_wp.mode & UFFDIO_WRITEPROTECT_MODE_DONTWAKE;
> +	ctx->async = uffdio_wp.mode & UFFDIO_WRITEPROTECT_MODE_ASYNC_WP;

Please no..  ctx attributes shouldn't be easily changed by a single ioctl.

I suggest we have a new feature bit as I mentioned above (say,
UFFD_FEATURE_WP_ASYNC), set it once with UFFDIO_API and it should apply to
the whole lifecycle of this uffd handle.  That flag should (something I can
quickly think of):

  - Have effect only if the uffd will be registered with WP mode (of
    course) or ignored in any other modes,

  - Should fail any attempts of UFFDIO_WRITEPROTECT with wp=false on this
    uffd handle because with async faults no page fault resolution needed
    from userspace,

  - Should apply to any region registered with this uffd ctx, so it's
    exclusively used with sync uffd-wp mode.

Then when the app wants to wr-protect in async mode, it simply goes ahead
with UFFDIO_WRITEPROTECT(wp=true), it'll happen exactly the same as when it
was sync mode.  It's only the pf handling procedure that's different (along
with how the fault is reported - rather than as a message but it'll be
consolidated into the soft-dirty bit).

>  
>  	if (mode_wp && mode_dontwake)
>  		return -EINVAL;
> @@ -2126,6 +2143,7 @@ static int new_userfaultfd(int flags)
>  	ctx->flags = flags;
>  	ctx->features = 0;
>  	ctx->released = false;
> +	ctx->async = false;
>  	atomic_set(&ctx->mmap_changing, 0);
>  	ctx->mm = current->mm;
>  	/* prevent the mm struct to be freed */
> diff --git a/include/uapi/linux/userfaultfd.h b/include/uapi/linux/userfaultfd.h
> index 005e5e306266..b89665653861 100644
> --- a/include/uapi/linux/userfaultfd.h
> +++ b/include/uapi/linux/userfaultfd.h
> @@ -284,6 +284,11 @@ struct uffdio_writeprotect {
>   * UFFDIO_WRITEPROTECT_MODE_DONTWAKE: set the flag to avoid waking up
>   * any wait thread after the operation succeeds.
>   *
> + * UFFDIO_WRITEPROTECT_MODE_ASYNC_WP: set the flag to write protect a
> + * range, the flag is unset automatically when the page is written.
> + * This is used to track which pages have been written to from the
> + * time the memory was write protected.
> + *
>   * NOTE: Write protecting a region (WP=1) is unrelated to page faults,
>   * therefore DONTWAKE flag is meaningless with WP=1.  Removing write
>   * protection (WP=0) in response to a page fault wakes the faulting
> @@ -291,6 +296,7 @@ struct uffdio_writeprotect {
>   */
>  #define UFFDIO_WRITEPROTECT_MODE_WP		((__u64)1<<0)
>  #define UFFDIO_WRITEPROTECT_MODE_DONTWAKE	((__u64)1<<1)
> +#define UFFDIO_WRITEPROTECT_MODE_ASYNC_WP	((__u64)1<<2)
>  	__u64 mode;
>  };
>  
> -- 
> 2.30.2
> 

-- 
Peter Xu