FTRACE_UPDATE_TRACE_FUNC        = (1 << 2),
        FTRACE_START_FUNC_RET           = (1 << 3),
        FTRACE_STOP_FUNC_RET            = (1 << 4),
+       FTRACE_MAY_SLEEP                = (1 << 5),
 };
 
 /*
 
 #define ASSIGN_OPS_HASH(opsname, val)
 #endif
 
+enum {
+       FTRACE_MODIFY_ENABLE_FL         = (1 << 0),
+       FTRACE_MODIFY_MAY_SLEEP_FL      = (1 << 1),
+};
+
 struct ftrace_ops ftrace_list_end __read_mostly = {
        .func           = ftrace_stub,
        .flags          = FTRACE_OPS_FL_RECURSION_SAFE | FTRACE_OPS_FL_STUB,
        return -1; /* unknow ftrace bug */
 }
 
-void __weak ftrace_replace_code(int enable)
+void __weak ftrace_replace_code(int mod_flags)
 {
        struct dyn_ftrace *rec;
        struct ftrace_page *pg;
+       int enable = mod_flags & FTRACE_MODIFY_ENABLE_FL;
+       int schedulable = mod_flags & FTRACE_MODIFY_MAY_SLEEP_FL;
        int failed;
 
        if (unlikely(ftrace_disabled))
                        /* Stop processing */
                        return;
                }
+               if (schedulable)
+                       cond_resched();
        } while_for_each_ftrace_rec();
 }
 
 void ftrace_modify_all_code(int command)
 {
        int update = command & FTRACE_UPDATE_TRACE_FUNC;
+       int mod_flags = 0;
        int err = 0;
 
+       if (command & FTRACE_MAY_SLEEP)
+               mod_flags = FTRACE_MODIFY_MAY_SLEEP_FL;
+
        /*
         * If the ftrace_caller calls a ftrace_ops func directly,
         * we need to make sure that it only traces functions it
        }
 
        if (command & FTRACE_UPDATE_CALLS)
-               ftrace_replace_code(1);
+               ftrace_replace_code(mod_flags | FTRACE_MODIFY_ENABLE_FL);
        else if (command & FTRACE_DISABLE_CALLS)
-               ftrace_replace_code(0);
+               ftrace_replace_code(mod_flags);
 
        if (update && ftrace_trace_function != ftrace_ops_list_func) {
                function_trace_op = set_function_trace_op;