KVM: Get reference to VM's address space in the async #PF worker

Get a reference to the target VM's address space in async_pf_execute()
instead of gifting a reference from kvm_setup_async_pf().  Keeping the
address space alive just to service an async #PF is counter-productive,
i.e. if the process is exiting and all vCPUs are dead, then NOT doing
get_user_pages_remote() and freeing the address space asap is desirable.

Handling the mm reference entirely within async_pf_execute() also
simplifies the async #PF flows as a whole, e.g. it's not immediately
obvious when the worker task vs. the vCPU task is responsible for putting
the gifted mm reference.

Reviewed-by: Vitaly Kuznetsov <vkuznets@redhat.com>
Reviewed-by: Xu Yilun <yilun.xu@intel.com>
Link: https://lore.kernel.org/r/20240110011533.503302-4-seanjc@google.com
Signed-off-by: Sean Christopherson <seanjc@google.com>

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 7e7fd25b09b3ebe3d81e30fb23f506a9ee5a6519..bbfefd7e612f7418a3fe77b281e6ed91b49eabab 100644 (file)
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -238,7 +238,6 @@ struct kvm_async_pf {
        struct list_head link;
        struct list_head queue;
        struct kvm_vcpu *vcpu;
-       struct mm_struct *mm;
        gpa_t cr2_or_gpa;
        unsigned long addr;
        struct kvm_arch_async_pf arch;

diff --git a/virt/kvm/async_pf.c b/virt/kvm/async_pf.c
index 85acc9e824b9136e9aa387dd259b1187eefc4cb5..628f6df7609f7186b82ae169b6f31e92679e8dfd 100644 (file)
--- a/virt/kvm/async_pf.c
+++ b/virt/kvm/async_pf.c
@@ -46,8 +46,8 @@ static void async_pf_execute(struct work_struct *work)
 {
        struct kvm_async_pf *apf =
                container_of(work, struct kvm_async_pf, work);
-       struct mm_struct *mm = apf->mm;
        struct kvm_vcpu *vcpu = apf->vcpu;
+       struct mm_struct *mm = vcpu->kvm->mm;
        unsigned long addr = apf->addr;
        gpa_t cr2_or_gpa = apf->cr2_or_gpa;
        int locked = 1;
@@ -56,16 +56,24 @@ static void async_pf_execute(struct work_struct *work)
        might_sleep();
 
        /*
-        * This work is run asynchronously to the task which owns
-        * mm and might be done in another context, so we must
-        * access remotely.
+        * Attempt to pin the VM's host address space, and simply skip gup() if
+        * acquiring a pin fail, i.e. if the process is exiting.  Note, KVM
+        * holds a reference to its associated mm_struct until the very end of
+        * kvm_destroy_vm(), i.e. the struct itself won't be freed before this
+        * work item is fully processed.
         */
-       mmap_read_lock(mm);
-       get_user_pages_remote(mm, addr, 1, FOLL_WRITE, NULL, &locked);
-       if (locked)
-               mmap_read_unlock(mm);
-       mmput(mm);
+       if (mmget_not_zero(mm)) {
+               mmap_read_lock(mm);
+               get_user_pages_remote(mm, addr, 1, FOLL_WRITE, NULL, &locked);
+               if (locked)
+                       mmap_read_unlock(mm);
+               mmput(mm);
+       }
 
+       /*
+        * Notify and kick the vCPU even if faulting in the page failed, e.g.
+        * so that the vCPU can retry the fault synchronously.
+        */
        if (IS_ENABLED(CONFIG_KVM_ASYNC_PF_SYNC))
                kvm_arch_async_page_present(vcpu, apf);
 
@@ -131,10 +139,8 @@ void kvm_clear_async_pf_completion_queue(struct kvm_vcpu *vcpu)
 #ifdef CONFIG_KVM_ASYNC_PF_SYNC
                flush_work(&work->work);
 #else
-               if (cancel_work_sync(&work->work)) {
-                       mmput(work->mm);
+               if (cancel_work_sync(&work->work))
                        kmem_cache_free(async_pf_cache, work);
-               }
 #endif
                spin_lock(&vcpu->async_pf.lock);
        }
@@ -205,8 +211,6 @@ bool kvm_setup_async_pf(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
        work->cr2_or_gpa = cr2_or_gpa;
        work->addr = hva;
        work->arch = *arch;
-       work->mm = current->mm;
-       mmget(work->mm);
 
        INIT_WORK(&work->work, async_pf_execute);