folio_end_dropbehind() uses in_task() to keep folio invalidation out of
interrupt context. Task context alone is not sufficient: preemption can
still be disabled, or the task can be in a preemptible RCU read-side
critical section, while filemap_end_dropbehind() may reach
folio_unmap_invalidate() and sleep.

Use the established conservative three-part atomic-context test: reject
preemptible RCU read-side sections, reject configurations without
PREEMPT_COUNT, and otherwise require a preemptible context. Unsafe
completions retain the existing best-effort behavior and skip invalidation.

Signed-off-by: Wenjie Qi <qiwenjie@xiaomi.com>
---
 mm/filemap.c | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/mm/filemap.c b/mm/filemap.c
index 6afec6368..0616057e6 100644
--- a/mm/filemap.c
+++ b/mm/filemap.c
@@ -45,6 +45,7 @@
 #include <linux/migrate.h>
 #include <linux/pipe_fs_i.h>
 #include <linux/splice.h>
+#include <linux/rcupdate.h>
 #include <linux/rcupdate_wait.h>
 #include <linux/sched/mm.h>
 #include <linux/sysctl.h>
@@ -1620,6 +1621,15 @@ static void filemap_end_dropbehind(struct folio *folio)
 		folio_unmap_invalidate(mapping, folio, 0);
 }
 
+static bool folio_dropbehind_in_atomic(void)
+{
+	if (IS_ENABLED(CONFIG_PREEMPTION) && rcu_preempt_depth())
+		return true;
+	if (!IS_ENABLED(CONFIG_PREEMPT_COUNT))
+		return true;
+	return !preemptible();
+}
+
 /*
  * If folio was marked as dropbehind, then pages should be dropped when writeback
  * completes. Do that now. If we fail, it's likely because of a big folio -
@@ -1631,13 +1641,12 @@ void folio_end_dropbehind(struct folio *folio)
 		return;
 
 	/*
-	 * Hitting !in_task() should not happen off RWF_DONTCACHE writeback,
-	 * but can happen if normal writeback just happens to find dirty folios
-	 * that were created as part of uncached writeback, and that writeback
-	 * would otherwise not need non-IRQ handling. Just skip the
+	 * Hitting an atomic context should not happen from RWF_DONTCACHE
+	 * writeback, but can happen if normal writeback just happens to find
+	 * dirty folios created as part of uncached writeback. Just skip the
 	 * invalidation in that case.
 	 */
-	if (in_task() && folio_trylock(folio)) {
+	if (!folio_dropbehind_in_atomic() && folio_trylock(folio)) {
 		filemap_end_dropbehind(folio);
 		folio_unlock(folio);
 	}
-- 
2.43.0