[PATCH bpf-next] bpf: Add necessary migrate_{disable,enable} in bpf arena

[Yonghong Song <yonghong.song@xxxxxxxxx>]

When running bpf selftest (./test_progs -j), the following warnings
showed up:

  $ ./test_progs -t arena_atomics
  ...
  BUG: using smp_processor_id() in preemptible [00000000] code: kworker/u19:0/12501
  caller is bpf_mem_free+0x128/0x330
  ...
  Call Trace:
   <TASK>
   dump_stack_lvl
   check_preemption_disabled
   bpf_mem_free
   range_tree_destroy
   arena_map_free
   bpf_map_free_deferred
   process_scheduled_works
   ...

For selftests arena_htab and arena_list, similar smp_process_id() BUGs are
dumped, and the following are two stack trace:

   <TASK>
   dump_stack_lvl
   check_preemption_disabled
   bpf_mem_alloc
   range_tree_set
   arena_map_alloc
   map_create
   ...

   <TASK>
   dump_stack_lvl
   check_preemption_disabled
   bpf_mem_alloc
   range_tree_clear
   arena_vm_fault
   do_pte_missing
   handle_mm_fault
   do_user_addr_fault
   ...

Adding migrate_{disable,enable}() around related arena_*() calls can fix the issue.

Fixes: b795379757eb ("bpf: Introduce range_tree data structure and use it in bpf arena")
Signed-off-by: Yonghong Song <yonghong.song@xxxxxxxxx>
---
 kernel/bpf/arena.c | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/kernel/bpf/arena.c b/kernel/bpf/arena.c
index 3e1dfe349ced..9a55d18032a4 100644
--- a/kernel/bpf/arena.c
+++ b/kernel/bpf/arena.c
@@ -134,7 +134,9 @@ static struct bpf_map *arena_map_alloc(union bpf_attr *attr)
 	INIT_LIST_HEAD(&arena->vma_list);
 	bpf_map_init_from_attr(&arena->map, attr);
 	range_tree_init(&arena->rt);
+	migrate_disable();
 	range_tree_set(&arena->rt, 0, attr->max_entries);
+	migrate_enable();
 	mutex_init(&arena->lock);
 
 	return &arena->map;
@@ -185,7 +187,9 @@ static void arena_map_free(struct bpf_map *map)
 	apply_to_existing_page_range(&init_mm, bpf_arena_get_kern_vm_start(arena),
 				     KERN_VM_SZ - GUARD_SZ, existing_page_cb, NULL);
 	free_vm_area(arena->kern_vm);
+	migrate_disable();
 	range_tree_destroy(&arena->rt);
+	migrate_enable();
 	bpf_map_area_free(arena);
 }
 
@@ -276,7 +280,9 @@ static vm_fault_t arena_vm_fault(struct vm_fault *vmf)
 		/* User space requested to segfault when page is not allocated by bpf prog */
 		return VM_FAULT_SIGSEGV;
 
+	migrate_disable();
 	ret = range_tree_clear(&arena->rt, vmf->pgoff, 1);
+	migrate_enable();
 	if (ret)
 		return VM_FAULT_SIGSEGV;
 
-- 
2.43.5