3 files changed, 13 insertions, 4 deletions

diff --git a/arch/x86/kernel/ftrace.c b/arch/x86/kernel/ftrace.c
index 1a5b8f8cb3cc..adba8e9a427c 100644
--- a/arch/x86/kernel/ftrace.c
+++ b/arch/x86/kernel/ftrace.c
@@ -484,14 +484,16 @@ void prepare_ftrace_return(unsigned long *parent, unsigned long self_addr)
                 : "memory"
         );
 
-        if (WARN_ON(faulted)) {
-                unregister_ftrace_graph();
+        if (unlikely(faulted)) {
+                ftrace_graph_stop();
+                WARN_ON(1);
                 return;
         }
 
-        if (WARN_ON(!__kernel_text_address(old))) {
-                unregister_ftrace_graph();
+        if (unlikely(!__kernel_text_address(old))) {
+                ftrace_graph_stop();
                 *parent = old;
+                WARN_ON(1);
                 return;
         }
 
diff --git a/include/linux/ftrace.h b/include/linux/ftrace.h
index afba918c623c..58ca1c3a3f4d 100644
--- a/include/linux/ftrace.h
+++ b/include/linux/ftrace.h
@@ -376,6 +376,8 @@ typedef void (*trace_func_graph_ent_t)(struct ftrace_graph_ent *); /* entry */
 extern int register_ftrace_graph(trace_func_graph_ret_t retfunc,
                                 trace_func_graph_ent_t entryfunc);
 
+extern void ftrace_graph_stop(void);
+
 /* The current handlers in use */
 extern trace_func_graph_ret_t ftrace_graph_return;
 extern trace_func_graph_ent_t ftrace_graph_entry;
diff --git a/kernel/trace/ftrace.c b/kernel/trace/ftrace.c
index 2e78628443e8..a44af05ae2d0 100644
--- a/kernel/trace/ftrace.c
+++ b/kernel/trace/ftrace.c
@@ -1769,5 +1769,10 @@ void ftrace_graph_exit_task(struct task_struct *t)
 
         kfree(ret_stack);
 }
+
+void ftrace_graph_stop(void)
+{
+        ftrace_stop();
+}
 #endif