mm: v2 drain TLB batching against CMA pages

MM has used TLB batching flush for memory unmapping. The batching
scheme does memory batch freeing as well as TLB batching flush.

Problem with CMA is process could schdule out once it kept holding
refcount of pages.

    CPU 0                                     CPU 1
do_madvise/munmap/exit_mmap
zap_pte_range
__tlb_remove_page

..                                          cma_alloc start
sched out

after 1sec                                  keep failing since

sched in
..
tlb_flush
free_pages_and_swap_cache
page_refcount is zero, finally.
					page migration succeded
					cma_alloc returns

If the process on CPU 0 is lower priority process, the CMA
allocation latency depends on the scheduler, sometimes, which
is priority inversion if process in CPU is higher priority.

This patch tries to fix it via using TLB draining right before
scheduling out(to release those pages immediately) if TLB has
CMA pages in the batch.

    CPU 0                                     CPU 1
do_madvise/munmap/exit_mmap
zap_pte_range
__tlb_remove_page

..                                          cma_alloc start
sched out

after 1sec                                  keep failing since

sched in
..
tlb_flush
free_pages_and_swap_cache
page_refcount is zero, finally.
					page migration succeded
					cma_alloc returns

If the process on CPU 0 is lower priority process, the CMA
allocation latency depends on the scheduler, sometimes, which
is priority inversion if process in CPU is higher priority.

This patch tries to fix it via using TLB draining right before
scheduling out(to release those pages immediately) if TLB has
CMA pages in the batch.

Due to vendor hook module load timing, following during boot can happens.

CPU A                                                               CPU B
zap_pte_range
..                                                  vendor_module registration
..                                                  register_trace_android_vh_zap_pte_range_tlb_start
trace_android_vh_zap_pte_range_tlb_start  ----->    preempt_disable
..
trace_android_vh_zap_pte_range_tlb_end    ----->    none
..                                                  register_trace_android_vh_zap_pte_range_tlb_end

To prevent it, we need to register register_trace_android_vh_zap_pte_range_tlb_end
first rather than register_trace_android_vh_zap_pte_range_tlb_start.

Bug: 238728493
Bug: 256549265
Change-Id: I7bdcf4f99f8a6ad79ba4e4bab7510549ad0b88bc
Signed-off-by: Minchan Kim <[email protected]>

drivers/soc/google/vh/include/mm.h

@@ -7,9 +7,9 @@ void vh_do_madvise_blk_plug(void *data, int behavior, bool *do_plug);
 void vh_shrink_inactive_list_blk_plug(void *data, bool *do_plug);
 void vh_reclaim_pages_plug(void *data, bool *do_plug);
 void vh_pagevec_drain(void *data, struct page *page, bool *ret);
-void vh_zap_pte_range_tlb_start(void *data, void *unused);
+void vh_zap_pte_range_tlb_start(void *data, void *preempt_off);
 void vh_zap_pte_range_tlb_force_flush(void *data, struct page *page, bool *flush);
-void vh_zap_pte_range_tlb_end(void *data, void *unused);
+void vh_zap_pte_range_tlb_end(void *data, void *preempt_off);
 void vh_skip_lru_disable(void *data, bool *skip);
 
 #endif

drivers/soc/google/vh/kernel/mm/memory.c

@@ -8,9 +8,10 @@
 
 #include <linux/mm.h>
 
-void vh_zap_pte_range_tlb_start(void *data, void *unused)
+void vh_zap_pte_range_tlb_start(void *data, void *preempt_off)
 {
 	preempt_disable();
+	*(int *)preempt_off = 1;
 }
 
 void vh_zap_pte_range_tlb_force_flush(void *data, struct page *page, bool *flush)
@@ -19,9 +20,12 @@ void vh_zap_pte_range_tlb_force_flush(void *data, struct page *page, bool *flush
 		*flush = true;
 }
 
-void vh_zap_pte_range_tlb_end(void *data, void *unused)
+void vh_zap_pte_range_tlb_end(void *data, void *preempt_off)
 {
-	preempt_enable();
+	if (likely(*(int *)preempt_off == 1)) {
+		preempt_enable();
+		*(int *)preempt_off = 0;
+	}
 }
 
 void vh_skip_lru_disable(void *data, bool *skip)

drivers/soc/google/vh/kernel/mm/vh_mm_init.c

@@ -63,6 +63,25 @@ static int vh_mm_init(void)
 			vh_pagevec_drain, NULL);
 	if (ret)
 		return ret;
+
+	/*
+	 * Do not reorder pte_range_tlb_end and pte_range_tlb_start
+	 * Otherwise, depending on module load timing, the pair can
+	 * be broken.
+	 */
+	ret = register_trace_android_vh_zap_pte_range_tlb_end(
+			vh_zap_pte_range_tlb_end, NULL);
+	if (ret)
+		return ret;
+	ret = register_trace_android_vh_zap_pte_range_tlb_force_flush(
+			vh_zap_pte_range_tlb_force_flush, NULL);
+	if (ret)
+		return ret;
+	ret = register_trace_android_vh_zap_pte_range_tlb_start(
+			vh_zap_pte_range_tlb_start, NULL);
+	if (ret)
+		return ret;
+
 	ret = register_trace_android_vh_skip_lru_disable(
 			vh_skip_lru_disable, NULL);
 	if (ret)