diff --git a/fs/nfs/inode.c b/fs/nfs/inode.c
index 134d9f560240..211a773cb39f 100644
--- a/fs/nfs/inode.c
+++ b/fs/nfs/inode.c
@@ -71,23 +71,28 @@ nfs_fattr_to_ino_t(struct nfs_fattr *fattr)
 	return nfs_fileid_to_ino_t(fattr->fileid);
 }
 
-static int nfs_wait_killable(int mode)
+static int nfs_wait_killable(int mode, struct wait_queue_entry *wait)
 {
-	freezable_schedule_unsafe();
+	freezer_do_not_count();
+	if (wait)
+		wait_woken(wait, mode, MAX_SCHEDULE_TIMEOUT);
+	else
+		schedule();
+	freezer_count_unsafe();
 	if (signal_pending_state(mode, current))
 		return -ERESTARTSYS;
 	return 0;
 }
 
-int nfs_wait_bit_killable(struct wait_bit_key *key, int mode)
+int nfs_wait_bit_killable(struct wait_bit_key *key, int mode, struct wait_queue_entry *wait)
 {
-	return nfs_wait_killable(mode);
+	return nfs_wait_killable(mode, wait);
 }
 EXPORT_SYMBOL_GPL(nfs_wait_bit_killable);
 
 int nfs_wait_atomic_killable(atomic_t *p)
 {
-	return nfs_wait_killable(TASK_KILLABLE);
+	return nfs_wait_killable(TASK_KILLABLE, NULL);
 }
 
 /**
diff --git a/fs/nfs/internal.h b/fs/nfs/internal.h
index 5bdf952f414b..a9144d311a50 100644
--- a/fs/nfs/internal.h
+++ b/fs/nfs/internal.h
@@ -386,7 +386,7 @@ extern void nfs_clear_inode(struct inode *);
 extern void nfs_evict_inode(struct inode *);
 void nfs_zap_acl_cache(struct inode *inode);
 extern bool nfs_check_cache_invalid(struct inode *, unsigned long);
-extern int nfs_wait_bit_killable(struct wait_bit_key *key, int mode);
+extern int nfs_wait_bit_killable(struct wait_bit_key *key, int mode, struct wait_queue_entry *wait);
 extern int nfs_wait_atomic_killable(atomic_t *p);
 
 /* super.c */
diff --git a/include/linux/freezer.h b/include/linux/freezer.h
index dd03e837ebb7..1d197994b181 100644
--- a/include/linux/freezer.h
+++ b/include/linux/freezer.h
@@ -272,6 +272,7 @@ static inline bool try_to_freeze(void) { return false; }
 
 static inline void freezer_do_not_count(void) {}
 static inline void freezer_count(void) {}
+static inline void freezer_count_unsafe(void) {}
 static inline int freezer_should_skip(struct task_struct *p) { return 0; }
 static inline void set_freezable(void) {}
 
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 26a7df4e558c..1fd387e6c554 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -184,6 +184,9 @@ extern void io_schedule_finish(int token);
 extern long io_schedule_timeout(long timeout);
 extern void io_schedule(void);
 
+struct wait_queue_entry;
+extern void io_wait_timeout(struct wait_queue_entry *wait, int mode, long timeout);
+
 /**
  * struct prev_cputime - snapshot of system and user cputime
  * @utime: time spent in user mode
diff --git a/include/linux/wait.h b/include/linux/wait.h
index 87c4641023fb..5a320e1c215e 100644
--- a/include/linux/wait.h
+++ b/include/linux/wait.h
@@ -987,6 +987,7 @@ void finish_wait(struct wait_queue_head *wq_head, struct wait_queue_entry *wq_en
 long wait_woken(struct wait_queue_entry *wq_entry, unsigned mode, long timeout);
 int woken_wake_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync, void *key);
 int autoremove_wake_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync, void *key);
+int autoremove_woken_wake_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync, void *key);
 
 #define DEFINE_WAIT_FUNC(name, function)					\
 	struct wait_queue_entry name = {					\
diff --git a/include/linux/wait_bit.h b/include/linux/wait_bit.h
index 12b26660d7e9..f7a98ba7100c 100644
--- a/include/linux/wait_bit.h
+++ b/include/linux/wait_bit.h
@@ -24,7 +24,7 @@ struct wait_bit_queue_entry {
 #define __WAIT_ATOMIC_T_KEY_INITIALIZER(p)					\
 	{ .flags = p, .bit_nr = WAIT_ATOMIC_T_BIT_NR, }
 
-typedef int wait_bit_action_f(struct wait_bit_key *key, int mode);
+typedef int wait_bit_action_f(struct wait_bit_key *key, int mode, struct wait_queue_entry *wait);
 void __wake_up_bit(struct wait_queue_head *wq_head, void *word, int bit);
 int __wait_on_bit(struct wait_queue_head *wq_head, struct wait_bit_queue_entry *wbq_entry, wait_bit_action_f *action, unsigned int mode);
 int __wait_on_bit_lock(struct wait_queue_head *wq_head, struct wait_bit_queue_entry *wbq_entry, wait_bit_action_f *action, unsigned int mode);
@@ -50,10 +50,10 @@ int wake_bit_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync
 		},								\
 	}
 
-extern int bit_wait(struct wait_bit_key *key, int bit);
-extern int bit_wait_io(struct wait_bit_key *key, int bit);
-extern int bit_wait_timeout(struct wait_bit_key *key, int bit);
-extern int bit_wait_io_timeout(struct wait_bit_key *key, int bit);
+extern int bit_wait(struct wait_bit_key *key, int bit, struct wait_queue_entry *wait);
+extern int bit_wait_io(struct wait_bit_key *key, int bit, struct wait_queue_entry *wait);
+extern int bit_wait_timeout(struct wait_bit_key *key, int bit, struct wait_queue_entry *wait);
+extern int bit_wait_io_timeout(struct wait_bit_key *key, int bit, struct wait_queue_entry *wait);
 
 /**
  * wait_on_bit - wait for a bit to be cleared
diff --git a/kernel/kthread.c b/kernel/kthread.c
index 1c19edf82427..41c9b4d901b8 100644
--- a/kernel/kthread.c
+++ b/kernel/kthread.c
@@ -90,7 +90,8 @@ void free_kthread_struct(struct task_struct *k)
  */
 bool kthread_should_stop(void)
 {
-	return test_bit(KTHREAD_SHOULD_STOP, &to_kthread(current)->flags);
+	return to_kthread(current) &&
+	       test_bit(KTHREAD_SHOULD_STOP, &to_kthread(current)->flags);
 }
 EXPORT_SYMBOL(kthread_should_stop);
 
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index d17c5da523a0..290322f2afaf 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5034,6 +5034,17 @@ void io_schedule(void)
 }
 EXPORT_SYMBOL(io_schedule);
 
+void io_wait_timeout(struct wait_queue_entry *wait, int mode, long timeout)
+{
+	int token;
+
+	token = io_schedule_prepare();
+	wait_woken(wait, mode, timeout);
+	io_schedule_finish(token);
+}
+EXPORT_SYMBOL(io_schedule);
+
+
 /**
  * sys_sched_get_priority_max - return maximum RT priority.
  * @policy: scheduling class.
diff --git a/kernel/sched/wait.c b/kernel/sched/wait.c
index 98feab7933c7..60e7f4a4f4f2 100644
--- a/kernel/sched/wait.c
+++ b/kernel/sched/wait.c
@@ -372,6 +372,16 @@ void finish_wait(struct wait_queue_head *wq_head, struct wait_queue_entry *wq_en
 }
 EXPORT_SYMBOL(finish_wait);
 
+int autoremove_woken_wake_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync, void *key)
+{
+	int ret = woken_wake_function(wq_entry, mode, sync, key);
+
+	if (ret)
+		list_del_init(&wq_entry->entry);
+	return ret;
+}
+EXPORT_SYMBOL(autoremove_wake_function);
+
 int autoremove_wake_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync, void *key)
 {
 	int ret = default_wake_function(wq_entry, mode, sync, key);
diff --git a/kernel/sched/wait_bit.c b/kernel/sched/wait_bit.c
index f8159698aa4d..cb13bac35e88 100644
--- a/kernel/sched/wait_bit.c
+++ b/kernel/sched/wait_bit.c
@@ -30,7 +30,7 @@ int wake_bit_function(struct wait_queue_entry *wq_entry, unsigned mode, int sync
 			test_bit(key->bit_nr, key->flags))
 		return 0;
 	else
-		return autoremove_wake_function(wq_entry, mode, sync, key);
+		return autoremove_woken_wake_function(wq_entry, mode, sync, key);
 }
 EXPORT_SYMBOL(wake_bit_function);
 
@@ -45,12 +45,12 @@ __wait_on_bit(struct wait_queue_head *wq_head, struct wait_bit_queue_entry *wbq_
 {
 	int ret = 0;
 
+	add_wait_queue(wq_head, &wbq_entry->wq_entry);
 	do {
-		prepare_to_wait(wq_head, &wbq_entry->wq_entry, mode);
 		if (test_bit(wbq_entry->key.bit_nr, wbq_entry->key.flags))
-			ret = (*action)(&wbq_entry->key, mode);
+			ret = (*action)(&wbq_entry->key, mode, &wbq_entry->wq_entry);
 	} while (test_bit(wbq_entry->key.bit_nr, wbq_entry->key.flags) && !ret);
-	finish_wait(wq_head, &wbq_entry->wq_entry);
+	remove_wait_queue(wq_head, &wbq_entry->wq_entry);
 	return ret;
 }
 EXPORT_SYMBOL(__wait_on_bit);
@@ -83,24 +83,22 @@ __wait_on_bit_lock(struct wait_queue_head *wq_head, struct wait_bit_queue_entry
 {
 	int ret = 0;
 
+	add_wait_queue_exclusive(wq_head, &wbq_entry->wq_entry);
 	for (;;) {
-		prepare_to_wait_exclusive(wq_head, &wbq_entry->wq_entry, mode);
 		if (test_bit(wbq_entry->key.bit_nr, wbq_entry->key.flags)) {
-			ret = action(&wbq_entry->key, mode);
+			ret = action(&wbq_entry->key, mode, &wbq_entry->wq_entry);
 			/*
 			 * See the comment in prepare_to_wait_event().
 			 * finish_wait() does not necessarily takes wwq_head->lock,
 			 * but test_and_set_bit() implies mb() which pairs with
 			 * smp_mb__after_atomic() before wake_up_page().
 			 */
-			if (ret)
-				finish_wait(wq_head, &wbq_entry->wq_entry);
 		}
 		if (!test_and_set_bit(wbq_entry->key.bit_nr, wbq_entry->key.flags)) {
-			if (!ret)
-				finish_wait(wq_head, &wbq_entry->wq_entry);
+			remove_wait_queue(wq_head, &wbq_entry->wq_entry);
 			return 0;
 		} else if (ret) {
+			remove_wait_queue(wq_head, &wbq_entry->wq_entry);
 			return ret;
 		}
 	}
@@ -235,42 +233,42 @@ void wake_up_atomic_t(atomic_t *p)
 }
 EXPORT_SYMBOL(wake_up_atomic_t);
 
-__sched int bit_wait(struct wait_bit_key *word, int mode)
+__sched int bit_wait(struct wait_bit_key *word, int mode, struct wait_queue_entry *wait)
 {
-	schedule();
+	wait_woken(wait, mode, MAX_SCHEDULE_TIMEOUT);
 	if (signal_pending_state(mode, current))
 		return -EINTR;
 	return 0;
 }
 EXPORT_SYMBOL(bit_wait);
 
-__sched int bit_wait_io(struct wait_bit_key *word, int mode)
+__sched int bit_wait_io(struct wait_bit_key *word, int mode, struct wait_queue_entry *wait)
 {
-	io_schedule();
+	io_wait_timeout(wait, mode, MAX_SCHEDULE_TIMEOUT);
 	if (signal_pending_state(mode, current))
 		return -EINTR;
 	return 0;
 }
 EXPORT_SYMBOL(bit_wait_io);
 
-__sched int bit_wait_timeout(struct wait_bit_key *word, int mode)
+__sched int bit_wait_timeout(struct wait_bit_key *word, int mode, struct wait_queue_entry *wait)
 {
 	unsigned long now = READ_ONCE(jiffies);
 	if (time_after_eq(now, word->timeout))
 		return -EAGAIN;
-	schedule_timeout(word->timeout - now);
+	wait_woken(wait, mode, word->timeout - now);
 	if (signal_pending_state(mode, current))
 		return -EINTR;
 	return 0;
 }
 EXPORT_SYMBOL_GPL(bit_wait_timeout);
 
-__sched int bit_wait_io_timeout(struct wait_bit_key *word, int mode)
+__sched int bit_wait_io_timeout(struct wait_bit_key *word, int mode, struct wait_queue_entry *wait)
 {
 	unsigned long now = READ_ONCE(jiffies);
 	if (time_after_eq(now, word->timeout))
 		return -EAGAIN;
-	io_schedule_timeout(word->timeout - now);
+	io_wait_timeout(wait, mode, word->timeout - now);
 	if (signal_pending_state(mode, current))
 		return -EINTR;
 	return 0;
diff --git a/net/sunrpc/sched.c b/net/sunrpc/sched.c
index 0cc83839c13c..ebd37d7d8096 100644
--- a/net/sunrpc/sched.c
+++ b/net/sunrpc/sched.c
@@ -251,9 +251,12 @@ void rpc_destroy_wait_queue(struct rpc_wait_queue *queue)
 }
 EXPORT_SYMBOL_GPL(rpc_destroy_wait_queue);
 
-static int rpc_wait_bit_killable(struct wait_bit_key *key, int mode)
+static int rpc_wait_bit_killable(struct wait_bit_key *key, int mode,
+				 struct wait_queue_entry *wait)
 {
-	freezable_schedule_unsafe();
+	freezer_do_not_count();
+	wait_woken(wait, mode, MAX_SCHEDULE_TIMEOUT);
+	freezer_count_unsafe();
 	if (signal_pending_state(mode, current))
 		return -ERESTARTSYS;
 	return 0;