Received-SPF: pass (google.com: domain of linux-kernel-owner@vger.kernel.org designates 2620:137:e000::1:20 as permitted sender) client-ip=2620:137:e000::1:20;
MIME-Version: 1.0
References: <20230712060144.3006358-1-fengwei.yin@intel.com> <20230712060144.3006358-4-fengwei.yin@intel.com>
In-Reply-To: <20230712060144.3006358-4-fengwei.yin@intel.com>
From:   Yu Zhao <yuzhao@google.com>
Date:   Wed, 12 Jul 2023 00:31:54 -0600
Message-ID: <CAOUHufYef--8MxFettL6fOGjVx2vyZHZQU6EEaTCoW0XBvuC8Q@mail.gmail.com>
Subject: Re: [RFC PATCH v2 3/3] mm: mlock: update mlock_pte_range to handle
 large folio
To:     Yin Fengwei <fengwei.yin@intel.com>
Cc:     linux-mm@kvack.org, linux-kernel@vger.kernel.org,
        akpm@linux-foundation.org, willy@infradead.org, david@redhat.com,
        ryan.roberts@arm.com, shy828301@gmail.com
Content-Type: text/plain; charset="UTF-8"
Content-Transfer-Encoding: quoted-printable
Precedence: bulk

On Wed, Jul 12, 2023 at 12:02=E2=80=AFAM Yin Fengwei <fengwei.yin@intel.com=
> wrote:
>
> Current kernel only lock base size folio during mlock syscall.
> Add large folio support with following rules:
>   - Only mlock large folio when it's in VM_LOCKED VMA range
>
>   - If there is cow folio, mlock the cow folio as cow folio
>     is also in VM_LOCKED VMA range.
>
>   - munlock will apply to the large folio which is in VMA range
>     or cross the VMA boundary.
>
> The last rule is used to handle the case that the large folio is
> mlocked, later the VMA is split in the middle of large folio
> and this large folio become cross VMA boundary.
>
> Signed-off-by: Yin Fengwei <fengwei.yin@intel.com>
> ---
>  mm/mlock.c | 104 ++++++++++++++++++++++++++++++++++++++++++++++++++---
>  1 file changed, 99 insertions(+), 5 deletions(-)
>
> diff --git a/mm/mlock.c b/mm/mlock.c
> index 0a0c996c5c214..f49e079066870 100644
> --- a/mm/mlock.c
> +++ b/mm/mlock.c
> @@ -305,6 +305,95 @@ void munlock_folio(struct folio *folio)
>         local_unlock(&mlock_fbatch.lock);
>  }
>
> +static inline bool should_mlock_folio(struct folio *folio,
> +                                       struct vm_area_struct *vma)
> +{
> +       if (vma->vm_flags & VM_LOCKED)
> +               return (!folio_test_large(folio) ||
> +                               folio_within_vma(folio, vma));
> +
> +       /*
> +        * For unlock, allow munlock large folio which is partially
> +        * mapped to VMA. As it's possible that large folio is
> +        * mlocked and VMA is split later.
> +        *
> +        * During memory pressure, such kind of large folio can
> +        * be split. And the pages are not in VM_LOCKed VMA
> +        * can be reclaimed.
> +        */
> +
> +       return true;

Looks good, or just

should_mlock_folio() // or whatever name you see fit, can_mlock_folio()?
{
  return !(vma->vm_flags & VM_LOCKED) || folio_within_vma();
}

> +}
> +
> +static inline unsigned int get_folio_mlock_step(struct folio *folio,
> +                       pte_t pte, unsigned long addr, unsigned long end)
> +{
> +       unsigned int nr;
> +
> +       nr =3D folio_pfn(folio) + folio_nr_pages(folio) - pte_pfn(pte);
> +       return min_t(unsigned int, nr, (end - addr) >> PAGE_SHIFT);
> +}
> +
> +void mlock_folio_range(struct folio *folio, struct vm_area_struct *vma,
> +               pte_t *pte, unsigned long addr, unsigned int nr)
> +{
> +       struct folio *cow_folio;
> +       unsigned int step =3D 1;
> +
> +       mlock_folio(folio);
> +       if (nr =3D=3D 1)
> +               return;
> +
> +       for (; nr > 0; pte +=3D step, addr +=3D (step << PAGE_SHIFT), nr =
-=3D step) {
> +               pte_t ptent;
> +
> +               step =3D 1;
> +               ptent =3D ptep_get(pte);
> +
> +               if (!pte_present(ptent))
> +                       continue;
> +
> +               cow_folio =3D vm_normal_folio(vma, addr, ptent);
> +               if (!cow_folio || cow_folio =3D=3D folio) {
> +                       continue;
> +               }
> +
> +               mlock_folio(cow_folio);
> +               step =3D get_folio_mlock_step(folio, ptent,
> +                               addr, addr + (nr << PAGE_SHIFT));
> +       }
> +}
> +
> +void munlock_folio_range(struct folio *folio, struct vm_area_struct *vma=
,
> +               pte_t *pte, unsigned long addr, unsigned int nr)
> +{
> +       struct folio *cow_folio;
> +       unsigned int step =3D 1;
> +
> +       munlock_folio(folio);
> +       if (nr =3D=3D 1)
> +               return;
> +
> +       for (; nr > 0; pte +=3D step, addr +=3D (step << PAGE_SHIFT), nr =
-=3D step) {
> +               pte_t ptent;
> +
> +               step =3D 1;
> +               ptent =3D ptep_get(pte);
> +
> +               if (!pte_present(ptent))
> +                       continue;
> +
> +               cow_folio =3D vm_normal_folio(vma, addr, ptent);
> +               if (!cow_folio || cow_folio =3D=3D folio) {
> +                       continue;
> +               }
> +
> +               munlock_folio(cow_folio);
> +               step =3D get_folio_mlock_step(folio, ptent,
> +                               addr, addr + (nr << PAGE_SHIFT));
> +       }
> +}

I'll finish the above later.

>  static int mlock_pte_range(pmd_t *pmd, unsigned long addr,
>                            unsigned long end, struct mm_walk *walk)
>
> @@ -314,6 +403,7 @@ static int mlock_pte_range(pmd_t *pmd, unsigned long =
addr,
>         pte_t *start_pte, *pte;
>         pte_t ptent;
>         struct folio *folio;
> +       unsigned int step =3D 1;
>
>         ptl =3D pmd_trans_huge_lock(pmd, vma);
>         if (ptl) {
> @@ -329,24 +419,28 @@ static int mlock_pte_range(pmd_t *pmd, unsigned lon=
g addr,
>                 goto out;
>         }
>
> -       start_pte =3D pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
> +       pte =3D start_pte =3D pte_offset_map_lock(vma->vm_mm, pmd, addr, =
&ptl);
>         if (!start_pte) {
>                 walk->action =3D ACTION_AGAIN;
>                 return 0;
>         }
> -       for (pte =3D start_pte; addr !=3D end; pte++, addr +=3D PAGE_SIZE=
) {
> +
> +       for (; addr !=3D end; pte +=3D step, addr +=3D (step << PAGE_SHIF=
T)) {
> +               step =3D 1;
>                 ptent =3D ptep_get(pte);
>                 if (!pte_present(ptent))
>                         continue;
>                 folio =3D vm_normal_folio(vma, addr, ptent);
>                 if (!folio || folio_is_zone_device(folio))
>                         continue;
> -               if (folio_test_large(folio))
> +               if (!should_mlock_folio(folio, vma))
>                         continue;
> +
> +               step =3D get_folio_mlock_step(folio, ptent, addr, end);
>                 if (vma->vm_flags & VM_LOCKED)
> -                       mlock_folio(folio);
> +                       mlock_folio_range(folio, vma, pte, addr, step);
>                 else
> -                       munlock_folio(folio);
> +                       munlock_folio_range(folio, vma, pte, addr, step);
>         }
>         pte_unmap(start_pte);
>  out:

Looks good.