#include <linux/mm.h>
 #include <linux/wait.h>
 #include <linux/hash.h>
+#include <linux/kthread.h>
 
 void __init_waitqueue_head(wait_queue_head_t *q, const char *name, struct lock_class_key *key)
 {
 }
 EXPORT_SYMBOL(autoremove_wake_function);
 
+static inline bool is_kthread_should_stop(void)
+{
+       return (current->flags & PF_KTHREAD) && kthread_should_stop();
+}
 
 /*
  * DEFINE_WAIT_FUNC(wait, woken_wake_func);
         * woken_wake_function() such that if we observe WQ_FLAG_WOKEN we must
         * also observe all state before the wakeup.
         */
-       if (!(wait->flags & WQ_FLAG_WOKEN))
+       if (!(wait->flags & WQ_FLAG_WOKEN) && !is_kthread_should_stop())
                timeout = schedule_timeout(timeout);
        __set_current_state(TASK_RUNNING);