1 files changed, 32 insertions, 14 deletions

diff --git a/kernel/trace/ftrace.c b/kernel/trace/ftrace.c
index ce8c3d68292f..98ca94a41819 100644
--- a/kernel/trace/ftrace.c
+++ b/kernel/trace/ftrace.c
@@ -3996,37 +3996,51 @@ static void ftrace_init_module(struct module *mod,
         ftrace_process_locs(mod, start, end);
 }
 
-static int ftrace_module_notify(struct notifier_block *self,
-                                unsigned long val, void *data)
+static int ftrace_module_notify_enter(struct notifier_block *self,
+                                      unsigned long val, void *data)
 {
         struct module *mod = data;
 
-        switch (val) {
-        case MODULE_STATE_COMING:
+        if (val == MODULE_STATE_COMING)
                 ftrace_init_module(mod, mod->ftrace_callsites,
                                    mod->ftrace_callsites +
                                    mod->num_ftrace_callsites);
-                break;
-        case MODULE_STATE_GOING:
+        return 0;
+}
+
+static int ftrace_module_notify_exit(struct notifier_block *self,
+                                     unsigned long val, void *data)
+{
+        struct module *mod = data;
+
+        if (val == MODULE_STATE_GOING)
                 ftrace_release_mod(mod);
-                break;
-        }
 
         return 0;
 }
 #else
-static int ftrace_module_notify(struct notifier_block *self,
-                                unsigned long val, void *data)
+static int ftrace_module_notify_enter(struct notifier_block *self,
+                                      unsigned long val, void *data)
+{
+        return 0;
+}
+static int ftrace_module_notify_exit(struct notifier_block *self,
+                                     unsigned long val, void *data)
 {
         return 0;
 }
 #endif /* CONFIG_MODULES */
 
-struct notifier_block ftrace_module_nb = {
-        .notifier_call = ftrace_module_notify,
+struct notifier_block ftrace_module_enter_nb = {
+        .notifier_call = ftrace_module_notify_enter,
         .priority = INT_MAX,    /* Run before anything that can use kprobes */
 };
 
+struct notifier_block ftrace_module_exit_nb = {
+        .notifier_call = ftrace_module_notify_exit,
+        .priority = INT_MIN,    /* Run after anything that can remove kprobes */
+};
+
 extern unsigned long __start_mcount_loc[];
 extern unsigned long __stop_mcount_loc[];
 
@@ -4058,9 +4072,13 @@ void __init ftrace_init(void)
                                   __start_mcount_loc,
                                   __stop_mcount_loc);
 
-        ret = register_module_notifier(&ftrace_module_nb);
+        ret = register_module_notifier(&ftrace_module_enter_nb);
+        if (ret)
+                pr_warning("Failed to register trace ftrace module enter notifier\n");
+
+        ret = register_module_notifier(&ftrace_module_exit_nb);
         if (ret)
-                pr_warning("Failed to register trace ftrace module notifier\n");
+                pr_warning("Failed to register trace ftrace module exit notifier\n");
 
         set_ftrace_early_filters();