diff --git a/net/bpf/test_run.c b/net/bpf/test_run.c
index 8b7d0b90fea7..1d0047073eda 100644
--- a/net/bpf/test_run.c
+++ b/net/bpf/test_run.c
@@ -36,7 +36,7 @@ struct bpf_test_timer {
 static void bpf_test_timer_enter(struct bpf_test_timer *t)
 	__acquires(rcu)
 {
-	rcu_read_lock_dont_migrate();
+	migrate_disable();
 	t->time_start = ktime_get_ns();
 }
 
@@ -44,7 +44,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,