2 files changed, 71 insertions, 38 deletions

diff --git a/kernel/module.c b/kernel/module.c
index a0394706f10c..2383e60fcf3f 100644
--- a/kernel/module.c
+++ b/kernel/module.c
@@ -1490,9 +1490,6 @@ static void free_module(struct module *mod)
         /* Free any allocated parameters. */
         destroy_params(mod->kp, mod->num_kp);
 
-        /* release any pointers to mcount in this module */
-        ftrace_release(mod->module_core, mod->core_size);
-
         /* This may be NULL, but that's OK */
         module_free(mod, mod->module_init);
         kfree(mod->args);
@@ -1893,11 +1890,9 @@ static noinline struct module *load_module(void __user *umod,
         unsigned int symindex = 0;
         unsigned int strindex = 0;
         unsigned int modindex, versindex, infoindex, pcpuindex;
-        unsigned int num_mcount;
         struct module *mod;
         long err = 0;
         void *percpu = NULL, *ptr = NULL; /* Stops spurious gcc warning */
-        unsigned long *mseg;
         mm_segment_t old_fs;
 
         DEBUGP("load_module: umod=%p, len=%lu, uargs=%p\n",
@@ -2179,7 +2174,13 @@ static noinline struct module *load_module(void __user *umod,
                                          sizeof(*mod->trace_events),
                                          &mod->num_trace_events);
 #endif
-
+#ifdef CONFIG_FTRACE_MCOUNT_RECORD
+        /* sechdrs[0].sh_size is always zero */
+        mod->ftrace_callsites = section_objs(hdr, sechdrs, secstrings,
+                                             "__mcount_loc",
+                                             sizeof(*mod->ftrace_callsites),
+                                             &mod->num_ftrace_callsites);
+#endif
 #ifdef CONFIG_MODVERSIONS
         if ((mod->num_syms && !mod->crcs)
             || (mod->num_gpl_syms && !mod->gpl_crcs)
@@ -2244,11 +2245,6 @@ static noinline struct module *load_module(void __user *umod,
                         dynamic_debug_setup(debug, num_debug);
         }
 
-        /* sechdrs[0].sh_size is always zero */
-        mseg = section_objs(hdr, sechdrs, secstrings, "__mcount_loc",
-                            sizeof(*mseg), &num_mcount);
-        ftrace_init_module(mod, mseg, mseg + num_mcount);
-
         err = module_finalize(hdr, sechdrs, mod);
         if (err < 0)
                 goto cleanup;
@@ -2309,7 +2305,6 @@ static noinline struct module *load_module(void __user *umod,
  cleanup:
         kobject_del(&mod->mkobj.kobj);
         kobject_put(&mod->mkobj.kobj);
-        ftrace_release(mod->module_core, mod->core_size);
  free_unload:
         module_unload_free(mod);
 #if defined(CONFIG_MODULE_UNLOAD) && defined(CONFIG_SMP)
diff --git a/kernel/trace/ftrace.c b/kernel/trace/ftrace.c
index a23488988581..5b606f45b6c4 100644
--- a/kernel/trace/ftrace.c
+++ b/kernel/trace/ftrace.c
@@ -916,30 +916,6 @@ static void ftrace_free_rec(struct dyn_ftrace *rec)
         rec->flags |= FTRACE_FL_FREE;
 }
 
-void ftrace_release(void *start, unsigned long size)
-{
-        struct dyn_ftrace *rec;
-        struct ftrace_page *pg;
-        unsigned long s = (unsigned long)start;
-        unsigned long e = s + size;
-
-        if (ftrace_disabled || !start)
-                return;
-
-        mutex_lock(&ftrace_lock);
-        do_for_each_ftrace_rec(pg, rec) {
-                if ((rec->ip >= s) && (rec->ip < e)) {
-                        /*
-                         * rec->ip is changed in ftrace_free_rec()
-                         * It should not between s and e if record was freed.
-                         */
-                        FTRACE_WARN_ON(rec->flags & FTRACE_FL_FREE);
-                        ftrace_free_rec(rec);
-                }
-        } while_for_each_ftrace_rec();
-        mutex_unlock(&ftrace_lock);
-}
-
 static struct dyn_ftrace *ftrace_alloc_dyn_node(unsigned long ip)
 {
         struct dyn_ftrace *rec;
@@ -2752,14 +2728,72 @@ static int ftrace_convert_nops(struct module *mod,
         return 0;
 }
 
-void ftrace_init_module(struct module *mod,
-                        unsigned long *start, unsigned long *end)
+#ifdef CONFIG_MODULES
+void ftrace_release(void *start, void *end)
+{
+        struct dyn_ftrace *rec;
+        struct ftrace_page *pg;
+        unsigned long s = (unsigned long)start;
+        unsigned long e = (unsigned long)end;
+
+        if (ftrace_disabled || !start || start == end)
+                return;
+
+        mutex_lock(&ftrace_lock);
+        do_for_each_ftrace_rec(pg, rec) {
+                if ((rec->ip >= s) && (rec->ip < e)) {
+                        /*
+                         * rec->ip is changed in ftrace_free_rec()
+                         * It should not between s and e if record was freed.
+                         */
+                        FTRACE_WARN_ON(rec->flags & FTRACE_FL_FREE);
+                        ftrace_free_rec(rec);
+                }
+        } while_for_each_ftrace_rec();
+        mutex_unlock(&ftrace_lock);
+}
+
+static void ftrace_init_module(struct module *mod,
+                               unsigned long *start, unsigned long *end)
 {
         if (ftrace_disabled || start == end)
                 return;
         ftrace_convert_nops(mod, start, end);
 }
 
+static int ftrace_module_notify(struct notifier_block *self,
+                                unsigned long val, void *data)
+{
+        struct module *mod = data;
+
+        switch (val) {
+        case MODULE_STATE_COMING:
+                ftrace_init_module(mod, mod->ftrace_callsites,
+                                   mod->ftrace_callsites +
+                                   mod->num_ftrace_callsites);
+                break;
+        case MODULE_STATE_GOING:
+                ftrace_release(mod->ftrace_callsites,
+                               mod->ftrace_callsites +
+                               mod->num_ftrace_callsites);
+                break;
+        }
+
+        return 0;
+}
+#else
+static int ftrace_module_notify(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,
+        .priority = 0,
+};
+
 extern unsigned long __start_mcount_loc[];
 extern unsigned long __stop_mcount_loc[];
 
@@ -2791,6 +2825,10 @@ void __init ftrace_init(void)
                                   __start_mcount_loc,
                                   __stop_mcount_loc);
 
+        ret = register_module_notifier(&ftrace_module_nb);
+        if (!ret)
+                pr_warning("Failed to register trace ftrace module notifier\n");
+
         return;
  failed:
         ftrace_disabled = 1;