*/
 # define might_sleep() \
        do { __might_sleep(__FILE__, __LINE__, 0); might_resched(); } while (0)
+# define sched_annotate_sleep()        __set_current_state(TASK_RUNNING)
 #else
   static inline void __might_sleep(const char *file, int line,
                                   int preempt_offset) { }
 # define might_sleep() do { might_resched(); } while (0)
+# define sched_annotate_sleep() do { } while (0)
 #endif
 
 #define might_sleep_if(cond) do { if (cond) might_sleep(); } while (0)
 
 
                get_task_struct(p);
                read_unlock(&tasklist_lock);
+               sched_annotate_sleep();
+
                if ((exit_code & 0x7f) == 0) {
                        why = CLD_EXITED;
                        status = exit_code >> 8;
         * thread can reap it because we its state == DEAD/TRACE.
         */
        read_unlock(&tasklist_lock);
+       sched_annotate_sleep();
 
        retval = wo->wo_rusage
                ? getrusage(p, RUSAGE_BOTH, wo->wo_rusage) : 0;
        pid = task_pid_vnr(p);
        why = ptrace ? CLD_TRAPPED : CLD_STOPPED;
        read_unlock(&tasklist_lock);
+       sched_annotate_sleep();
 
        if (unlikely(wo->wo_flags & WNOWAIT))
                return wait_noreap_copyout(wo, p, pid, uid, why, exit_code);
        pid = task_pid_vnr(p);
        get_task_struct(p);
        read_unlock(&tasklist_lock);
+       sched_annotate_sleep();
 
        if (!wo->wo_info) {
                retval = wo->wo_rusage