struct module;
 struct ftrace_hash;
+struct ftrace_direct_func;
 
 #if defined(CONFIG_FUNCTION_TRACER) && defined(CONFIG_MODULES) && \
        defined(CONFIG_DYNAMIC_FTRACE)
 #ifdef CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
 int register_ftrace_direct(unsigned long ip, unsigned long addr);
 int unregister_ftrace_direct(unsigned long ip, unsigned long addr);
+struct ftrace_direct_func *ftrace_find_direct_func(unsigned long addr);
 #else
 static inline int register_ftrace_direct(unsigned long ip, unsigned long addr)
 {
 {
        return -ENODEV;
 }
+static inline struct ftrace_direct_func *ftrace_find_direct_func(unsigned long addr)
+{
+       return NULL;
+}
 #endif /* CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS */
 
 #ifndef CONFIG_HAVE_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
 
 }
 
 #ifdef CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
+
+struct ftrace_direct_func {
+       struct list_head        next;
+       unsigned long           addr;
+       int                     count;
+};
+
+static LIST_HEAD(ftrace_direct_funcs);
+
+/**
+ * ftrace_find_direct_func - test an address if it is a registered direct caller
+ * @addr: The address of a registered direct caller
+ *
+ * This searches to see if a ftrace direct caller has been registered
+ * at a specific address, and if so, it returns a descriptor for it.
+ *
+ * This can be used by architecture code to see if an address is
+ * a direct caller (trampoline) attached to a fentry/mcount location.
+ * This is useful for the function_graph tracer, as it may need to
+ * do adjustments if it traced a location that also has a direct
+ * trampoline attached to it.
+ */
+struct ftrace_direct_func *ftrace_find_direct_func(unsigned long addr)
+{
+       struct ftrace_direct_func *entry;
+       bool found = false;
+
+       /* May be called by fgraph trampoline (protected by rcu tasks) */
+       list_for_each_entry_rcu(entry, &ftrace_direct_funcs, next) {
+               if (entry->addr == addr) {
+                       found = true;
+                       break;
+               }
+       }
+       if (found)
+               return entry;
+
+       return NULL;
+}
+
 /**
  * register_ftrace_direct - Call a custom trampoline directly
  * @ip: The address of the nop at the beginning of a function
  */
 int register_ftrace_direct(unsigned long ip, unsigned long addr)
 {
+       struct ftrace_direct_func *direct;
        struct ftrace_func_entry *entry;
        struct ftrace_hash *free_hash = NULL;
        struct dyn_ftrace *rec;
        if (!entry)
                goto out_unlock;
 
+       direct = ftrace_find_direct_func(addr);
+       if (!direct) {
+               direct = kmalloc(sizeof(*direct), GFP_KERNEL);
+               if (!direct) {
+                       kfree(entry);
+                       goto out_unlock;
+               }
+               direct->addr = addr;
+               direct->count = 0;
+               list_add_rcu(&direct->next, &ftrace_direct_funcs);
+       }
+
        entry->ip = ip;
        entry->direct = addr;
        __add_hash_entry(direct_functions, entry);
                        ftrace_set_filter_ip(&direct_ops, ip, 1, 0);
        }
 
-       if (ret)
+       if (ret) {
                kfree(entry);
+               if (!direct->count) {
+                       list_del_rcu(&direct->next);
+                       synchronize_rcu_tasks();
+                       kfree(direct);
+                       if (free_hash)
+                               free_ftrace_hash(free_hash);
+                       free_hash = NULL;
+               }
+       } else {
+               if (!direct->count)
+                       direct->count++;
+       }
  out_unlock:
        mutex_unlock(&direct_mutex);
 
 int unregister_ftrace_direct(unsigned long ip, unsigned long addr)
 {
        struct ftrace_func_entry *entry;
+       struct ftrace_direct_func *direct;
        struct dyn_ftrace *rec;
        int ret = -ENODEV;
 
 
        remove_hash_entry(direct_functions, entry);
 
+       direct = ftrace_find_direct_func(addr);
+       if (!WARN_ON(!direct)) {
+               /* This is the good path (see the ! before WARN) */
+               direct->count--;
+               WARN_ON(direct->count < 0);
+               if (!direct->count) {
+                       list_del_rcu(&direct->next);
+                       synchronize_rcu_tasks();
+                       kfree(direct);
+               }
+       }
  out_unlock:
        mutex_unlock(&direct_mutex);