Commit 49147beb0ccb ("x86/xen: allow nesting of same lazy mode")
originally introduced support for nested lazy sections (LAZY_MMU and
LAZY_CPU). It later got reverted by commit c36549ff8d84 as its
implementation turned out to be intolerant to preemption.

Now that the lazy_mmu API allows enter() to pass through a state to
the matching leave() call, we can support nesting again for the
LAZY_MMU mode in a preemption-safe manner. If xen_enter_lazy_mmu() is
called inside an active lazy_mmu section, xen_lazy_mode will already
be set to XEN_LAZY_MMU and we can then return LAZY_MMU_NESTED to
instruct the matching xen_leave_lazy_mmu() call to leave
xen_lazy_mode unchanged.

The only effect of this patch is to ensure that xen_lazy_mode
remains set to XEN_LAZY_MMU until the outermost lazy_mmu section
ends. xen_leave_lazy_mmu() still calls xen_mc_flush()
unconditionally.

Signed-off-by: Kevin Brodsky <kevin.brodsky@arm.com>
---
 arch/x86/include/asm/paravirt.h       |  6 ++----
 arch/x86/include/asm/paravirt_types.h |  4 ++--
 arch/x86/xen/mmu_pv.c                 | 11 ++++++++---
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/arch/x86/include/asm/paravirt.h b/arch/x86/include/asm/paravirt.h
index 65a0d394fba1..4ecd3a6b1dea 100644
--- a/arch/x86/include/asm/paravirt.h
+++ b/arch/x86/include/asm/paravirt.h
@@ -529,14 +529,12 @@ static inline void arch_end_context_switch(struct task_struct *next)
 #define  __HAVE_ARCH_ENTER_LAZY_MMU_MODE
 static inline lazy_mmu_state_t arch_enter_lazy_mmu_mode(void)
 {
-	PVOP_VCALL0(mmu.lazy_mode.enter);
-
-	return LAZY_MMU_DEFAULT;
+	return PVOP_CALL0(lazy_mmu_state_t, mmu.lazy_mode.enter);
 }
 
 static inline void arch_leave_lazy_mmu_mode(lazy_mmu_state_t state)
 {
-	PVOP_VCALL0(mmu.lazy_mode.leave);
+	PVOP_VCALL1(mmu.lazy_mode.leave, state);
 }
 
 static inline void arch_flush_lazy_mmu_mode(void)
diff --git a/arch/x86/include/asm/paravirt_types.h b/arch/x86/include/asm/paravirt_types.h
index bc1af86868a3..b7c567ccbf32 100644
--- a/arch/x86/include/asm/paravirt_types.h
+++ b/arch/x86/include/asm/paravirt_types.h
@@ -45,8 +45,8 @@ typedef int lazy_mmu_state_t;
 
 struct pv_lazy_ops {
 	/* Set deferred update mode, used for batching operations. */
-	void (*enter)(void);
-	void (*leave)(void);
+	lazy_mmu_state_t (*enter)(void);
+	void (*leave)(lazy_mmu_state_t);
 	void (*flush)(void);
 } __no_randomize_layout;
 #endif
diff --git a/arch/x86/xen/mmu_pv.c b/arch/x86/xen/mmu_pv.c
index 2039d5132ca3..6e5390ff06a5 100644
--- a/arch/x86/xen/mmu_pv.c
+++ b/arch/x86/xen/mmu_pv.c
@@ -2130,9 +2130,13 @@ static void xen_set_fixmap(unsigned idx, phys_addr_t phys, pgprot_t prot)
 #endif
 }
 
-static void xen_enter_lazy_mmu(void)
+static lazy_mmu_state_t xen_enter_lazy_mmu(void)
 {
+	if (this_cpu_read(xen_lazy_mode) == XEN_LAZY_MMU)
+		return LAZY_MMU_NESTED;
+
 	enter_lazy(XEN_LAZY_MMU);
+	return LAZY_MMU_DEFAULT;
 }
 
 static void xen_flush_lazy_mmu(void)
@@ -2167,11 +2171,12 @@ static void __init xen_post_allocator_init(void)
 	pv_ops.mmu.write_cr3 = &xen_write_cr3;
 }
 
-static void xen_leave_lazy_mmu(void)
+static void xen_leave_lazy_mmu(lazy_mmu_state_t state)
 {
 	preempt_disable();
 	xen_mc_flush();
-	leave_lazy(XEN_LAZY_MMU);
+	if (state != LAZY_MMU_NESTED)
+		leave_lazy(XEN_LAZY_MMU);
 	preempt_enable();
 }
 
-- 
2.47.0

On Thu, Sep 04, 2025 at 01:57:33PM +0100, Kevin Brodsky wrote:
...
> -static void xen_enter_lazy_mmu(void)
> +static lazy_mmu_state_t xen_enter_lazy_mmu(void)
>  {
> +	if (this_cpu_read(xen_lazy_mode) == XEN_LAZY_MMU)
> +		return LAZY_MMU_NESTED;
> +
>  	enter_lazy(XEN_LAZY_MMU);
> +	return LAZY_MMU_DEFAULT;
>  }
>  
>  static void xen_flush_lazy_mmu(void)
> @@ -2167,11 +2171,12 @@ static void __init xen_post_allocator_init(void)
>  	pv_ops.mmu.write_cr3 = &xen_write_cr3;
>  }
>  
> -static void xen_leave_lazy_mmu(void)
> +static void xen_leave_lazy_mmu(lazy_mmu_state_t state)
>  {
>  	preempt_disable();
>  	xen_mc_flush();
> -	leave_lazy(XEN_LAZY_MMU);
> +	if (state != LAZY_MMU_NESTED)
> +		leave_lazy(XEN_LAZY_MMU);

Based on xen_enter_lazy_mmu(), whether this condition needs to be
executed with the preemption disabled?

Or may be this_cpu_read(xen_lazy_mode) + enter_lazy(XEN_LAZY_MMU)
should be executed with the preemption disabled?

>  	preempt_enable();
>  }

Thanks!

On 05/09/2025 17:48, Alexander Gordeev wrote:
> On Thu, Sep 04, 2025 at 01:57:33PM +0100, Kevin Brodsky wrote:
> ...
>> -static void xen_enter_lazy_mmu(void)
>> +static lazy_mmu_state_t xen_enter_lazy_mmu(void)
>>  {
>> +	if (this_cpu_read(xen_lazy_mode) == XEN_LAZY_MMU)
>> +		return LAZY_MMU_NESTED;
>> +
>>  	enter_lazy(XEN_LAZY_MMU);
>> +	return LAZY_MMU_DEFAULT;
>>  }
>>  
>>  static void xen_flush_lazy_mmu(void)
>> @@ -2167,11 +2171,12 @@ static void __init xen_post_allocator_init(void)
>>  	pv_ops.mmu.write_cr3 = &xen_write_cr3;
>>  }
>>  
>> -static void xen_leave_lazy_mmu(void)
>> +static void xen_leave_lazy_mmu(lazy_mmu_state_t state)
>>  {
>>  	preempt_disable();
>>  	xen_mc_flush();
>> -	leave_lazy(XEN_LAZY_MMU);
>> +	if (state != LAZY_MMU_NESTED)
>> +		leave_lazy(XEN_LAZY_MMU);
> Based on xen_enter_lazy_mmu(), whether this condition needs to be
> executed with the preemption disabled?

AFAIU xen_mc_flush() needs preemption to be disabled. I don't think
{enter,leave}_lazy() do, but this patch doesn't introduce any change
from that perspective. I suppose it doesn't hurt that
xen_leave_lazy_mmu() calls leave_lazy() with preemption disabled.

> Or may be this_cpu_read(xen_lazy_mode) + enter_lazy(XEN_LAZY_MMU)
> should be executed with the preemption disabled?

Adding another this_cpu_read(xen_lazy_mode) in xen_enter_lazy_mmu()
shouldn't change the situation, i.e. preemption should still be safe. If
preemption occurs in the middle of that function,
xen_{start,end}_context_switch() will do the right thing to save/restore
xen_lazy_mode.

- Kevin