diff --git a/net/bpf/test_run.c b/net/bpf/test_run.c
index 8b7d0b90fea7..cc4288f408f2 100644
--- a/net/bpf/test_run.c
+++ b/net/bpf/test_run.c
@@ -36,7 +36,8 @@ struct bpf_test_timer {
 static void bpf_test_timer_enter(struct bpf_test_timer *t)
 	__acquires(rcu)
 {
-	rcu_read_lock_dont_migrate();
+	rcu_read_lock();
+	migrate_disable();
 	t->time_start = ktime_get_ns();
 }
 
@@ -44,7 +45,8 @@ static void bpf_test_timer_leave(struct bpf_test_timer *t)
 	__releases(rcu)
 {
 	t->time_start = 0;
-	rcu_read_unlock_migrate();
+	migrate_enable();
+	rcu_read_unlock();
 }
 
 static bool bpf_test_timer_continue(struct bpf_test_timer *t, int iterations,