diff --git a/include/linux/interval_tree_generic.h b/include/linux/interval_tree_generic.h
index aaa8a0767aa3a512c978d047556af3cee07af66f..6e27d315e6cfa16950c60045a509d23431f16b7c 100644
--- a/include/linux/interval_tree_generic.h
+++ b/include/linux/interval_tree_generic.h
@@ -85,6 +85,7 @@ ITPREFIX ## _subtree_search(ITSTRUCT *node, ITTYPE start, ITTYPE last)	      \
 		 * Loop invariant: start <= node->ITSUBTREE		      \
 		 * (Cond2 is satisfied by one of the subtree nodes)	      \
 		 */							      \
+		cond_resched();						      \
 		if (node->ITRB.rb_left) {				      \
 			ITSTRUCT *left = rb_entry(node->ITRB.rb_left,	      \
 						  ITSTRUCT, ITRB);	      \
diff --git a/lib/interval_tree.c b/lib/interval_tree.c
index 593ce56ece50500b63f9fa5cdb93f1fa09697f1a..77864645523ed4bdea1556b77b618fcaa7f1bad2 100644
--- a/lib/interval_tree.c
+++ b/lib/interval_tree.c
@@ -3,6 +3,7 @@
 #include <linux/interval_tree_generic.h>
 #include <linux/compiler.h>
 #include <linux/export.h>
+#include <linux/sched.h>
 
 #define START(node) ((node)->start)
 #define LAST(node)  ((node)->last)