FTRACE_WARN_ON(ret);
 }
 
-static void ftrace_run_modify_code(struct ftrace_ops *ops, int command)
+static void ftrace_run_modify_code(struct ftrace_ops *ops, int command,
+                                  struct ftrace_hash *old_hash)
 {
        ops->flags |= FTRACE_OPS_FL_MODIFYING;
+       ops->old_hash.filter_hash = old_hash;
        ftrace_run_update_code(command);
+       ops->old_hash.filter_hash = NULL;
        ops->flags &= ~FTRACE_OPS_FL_MODIFYING;
 }
 
 
 static int ftrace_probe_registered;
 
-static void __enable_ftrace_function_probe(void)
+static void __enable_ftrace_function_probe(struct ftrace_hash *old_hash)
 {
        int ret;
        int i;
        if (ftrace_probe_registered) {
                /* still need to update the function call sites */
                if (ftrace_enabled)
-                       ftrace_run_modify_code(&trace_probe_ops, FTRACE_UPDATE_CALLS);
+                       ftrace_run_modify_code(&trace_probe_ops, FTRACE_UPDATE_CALLS,
+                                              old_hash);
                return;
        }
 
        } while_for_each_ftrace_rec();
 
        ret = ftrace_hash_move(&trace_probe_ops, 1, orig_hash, hash);
+
+       __enable_ftrace_function_probe(old_hash);
+
        if (!ret)
                free_ftrace_hash_rcu(old_hash);
        else
                count = ret;
 
-       __enable_ftrace_function_probe();
-
  out_unlock:
        mutex_unlock(&ftrace_lock);
  out:
        return add_hash_entry(hash, ip);
 }
 
-static void ftrace_ops_update_code(struct ftrace_ops *ops)
+static void ftrace_ops_update_code(struct ftrace_ops *ops,
+                                  struct ftrace_hash *old_hash)
 {
        if (ops->flags & FTRACE_OPS_FL_ENABLED && ftrace_enabled)
-               ftrace_run_modify_code(ops, FTRACE_UPDATE_CALLS);
+               ftrace_run_modify_code(ops, FTRACE_UPDATE_CALLS, old_hash);
 }
 
 static int
        old_hash = *orig_hash;
        ret = ftrace_hash_move(ops, enable, orig_hash, hash);
        if (!ret) {
-               ftrace_ops_update_code(ops);
+               ftrace_ops_update_code(ops, old_hash);
                free_ftrace_hash_rcu(old_hash);
        }
        mutex_unlock(&ftrace_lock);
                ret = ftrace_hash_move(iter->ops, filter_hash,
                                       orig_hash, iter->hash);
                if (!ret) {
-                       ftrace_ops_update_code(iter->ops);
+                       ftrace_ops_update_code(iter->ops, old_hash);
                        free_ftrace_hash_rcu(old_hash);
                }
                mutex_unlock(&ftrace_lock);