hash->count++;
 }
 
-static int add_hash_entry(struct ftrace_hash *hash, unsigned long ip)
+static struct ftrace_func_entry *
+add_hash_entry(struct ftrace_hash *hash, unsigned long ip)
 {
        struct ftrace_func_entry *entry;
 
        entry = kmalloc(sizeof(*entry), GFP_KERNEL);
        if (!entry)
-               return -ENOMEM;
+               return NULL;
 
        entry->ip = ip;
        __add_hash_entry(hash, entry);
 
-       return 0;
+       return entry;
 }
 
 static void
        struct ftrace_func_entry *entry;
        struct ftrace_hash *new_hash;
        int size;
-       int ret;
        int i;
 
        new_hash = alloc_ftrace_hash(size_bits);
        size = 1 << hash->size_bits;
        for (i = 0; i < size; i++) {
                hlist_for_each_entry(entry, &hash->buckets[i], hlist) {
-                       ret = add_hash_entry(new_hash, entry->ip);
-                       if (ret < 0)
+                       if (add_hash_entry(new_hash, entry->ip) == NULL)
                                goto free_hash;
                }
        }
 
 #ifdef CONFIG_DYNAMIC_FTRACE_WITH_DIRECT_CALLS
 /* Protected by rcu_tasks for reading, and direct_mutex for writing */
-static struct ftrace_hash *direct_functions = EMPTY_HASH;
+static struct ftrace_hash __rcu *direct_functions = EMPTY_HASH;
 static DEFINE_MUTEX(direct_mutex);
 int ftrace_direct_func_count;
 
        return entry->direct;
 }
 
-static struct ftrace_func_entry*
-ftrace_add_rec_direct(unsigned long ip, unsigned long addr,
-                     struct ftrace_hash **free_hash)
-{
-       struct ftrace_func_entry *entry;
-
-       if (ftrace_hash_empty(direct_functions) ||
-           direct_functions->count > 2 * (1 << direct_functions->size_bits)) {
-               struct ftrace_hash *new_hash;
-               int size = ftrace_hash_empty(direct_functions) ? 0 :
-                       direct_functions->count + 1;
-
-               if (size < 32)
-                       size = 32;
-
-               new_hash = dup_hash(direct_functions, size);
-               if (!new_hash)
-                       return NULL;
-
-               *free_hash = direct_functions;
-               direct_functions = new_hash;
-       }
-
-       entry = kmalloc(sizeof(*entry), GFP_KERNEL);
-       if (!entry)
-               return NULL;
-
-       entry->ip = ip;
-       entry->direct = addr;
-       __add_hash_entry(direct_functions, entry);
-       return entry;
-}
-
 static void call_direct_funcs(unsigned long ip, unsigned long pip,
                              struct ftrace_ops *ops, struct ftrace_regs *fregs)
 {
                /* Do nothing if it exists */
                if (entry)
                        return 0;
-
-               ret = add_hash_entry(hash, rec->ip);
+               if (add_hash_entry(hash, rec->ip) == NULL)
+                       ret = -ENOMEM;
        }
        return ret;
 }
                return 0;
        }
 
-       return add_hash_entry(hash, ip);
+       entry = add_hash_entry(hash, ip);
+       return entry ? 0 :  -ENOMEM;
 }
 
 static int
  */
 int register_ftrace_direct(struct ftrace_ops *ops, unsigned long addr)
 {
-       struct ftrace_hash *hash, *free_hash = NULL;
+       struct ftrace_hash *hash, *new_hash = NULL, *free_hash = NULL;
        struct ftrace_func_entry *entry, *new;
        int err = -EBUSY, size, i;
 
                }
        }
 
-       /* ... and insert them to direct_functions hash. */
        err = -ENOMEM;
+
+       /* Make a copy hash to place the new and the old entries in */
+       size = hash->count + direct_functions->count;
+       if (size > 32)
+               size = 32;
+       new_hash = alloc_ftrace_hash(fls(size));
+       if (!new_hash)
+               goto out_unlock;
+
+       /* Now copy over the existing direct entries */
+       size = 1 << direct_functions->size_bits;
+       for (i = 0; i < size; i++) {
+               hlist_for_each_entry(entry, &direct_functions->buckets[i], hlist) {
+                       new = add_hash_entry(new_hash, entry->ip);
+                       if (!new)
+                               goto out_unlock;
+                       new->direct = entry->direct;
+               }
+       }
+
+       /* ... and add the new entries */
+       size = 1 << hash->size_bits;
        for (i = 0; i < size; i++) {
                hlist_for_each_entry(entry, &hash->buckets[i], hlist) {
-                       new = ftrace_add_rec_direct(entry->ip, addr, &free_hash);
+                       new = add_hash_entry(new_hash, entry->ip);
                        if (!new)
-                               goto out_remove;
+                               goto out_unlock;
+                       /* Update both the copy and the hash entry */
+                       new->direct = addr;
                        entry->direct = addr;
                }
        }
 
+       free_hash = direct_functions;
+       rcu_assign_pointer(direct_functions, new_hash);
+       new_hash = NULL;
+
        ops->func = call_direct_funcs;
        ops->flags = MULTI_FLAGS;
        ops->trampoline = FTRACE_REGS_ADDR;
 
        err = register_ftrace_function_nolock(ops);
 
- out_remove:
-       if (err)
-               remove_direct_functions_hash(hash, addr);
-
  out_unlock:
        mutex_unlock(&direct_mutex);
 
-       if (free_hash) {
+       if (free_hash && free_hash != EMPTY_HASH) {
                synchronize_rcu_tasks();
                free_ftrace_hash(free_hash);
        }
+
+       if (new_hash)
+               free_ftrace_hash(new_hash);
+
        return err;
 }
 EXPORT_SYMBOL_GPL(register_ftrace_direct);
 
                                if (entry)
                                        continue;
-                               if (add_hash_entry(hash, rec->ip) < 0)
+                               if (add_hash_entry(hash, rec->ip) == NULL)
                                        goto out;
                        } else {
                                if (entry) {