1 files changed, 2 insertions, 6 deletions

diff --git a/kernel/trace/trace.c b/kernel/trace/trace.c
index e4c40c868d67..3e7bf5eb9007 100644
--- a/kernel/trace/trace.c
+++ b/kernel/trace/trace.c
@@ -904,8 +904,7 @@ function_trace_call(unsigned long ip, unsigned long parent_ip)
                 return;
 
         pc = preempt_count();
-        resched = need_resched();
-        preempt_disable_notrace();
+        resched = ftrace_preempt_disable();
         local_save_flags(flags);
         cpu = raw_smp_processor_id();
         data = tr->data[cpu];
@@ -915,10 +914,7 @@ function_trace_call(unsigned long ip, unsigned long parent_ip)
                 trace_function(tr, data, ip, parent_ip, flags, pc);
 
         atomic_dec(&data->disabled);
-        if (resched)
-                preempt_enable_no_resched_notrace();
-        else
-                preempt_enable_notrace();
+        ftrace_preempt_enable(resched);
 }
 
 static struct ftrace_ops trace_ops __read_mostly =