[RFC PATCH 0/6] fs: range-based admission for generation-drained DIO

[RFC PATCH 0/6] fs: range-based admission for generation-drained DIO

[Vitali Lovich]

Hi all,

This patch series attempts to resolve a major source of contention between
concurrent readers and writers on the same inode, specifically when
those writers are punching holes in the file or gapped appends provided
that readers aren't overlapping these operations. There's some interface
changes I needed just for this initial exclusive IOLOCK removal because
DIO completion needs to cooperate and that lives in generic layers.

I believe the patch set to be complete and well-tested. The details
performance numbers are available in my earlier XFS post:
https://lore.kernel.org/linux-xfs/20261019182620.1527591-1-vlovich@gmail.com/T/#t
This patch series is identical, just widening the group.

Performance measurements on XFS available in the previous posting indicate
substantial read throughput + P99 latency improvements under concurrent
read/write contention with no negative impact to pure read scenarios.
Fixing all sources of contention even manages to improve pure read
workloads by ~2-5%.

Technical overview
------------------

To opt-in, the filesystem serializes compatible punch/gapped operations
with a per-inode drain mutex (only for filesystems that opt-in). It
publishes an exclusion range through a seqcount latch and flips the DIO
generation counter, capturing the outstanding request count. Completion of
old-generation requests retires that finite cohort to wake the drainer;
this is similar semantics to today with the exception that reads can
continue during this exclusive drain mutex hold provided they don't
overlap the range (the writer doesn't wait for them to drain either
since it's guaranteed they don't overlap). Overlapping reads queue up
behind the drain mutex.

Filesystems not using the new interface don't pay a cost.

Not all cases are eligible for this fast-path and those degrade to the old
exclusive IOLOCK + drain machinery, including DAX, realtime and unaligned
punches; the goal is to remain conservative for now while correctly solving
scaling bottlenecks for the common cases. I've tried to make sure there's
good additional test coverage for the new code paths.

I expect the same concept would also benefit ext4 and Btrfs which
similarly hold exclusive IOLOCKs; I have only double-checked that they
broadly follow the similar pattern but I'm not 100% sure if the
interface introduced here strictly applies the same to solve the same
scaling problem. I could imagine small adjustments being necessary. Like
XFS they have additional sources of contention present beyond what this
specific patch series solves and like XFS I expect filesystem-specific
optimizations will need to be written; it's possible there could be a
shared mechanism to encode some of this but I think all contention sources
for at least two or three filesystems need to be solved before generic
interfaces are codified.

In particular, I would appreciate feedback on:

* Whether the generic inode DIO accounting and iomap interfaces, with
  filesystem-private lifetime management through ->dio_drain_complete is
  correct.
* The generation/count encoding and ordering guarantees.
* The XFS lock/revalidation contract and whether my analysis is correct that
  ext4/Btrfs would benefit from the same interface or if different hooks
  are needed.

The base is 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4. The patch ordering is
generic accounting, iomap ownership, KUnit coverage, XFS admission, aligned
hole punching, then post-EOF zeroing.

Thanks,
Vitali

Vitali Lovich (6):
  fs: implement range-based generation-coordinated DIO draining
  iomap: support generation accounting and caller-registered DIO
  fs: add KUnit coverage for generation-based DIO draining
  xfs: add range-drain state and latch-based DIO admission
  xfs: use range-based DIO draining for aligned hole punches
  xfs: use range-based DIO draining for post-EOF zeroing

 Documentation/filesystems/locking.rst         |  89 +-
 Documentation/filesystems/vfs.rst             |  41 +
 fs/Kconfig                                    |  13 +
 fs/Makefile                                   |   1 +
 fs/ext4/file.c                                |   2 +-
 fs/inode.c                                    | 120 ++-
 fs/iomap/direct-io.c                          |  56 +-
 fs/tests/dio_drain_kunit.c                    | 893 ++++++++++++++++++
 fs/xfs/Kconfig                                |   9 +
 fs/xfs/Makefile                               |   2 +
 fs/xfs/tests/xfs_dio_test.c                   | 190 ++++
 fs/xfs/xfs_bmap_util.c                        |   9 +-
 fs/xfs/xfs_dio_test.h                         |  32 +
 fs/xfs/xfs_file.c                             | 357 ++++++-
 fs/xfs/xfs_icache.c                           |   2 +
 fs/xfs/xfs_inode.c                            |  10 +
 fs/xfs/xfs_inode.h                            |   1 +
 fs/xfs/xfs_iops.c                             |   9 +
 fs/xfs/xfs_reflink.c                          |   2 +
 fs/xfs/xfs_reflink.h                          |   2 +-
 fs/xfs/xfs_super.c                            |   2 +
 include/linux/fs.h                            |  97 +-
 include/linux/iomap.h                         |  60 +-
 .../litmus-tests/DIO-register-first.litmus    |  36 +
 tools/testing/selftests/filesystems/Makefile  |   6 +
 .../selftests/filesystems/xfs_dio_admission.c | 380 ++++++++
 .../selftests/filesystems/xfs_dio_drain.c     | 392 ++++++++
 .../selftests/filesystems/xfs_dio_remap.c     | 391 ++++++++
 28 files changed, 3116 insertions(+), 88 deletions(-)
 create mode 100644 fs/tests/dio_drain_kunit.c
 create mode 100644 fs/xfs/tests/xfs_dio_test.c
 create mode 100644 fs/xfs/xfs_dio_test.h
 create mode 100644 tools/memory-model/litmus-tests/DIO-register-first.litmus
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_admission.c
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_drain.c
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_remap.c


base-commit: 6628e7d4c6b25c8d18f1345e3a20d21829aae7e4
-- 
2.56.0

[RFC PATCH 1/6] fs: implement range-based generation-coordinated DIO draining

[Vitali Lovich]

Whole-inode direct I/O (DIO) draining currently requires filesystems to
exclude all new submission requests, even those completely disjoint from
the range being modified.

Add generation accounting and range admission to allow filesystems to
drain a finite cohort of in-flight DIO while concurrently admitting
new, disjoint DIO under a shared i_rwsem.

This reserves the top two bits of i_inode->i_dio_count for generation
and draining state, utilizes seqcount_latch_t for lockless range snapshot
validation, and introduces a new ->dio_drain_complete VFS operation to
resolve private drain state lazily without bloating struct inode or
per-request storage.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 Documentation/filesystems/locking.rst         |  89 ++++++++++-----
 Documentation/filesystems/vfs.rst             |  41 +++++++
 fs/inode.c                                    | 107 +++++++++++++++++-
 fs/xfs/xfs_reflink.h                          |   2 +-
 include/linux/fs.h                            |  92 ++++++++++++++-
 .../litmus-tests/DIO-register-first.litmus    |  36 ++++++
 6 files changed, 334 insertions(+), 33 deletions(-)
 create mode 100644 tools/memory-model/litmus-tests/DIO-register-first.litmus

diff --git a/Documentation/filesystems/locking.rst b/Documentation/filesystems/locking.rst
index f7363ad770c7..cfd4ce7ed563 100644
--- a/Documentation/filesystems/locking.rst
+++ b/Documentation/filesystems/locking.rst
@@ -93,40 +93,44 @@ prototypes::
 	int (*fileattr_get)(struct dentry *dentry, struct file_kattr *fa);
 	struct posix_acl * (*get_acl)(const struct mnt_idmap *, struct dentry *, int);
 	struct offset_ctx *(*get_offset_ctx)(struct inode *inode);
+	void (*dio_drain_complete)(struct inode *inode);
 
 locking rules:
-	all may block
+	all may block except dio_drain_complete
 
-==============	==================================================
-ops		i_rwsem(inode)
-==============	==================================================
-lookup:		shared
-create:		exclusive
-link:		exclusive (both)
-mknod:		exclusive
-symlink:	exclusive
-mkdir:		exclusive
-unlink:		exclusive (both)
-rmdir:		exclusive (both)(see below)
-rename:		exclusive (both parents, some children)	(see below)
-readlink:	no
-get_link:	no
-setattr:	exclusive
-permission:	no (may not block if called in rcu-walk mode)
-get_inode_acl:	no
-get_acl:	no
-getattr:	no
-listxattr:	no
-fiemap:		no
-update_time:	no
-sync_lazytime:	no
-atomic_open:	shared (exclusive if O_CREAT is set in open flags)
-tmpfile:	no
-fileattr_get:	no or exclusive
-fileattr_set:	exclusive
-get_offset_ctx  no
-==============	==================================================
+==================  ==================================================
+ops                 i_rwsem(inode)
+==================  ==================================================
+lookup:             shared
+create:             exclusive
+link:               exclusive (both)
+mknod:              exclusive
+symlink:            exclusive
+mkdir:              exclusive
+unlink:             exclusive (both)
+rmdir:              exclusive (both)(see below)
+rename:             exclusive (both parents, some children) (see below)
+readlink:           no
+get_link:           no
+setattr:            exclusive
+permission:         no (may not block if called in rcu-walk mode)
+get_inode_acl:      no
+get_acl:            no
+getattr:            no
+listxattr:          no
+fiemap:             no
+update_time:        no
+sync_lazytime:      no
+atomic_open:        shared (exclusive if O_CREAT is set in open flags)
+tmpfile:            no
+fileattr_get:       no or exclusive
+fileattr_set:       exclusive
+get_offset_ctx      no
+dio_drain_complete  no
+==================  ==================================================
 
+->dio_drain_complete() does not require i_rwsem and must not block. It may
+run in interrupt context.
 
 	Additionally, ->rmdir(), ->unlink() and ->rename() have ->i_rwsem
 	exclusive on victim.
@@ -138,6 +142,31 @@ get_offset_ctx  no
 See Documentation/filesystems/directory-locking.rst for more detailed discussion
 of the locking scheme for directory operations.
 
+Direct I/O range draining
+=========================
+
+Range drainers acquire ``drain->lock`` before ``inode->i_rwsem`` and any
+filesystem locks needed by the protected operation. Hold the drain mutex
+through inode_dio_drain_finish(): it serializes range publication, cohort
+draining, the protected operation, and exclusion removal. Never wait for
+this mutex while holding i_rwsem, even shared.
+
+A drainer can drop i_rwsem and reacquire it exclusively during write
+checks, for example to strip privileges. Acquiring the mutex inside shared
+i_rwsem would permit this deadlock:
+
+1. Drainer A holds the drain mutex and shared i_rwsem.
+2. Drainer B holds shared i_rwsem and waits for A's drain mutex.
+3. A drops shared i_rwsem and tries to acquire it exclusively. B's shared
+   lock prevents A from proceeding, while B cannot release it until A
+   releases the mutex.
+
+Callers that already hold i_rwsem must drop it before waiting for the drain
+mutex, then reacquire i_rwsem and repeat their checks. DIO admission and
+completion do not acquire the drain mutex, so the old cohort can complete
+while a drainer holds it. Drain the old cohort before taking any filesystem
+lock that its I/O preparation or completion might need.
+
 xattr_handler operations
 ========================
 
diff --git a/Documentation/filesystems/vfs.rst b/Documentation/filesystems/vfs.rst
index 1824462a4371..60fa5c12e451 100644
--- a/Documentation/filesystems/vfs.rst
+++ b/Documentation/filesystems/vfs.rst
@@ -445,6 +445,7 @@ As of kernel 2.6.22, the following members are defined:
 				    struct dentry *dentry, struct file_kattr *fa);
 		int (*fileattr_get)(struct dentry *dentry, struct file_kattr *fa);
 	        struct offset_ctx *(*get_offset_ctx)(struct inode *inode);
+		void (*dio_drain_complete)(struct inode *inode);
 	};
 
 Again, all methods are called without any locks being held, unless
@@ -634,6 +635,46 @@ otherwise noted.
         filesystem must define this operation to use
         simple_offset_dir_operations.
 
+``dio_drain_complete``
+	called by inode_dio_end_generation() when a direct I/O request
+	from the old generation completes while a drain is active.
+	Filesystems that use inode_dio_drain_start() must implement this
+	method to look up their private struct inode_dio_drain pointer
+	(allocated via inode_dio_drain_get() and retained until inode
+	reclamation) and pass it to inode_dio_drain_complete(). Because
+	a drain may be started while earlier requests are already in
+	flight, the method must locate the drain state from the inode
+	itself rather than from per-request submission state. May be
+	called in interrupt context and must not sleep.
+
+	inode_dio_drain_start() returns -EOPNOTSUPP without changing drain
+	state if this method is absent. Callers must check for success
+	before waiting for or finishing the drain.
+
+	Drain operations serialize with the drain mutex, then hold i_rwsem
+	shared to exclude truncate and buffered writes. Never wait for the
+	drain mutex while holding i_rwsem: its owner may need to drop shared
+	i_rwsem and reacquire it exclusively during write checks. See the
+	deadlock example in Documentation/filesystems/locking.rst.
+
+	Every DIO admission must use inode_dio_try_begin() to register
+	provisionally with a fully ordered atomic increment before reading the
+	range snapshot and preparing or submitting I/O. A failed admission
+	withdraws its token with generation-aware completion accounting before
+	returning. The caller must drop its submission lock before waiting for
+	the exclusion to clear, then reacquire the lock and recheck the request.
+
+	inode_dio_drain_start() publishes both copies of a seqcount_latch_t
+	range snapshot before the fully ordered generation flip and DIO count
+	snapshot. Earlier registrations belong to the finite old cohort;
+	later registrations observe the exclusion or a later snapshot.
+	Nonoverlapping new-generation I/O can proceed without extending the
+	old cohort. Readers only load and validate scalar snapshots, and do
+	not acquire the drain mutex, including on PREEMPT_RT. The lazy drain
+	allocation is acquired through its published pointer and lives until
+	inode reclaim. inode_dio_drain_finish() publishes an inactive snapshot
+	and wakes blocked admissions. Range reuse requires no RCU grace period.
+
 The Address Space Object
 ========================
 
diff --git a/fs/inode.c b/fs/inode.c
index f63535ee3629..7b027547b305 100644
--- a/fs/inode.c
+++ b/fs/inode.c
@@ -358,6 +358,7 @@ struct inode *alloc_inode(struct super_block *sb)
 
 void __destroy_inode(struct inode *inode)
 {
+	WARN_ON_ONCE(atomic_read(&inode->i_dio_count) & INODE_DIO_DRAINING);
 	inode_detach_wb(inode);
 	security_inode_free(inode);
 	fsnotify_inode_delete(inode);
@@ -2778,10 +2779,114 @@ EXPORT_SYMBOL(inode_owner_or_capable);
  */
 bool inode_dio_finished(const struct inode *inode)
 {
-	return atomic_read(&inode->i_dio_count) == 0;
+	return !(atomic_read(&inode->i_dio_count) & INODE_DIO_COUNT_MASK);
 }
 EXPORT_SYMBOL(inode_dio_finished);
 
+/*
+ * Allocated on the first draining operation, never on ordinary DIO. The
+ * filesystem owns the pointer and retains the allocation until inode reclaim.
+ */
+struct inode_dio_drain *inode_dio_drain_get(struct inode_dio_drain **drainp)
+{
+	struct inode_dio_drain *drain, *existing;
+
+	/* Acquire the initialized metadata published by the successful cmpxchg. */
+	drain = smp_load_acquire(drainp);
+	if (drain)
+		return drain;
+	drain = kzalloc_obj(*drain, GFP_NOFS);
+	if (!drain)
+		return ERR_PTR(-ENOMEM);
+	mutex_init(&drain->lock);
+	seqcount_latch_init(&drain->range_seq);
+	existing = cmpxchg(drainp, NULL, drain);
+	if (existing) {
+		kfree(drain);
+		/* A failed cmpxchg does not acquire the winner's initialization. */
+		drain = smp_load_acquire(drainp);
+	}
+	return drain;
+}
+EXPORT_SYMBOL_GPL(inode_dio_drain_get);
+
+/* Writers hold drain->lock; readers retain scalars, never copy pointers. */
+static void inode_dio_publish_range(struct inode_dio_drain *drain,
+		bool active, loff_t start, loff_t last)
+{
+	struct inode_dio_range range = {
+		.active = active, .start = start, .last = last,
+	};
+
+	lockdep_assert_held(&drain->lock);
+	write_seqcount_latch_begin(&drain->range_seq);
+	drain->range[0] = range;
+	write_seqcount_latch(&drain->range_seq);
+	drain->range[1] = range;
+	write_seqcount_latch_end(&drain->range_seq);
+}
+
+/*
+ * The drain mutex serializes drain operations before they acquire the shared
+ * inode lock. Never wait for it while holding i_rwsem: its owner may need to
+ * drop shared i_rwsem and reacquire it exclusively during write checks.
+ * Keep the mutex through the protected operation and inode_dio_drain_finish().
+ * Every DIO on this inode must use inode_dio_try_begin() for admission and
+ * generation-aware completion accounting, including pre-submission failures.
+ * The filesystem must publish the drain pointer before calling this function
+ * and provide ->dio_drain_complete to find it at completion, including for
+ * requests admitted before the first drain allocation.
+ * Return -EOPNOTSUPP without publishing a drain if the callback is absent.
+ */
+int inode_dio_drain_start(struct inode *inode, struct inode_dio_drain *drain,
+			 loff_t start, loff_t last)
+{
+	unsigned int old;
+
+	lockdep_assert_held(&drain->lock);
+	lockdep_assert_held(&inode->i_rwsem);
+	if (!inode->i_op->dio_drain_complete)
+		return -EOPNOTSUPP;
+	WARN_ON_ONCE(atomic_read(&inode->i_dio_count) & INODE_DIO_DRAINING);
+	WARN_ON_ONCE(atomic_read(&drain->pending));
+	WARN_ON_ONCE(drain->range[0].active);
+	atomic_set(&drain->pending, 0);
+	/* Both latch copies are published before the fully ordered cohort cutover. */
+	inode_dio_publish_range(drain, true, start, last);
+	old = atomic_fetch_xor(INODE_DIO_DRAINING | INODE_DIO_GENERATION,
+			&inode->i_dio_count);
+	/* Early completions subtract from zero; adding the snapshot reconciles it. */
+	atomic_add(old & INODE_DIO_COUNT_MASK, &drain->pending);
+	return 0;
+}
+EXPORT_SYMBOL_GPL(inode_dio_drain_start);
+
+void inode_dio_drain_wait(struct inode *inode, struct inode_dio_drain *drain)
+{
+	lockdep_assert_held(&drain->lock);
+	wait_var_event(&inode->i_dio_count, !atomic_read(&drain->pending));
+}
+EXPORT_SYMBOL_GPL(inode_dio_drain_wait);
+
+void inode_dio_drain_finish(struct inode *inode, struct inode_dio_drain *drain)
+{
+	lockdep_assert_held(&drain->lock);
+	WARN_ON_ONCE(atomic_read(&drain->pending));
+	atomic_fetch_and(~INODE_DIO_DRAINING, &inode->i_dio_count);
+	inode_dio_publish_range(drain, false, 0, 0);
+	/* Order exclusion clearing before wake_up_var's waiter check. */
+	smp_mb();
+	wake_up_var(&inode->i_dio_count);
+}
+EXPORT_SYMBOL_GPL(inode_dio_drain_finish);
+
+void inode_dio_drain_complete(struct inode *inode, struct inode_dio_drain *drain)
+{
+	if (atomic_dec_and_test(&drain->pending))
+		wake_up_var(&inode->i_dio_count);
+}
+EXPORT_SYMBOL_GPL(inode_dio_drain_complete);
+
 /**
  * inode_dio_wait - wait for outstanding DIO requests to finish
  * @inode: inode to wait for
diff --git a/fs/xfs/xfs_reflink.h b/fs/xfs/xfs_reflink.h
index 683c1841e640..c862b9a01c78 100644
--- a/fs/xfs/xfs_reflink.h
+++ b/fs/xfs/xfs_reflink.h
@@ -20,7 +20,7 @@ xfs_can_free_cowblocks(struct xfs_inode *ip)
 	if ((inode_state_read_once(inode) & I_DIRTY_PAGES) ||
 	    mapping_tagged(inode->i_mapping, PAGECACHE_TAG_DIRTY) ||
 	    mapping_tagged(inode->i_mapping, PAGECACHE_TAG_WRITEBACK) ||
-	    atomic_read(&inode->i_dio_count))
+	    !inode_dio_finished(inode))
 		return false;
 	return true;
 }
diff --git a/include/linux/fs.h b/include/linux/fs.h
index 99caef43ff3c..016f621345af 100644
--- a/include/linux/fs.h
+++ b/include/linux/fs.h
@@ -24,6 +24,7 @@
 #include <linux/mm_types.h>
 #include <linux/capability.h>
 #include <linux/semaphore.h>
+#include <linux/seqlock.h>
 #include <linux/fcntl.h>
 #include <linux/rculist_bl.h>
 #include <linux/atomic.h>
@@ -2040,6 +2041,8 @@ struct inode_operations {
 			    struct dentry *dentry, struct file_kattr *fa);
 	int (*fileattr_get)(struct dentry *dentry, struct file_kattr *fa);
 	struct offset_ctx *(*get_offset_ctx)(struct inode *inode);
+	/* Complete an old-generation DIO during a drain; must not sleep. */
+	void (*dio_drain_complete)(struct inode *inode);
 } ____cacheline_aligned;
 
 /* Did the driver provide valid mmap hook configuration? */
@@ -3170,6 +3173,92 @@ bool inode_dio_finished(const struct inode *inode);
 void inode_dio_wait(struct inode *inode);
 void inode_dio_wait_interruptible(struct inode *inode);
 
+/*
+ * Completion and drain publication share this single atomic cutover. Only
+ * inodes whose DIO always uses the generation-aware helpers may arm a drain.
+ * The low bits retain the normal DIO count, including newer generations.
+ */
+#define INODE_DIO_COUNT_MASK	((1U << 30) - 1)
+#define INODE_DIO_GENERATION	(1U << 30)
+#define INODE_DIO_DRAINING	(1U << 31)
+
+struct inode_dio_range {
+	bool		active;
+	loff_t		start;
+	loff_t		last;
+};
+
+struct inode_dio_drain {
+	struct mutex	lock;
+	atomic_t	pending;
+	seqcount_latch_t range_seq;
+	struct inode_dio_range range[2];
+};
+
+struct inode_dio_drain *inode_dio_drain_get(struct inode_dio_drain **drainp);
+int __must_check inode_dio_drain_start(struct inode *inode,
+		struct inode_dio_drain *drain, loff_t start, loff_t last);
+void inode_dio_drain_wait(struct inode *inode, struct inode_dio_drain *drain);
+void inode_dio_drain_finish(struct inode *inode, struct inode_dio_drain *drain);
+void inode_dio_drain_complete(struct inode *inode, struct inode_dio_drain *drain);
+
+/* The allocation lives until inode reclaim; readers need no RCU section. */
+static inline bool inode_dio_drain_overlap(struct inode *inode,
+		struct inode_dio_drain *drain, loff_t pos, loff_t len)
+{
+	struct inode_dio_range range;
+	unsigned int seq;
+	loff_t last;
+
+	if (!len || !drain)
+		return false;
+	do {
+		seq = read_seqcount_latch(&drain->range_seq);
+		range.active = READ_ONCE(drain->range[seq & 1].active);
+		range.start = READ_ONCE(drain->range[seq & 1].start);
+		range.last = READ_ONCE(drain->range[seq & 1].last);
+	} while (read_seqcount_latch_retry(&drain->range_seq, seq));
+	last = min_t(u64, LLONG_MAX, (u64)pos + len - 1);
+	return range.active && last >= range.start && pos <= range.last;
+}
+
+static inline bool inode_dio_begin_generation(struct inode *inode)
+{
+	return atomic_fetch_inc(&inode->i_dio_count) &
+		INODE_DIO_GENERATION;
+}
+
+static inline void inode_dio_end_generation(struct inode *inode, bool generation)
+{
+	unsigned int old = atomic_fetch_sub(1, &inode->i_dio_count);
+
+	/* A later load of the generation would race with drain publication. */
+	if (unlikely((old & INODE_DIO_DRAINING) &&
+		     !!(old & INODE_DIO_GENERATION) != generation))
+		inode->i_op->dio_drain_complete(inode);
+	if ((old & INODE_DIO_COUNT_MASK) == 1)
+		wake_up_var(&inode->i_dio_count);
+}
+
+/*
+ * Register before inspecting the exclusion. The fully ordered increment pairs
+ * with the publisher's fully ordered generation cutover: a request either
+ * belongs to the finite old cohort or observes the published range (or a
+ * later snapshot). No I/O preparation may precede this admission helper.
+ * Rejected requests withdraw their provisional token before waiting/returning.
+ */
+static inline bool inode_dio_try_begin(struct inode *inode,
+		struct inode_dio_drain **drainp, loff_t pos, loff_t len,
+		bool *generation)
+{
+	*generation = inode_dio_begin_generation(inode);
+	/* Pair with the first allocation's publication in inode_dio_drain_get(). */
+	if (!inode_dio_drain_overlap(inode, smp_load_acquire(drainp), pos, len))
+		return true;
+	inode_dio_end_generation(inode, *generation);
+	return false;
+}
+
 /**
  * inode_dio_begin - signal start of a direct I/O requests
  * @inode: inode the direct I/O happens on
@@ -3191,7 +3280,8 @@ static inline void inode_dio_begin(struct inode *inode)
  */
 static inline void inode_dio_end(struct inode *inode)
 {
-	if (atomic_dec_and_test(&inode->i_dio_count))
+	if ((atomic_fetch_sub(1, &inode->i_dio_count) &
+	     INODE_DIO_COUNT_MASK) == 1)
 		wake_up_var(&inode->i_dio_count);
 }
 
diff --git a/tools/memory-model/litmus-tests/DIO-register-first.litmus b/tools/memory-model/litmus-tests/DIO-register-first.litmus
new file mode 100644
index 000000000000..a7bca403bdfa
--- /dev/null
+++ b/tools/memory-model/litmus-tests/DIO-register-first.litmus
@@ -0,0 +1,36 @@
+C DIO-register-first
+
+(*
+ * Result: Never
+ *
+ * Model one drain cutover, with count bit 0 and generation bit 1.
+ * Publishing both latch copies (or the first initialized drain pointer)
+ * precedes the fully ordered generation RMW. An admission increment which
+ * follows the cutover must see the exclusion. A pre-cutover registration is
+ * instead accounted in the finite old cohort. Completion and subsequent
+ * latch reuse are deliberately outside this single-cutover ordering model.
+ *)
+
+{
+	atomic_t count = 0;
+	int active = 0;
+}
+
+P0(atomic_t *count, int *active)
+{
+	int r0;
+
+	WRITE_ONCE(*active, 1);
+	r0 = atomic_fetch_add(2, count);
+}
+
+P1(atomic_t *count, int *active)
+{
+	int r0;
+	int r1;
+
+	r0 = atomic_fetch_add(1, count);
+	r1 = READ_ONCE(*active);
+}
+
+exists (1:r0=2 /\ 1:r1=0)
-- 
2.56.0

[RFC PATCH 2/6] iomap: support generation accounting and caller-registered DIO

[Vitali Lovich]

Generation draining requires every DIO request on an opted-in inode to
preserve its admission generation until completion. Convert both generic
iomap DIO and the simple-read path to generation-aware accounting,
reusing their existing internal flag fields.

Allow callers to register DIO before entering iomap with
IOMAP_DIO_ACCOUNTED and pass the saved token in IOMAP_DIO_GENERATION.
Iomap owns that registration and releases it on completion or any early
return, including allocation failures, empty requests and -ENOTBLK.
Callers retrying through the generic path must acquire a fresh admission.
Callers without a registration retain internal iomap accounting.

Pass the flags into iomap_dio_read_simple() and update existing XFS and
ext4 callers to pass zero. Preserve the token through mapping and
submission errors, and read it before bio_put() frees the bio-backed
simple-read struct. The completion lifetime still includes filesystem
end_io work and post-write invalidation.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 fs/ext4/file.c        |  2 +-
 fs/iomap/direct-io.c  | 36 +++++++++++++++++++++-------
 fs/xfs/xfs_file.c     |  2 +-
 include/linux/iomap.h | 56 +++++++++++++++++++++++++++++--------------
 4 files changed, 67 insertions(+), 29 deletions(-)

diff --git a/fs/ext4/file.c b/fs/ext4/file.c
index 50d3c92709c8..cb6d6722c18e 100644
--- a/fs/ext4/file.c
+++ b/fs/ext4/file.c
@@ -91,7 +91,7 @@ static ssize_t ext4_dio_read_iter(struct kiocb *iocb, struct iov_iter *to)
 		return generic_file_read_iter(iocb, to);
 	}
 
-	ret = iomap_dio_read_simple(iocb, to, ext4_iomap_begin);
+	ret = iomap_dio_read_simple(iocb, to, ext4_iomap_begin, 0);
 	if (ret == -ENOTBLK)
 		ret = iomap_dio_rw(iocb, to, &ext4_iomap_ops, NULL, 0, NULL, 0);
 	inode_unlock_shared(inode);
diff --git a/fs/iomap/direct-io.c b/fs/iomap/direct-io.c
index ad041ac3252a..fffe6702c42f 100644
--- a/fs/iomap/direct-io.c
+++ b/fs/iomap/direct-io.c
@@ -149,7 +149,8 @@ ssize_t iomap_dio_complete(struct iomap_dio *dio)
 	    !(dio->flags & IOMAP_DIO_NO_INVALIDATE))
 		kiocb_invalidate_post_direct_write(iocb, dio->size);
 
-	inode_dio_end(file_inode(iocb->ki_filp));
+	inode_dio_end_generation(file_inode(iocb->ki_filp),
+				dio->flags & IOMAP_DIO_GENERATION);
 
 	if (ret > 0) {
 		iocb->ki_pos += ret;
@@ -690,6 +691,12 @@ static int iomap_dio_iter(struct iomap_iter *iter, struct iomap_dio *dio)
  * the correct total number of bytes transferred; this is essential for
  * completing partial requests asynchronously.
  *
+ * When @dio_flags includes IOMAP_DIO_ACCOUNTED, iomap takes ownership of the
+ * caller's registration from inode_dio_try_begin() on entry, with its saved
+ * generation token supplied in IOMAP_DIO_GENERATION. Iomap releases the
+ * registration on completion or any early error or NULL return, including
+ * -ENOTBLK. Callers retrying the request must acquire a fresh registration.
+ *
  * Returns -ENOTBLK In case of a page invalidation invalidation failure for
  * writes.  The callers needs to fall back to buffered I/O in this case.
  */
@@ -708,6 +715,8 @@ __iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 	};
 	bool wait_for_completion =
 		is_sync_kiocb(iocb) || (dio_flags & IOMAP_DIO_FORCE_WAIT);
+	bool accounted = dio_flags & IOMAP_DIO_ACCOUNTED;
+	bool generation = dio_flags & IOMAP_DIO_GENERATION;
 	struct blk_plug plug;
 	struct iomap_dio *dio;
 	loff_t ret = 0;
@@ -715,11 +724,13 @@ __iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 	trace_iomap_dio_rw_begin(iocb, iter, dio_flags, done_before);
 
 	if (!iomi.len)
-		return NULL;
+		goto out_account;
 
 	dio = kmalloc_obj(*dio);
-	if (!dio)
-		return ERR_PTR(-ENOMEM);
+	if (!dio) {
+		ret = -ENOMEM;
+		goto out_account;
+	}
 
 	dio->iocb = iocb;
 	atomic_set(&dio->ref, 1);
@@ -816,7 +827,10 @@ __iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 			goto out_free_dio;
 	}
 
-	inode_dio_begin(inode);
+	if (!accounted)
+		generation = inode_dio_begin_generation(inode);
+	if (generation)
+		dio->flags |= IOMAP_DIO_GENERATION;
 
 	blk_start_plug(&plug);
 	while ((ret = iomap_iter(&iomi, ops)) > 0) {
@@ -901,6 +915,9 @@ __iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 
 out_free_dio:
 	kfree(dio);
+out_account:
+	if (accounted)
+		inode_dio_end_generation(inode, generation);
 	if (ret)
 		return ERR_PTR(ret);
 	return NULL;
@@ -943,6 +960,7 @@ static ssize_t iomap_dio_simple_complete(struct iomap_dio_simple *sr)
 	struct bio *bio = &sr->bio;
 	struct kiocb *iocb = sr->iocb;
 	struct inode *inode = file_inode(iocb->ki_filp);
+	bool generation = sr->dio_flags & IOMAP_DIO_GENERATION;
 	ssize_t ret;
 
 	if (unlikely(bio->bi_status)) {
@@ -962,7 +980,7 @@ static ssize_t iomap_dio_simple_complete(struct iomap_dio_simple *sr)
 		bio_release_pages(bio, false);
 		bio_put(bio);
 	}
-	inode_dio_end(inode);
+	inode_dio_end_generation(inode, generation);
 	trace_iomap_dio_complete(iocb, ret < 0 ? ret : 0, ret);
 	return ret;
 }
@@ -980,7 +998,7 @@ static void iomap_dio_simple_end_io(struct bio *bio)
 }
 
 ssize_t __iomap_dio_read_simple(struct kiocb *iocb, struct iov_iter *iter,
-		struct iomap_iter *iomi)
+		struct iomap_iter *iomi, bool generation)
 {
 	gfp_t gfp = (iomi->flags & IOMAP_NOWAIT) ? GFP_NOWAIT : GFP_KERNEL;
 	struct iomap_dio_simple *sr;
@@ -1019,7 +1037,7 @@ ssize_t __iomap_dio_read_simple(struct kiocb *iocb, struct iov_iter *iter,
 	}
 	sr = container_of(bio, struct iomap_dio_simple, bio);
 	sr->iocb = iocb;
-	sr->dio_flags = 0;
+	sr->dio_flags = generation ? IOMAP_DIO_GENERATION : 0;
 
 	bio->bi_iter.bi_sector = iomap_sector(&iomi->iomap, iomi->pos);
 	bio->bi_ioprio = iocb->ki_ioprio;
@@ -1064,7 +1082,7 @@ ssize_t __iomap_dio_read_simple(struct kiocb *iocb, struct iov_iter *iter,
 out_bio_put:
 	bio_put(bio);
 out_dio_end:
-	inode_dio_end(iomi->inode);
+	inode_dio_end_generation(iomi->inode, generation);
 	return ret;
 }
 EXPORT_SYMBOL_GPL(__iomap_dio_read_simple);
diff --git a/fs/xfs/xfs_file.c b/fs/xfs/xfs_file.c
index 7d0c8a197575..124a037cfa9f 100644
--- a/fs/xfs/xfs_file.c
+++ b/fs/xfs/xfs_file.c
@@ -253,7 +253,7 @@ xfs_file_dio_read(
 		ret = iomap_dio_rw(iocb, to, &xfs_read_iomap_ops,
 				&xfs_dio_read_bounce_ops, 0, NULL, 0);
 	} else {
-		ret = iomap_dio_read_simple(iocb, to, xfs_read_iomap_begin);
+		ret = iomap_dio_read_simple(iocb, to, xfs_read_iomap_begin, 0);
 		if (ret == -ENOTBLK)
 			ret = iomap_dio_rw(iocb, to, &xfs_read_iomap_ops, NULL,
 					0, NULL, 0);
diff --git a/include/linux/iomap.h b/include/linux/iomap.h
index 5254430844b2..611c37c19f82 100644
--- a/include/linux/iomap.h
+++ b/include/linux/iomap.h
@@ -710,6 +710,16 @@ struct iomap_dio_ops {
  */
 #define IOMAP_DIO_BOUNCE		(1 << 4)
 
+/*
+ * IOMAP_DIO_ACCOUNTED indicates the caller has already registered this
+ * request with inode_dio_try_begin(); iomap takes ownership of the
+ * registration and releases it on completion or early return (including
+ * -ENOTBLK). IOMAP_DIO_GENERATION carries the generation token returned
+ * by that admission and is only meaningful with IOMAP_DIO_ACCOUNTED.
+ */
+#define IOMAP_DIO_ACCOUNTED	(1 << 5)
+#define IOMAP_DIO_GENERATION	(1 << 6)
+
 ssize_t iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 		const struct iomap_ops *ops, const struct iomap_dio_ops *dops,
 		unsigned int dio_flags, void *private, size_t done_before);
@@ -727,7 +737,8 @@ void iomap_dio_bio_end_io(struct bio *bio);
  * writes, zero-length I/O, and larger requests need the generic iomap direct
  * I/O path.
  *
- * Does not support iomap_dio_ops, dio_flags, done_before or private data.
+ * Does not support iomap_dio_ops, done_before, private data, or dio_flags
+ * other than IOMAP_DIO_ACCOUNTED and IOMAP_DIO_GENERATION.
  * The range must also stay within i_size and encrypted inodes must use the
  * generic iomap direct I/O path.
  *
@@ -739,9 +750,10 @@ void iomap_dio_bio_end_io(struct bio *bio);
  * is never advanced.
  */
 ssize_t __iomap_dio_read_simple(struct kiocb *iocb, struct iov_iter *iter,
-		struct iomap_iter *iomi);
+		struct iomap_iter *iomi, bool generation);
 static __always_inline ssize_t iomap_dio_read_simple(struct kiocb *iocb,
-		struct iov_iter *iter, iomap_iter_begin_fn begin)
+		struct iov_iter *iter, iomap_iter_begin_fn begin,
+		unsigned int dio_flags)
 {
 	struct iomap_iter iomi = {
 		.inode		= file_inode(iocb->ki_filp),
@@ -749,39 +761,47 @@ static __always_inline ssize_t iomap_dio_read_simple(struct kiocb *iocb,
 		.len		= iov_iter_count(iter),
 		.flags		= IOMAP_DIRECT,
 	};
+	bool accounted = dio_flags & IOMAP_DIO_ACCOUNTED;
+	bool generation = dio_flags & IOMAP_DIO_GENERATION;
 	ssize_t ret;
 
-	if (!iomi.len)
-		return 0;
+	if (!iomi.len) {
+		ret = 0;
+		goto out_account;
+	}
 
 	/*
 	 * Simple dio is an optimization for small IO. Filter out large IO
 	 * early as it's the most common case to fail for typical direct IO
 	 * workloads.
 	 */
-	if (iomi.len > iomi.inode->i_sb->s_blocksize)
-		return -ENOTBLK;
-	if (iocb->ki_pos + iomi.len > i_size_read(iomi.inode))
-		return -ENOTBLK;
-	if (IS_ENCRYPTED(iomi.inode))
-		return -ENOTBLK;
+	ret = -ENOTBLK;
+	if (iomi.len > iomi.inode->i_sb->s_blocksize ||
+	    iocb->ki_pos + iomi.len > i_size_read(iomi.inode) ||
+	    IS_ENCRYPTED(iomi.inode))
+		goto out_account;
 
 	ret = kiocb_write_and_wait(iocb, iomi.len);
 	if (ret)
-		return ret;
+		goto out_account;
 
 	if (iocb->ki_flags & IOCB_NOWAIT)
 		iomi.flags |= IOMAP_NOWAIT;
 
-	inode_dio_begin(iomi.inode);
+	if (!accounted) {
+		generation = inode_dio_begin_generation(iomi.inode);
+		accounted = true;
+	}
 	ret = begin(iomi.inode, iomi.pos, iomi.len, iomi.flags, &iomi.iomap,
 			&iomi.srcmap);
-	if (ret) {
-		inode_dio_end(iomi.inode);
-		return ret;
-	}
+	if (ret)
+		goto out_account;
 
-	return __iomap_dio_read_simple(iocb, iter, &iomi);
+	return __iomap_dio_read_simple(iocb, iter, &iomi, generation);
+out_account:
+	if (accounted)
+		inode_dio_end_generation(iomi.inode, generation);
+	return ret;
 }
 
 #ifdef CONFIG_SWAP
-- 
2.56.0

[RFC PATCH 3/6] fs: add KUnit coverage for generation-based DIO draining

[Vitali Lovich]

Exercise admission before lazy allocation, rejection of overlapping DIO,
disjoint new-generation requests and cohort rollover. Check that iomap
consumes caller registrations on simple-read fallback and generic early
returns, including an old-generation completion during an active drain.

Keep generation registration inline in normal builds. When
DIO_DRAIN_KUNIT_TEST is enabled, provide an out-of-line definition in
fs/inode.c with a static stub and one function address across translation
units. Keep the KUnit stub headers out of linux/fs.h.

Use the stub to pause the real admission helper after its provisional
increment and before its first pointer load. Publish the first drain
allocation and verify that cutover completes while the reader is paused.
The resumed overlapping admission must reject the request and withdraw
its old-generation registration. Release both threads and any outstanding
DIO count before reporting failures. Skip this race test without two
online CPUs.

Check snapshot reads while a latch writer is interrupted at either copy
update, and exercise repeated cutover and range reuse across both
generations.

Inject iomap allocation and mapping failures with both saved generation
values. Pause filesystem end_io in the real iomap completion path and
verify that it remains in the old cohort while newer disjoint I/O finishes.
Check successful short-read and failed completion accounting.

Exercise real bio-based writes on a hidden test block device, including
a successful partial request, driver completion errors and block-layer
submission rejection. Check both saved generation values, delayed
filesystem completion, newer disjoint progress and exact token release.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 fs/Kconfig                 |  13 +
 fs/Makefile                |   1 +
 fs/inode.c                 |  13 +
 fs/iomap/direct-io.c       |  20 +-
 fs/tests/dio_drain_kunit.c | 893 +++++++++++++++++++++++++++++++++++++
 include/linux/fs.h         |   5 +
 include/linux/iomap.h      |   4 +
 7 files changed, 948 insertions(+), 1 deletion(-)
 create mode 100644 fs/tests/dio_drain_kunit.c

diff --git a/fs/Kconfig b/fs/Kconfig
index 70a1d0760d28..e004f5a0d6f0 100644
--- a/fs/Kconfig
+++ b/fs/Kconfig
@@ -428,4 +428,17 @@ config FDTABLE_KUNIT_TEST
 	  This builds the fdtable KUnit tests, which tests various aspects
 	  of the fdtable structure and allocation.
 
+config DIO_DRAIN_KUNIT_TEST
+	tristate "KUnit tests for direct I/O range draining" if !KUNIT_ALL_TESTS
+	depends on KUNIT && BLOCK
+	select FS_IOMAP
+	default KUNIT_ALL_TESTS
+	help
+	  Test direct I/O range admission, generation accounting, and
+	  seqcount_latch range snapshots across cohort cutovers, including
+	  iomap early-return ownership and delayed-admission races.
+	  The concurrent admission test requires at least two online CPUs.
+
+	  If unsure, say N.
+
 endmenu
diff --git a/fs/Makefile b/fs/Makefile
index 16f1108b64e1..8974a3bba893 100644
--- a/fs/Makefile
+++ b/fs/Makefile
@@ -28,6 +28,7 @@ obj-$(CONFIG_SIGNALFD)		+= signalfd.o
 obj-$(CONFIG_TIMERFD)		+= timerfd.o
 obj-$(CONFIG_EVENTFD)		+= eventfd.o
 obj-$(CONFIG_AIO)               += aio.o
+obj-$(CONFIG_DIO_DRAIN_KUNIT_TEST) += tests/dio_drain_kunit.o
 obj-$(CONFIG_FS_DAX)		+= dax.o
 obj-$(CONFIG_FS_ENCRYPTION)	+= crypto/
 obj-$(CONFIG_FS_VERITY)		+= verity/
diff --git a/fs/inode.c b/fs/inode.c
index 7b027547b305..b31d5ac5bb96 100644
--- a/fs/inode.c
+++ b/fs/inode.c
@@ -23,6 +23,9 @@
 #include <linux/rw_hint.h>
 #include <linux/seq_file.h>
 #include <linux/debugfs.h>
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+#include <kunit/static_stub.h>
+#endif
 #include <trace/events/writeback.h>
 #define CREATE_TRACE_POINTS
 #include <trace/events/timestamp.h>
@@ -2783,6 +2786,16 @@ bool inode_dio_finished(const struct inode *inode)
 }
 EXPORT_SYMBOL(inode_dio_finished);
 
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+bool inode_dio_begin_generation(struct inode *inode)
+{
+	KUNIT_STATIC_STUB_REDIRECT(inode_dio_begin_generation, inode);
+	return atomic_fetch_inc(&inode->i_dio_count) &
+		INODE_DIO_GENERATION;
+}
+EXPORT_SYMBOL_GPL(inode_dio_begin_generation);
+#endif
+
 /*
  * Allocated on the first draining operation, never on ordinary DIO. The
  * filesystem owns the pointer and retains the allocation until inode reclaim.
diff --git a/fs/iomap/direct-io.c b/fs/iomap/direct-io.c
index fffe6702c42f..2c840358abc8 100644
--- a/fs/iomap/direct-io.c
+++ b/fs/iomap/direct-io.c
@@ -13,6 +13,9 @@
 #include <linux/init.h>
 #include "internal.h"
 #include "trace.h"
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+#include <kunit/static_stub.h>
+#endif
 
 #include "../internal.h"
 
@@ -674,6 +677,21 @@ static int iomap_dio_iter(struct iomap_iter *iter, struct iomap_dio *dio)
 	}
 }
 
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+/* Fault injection is scoped to the task running the ownership test. */
+struct iomap_dio *iomap_dio_alloc(void)
+{
+	KUNIT_STATIC_STUB_REDIRECT(iomap_dio_alloc);
+	return kmalloc_obj(struct iomap_dio);
+}
+EXPORT_SYMBOL_GPL(iomap_dio_alloc);
+#else
+static inline struct iomap_dio *iomap_dio_alloc(void)
+{
+	return kmalloc_obj(struct iomap_dio);
+}
+#endif
+
 /*
  * iomap_dio_rw() always completes O_[D]SYNC writes regardless of whether the IO
  * is being issued as AIO or not.  This allows us to optimise pure data writes
@@ -726,7 +744,7 @@ __iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 	if (!iomi.len)
 		goto out_account;
 
-	dio = kmalloc_obj(*dio);
+	dio = iomap_dio_alloc();
 	if (!dio) {
 		ret = -ENOMEM;
 		goto out_account;
diff --git a/fs/tests/dio_drain_kunit.c b/fs/tests/dio_drain_kunit.c
new file mode 100644
index 000000000000..a99581f9edc4
--- /dev/null
+++ b/fs/tests/dio_drain_kunit.c
@@ -0,0 +1,893 @@
+// SPDX-License-Identifier: GPL-2.0-only
+#include <kunit/static_stub.h>
+#include <kunit/test.h>
+#include <linux/blkdev.h>
+#include <linux/completion.h>
+#include <linux/highmem.h>
+#include <linux/delay.h>
+#include <linux/fs.h>
+#include <linux/iomap.h>
+#include <linux/kthread.h>
+#include <linux/module.h>
+#include <linux/uio.h>
+
+struct dio_drain_race;
+struct dio_test_device;
+
+struct dio_drain_test {
+	struct inode inode;
+	struct super_block sb;
+	struct inode_dio_drain *drain;
+	struct dio_drain_race *race;
+	struct kunit *test;
+	struct completion endio_entered;
+	struct completion endio_release;
+	struct iomap_dio *dio;
+	ssize_t result;
+	int endio_error;
+	struct dio_test_device *device;
+};
+
+static void dio_drain_test_complete(struct inode *inode)
+{
+	struct dio_drain_test *ctx = inode->i_private;
+
+	inode_dio_drain_complete(inode, ctx->drain);
+}
+
+static const struct inode_operations dio_drain_test_iops = {
+	.dio_drain_complete = dio_drain_test_complete,
+};
+
+static int dio_drain_test_init(struct kunit *test)
+{
+	struct dio_drain_test *ctx;
+
+	ctx = kunit_kzalloc(test, sizeof(*ctx), GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, ctx);
+	inode_init_once(&ctx->inode);
+	init_rwsem(&ctx->inode.i_rwsem);
+	ctx->inode.i_mapping = &ctx->inode.i_data;
+	ctx->inode.i_data.host = &ctx->inode;
+	ctx->inode.i_sb = &ctx->sb;
+	ctx->inode.i_op = &dio_drain_test_iops;
+	ctx->inode.i_private = ctx;
+	ctx->sb.s_blocksize = 4096;
+	ctx->test = test;
+	init_completion(&ctx->endio_entered);
+	init_completion(&ctx->endio_release);
+	test->priv = ctx;
+	return 0;
+}
+
+static void dio_drain_test_exit(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+
+	if (!IS_ERR(ctx->drain))
+		kfree(ctx->drain);
+}
+
+static void dio_drain_cohort_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	bool old, newer, resumed;
+	int error;
+
+	/* Count a request admitted before the first drain allocation. */
+	KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			4096, 4096, &old));
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	error = inode_dio_drain_start(&ctx->inode, ctx->drain, 4096, 8191);
+	KUNIT_EXPECT_EQ(test, error, 0);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+	KUNIT_EXPECT_FALSE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			4096, 4096, &resumed));
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+			INODE_DIO_COUNT_MASK, 1);
+	KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			8192, 4096, &newer));
+	KUNIT_EXPECT_NE(test, newer, old);
+	inode_dio_end_generation(&ctx->inode, old);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+	/* The newer request must not extend this wait. */
+	inode_dio_drain_wait(&ctx->inode, ctx->drain);
+	inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			4096, 4096, &resumed));
+	inode_dio_end_generation(&ctx->inode, resumed);
+
+	/* Start a second drain covering the in-flight request from the previous cohort. */
+	error = inode_dio_drain_start(&ctx->inode, ctx->drain, 8192, 12287);
+	KUNIT_EXPECT_EQ(test, error, 0);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+	inode_dio_end_generation(&ctx->inode, newer);
+	inode_dio_drain_wait(&ctx->inode, ctx->drain);
+	inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	KUNIT_EXPECT_FALSE(test, ctx->drain->range[0].active);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+			INODE_DIO_COUNT_MASK, 0);
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+}
+
+static void dio_drain_iomap_ownership_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct file file = {
+		.f_inode = &ctx->inode,
+		.f_mapping = ctx->inode.i_mapping,
+	};
+	struct kiocb iocb = { .ki_filp = &file };
+	struct kvec vec = { .iov_len = 8192 };
+	struct iov_iter iter;
+	unsigned int flags = IOMAP_DIO_ACCOUNTED;
+	bool generation;
+	ssize_t ret;
+
+	KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			0, vec.iov_len, &generation));
+	iov_iter_kvec(&iter, ITER_DEST, &vec, 1, vec.iov_len);
+	ret = iomap_dio_read_simple(&iocb, &iter, NULL, flags);
+	KUNIT_EXPECT_EQ(test, ret, -ENOTBLK);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count), 0);
+
+	/* Generic fallback re-admits; its EOF early return consumes the token. */
+	KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			0, vec.iov_len, &generation));
+	ret = iomap_dio_rw(&iocb, &iter, NULL, NULL, flags, NULL, 0);
+	KUNIT_EXPECT_EQ(test, ret, 0);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count), 0);
+
+	/* Exercise an early failure carrying a nonzero old-generation token. */
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode, ctx->drain,
+			4096, 8191), 0);
+	inode_dio_drain_wait(&ctx->inode, ctx->drain);
+	inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			4096, vec.iov_len, &generation));
+	KUNIT_EXPECT_TRUE(test, generation);
+	flags |= IOMAP_DIO_GENERATION;
+	KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode, ctx->drain,
+			4096, 8191), 0);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+	iov_iter_kvec(&iter, ITER_DEST, &vec, 1, vec.iov_len);
+	ret = iomap_dio_read_simple(&iocb, &iter, NULL, flags);
+	KUNIT_EXPECT_EQ(test, ret, -ENOTBLK);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+	inode_dio_drain_wait(&ctx->inode, ctx->drain);
+	inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+
+	/* An empty request must also release a pre-acquired registration. */
+	iov_iter_kvec(&iter, ITER_DEST, &vec, 1, 0);
+	KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			0, 0, &generation));
+	flags = IOMAP_DIO_ACCOUNTED |
+		(generation ? IOMAP_DIO_GENERATION : 0);
+	ret = iomap_dio_rw(&iocb, &iter, NULL, NULL, flags, NULL, 0);
+	KUNIT_EXPECT_EQ(test, ret, 0);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count), 0);
+}
+
+/* Use real iomap submission/completion with a hole, without a block device. */
+static int dio_drain_map(struct inode *inode, loff_t offset, loff_t length,
+		unsigned int flags, struct iomap *iomap, struct iomap *srcmap)
+{
+	iomap->type = IOMAP_HOLE;
+	iomap->addr = IOMAP_NULL_ADDR;
+	iomap->offset = offset;
+	iomap->length = length;
+	return 0;
+}
+
+static int dio_drain_map_fail(struct inode *inode, loff_t offset, loff_t length,
+		unsigned int flags, struct iomap *iomap, struct iomap *srcmap)
+{
+	return -EAGAIN;
+}
+
+static struct iomap_dio *dio_drain_alloc_fail(void)
+{
+	return NULL;
+}
+
+static const struct iomap_ops dio_drain_map_ops = {
+	.iomap_begin = dio_drain_map,
+};
+
+static const struct iomap_ops dio_drain_map_fail_ops = {
+	.iomap_begin = dio_drain_map_fail,
+};
+
+static void dio_drain_iomap_errors_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct file file = { .f_inode = &ctx->inode,
+		.f_mapping = ctx->inode.i_mapping };
+	struct kiocb iocb = { .ki_filp = &file };
+	struct kvec vec = { .iov_len = 4096 };
+	struct iov_iter iter;
+	unsigned int flags;
+	bool generation;
+	int pass;
+
+	vec.iov_base = kunit_kzalloc(test, vec.iov_len, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, vec.iov_base);
+	i_size_write(&ctx->inode, vec.iov_len);
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	/* Repeat with both saved generation values and an armed old cohort. */
+	for (pass = 0; pass < 2; pass++) {
+		KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, 0, vec.iov_len, &generation));
+		flags = IOMAP_DIO_ACCOUNTED |
+			(generation ? IOMAP_DIO_GENERATION : 0);
+		KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode,
+				ctx->drain, 0, vec.iov_len - 1), 0);
+		iov_iter_kvec(&iter, ITER_DEST, &vec, 1, vec.iov_len);
+		KUNIT_EXPECT_EQ(test, iomap_dio_rw(&iocb, &iter,
+				&dio_drain_map_fail_ops, NULL, flags, NULL, 0),
+				(ssize_t)-EAGAIN);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+
+		KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, 0, vec.iov_len, &generation));
+		flags = IOMAP_DIO_ACCOUNTED |
+			(generation ? IOMAP_DIO_GENERATION : 0);
+		KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode,
+				ctx->drain, 0, vec.iov_len - 1), 0);
+		iov_iter_kvec(&iter, ITER_DEST, &vec, 1, vec.iov_len);
+		KUNIT_EXPECT_EQ(test, iomap_dio_read_simple(&iocb, &iter,
+				dio_drain_map_fail, flags), (ssize_t)-EAGAIN);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+
+		KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, 0, vec.iov_len, &generation));
+		flags = IOMAP_DIO_ACCOUNTED |
+			(generation ? IOMAP_DIO_GENERATION : 0);
+		KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode,
+				ctx->drain, 0, vec.iov_len - 1), 0);
+		kunit_activate_static_stub(test, iomap_dio_alloc,
+				dio_drain_alloc_fail);
+		KUNIT_EXPECT_EQ(test, iomap_dio_rw(&iocb, &iter,
+				&dio_drain_map_ops, NULL, flags, NULL, 0),
+				(ssize_t)-ENOMEM);
+		kunit_deactivate_static_stub(test, iomap_dio_alloc);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	}
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+			INODE_DIO_COUNT_MASK, 0);
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+}
+
+static int dio_drain_endio(struct kiocb *iocb, ssize_t size, int error,
+		unsigned int flags)
+{
+	struct dio_drain_test *ctx = file_inode(iocb->ki_filp)->i_private;
+
+	/* Filesystem completion must still belong to the old cohort. */
+	KUNIT_EXPECT_EQ(ctx->test, atomic_read(&ctx->drain->pending), 1);
+	complete(&ctx->endio_entered);
+	if (!wait_for_completion_timeout(&ctx->endio_release, 5 * HZ))
+		KUNIT_FAIL(ctx->test, "filesystem completion was not released");
+	return ctx->endio_error;
+}
+
+static const struct iomap_dio_ops dio_drain_endio_ops = {
+	.end_io = dio_drain_endio,
+};
+
+static int dio_drain_finish_io(void *data)
+{
+	struct dio_drain_test *ctx = data;
+
+	ctx->result = iomap_dio_complete(ctx->dio);
+	return 0;
+}
+
+static void dio_drain_iomap_endio_check(struct kunit *test, int error)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct file file = { .f_inode = &ctx->inode,
+		.f_mapping = ctx->inode.i_mapping };
+	struct kiocb iocb = { .ki_filp = &file };
+	struct kvec vec = { .iov_len = 8192 };
+	struct iov_iter iter;
+	struct task_struct *task;
+	unsigned int flags;
+	bool generation, newer;
+	bool entered;
+
+	vec.iov_base = kunit_kmalloc(test, vec.iov_len, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, vec.iov_base);
+	memset(vec.iov_base, 0x5a, vec.iov_len);
+	/* Also check a short read and the iterator/position accounting. */
+	i_size_write(&ctx->inode, 4096);
+	KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			0, vec.iov_len, &generation));
+	flags = IOMAP_DIO_ACCOUNTED |
+		(generation ? IOMAP_DIO_GENERATION : 0);
+	iov_iter_kvec(&iter, ITER_DEST, &vec, 1, vec.iov_len);
+	ctx->dio = __iomap_dio_rw(&iocb, &iter, &dio_drain_map_ops,
+			&dio_drain_endio_ops, flags, NULL, 0);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->dio);
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	ctx->endio_error = error;
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode, ctx->drain,
+			0, 4095), 0);
+	task = kthread_create(dio_drain_finish_io, ctx, "dio-endio-test");
+	if (IS_ERR(task)) {
+		complete(&ctx->endio_release);
+		ctx->result = iomap_dio_complete(ctx->dio);
+		KUNIT_FAIL(test, "cannot start filesystem completion");
+	} else {
+		get_task_struct(task);
+		wake_up_process(task);
+		entered = wait_for_completion_timeout(&ctx->endio_entered, HZ);
+		KUNIT_EXPECT_TRUE(test, entered);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+		KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, 8192, 4096, &newer));
+		/* New disjoint I/O can finish while old end_io remains blocked. */
+		inode_dio_end_generation(&ctx->inode, newer);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+		complete(&ctx->endio_release);
+		kthread_stop_put(task);
+	}
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+	inode_dio_drain_wait(&ctx->inode, ctx->drain);
+	inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+	KUNIT_EXPECT_EQ(test, ctx->result, error ? (ssize_t)error : (ssize_t)4096);
+	KUNIT_EXPECT_EQ(test, iocb.ki_pos, error ? (loff_t)0 : (loff_t)4096);
+	KUNIT_EXPECT_EQ(test, iov_iter_count(&iter), (size_t)4096);
+	KUNIT_EXPECT_PTR_EQ(test, memchr_inv(vec.iov_base, 0, 4096), NULL);
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+			INODE_DIO_COUNT_MASK, 0);
+}
+
+static void dio_drain_iomap_endio_test(struct kunit *test)
+{
+	dio_drain_iomap_endio_check(test, 0);
+}
+
+static void dio_drain_iomap_endio_error_test(struct kunit *test)
+{
+	dio_drain_iomap_endio_check(test, -EIO);
+}
+
+/*
+ * A real bio-based block device: iomap allocates and submits the bio through
+ * the block layer.  Keeping the bio in the driver gives the test control over
+ * device completion without replacing iomap's submit or end_io callbacks.
+ */
+struct dio_test_device {
+	struct dio_drain_test *ctx;
+	struct gendisk *disk;
+	struct bio *bio;
+	struct completion submitted;
+	struct completion done;
+	struct kiocb iocb;
+	struct file file;
+	struct iov_iter iter;
+	struct kvec vec;
+	u8 *data;
+	unsigned int dio_flags;
+	unsigned int submissions;
+	unsigned int maps;
+	unsigned int endios;
+	ssize_t result;
+	ssize_t endio_size;
+	int endio_error;
+	int pending_at_endio;
+	bool partial;
+	bool beyond_end;
+};
+
+static void dio_test_submit_bio(struct bio *bio)
+{
+	struct dio_test_device *dev = bio->bi_bdev->bd_disk->private_data;
+
+	dev->submissions++;
+	dev->bio = bio;
+	complete(&dev->submitted);
+}
+
+static const struct block_device_operations dio_test_bdev_ops = {
+	.owner = THIS_MODULE,
+	.submit_bio = dio_test_submit_bio,
+};
+
+static void dio_test_free_disk(void *data)
+{
+	struct gendisk *disk = data;
+
+	del_gendisk(disk);
+	put_disk(disk);
+}
+
+static struct dio_test_device *dio_test_alloc_device(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct queue_limits lim = { .logical_block_size = 512 };
+	struct dio_test_device *dev;
+	struct gendisk *disk;
+	int error;
+
+	dev = kunit_kzalloc(test, sizeof(*dev), GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, dev);
+	dev->ctx = ctx;
+	dev->data = kunit_kmalloc(test, 2 * PAGE_SIZE, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, dev->data);
+	dev->vec.iov_base = kunit_kmalloc(test, 2 * PAGE_SIZE, GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, dev->vec.iov_base);
+	init_completion(&dev->submitted);
+	init_completion(&dev->done);
+	disk = blk_alloc_disk(&lim, NUMA_NO_NODE);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, disk);
+	disk->fops = &dio_test_bdev_ops;
+	disk->private_data = dev;
+	disk->flags |= GENHD_FL_HIDDEN | GENHD_FL_NO_PART;
+	strscpy(disk->disk_name, "dio-kunit", DISK_NAME_LEN);
+	set_capacity(disk, 2 * PAGE_SIZE >> SECTOR_SHIFT);
+	error = add_disk(disk);
+	if (error) {
+		put_disk(disk);
+		KUNIT_FAIL(test, "cannot add block device: %d", error);
+		return NULL;
+	}
+	error = kunit_add_action_or_reset(test, dio_test_free_disk, disk);
+	KUNIT_ASSERT_EQ(test, error, 0);
+	dev->disk = disk;
+	ctx->device = dev;
+	dev->file.f_inode = &ctx->inode;
+	dev->file.f_mapping = ctx->inode.i_mapping;
+	dev->iocb.ki_filp = &dev->file;
+	return dev;
+}
+
+static int dio_test_map(struct inode *inode, loff_t offset, loff_t length,
+		unsigned int flags, struct iomap *iomap, struct iomap *srcmap)
+{
+	struct dio_drain_test *ctx = inode->i_private;
+	struct dio_test_device *dev = ctx->device;
+
+	dev->maps++;
+	/* A second-extent fault after the first page has really been issued. */
+	if (dev->partial && offset >= PAGE_SIZE)
+		return -EFAULT;
+	iomap->type = IOMAP_MAPPED;
+	iomap->offset = offset;
+	iomap->length = PAGE_SIZE;
+	iomap->addr = dev->beyond_end ? 2 * PAGE_SIZE : offset;
+	iomap->bdev = dev->disk->part0;
+	return 0;
+}
+
+static int dio_test_endio(struct kiocb *iocb, ssize_t size, int error,
+		unsigned int flags)
+{
+	struct dio_drain_test *ctx = file_inode(iocb->ki_filp)->i_private;
+	struct dio_test_device *dev = ctx->device;
+
+	dev->endios++;
+	dev->endio_size = size;
+	dev->endio_error = error;
+	dev->pending_at_endio = atomic_read(&ctx->drain->pending);
+	return error;
+}
+
+static const struct iomap_ops dio_test_map_ops = {
+	.iomap_begin = dio_test_map,
+};
+
+static const struct iomap_dio_ops dio_test_endio_ops = {
+	.end_io = dio_test_endio,
+};
+
+static int dio_test_write(void *data)
+{
+	struct dio_test_device *dev = data;
+
+	inode_lock_shared(&dev->ctx->inode);
+	dev->result = iomap_dio_rw(&dev->iocb, &dev->iter,
+			&dio_test_map_ops, &dio_test_endio_ops,
+			dev->dio_flags, NULL, 0);
+	inode_unlock_shared(&dev->ctx->inode);
+	complete(&dev->done);
+	return 0;
+}
+
+static void dio_test_complete_bio(struct dio_test_device *dev, bool error)
+{
+	struct bio *bio = dev->bio;
+	struct bio_vec bvec;
+	struct bvec_iter iter;
+	size_t offset = bio->bi_iter.bi_sector << SECTOR_SHIFT;
+
+	if (!error) {
+		bio_for_each_segment(bvec, bio, iter) {
+			void *source = bvec_kmap_local(&bvec);
+
+			memcpy(dev->data + offset, source, bvec.bv_len);
+			kunmap_local(source);
+			offset += bvec.bv_len;
+		}
+	}
+	dev->bio = NULL;
+	bio_endio_status(bio, error ? BLK_STS_IOERR : BLK_STS_OK);
+}
+
+static void dio_drain_bio_write_check(struct kunit *test, bool partial,
+		bool beyond_end)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct dio_test_device *dev = dio_test_alloc_device(test);
+	struct task_struct *task;
+	bool generation, newer, admitted;
+	unsigned long waited;
+	int pass;
+
+	KUNIT_ASSERT_NOT_NULL(test, dev);
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	i_size_write(&ctx->inode, 2 * PAGE_SIZE);
+	dev->partial = partial;
+	dev->beyond_end = beyond_end;
+	/* Both generation tokens must survive successful and failed writes. */
+	for (pass = 0; pass < 2; pass++) {
+		dev->vec.iov_len = partial ? 2 * PAGE_SIZE : PAGE_SIZE;
+		memset(dev->vec.iov_base, 0x5a, dev->vec.iov_len);
+		memset(dev->data, 0xa5, 2 * PAGE_SIZE);
+		dev->iocb.ki_pos = 0;
+		dev->submissions = dev->maps = dev->endios = 0;
+		reinit_completion(&dev->submitted);
+		reinit_completion(&dev->done);
+		iov_iter_kvec(&dev->iter, ITER_SOURCE, &dev->vec, 1,
+				dev->vec.iov_len);
+		KUNIT_ASSERT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, 0, dev->vec.iov_len, &generation));
+		KUNIT_EXPECT_EQ(test, generation, (bool)pass);
+		dev->dio_flags = IOMAP_DIO_ACCOUNTED |
+			(generation ? IOMAP_DIO_GENERATION : 0) |
+			(partial ? IOMAP_DIO_PARTIAL : 0);
+		mutex_lock(&ctx->drain->lock);
+		inode_lock_shared(&ctx->inode);
+		KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode,
+				ctx->drain, 0, 2 * PAGE_SIZE - 1), 0);
+		task = kthread_create(dio_test_write, dev, "dio-bio-test");
+		if (IS_ERR(task)) {
+			inode_dio_end_generation(&ctx->inode, generation);
+			inode_dio_drain_finish(&ctx->inode, ctx->drain);
+			inode_unlock_shared(&ctx->inode);
+			mutex_unlock(&ctx->drain->lock);
+			KUNIT_FAIL(test, "cannot start direct write");
+			break;
+		}
+		get_task_struct(task);
+		wake_up_process(task);
+		if (!beyond_end) {
+			waited = wait_for_completion_timeout(&dev->submitted,
+					5 * HZ);
+			KUNIT_EXPECT_NE(test, waited, 0UL);
+			KUNIT_EXPECT_EQ(test, dev->submissions, 1U);
+			KUNIT_EXPECT_FALSE(test, completion_done(&dev->done));
+			KUNIT_EXPECT_EQ(test,
+					atomic_read(&ctx->drain->pending), 1);
+			admitted = inode_dio_try_begin(&ctx->inode,
+					&ctx->drain, 2 * PAGE_SIZE, PAGE_SIZE,
+					&newer);
+			KUNIT_EXPECT_TRUE(test, admitted);
+			if (admitted)
+				inode_dio_end_generation(&ctx->inode, newer);
+			KUNIT_EXPECT_EQ(test,
+					atomic_read(&ctx->drain->pending), 1);
+			if (dev->bio) {
+				KUNIT_EXPECT_EQ(test, bio_op(dev->bio),
+						REQ_OP_WRITE);
+				KUNIT_EXPECT_EQ(test, dev->bio->bi_iter.bi_size,
+						(unsigned int)PAGE_SIZE);
+				dio_test_complete_bio(dev, !partial);
+			}
+		}
+		waited = wait_for_completion_timeout(&dev->done, 5 * HZ);
+		KUNIT_EXPECT_NE(test, waited, 0UL);
+		kthread_stop_put(task);
+		KUNIT_EXPECT_EQ(test, dev->submissions, beyond_end ? 0U : 1U);
+		KUNIT_EXPECT_EQ(test, dev->maps, partial ? 2U : 1U);
+		KUNIT_EXPECT_EQ(test, dev->endios, 1U);
+		KUNIT_EXPECT_EQ(test, dev->endio_size, (ssize_t)PAGE_SIZE);
+		KUNIT_EXPECT_EQ(test, dev->endio_error, partial ? 0 : -EIO);
+		KUNIT_EXPECT_EQ(test, dev->pending_at_endio, 1);
+		KUNIT_EXPECT_EQ(test, dev->result,
+				partial ? (ssize_t)PAGE_SIZE : (ssize_t)-EIO);
+		KUNIT_EXPECT_EQ(test, dev->iocb.ki_pos,
+				partial ? (loff_t)PAGE_SIZE : (loff_t)0);
+		KUNIT_EXPECT_EQ(test, iov_iter_count(&dev->iter),
+				partial ? (size_t)PAGE_SIZE : (size_t)0);
+		KUNIT_EXPECT_PTR_EQ(test, memchr_inv(dev->data,
+				partial ? 0x5a : 0xa5, PAGE_SIZE), NULL);
+		KUNIT_EXPECT_PTR_EQ(test, memchr_inv(dev->data + PAGE_SIZE,
+				0xa5, PAGE_SIZE), NULL);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+		/* No I/O remains; let accounting regressions fail without hanging. */
+		if (atomic_read(&ctx->drain->pending)) {
+			atomic_set(&ctx->drain->pending, 0);
+			wake_up_var(&ctx->inode.i_dio_count);
+		}
+		inode_dio_drain_wait(&ctx->inode, ctx->drain);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+				INODE_DIO_COUNT_MASK, 0);
+		inode_unlock_shared(&ctx->inode);
+		mutex_unlock(&ctx->drain->lock);
+	}
+}
+
+static void dio_drain_partial_write_test(struct kunit *test)
+{
+	dio_drain_bio_write_check(test, true, false);
+}
+
+static void dio_drain_bio_completion_error_test(struct kunit *test)
+{
+	dio_drain_bio_write_check(test, false, false);
+}
+
+static void dio_drain_bio_submission_error_test(struct kunit *test)
+{
+	/* The block layer rejects this before the device can see the bio. */
+	dio_drain_bio_write_check(test, false, true);
+}
+
+struct dio_drain_race {
+	struct dio_drain_test *ctx;
+	struct kunit *test;
+	struct completion checked;
+	struct completion reader_done;
+	struct completion started;
+	struct completion finish;
+	bool release;
+	bool admitted;
+	bool registered;
+	bool generation;
+	bool timed_out;
+	int error;
+};
+
+static bool dio_drain_delayed_registration(struct inode *inode)
+{
+	struct dio_drain_test *ctx = inode->i_private;
+	struct dio_drain_race *race = ctx->race;
+	unsigned long deadline = jiffies + 5 * HZ;
+	bool generation;
+
+	/* Pause the real helper after registration, before its first pointer load. */
+	generation = atomic_fetch_inc(&inode->i_dio_count) &
+		INODE_DIO_GENERATION;
+
+	race->registered = true;
+	complete(&race->checked);
+	while (!READ_ONCE(race->release) && time_before(jiffies, deadline))
+		cpu_relax();
+	race->timed_out = !READ_ONCE(race->release);
+	return generation;
+}
+
+static int dio_drain_delayed_admission(void *data)
+{
+	struct dio_drain_race *race = data;
+	struct dio_drain_test *ctx = race->ctx;
+
+	/* Static stubs are scoped to the test associated with the current task. */
+	current->kunit_test = race->test;
+	race->admitted = inode_dio_try_begin(&ctx->inode, &ctx->drain,
+			4096, 4096, &race->generation);
+	current->kunit_test = NULL;
+	complete(&race->reader_done);
+	return 0;
+}
+
+static int dio_drain_race_publisher(void *data)
+{
+	struct dio_drain_race *race = data;
+	struct dio_drain_test *ctx = race->ctx;
+
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	race->error = inode_dio_drain_start(&ctx->inode, ctx->drain, 4096, 8191);
+	complete(&race->started);
+	wait_for_completion(&race->finish);
+	if (!race->error) {
+		inode_dio_drain_wait(&ctx->inode, ctx->drain);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+	}
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+	return 0;
+}
+
+static void dio_drain_delayed_admission_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct dio_drain_race race = { .ctx = ctx, .test = test };
+	struct task_struct *reader, *publisher = NULL;
+	struct inode_dio_drain *drain;
+	unsigned long deadline;
+	bool checked, published = false;
+	int cpu, pending = 0, error = 0;
+
+	if (num_online_cpus() < 2) {
+		kunit_skip(test, "requires two online CPUs");
+		return;
+	}
+	init_completion(&race.checked);
+	init_completion(&race.reader_done);
+	init_completion(&race.started);
+	init_completion(&race.finish);
+	ctx->race = &race;
+	kunit_activate_static_stub(test, inode_dio_begin_generation,
+			dio_drain_delayed_registration);
+	reader = kthread_create(dio_drain_delayed_admission, &race, "dio-admit-test");
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, reader);
+	get_task_struct(reader);
+	cpu = cpumask_any_but(cpu_online_mask, raw_smp_processor_id());
+	kthread_bind(reader, cpu);
+	wake_up_process(reader);
+	checked = wait_for_completion_timeout(&race.checked, HZ);
+	if (!checked)
+		goto out_reader;
+	/* First publication occurs while a registered admission is paused. */
+	drain = inode_dio_drain_get(&ctx->drain);
+	if (IS_ERR(drain)) {
+		error = PTR_ERR(drain);
+		goto out_reader;
+	}
+	publisher = kthread_create(dio_drain_race_publisher, &race, "dio-drain-test");
+	if (IS_ERR(publisher)) {
+		error = PTR_ERR(publisher);
+		publisher = NULL;
+		goto out_reader;
+	}
+	get_task_struct(publisher);
+	wake_up_process(publisher);
+	deadline = jiffies + HZ;
+	/* Cutover must complete while the old registration remains paused. */
+	while (!completion_done(&race.started) &&
+	       time_before(jiffies, deadline))
+		usleep_range(1000, 2000);
+	/* A timeout must not turn this into an admission before publication. */
+	published = completion_done(&race.started);
+out_reader:
+	WRITE_ONCE(race.release, true);
+	wait_for_completion(&race.reader_done);
+	kthread_stop_put(reader);
+	if (publisher) {
+		wait_for_completion(&race.started);
+		pending = atomic_read(&ctx->drain->pending);
+	}
+	if (race.admitted)
+		inode_dio_end_generation(&ctx->inode, race.generation);
+	if (publisher) {
+		complete(&race.finish);
+		kthread_stop_put(publisher);
+	}
+	kunit_deactivate_static_stub(test, inode_dio_begin_generation);
+	ctx->race = NULL;
+
+	/* Report failures only after releasing both tasks and their DIO count. */
+	KUNIT_ASSERT_TRUE_MSG(test, checked, "reader did not reach registration");
+	KUNIT_ASSERT_EQ_MSG(test, error, 0, "cannot start publisher");
+	KUNIT_ASSERT_TRUE_MSG(test, published, "publisher did not publish exclusion");
+	KUNIT_EXPECT_FALSE(test, race.admitted);
+	KUNIT_EXPECT_TRUE(test, race.registered);
+	KUNIT_EXPECT_FALSE(test, race.timed_out);
+	KUNIT_EXPECT_FALSE(test, race.generation);
+	KUNIT_EXPECT_EQ(test, race.error, 0);
+	KUNIT_EXPECT_EQ(test, pending, 0);
+}
+
+/* Even when a publisher is suspended, the selected copy remains readable. */
+static void dio_drain_latch_interrupted_writer_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	struct inode_dio_drain *drain;
+
+	drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, drain);
+	mutex_lock(&drain->lock);
+	write_seqcount_latch_begin(&drain->range_seq);
+	drain->range[0].active = true;
+	drain->range[0].start = 4096;
+	drain->range[0].last = 8191;
+	/* Odd copy is still inactive while the even copy is being replaced. */
+	KUNIT_EXPECT_FALSE(test, inode_dio_drain_overlap(&ctx->inode, drain,
+			4096, 4096));
+	write_seqcount_latch(&drain->range_seq);
+	/* Even copy is now complete while the odd copy is being replaced. */
+	KUNIT_EXPECT_TRUE(test, inode_dio_drain_overlap(&ctx->inode, drain,
+			4096, 4096));
+	KUNIT_EXPECT_FALSE(test, inode_dio_drain_overlap(&ctx->inode, drain,
+			8192, 4096));
+	drain->range[1] = drain->range[0];
+	write_seqcount_latch_end(&drain->range_seq);
+	mutex_unlock(&drain->lock);
+}
+
+static void dio_drain_repeated_cutover_test(struct kunit *test)
+{
+	struct dio_drain_test *ctx = test->priv;
+	bool old, rejected;
+	int i;
+
+	ctx->drain = inode_dio_drain_get(&ctx->drain);
+	KUNIT_ASSERT_NOT_ERR_OR_NULL(test, ctx->drain);
+	mutex_lock(&ctx->drain->lock);
+	inode_lock_shared(&ctx->inode);
+	for (i = 0; i < 128; i++) {
+		loff_t pos = (i & 1) ? 4096 : 8192;
+
+		KUNIT_EXPECT_TRUE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, pos, 4096, &old));
+		KUNIT_EXPECT_EQ(test, inode_dio_drain_start(&ctx->inode,
+				ctx->drain, pos, pos + 4095), 0);
+		KUNIT_EXPECT_FALSE(test, inode_dio_try_begin(&ctx->inode,
+				&ctx->drain, pos, 4096, &rejected));
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 1);
+		inode_dio_end_generation(&ctx->inode, old);
+		KUNIT_EXPECT_EQ(test, atomic_read(&ctx->drain->pending), 0);
+		inode_dio_drain_wait(&ctx->inode, ctx->drain);
+		inode_dio_drain_finish(&ctx->inode, ctx->drain);
+		KUNIT_EXPECT_FALSE(test, inode_dio_drain_overlap(&ctx->inode,
+				ctx->drain, pos, 4096));
+	}
+	KUNIT_EXPECT_EQ(test, atomic_read(&ctx->inode.i_dio_count) &
+			INODE_DIO_COUNT_MASK, 0);
+	inode_unlock_shared(&ctx->inode);
+	mutex_unlock(&ctx->drain->lock);
+}
+
+static struct kunit_case dio_drain_test_cases[] = {
+	KUNIT_CASE(dio_drain_cohort_test),
+	KUNIT_CASE(dio_drain_iomap_ownership_test),
+	KUNIT_CASE(dio_drain_iomap_errors_test),
+	KUNIT_CASE(dio_drain_iomap_endio_test),
+	KUNIT_CASE(dio_drain_iomap_endio_error_test),
+	KUNIT_CASE(dio_drain_partial_write_test),
+	KUNIT_CASE(dio_drain_bio_completion_error_test),
+	KUNIT_CASE(dio_drain_bio_submission_error_test),
+	KUNIT_CASE(dio_drain_delayed_admission_test),
+	KUNIT_CASE(dio_drain_latch_interrupted_writer_test),
+	KUNIT_CASE(dio_drain_repeated_cutover_test),
+	{}
+};
+
+static struct kunit_suite dio_drain_test_suite = {
+	.name = "inode-dio-drain",
+	.init = dio_drain_test_init,
+	.exit = dio_drain_test_exit,
+	.test_cases = dio_drain_test_cases,
+};
+
+kunit_test_suite(dio_drain_test_suite);
+MODULE_DESCRIPTION("KUnit tests for generation-based direct I/O draining");
+MODULE_LICENSE("GPL");
diff --git a/include/linux/fs.h b/include/linux/fs.h
index 016f621345af..3eb62a40ea21 100644
--- a/include/linux/fs.h
+++ b/include/linux/fs.h
@@ -3222,11 +3222,16 @@ static inline bool inode_dio_drain_overlap(struct inode *inode,
 	return range.active && last >= range.start && pos <= range.last;
 }
 
+/* Test builds need a single function address for the registration stub. */
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+bool inode_dio_begin_generation(struct inode *inode);
+#else
 static inline bool inode_dio_begin_generation(struct inode *inode)
 {
 	return atomic_fetch_inc(&inode->i_dio_count) &
 		INODE_DIO_GENERATION;
 }
+#endif
 
 static inline void inode_dio_end_generation(struct inode *inode, bool generation)
 {
diff --git a/include/linux/iomap.h b/include/linux/iomap.h
index 611c37c19f82..f6652f5d8ac5 100644
--- a/include/linux/iomap.h
+++ b/include/linux/iomap.h
@@ -723,6 +723,10 @@ struct iomap_dio_ops {
 ssize_t iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 		const struct iomap_ops *ops, const struct iomap_dio_ops *dops,
 		unsigned int dio_flags, void *private, size_t done_before);
+#if IS_ENABLED(CONFIG_DIO_DRAIN_KUNIT_TEST)
+struct iomap_dio *iomap_dio_alloc(void);
+#endif
+
 struct iomap_dio *__iomap_dio_rw(struct kiocb *iocb, struct iov_iter *iter,
 		const struct iomap_ops *ops, const struct iomap_dio_ops *dops,
 		unsigned int dio_flags, void *private, size_t done_before);
-- 
2.56.0

[RFC PATCH 4/6] xfs: add range-drain state and latch-based DIO admission

[Vitali Lovich]

Prepare XFS for range draining with lazily allocated inode-private state,
retained across VFS inode recycling and freed at XFS inode reclamation.
Provide ->dio_drain_complete to resolve the state for older requests.

Register provisionally before the final latch range snapshot for every
direct read and write submission. Transfer the token to iomap, and obtain
a fresh admission for simple-read fallback or write-check retries.

Overlapping IOCB_NOWAIT requests return -EAGAIN. Blocking requests drop
IOLOCK before waiting for a nonoverlapping snapshot and contribute no DIO
count while asleep. Writers reacquire through
xfs_ilock_iocb_for_write(), repeat append and limit checks, and recheck
unaligned COW eligibility because a clone can share blocks while IOLOCK
is dropped.

Establish the final admission and ownership rules before enabling either
range-drain consumer.

Provide optional root-only debugfs gates and counters for deterministic
filesystem selftests of request revalidation. The controls pin the selected
file and release paused tasks on close or timeout; normal builds omit them.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 fs/xfs/Kconfig              |   9 ++
 fs/xfs/Makefile             |   2 +
 fs/xfs/tests/xfs_dio_test.c | 190 ++++++++++++++++++++++++++++++++++++
 fs/xfs/xfs_dio_test.h       |  32 ++++++
 fs/xfs/xfs_file.c           | 162 ++++++++++++++++++++++++++++--
 fs/xfs/xfs_icache.c         |   2 +
 fs/xfs/xfs_inode.h          |   1 +
 fs/xfs/xfs_iops.c           |   9 ++
 fs/xfs/xfs_reflink.c        |   2 +
 fs/xfs/xfs_super.c          |   2 +
 10 files changed, 401 insertions(+), 10 deletions(-)
 create mode 100644 fs/xfs/tests/xfs_dio_test.c
 create mode 100644 fs/xfs/xfs_dio_test.h

diff --git a/fs/xfs/Kconfig b/fs/xfs/Kconfig
index b99da294e9a3..9d36201b5cdc 100644
--- a/fs/xfs/Kconfig
+++ b/fs/xfs/Kconfig
@@ -204,6 +204,15 @@ config XFS_WARN
 	  However, similar to XFS_DEBUG, it is only advisable to use this if you
 	  are debugging a particular problem.
 
+config XFS_DIO_TEST
+	bool "XFS direct I/O filesystem selftest synchronization"
+	depends on XFS_FS && DEBUG_FS
+	help
+	  Enable root-only debugfs synchronization and event counters used by
+	  direct I/O selftests on a mounted XFS filesystem. Tests can pause a
+	  selected inode at explicit events; closing the control file or a
+	  ten-second timeout releases paused tasks. Say N for production kernels.
+
 config XFS_DEBUG
 	bool "XFS Debugging support"
 	depends on XFS_FS
diff --git a/fs/xfs/Makefile b/fs/xfs/Makefile
index 399a207f2d0e..2c31202b6ba1 100644
--- a/fs/xfs/Makefile
+++ b/fs/xfs/Makefile
@@ -110,6 +110,8 @@ xfs-y				+= xfs_aops.o \
 				   xfs_verify_media.o \
 				   xfs_xattr.o
 
+xfs-$(CONFIG_XFS_DIO_TEST) += tests/xfs_dio_test.o
+
 # low-level transaction/log code
 xfs-y				+= xfs_log.o \
 				   xfs_log_cil.o \
diff --git a/fs/xfs/tests/xfs_dio_test.c b/fs/xfs/tests/xfs_dio_test.c
new file mode 100644
index 000000000000..5d7df903198d
--- /dev/null
+++ b/fs/xfs/tests/xfs_dio_test.c
@@ -0,0 +1,190 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Test-only synchronization for real filesystem DIO selftests. */
+#include <linux/debugfs.h>
+#include <linux/file.h>
+#include <linux/fs.h>
+#include <linux/magic.h>
+#include <linux/module.h>
+#include <linux/mutex.h>
+#include <linux/spinlock.h>
+#include <linux/uaccess.h>
+#include <linux/wait.h>
+#include "xfs_dio_test.h"
+
+static DEFINE_MUTEX(control_lock);
+static DEFINE_SPINLOCK(state_lock);
+static DECLARE_WAIT_QUEUE_HEAD(event_wait);
+static struct inode *target;
+static struct file *target_file;
+static unsigned long generation;
+static bool opened;
+static unsigned long gates, timedout;
+static unsigned long hits[XFS_DIO_TEST_NR_EVENTS];
+static unsigned long values[XFS_DIO_TEST_NR_EVENTS];
+
+void xfs_dio_test_event(struct inode *inode, enum xfs_dio_test_event event,
+		unsigned long value)
+{
+	bool pause;
+	unsigned long current_generation;
+
+	spin_lock(&state_lock);
+	if (target != inode) {
+		spin_unlock(&state_lock);
+		return;
+	}
+	hits[event]++;
+	values[event] = value;
+	pause = gates & BIT(event);
+	current_generation = generation;
+	spin_unlock(&state_lock);
+	wake_up_all(&event_wait);
+	if (pause && !wait_event_timeout(event_wait,
+			READ_ONCE(generation) != current_generation ||
+			!(READ_ONCE(gates) & BIT(event)), 10 * HZ)) {
+		spin_lock(&state_lock);
+		if (generation == current_generation) {
+			timedout++;
+			gates = 0;
+		}
+		spin_unlock(&state_lock);
+		wake_up_all(&event_wait);
+	}
+}
+
+static int dio_test_open(struct inode *inode, struct file *file)
+{
+	int error = 0;
+
+	mutex_lock(&control_lock);
+	if (opened)
+		error = -EBUSY;
+	else
+		opened = true;
+	mutex_unlock(&control_lock);
+	return error;
+}
+
+static int dio_test_release(struct inode *inode, struct file *file)
+{
+	struct file *old;
+
+	mutex_lock(&control_lock);
+	spin_lock(&state_lock);
+	old = target_file;
+	target_file = NULL;
+	target = NULL;
+	generation++;
+	gates = 0;
+	spin_unlock(&state_lock);
+	wake_up_all(&event_wait);
+	if (old)
+		fput(old);
+	opened = false;
+	mutex_unlock(&control_lock);
+	return 0;
+}
+
+static ssize_t dio_test_read(struct file *file, char __user *buf, size_t len,
+		loff_t *ppos)
+{
+	char result[512];
+	unsigned int i, nr;
+
+	spin_lock(&state_lock);
+	nr = scnprintf(result, sizeof(result), "%lu %lu\n", gates, timedout);
+	for (i = 0; i < XFS_DIO_TEST_NR_EVENTS; i++)
+		nr += scnprintf(result + nr, sizeof(result) - nr, "%lu %lu\n",
+				hits[i], values[i]);
+	spin_unlock(&state_lock);
+	return simple_read_from_buffer(buf, len, ppos, result, nr);
+}
+
+static ssize_t dio_test_write(struct file *file, const char __user *buf,
+		size_t len, loff_t *ppos)
+{
+	char command[80];
+	struct inode *inode;
+	unsigned long mask;
+	int fd, error = 0;
+
+	if (!len || len >= sizeof(command))
+		return -EINVAL;
+	if (copy_from_user(command, buf, len))
+		return -EFAULT;
+	command[len] = 0;
+	mutex_lock(&control_lock);
+	if (sysfs_streq(command, "release")) {
+		spin_lock(&state_lock);
+		gates = 0;
+		spin_unlock(&state_lock);
+		wake_up_all(&event_wait);
+		goto out;
+	}
+	if (sscanf(command, "gate %lu", &mask) == 1) {
+		if (mask >= BIT(XFS_DIO_TEST_NR_EVENTS)) {
+			error = -EINVAL;
+			goto out;
+		}
+		spin_lock(&state_lock);
+		gates = mask;
+		spin_unlock(&state_lock);
+		wake_up_all(&event_wait);
+		goto out;
+	}
+	if (sscanf(command, "release %lu", &mask) == 1) {
+		spin_lock(&state_lock);
+		gates &= ~mask;
+		spin_unlock(&state_lock);
+		wake_up_all(&event_wait);
+		goto out;
+	}
+	if (sscanf(command, "arm %d %lu", &fd, &mask) != 2 ||
+	    mask >= BIT(XFS_DIO_TEST_NR_EVENTS)) {
+		error = -EINVAL;
+		goto out;
+	}
+	{
+		CLASS(fd, f)(fd);
+
+		if (fd_empty(f)) {
+			error = -EBADF;
+			goto out;
+		}
+		inode = file_inode(fd_file(f));
+		if (inode->i_sb->s_magic != XFS_SUPER_MAGIC || !S_ISREG(inode->i_mode)) {
+			error = -EINVAL;
+			goto out;
+		}
+		spin_lock(&state_lock);
+		if (target) {
+			error = -EBUSY;
+		} else {
+			target_file = get_file(fd_file(f));
+			target = inode;
+			generation++;
+			gates = mask;
+			timedout = 0;
+			memset(hits, 0, sizeof(hits));
+			memset(values, 0, sizeof(values));
+		}
+		spin_unlock(&state_lock);
+	}
+out:
+	mutex_unlock(&control_lock);
+	return error ? error : len;
+}
+
+static const struct file_operations dio_test_fops = {
+	.owner = THIS_MODULE,
+	.open = dio_test_open,
+	.release = dio_test_release,
+	.read = dio_test_read,
+	.write = dio_test_write,
+	.llseek = default_llseek,
+};
+
+void xfs_dio_test_init(struct dentry *parent)
+{
+	debugfs_create_file("dio_test", 0600, parent, NULL, &dio_test_fops);
+}
diff --git a/fs/xfs/xfs_dio_test.h b/fs/xfs/xfs_dio_test.h
new file mode 100644
index 000000000000..178b668580d8
--- /dev/null
+++ b/fs/xfs/xfs_dio_test.h
@@ -0,0 +1,32 @@
+/* SPDX-License-Identifier: GPL-2.0 */
+#ifndef __XFS_DIO_TEST_H__
+#define __XFS_DIO_TEST_H__
+
+struct inode;
+struct dentry;
+
+enum xfs_dio_test_event {
+	XFS_DIO_TEST_BATCH_LEADER,
+	XFS_DIO_TEST_BATCH_FOLLOWER,
+	XFS_DIO_TEST_BATCH_COLLECT,
+	XFS_DIO_TEST_BATCH_ZERO,
+	XFS_DIO_TEST_WAIT_UNLOCKED,
+	XFS_DIO_TEST_WRITE_CHECKS,
+	XFS_DIO_TEST_WRITE_ADMIT,
+	XFS_DIO_TEST_WRITE_REJECT,
+	XFS_DIO_TEST_UNALIGNED_RECHECK,
+	XFS_DIO_TEST_REMAP_ACTIVE,
+	XFS_DIO_TEST_REMAP_UPGRADE,
+	XFS_DIO_TEST_NR_EVENTS,
+};
+
+#ifdef CONFIG_XFS_DIO_TEST
+void xfs_dio_test_init(struct dentry *parent);
+void xfs_dio_test_event(struct inode *inode, enum xfs_dio_test_event event,
+		unsigned long value);
+#else
+static inline void xfs_dio_test_init(struct dentry *parent) { }
+static inline void xfs_dio_test_event(struct inode *inode,
+		enum xfs_dio_test_event event, unsigned long value) { }
+#endif
+#endif
diff --git a/fs/xfs/xfs_file.c b/fs/xfs/xfs_file.c
index 124a037cfa9f..e7452343671b 100644
--- a/fs/xfs/xfs_file.c
+++ b/fs/xfs/xfs_file.c
@@ -11,6 +11,7 @@
 #include "xfs_trans_resv.h"
 #include "xfs_mount.h"
 #include "xfs_inode.h"
+#include "xfs_dio_test.h"
 #include "xfs_trans.h"
 #include "xfs_inode_item.h"
 #include "xfs_bmap.h"
@@ -38,6 +39,7 @@
 #include <linux/mount.h>
 #include <linux/filelock.h>
 #include <linux/bio-integrity.h>
+#include <linux/wait_bit.h>
 
 static const struct vm_operations_struct xfs_file_vm_ops;
 
@@ -206,6 +208,8 @@ xfs_ilock_iocb_for_write(
 	    xfs_iflags_test(ip, XFS_IREMAPPING)) {
 		xfs_iunlock(ip, *lock_mode);
 		*lock_mode = XFS_IOLOCK_EXCL;
+		xfs_dio_test_event(VFS_I(ip), XFS_DIO_TEST_REMAP_UPGRADE,
+				*lock_mode);
 		return xfs_ilock_iocb(iocb, *lock_mode);
 	}
 
@@ -226,6 +230,80 @@ xfs_dio_read_bounce_submit_io(
 			iomap_ioend_flags(&iter->iomap) | IOMAP_IOEND_DIRECT);
 }
 
+static bool
+xfs_dio_range_overlap(
+	struct xfs_inode	*ip,
+	loff_t			pos,
+	size_t			len)
+{
+	struct inode_dio_drain	*drain;
+	bool			overlap;
+
+	/* Acquire the initialized allocation published by inode_dio_drain_get(). */
+	drain = smp_load_acquire(&ip->i_dio_drain);
+	overlap = inode_dio_drain_overlap(VFS_I(ip), drain, pos, len);
+	return overlap;
+}
+
+/* Register provisionally before the final exclusion snapshot. */
+static bool
+xfs_dio_try_begin(
+	struct xfs_inode	*ip,
+	loff_t			pos,
+	size_t			len,
+	unsigned int		*dio_flags)
+{
+	bool			generation;
+
+	if (!inode_dio_try_begin(VFS_I(ip), &ip->i_dio_drain, pos, len,
+			&generation))
+		return false;
+	*dio_flags &= ~IOMAP_DIO_GENERATION;
+	*dio_flags |= IOMAP_DIO_ACCOUNTED;
+	if (generation)
+		*dio_flags |= IOMAP_DIO_GENERATION;
+	return true;
+}
+
+/* No request is counted while waiting for an overlapping exclusion. */
+static int
+xfs_dio_range_wait(
+	struct kiocb		*iocb,
+	size_t			len,
+	unsigned int		*iolock,
+	bool			write)
+{
+	struct inode		*inode = file_inode(iocb->ki_filp);
+	struct xfs_inode	*ip = XFS_I(inode);
+	struct inode_dio_drain	*drain;
+	int			waited = 0;
+	int			error;
+
+	while (xfs_dio_range_overlap(ip, iocb->ki_pos, len)) {
+		if (iocb->ki_flags & IOCB_NOWAIT)
+			return -EAGAIN;
+		/* The successful overlap check acquired this permanent allocation. */
+		drain = READ_ONCE(ip->i_dio_drain);
+		xfs_iunlock(ip, *iolock);
+		xfs_dio_test_event(inode, XFS_DIO_TEST_WAIT_UNLOCKED,
+				iocb->ki_pos);
+		wait_var_event(&inode->i_dio_count,
+			!inode_dio_drain_overlap(inode, drain, iocb->ki_pos, len));
+		/* Re-evaluate the required write lock mode after dropping IOLOCK. */
+		if (write) {
+			error = xfs_ilock_iocb_for_write(iocb, iolock);
+			if (error) {
+				*iolock = 0;
+				return error;
+			}
+		} else {
+			xfs_ilock(ip, *iolock);
+		}
+		waited = 1;
+	}
+	return waited;
+}
+
 static const struct iomap_dio_ops xfs_dio_read_bounce_ops = {
 	.submit_io	= xfs_dio_read_bounce_submit_io,
 	.bio_set	= &iomap_ioend_bioset,
@@ -237,6 +315,10 @@ xfs_file_dio_read(
 	struct iov_iter		*to)
 {
 	struct xfs_inode	*ip = XFS_I(file_inode(iocb->ki_filp));
+	unsigned int		iolock = XFS_IOLOCK_SHARED;
+	unsigned int		dio_flags = 0;
+	bool			bounce = mapping_stable_writes(iocb->ki_filp->f_mapping);
+	bool			simple = !bounce;
 	ssize_t			ret;
 
 	trace_xfs_file_direct_read(iocb, to);
@@ -249,15 +331,27 @@ xfs_file_dio_read(
 	ret = xfs_ilock_iocb(iocb, XFS_IOLOCK_SHARED);
 	if (ret)
 		return ret;
-	if (mapping_stable_writes(iocb->ki_filp->f_mapping)) {
+retry:
+	ret = xfs_dio_range_wait(iocb, iov_iter_count(to), &iolock, false);
+	if (ret < 0)
+		goto out_unlock;
+	if (!xfs_dio_try_begin(ip, iocb->ki_pos, iov_iter_count(to), &dio_flags))
+		goto retry;
+	if (!simple) {
 		ret = iomap_dio_rw(iocb, to, &xfs_read_iomap_ops,
-				&xfs_dio_read_bounce_ops, 0, NULL, 0);
+					bounce ? &xfs_dio_read_bounce_ops : NULL,
+					dio_flags,
+					NULL, 0);
 	} else {
-		ret = iomap_dio_read_simple(iocb, to, xfs_read_iomap_begin, 0);
-		if (ret == -ENOTBLK)
-			ret = iomap_dio_rw(iocb, to, &xfs_read_iomap_ops, NULL,
-					0, NULL, 0);
+		ret = iomap_dio_read_simple(iocb, to, xfs_read_iomap_begin,
+				dio_flags);
+		if (ret == -ENOTBLK) {
+			/* Fallback needs a fresh admission and registration. */
+			simple = false;
+			goto retry;
+		}
 	}
+out_unlock:
 	xfs_iunlock(ip, XFS_IOLOCK_SHARED);
 
 	return ret;
@@ -436,9 +530,10 @@ xfs_file_write_zero_eof(
 /*
  * Common pre-write limit and setup checks.
  *
- * Called with the iolock held either shared and exclusive according to
- * @iolock, and returns with it held.  Might upgrade the iolock to exclusive
- * if called for a direct write beyond i_size.
+ * Called with the iolock held shared or exclusive according to @iolock.
+ * Returns with it held unless reacquisition fails, in which case *@iolock
+ * is set to zero.  Might upgrade the iolock to exclusive if called for a
+ * direct write beyond i_size.
  */
 STATIC ssize_t
 xfs_file_write_checks(
@@ -453,10 +548,25 @@ xfs_file_write_checks(
 	ssize_t			error;
 
 restart:
+	/* Normalize append position and enforce limits before range admission. */
 	error = generic_write_checks(iocb, from);
+	xfs_dio_test_event(inode, XFS_DIO_TEST_WRITE_CHECKS,
+			iov_iter_count(from));
 	if (error <= 0)
 		return error;
 
+	if (iocb->ki_flags & IOCB_DIRECT) {
+		error = xfs_dio_range_wait(iocb, iov_iter_count(from), iolock,
+				true);
+		if (error < 0)
+			return error;
+		if (error) {
+			iov_iter_reexpand(from, count);
+			/* EOF and limits can change while IOLOCK is dropped. */
+			goto restart;
+		}
+	}
+
 	if (iocb->ki_flags & IOCB_NOWAIT) {
 		error = break_layout(inode, false);
 		if (error == -EWOULDBLOCK)
@@ -735,6 +845,7 @@ xfs_file_dio_write_aligned(
 {
 	unsigned int		iolock = XFS_IOLOCK_SHARED;
 	unsigned int		dio_flags = 0;
+	size_t			count = iov_iter_count(from);
 	ssize_t			ret;
 
 	/*
@@ -748,6 +859,7 @@ xfs_file_dio_write_aligned(
 	ret = xfs_ilock_iocb_for_write(iocb, &iolock);
 	if (ret)
 		return ret;
+retry_checks:
 	ret = xfs_file_write_checks(iocb, from, &iolock, ac);
 	if (ret)
 		goto out_unlock;
@@ -763,10 +875,18 @@ xfs_file_dio_write_aligned(
 	}
 	if (mapping_stable_writes(iocb->ki_filp->f_mapping))
 		dio_flags |= IOMAP_DIO_BOUNCE;
+	xfs_dio_test_event(VFS_I(ip), XFS_DIO_TEST_WRITE_ADMIT, iocb->ki_pos);
+	if (!xfs_dio_try_begin(ip, iocb->ki_pos, iov_iter_count(from), &dio_flags)) {
+		xfs_dio_test_event(VFS_I(ip), XFS_DIO_TEST_WRITE_REJECT,
+				iocb->ki_pos);
+		iov_iter_reexpand(from, count);
+		goto retry_checks;
+	}
 	trace_xfs_file_direct_write(iocb, from);
 	ret = iomap_dio_rw(iocb, from, ops, dops, dio_flags, ac, 0);
 out_unlock:
-	xfs_iunlock(ip, iolock);
+	if (iolock)
+		xfs_iunlock(ip, iolock);
 	return ret;
 }
 
@@ -827,6 +947,7 @@ xfs_file_dio_write_atomic(
 	if (ret)
 		return ret;
 
+retry_checks:
 	ret = xfs_file_write_checks(iocb, from, &iolock, NULL);
 	if (ret)
 		goto out_unlock;
@@ -840,6 +961,10 @@ xfs_file_dio_write_atomic(
 	trace_xfs_file_direct_write(iocb, from);
 	if (mapping_stable_writes(iocb->ki_filp->f_mapping))
 		dio_flags |= IOMAP_DIO_BOUNCE;
+	if (!xfs_dio_try_begin(ip, iocb->ki_pos, iov_iter_count(from), &dio_flags)) {
+		iov_iter_reexpand(from, ocount);
+		goto retry_checks;
+	}
 	ret = iomap_dio_rw(iocb, from, dops, &xfs_dio_write_ops, dio_flags,
 			NULL, 0);
 
@@ -917,10 +1042,22 @@ xfs_file_dio_write_unaligned(
 		goto out_unlock;
 	}
 
+retry_checks:
 	ret = xfs_file_write_checks(iocb, from, &iolock, NULL);
 	if (ret)
 		goto out_unlock;
 
+	if (IS_ENABLED(CONFIG_XFS_DIO_TEST))
+		xfs_dio_test_event(VFS_I(ip), XFS_DIO_TEST_UNALIGNED_RECHECK,
+				(iolock << 8) | xfs_is_cow_inode(ip));
+
+	/* The checks can drop IOLOCK, allowing a clone to share this block. */
+	if (xfs_is_cow_inode(ip)) {
+		trace_xfs_reflink_bounce_dio_write(iocb, from);
+		ret = -ENOTBLK;
+		goto out_unlock;
+	}
+
 	/*
 	 * If we are doing exclusive unaligned I/O, this must be the only I/O
 	 * in-flight.  Otherwise we risk data corruption due to unwritten extent
@@ -932,6 +1069,11 @@ xfs_file_dio_write_unaligned(
 
 	if (mapping_stable_writes(iocb->ki_filp->f_mapping))
 		flags |= IOMAP_DIO_BOUNCE;
+	xfs_dio_test_event(VFS_I(ip), XFS_DIO_TEST_WRITE_ADMIT, iocb->ki_pos);
+	if (!xfs_dio_try_begin(ip, iocb->ki_pos, iov_iter_count(from), &flags)) {
+		iov_iter_reexpand(from, count);
+		goto retry_checks;
+	}
 
 	trace_xfs_file_direct_write(iocb, from);
 	ret = iomap_dio_rw(iocb, from, &xfs_direct_write_iomap_ops,
diff --git a/fs/xfs/xfs_icache.c b/fs/xfs/xfs_icache.c
index de8be344e987..e26df6894618 100644
--- a/fs/xfs/xfs_icache.c
+++ b/fs/xfs/xfs_icache.c
@@ -110,6 +110,7 @@ xfs_inode_alloc(
 	ip->i_mount = mp;
 	memset(&ip->i_imap, 0, sizeof(struct xfs_imap));
 	ip->i_cowfp = NULL;
+	ip->i_dio_drain = NULL;
 	memset(&ip->i_af, 0, sizeof(ip->i_af));
 	ip->i_af.if_format = XFS_DINODE_FMT_EXTENTS;
 	memset(&ip->i_df, 0, sizeof(ip->i_df));
@@ -156,6 +157,7 @@ xfs_inode_free_callback(
 		xfs_inode_item_destroy(ip);
 	}
 
+	kfree(ip->i_dio_drain);
 	kmem_cache_free(xfs_inode_cache, ip);
 }
 
diff --git a/fs/xfs/xfs_inode.h b/fs/xfs/xfs_inode.h
index 75eff68c1e39..04907a7b431f 100644
--- a/fs/xfs/xfs_inode.h
+++ b/fs/xfs/xfs_inode.h
@@ -41,6 +41,7 @@ typedef struct xfs_inode {
 	struct xfs_inode_log_item *i_itemp;	/* logging information */
 	struct rw_semaphore	i_lock;		/* inode lock */
 	atomic_t		i_pincount;	/* inode pin count */
+	struct inode_dio_drain	*i_dio_drain; /* lazily allocated DIO drain */
 	struct llist_node	i_gclist;	/* deferred inactivation list */
 
 	/*
diff --git a/fs/xfs/xfs_iops.c b/fs/xfs/xfs_iops.c
index d307ec97fac8..dec2139c0c6f 100644
--- a/fs/xfs/xfs_iops.c
+++ b/fs/xfs/xfs_iops.c
@@ -1278,6 +1278,14 @@ xfs_vn_tmpfile(
 	return finish_open_simple(file, err);
 }
 
+static void
+xfs_vn_dio_drain_complete(
+	struct inode		*inode)
+{
+	/* The generation cutover orders i_dio_drain initialization before this callback. */
+	inode_dio_drain_complete(inode, READ_ONCE(XFS_I(inode)->i_dio_drain));
+}
+
 static const struct inode_operations xfs_inode_operations = {
 	.get_inode_acl		= xfs_get_acl,
 	.set_acl		= xfs_set_acl,
@@ -1289,6 +1297,7 @@ static const struct inode_operations xfs_inode_operations = {
 	.sync_lazytime		= xfs_vn_sync_lazytime,
 	.fileattr_get		= xfs_fileattr_get,
 	.fileattr_set		= xfs_fileattr_set,
+	.dio_drain_complete	= xfs_vn_dio_drain_complete,
 };
 
 static const struct inode_operations xfs_dir_inode_operations = {
diff --git a/fs/xfs/xfs_reflink.c b/fs/xfs/xfs_reflink.c
index 104542b10d9e..3e1b7d628477 100644
--- a/fs/xfs/xfs_reflink.c
+++ b/fs/xfs/xfs_reflink.c
@@ -12,6 +12,7 @@
 #include "xfs_mount.h"
 #include "xfs_defer.h"
 #include "xfs_inode.h"
+#include "xfs_dio_test.h"
 #include "xfs_trans.h"
 #include "xfs_bmap.h"
 #include "xfs_bmap_util.h"
@@ -1763,6 +1764,7 @@ xfs_reflink_remap_prep(
 	xfs_iflags_set(src, XFS_IREMAPPING);
 	if (inode_in != inode_out)
 		xfs_ilock_demote(src, XFS_IOLOCK_EXCL | XFS_MMAPLOCK_EXCL);
+	xfs_dio_test_event(inode_in, XFS_DIO_TEST_REMAP_ACTIVE, 0);
 
 	return 0;
 out_unlock:
diff --git a/fs/xfs/xfs_super.c b/fs/xfs/xfs_super.c
index 5a06132aa384..a7e84ddd145b 100644
--- a/fs/xfs/xfs_super.c
+++ b/fs/xfs/xfs_super.c
@@ -12,6 +12,7 @@
 #include "xfs_sb.h"
 #include "xfs_mount.h"
 #include "xfs_inode.h"
+#include "xfs_dio_test.h"
 #include "xfs_btree.h"
 #include "xfs_bmap.h"
 #include "xfs_alloc.h"
@@ -2677,6 +2678,7 @@ init_xfs_fs(void)
 		goto out_cleanup_procfs;
 
 	xfs_debugfs = xfs_debugfs_mkdir("xfs", NULL);
+	xfs_dio_test_init(xfs_debugfs);
 
 	xfs_kset = kset_create_and_add("xfs", NULL, fs_kobj);
 	if (!xfs_kset) {
-- 
2.56.0

[RFC PATCH 5/6] xfs: use range-based DIO draining for aligned hole punches

[Vitali Lovich]

Hole punching holds exclusive IOLOCK while draining DIO and modifying
the range, blocking new submissions throughout the inode. Use shared
IOLOCK and a latch-published range exclusion for eligible aligned
data-device punches so disjoint new DIO can proceed.

Reject obvious ineligible requests before allocating drain state or
taking the drain mutex. Obtain the drain mutex before IOLOCK and recheck
eligibility under shared IOLOCK. A waiter must not hold shared IOLOCK
while the mutex owner may drop and reacquire IOLOCK exclusively during
write checks. Retain the existing path for DAX, realtime,
privilege/capability stripping (!IS_NOSEC), reflink remapping, unaligned
punches, layout breaking, or an unrepresentable rounded endpoint. Round
the exclusion to the larger of the allocation unit and PAGE_SIZE,
matching cache invalidation.

Drain the pre-cutover cohort before taking MMAPLOCK: an in-flight DIO
can fault in a buffer backed by the same file. Skip unnecessary partial
block zeroing for aligned data-device punches so the operation can retain
shared IOLOCK. Keep the exclusion through extent removal and any required
synchronous log force, and report allocation or drain-start errors.

Add a filesystem selftest that parks registered DIO in a userfaultfd
fault, observes the active exclusion, and checks disjoint progress and
overlapping NOWAIT write rejection. Race faulting same-inode buffers
with punching, using synchronous and AIO extending and COW writes, and
verify clone isolation.

Exercise repeated exclusion retries and append-position/limit changes
while IOLOCK is dropped, including iterator expansion and a limit that
changes DIO alignment.

Verify completed and in-progress reflink remaps during the unlocked
window, including exact shared-extent evidence, unaligned COW fallback,
exclusive lock reacquisition and clone isolation.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 fs/xfs/xfs_bmap_util.c                        |   9 +-
 fs/xfs/xfs_file.c                             |  88 ++++
 fs/xfs/xfs_inode.c                            |  10 +
 tools/testing/selftests/filesystems/Makefile  |   6 +
 .../selftests/filesystems/xfs_dio_admission.c | 380 +++++++++++++++++
 .../selftests/filesystems/xfs_dio_drain.c     | 371 +++++++++++++++++
 .../selftests/filesystems/xfs_dio_remap.c     | 391 ++++++++++++++++++
 7 files changed, 1253 insertions(+), 2 deletions(-)
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_admission.c
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_drain.c
 create mode 100644 tools/testing/selftests/filesystems/xfs_dio_remap.c

diff --git a/fs/xfs/xfs_bmap_util.c b/fs/xfs/xfs_bmap_util.c
index 865aec3ae5ab..30759c6f94d3 100644
--- a/fs/xfs/xfs_bmap_util.c
+++ b/fs/xfs/xfs_bmap_util.c
@@ -889,8 +889,8 @@ xfs_free_file_space(
 		return 0;
 
 	/*
-	 * Now AIO and DIO has drained we flush and (if necessary) invalidate
-	 * the cached range over the first operation we are about to run.
+	 * The caller has excluded and drained DIO over the affected range.
+	 * Flush and invalidate the cached range before removing its mappings.
 	 */
 	error = xfs_flush_unmap_range(ip, offset, len);
 	if (error)
@@ -917,6 +917,11 @@ xfs_free_file_space(
 		}
 	}
 
+	/* An aligned data-device punch needs no partial-block zeroing. */
+	if (!XFS_IS_REALTIME_INODE(ip) &&
+	    !((offset | len) & mp->m_blockmask))
+		return 0;
+
 	/*
 	 * Now that we've unmap all full blocks we'll have to zero out any
 	 * partial block at the beginning and/or end.  xfs_zero_range is smart
diff --git a/fs/xfs/xfs_file.c b/fs/xfs/xfs_file.c
index e7452343671b..ad83328b7350 100644
--- a/fs/xfs/xfs_file.c
+++ b/fs/xfs/xfs_file.c
@@ -1682,6 +1682,90 @@ xfs_falloc_allocate_range(
 		 FALLOC_FL_ZERO_RANGE |	FALLOC_FL_INSERT_RANGE |	\
 		 FALLOC_FL_UNSHARE_RANGE | FALLOC_FL_WRITE_ZEROES)
 
+/*
+ * Aligned punches do not change the file size or zero partial
+ * blocks.  Exclude only the range that cache invalidation can touch, while
+ * retaining the shared IOLOCK to exclude truncate and buffered writes.
+ */
+STATIC bool
+xfs_file_punch_hole_range(
+	struct file		*file,
+	loff_t			offset,
+	loff_t			len,
+	long			*result)
+{
+	struct inode		*inode = file_inode(file);
+	struct xfs_inode	*ip = XFS_I(inode);
+	struct inode_dio_drain	*drain;
+	loff_t			rounding, start, last;
+	unsigned int		inode_flags = READ_ONCE(inode->i_flags);
+	unsigned int		iolock = XFS_IOLOCK_SHARED;
+	bool			handled = false;
+	int			error;
+
+	/*
+	 * Avoid drain allocation and serialization for obvious fallbacks.
+	 * Only data-device files qualify, so check alignment to the filesystem
+	 * block size. Recheck mutable inode state after acquiring IOLOCK.
+	 */
+	if ((inode_flags & S_DAX) || !(inode_flags & S_NOSEC) ||
+	    (READ_ONCE(ip->i_diflags) & XFS_DIFLAG_REALTIME) ||
+	    (READ_ONCE(ip->i_flags) & XFS_IREMAPPING) ||
+	    !IS_ALIGNED(offset | len, i_blocksize(inode)))
+		return false;
+
+	drain = inode_dio_drain_get(&ip->i_dio_drain);
+	if (IS_ERR(drain)) {
+		*result = PTR_ERR(drain);
+		return true;
+	}
+	mutex_lock(&drain->lock);
+	xfs_ilock(ip, iolock);
+	if (IS_DAX(inode) || XFS_IS_REALTIME_INODE(ip) || !IS_NOSEC(inode) ||
+	    xfs_iflags_test(ip, XFS_IREMAPPING) ||
+	    !xfs_is_falloc_aligned(ip, offset, len))
+		goto out_unlock;
+
+	rounding = max_t(loff_t, xfs_inode_alloc_unitsize(ip), PAGE_SIZE);
+	/* The invalidation range must remain representable even beyond EOF. */
+	if (offset + len > rounddown_64(LLONG_MAX, rounding))
+		goto out_unlock;
+	start = rounddown_64(offset, rounding);
+	last = roundup_64(offset + len, rounding) - 1;
+
+	/* Fall back to the exclusive path if breaking a layout would block or fail. */
+	if (break_layout(inode, false))
+		goto out_unlock;
+
+	handled = true;
+	/*
+	 * IS_NOSEC was checked under shared IOLOCK, which excludes chmod and
+	 * security xattr changes that could clear S_NOSEC. file_modified() can
+	 * update timestamps without stripping privileges under shared IOLOCK.
+	 */
+	error = file_modified(file);
+	if (error)
+		goto out_result;
+
+	/* DIO can fault in a buffer from this file. Drain before MMAPLOCK. */
+	error = inode_dio_drain_start(inode, drain, start, last);
+	if (error)
+		goto out_result;
+	inode_dio_drain_wait(inode, drain);
+	xfs_ilock(ip, XFS_MMAPLOCK_EXCL);
+	error = xfs_free_file_space(ip, offset, len, NULL);
+	if (!error && xfs_file_sync_writes(file))
+		error = xfs_log_force_inode(ip);
+	xfs_iunlock(ip, XFS_MMAPLOCK_EXCL);
+	inode_dio_drain_finish(inode, drain);
+out_result:
+	*result = error;
+out_unlock:
+	xfs_iunlock(ip, iolock);
+	mutex_unlock(&drain->lock);
+	return handled;
+}
+
 STATIC long
 __xfs_file_fallocate(
 	struct file		*file,
@@ -1695,6 +1779,10 @@ __xfs_file_fallocate(
 	long			error;
 	uint			iolock = XFS_IOLOCK_EXCL | XFS_MMAPLOCK_EXCL;
 
+	if (mode == (FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE) &&
+	    xfs_file_punch_hole_range(file, offset, len, &error))
+		return error;
+
 	xfs_ilock(ip, iolock);
 	error = xfs_break_layouts(inode, &iolock, BREAK_UNMAP);
 	if (error)
diff --git a/fs/xfs/xfs_inode.c b/fs/xfs/xfs_inode.c
index 05a14da28031..b8119abab30a 100644
--- a/fs/xfs/xfs_inode.c
+++ b/fs/xfs/xfs_inode.c
@@ -118,6 +118,16 @@ xfs_lock_flags_assert(
  *
  * i_rwsem -> invalidate_lock -> page_lock -> i_ilock
  *
+ * DIO range drainers add an outer mutex:
+ *
+ * i_dio_drain->lock -> i_rwsem -> invalidate_lock -> page_lock -> i_ilock
+ *
+ * Never wait for the drain mutex while holding IOLOCK: its owner may drop
+ * shared IOLOCK and reacquire it exclusively during write checks. A waiter
+ * holding shared IOLOCK would prevent that upgrade and deadlock both drainers.
+ * The mutex stays held through exclusion removal, including extent freeing
+ * and any synchronous log force. DIO admission and completion do not take it.
+ *
  * mmap_lock locking order:
  *
  * i_rwsem -> page lock -> mmap_lock
diff --git a/tools/testing/selftests/filesystems/Makefile b/tools/testing/selftests/filesystems/Makefile
index fbd5c27505ef..0c572f10fa71 100644
--- a/tools/testing/selftests/filesystems/Makefile
+++ b/tools/testing/selftests/filesystems/Makefile
@@ -5,7 +5,13 @@ TEST_GEN_PROGS := devpts_pts anon_inode_test kernfs_test fclog ustat_test rw_hin
 TEST_GEN_PROGS += idmapped_tmpfile
 TEST_GEN_PROGS_EXTENDED := dnotify_test
 
+TEST_GEN_PROGS_EXTENDED += xfs_dio_drain xfs_dio_admission xfs_dio_remap
+
 include ../lib.mk
 
+$(OUTPUT)/xfs_dio_drain: LDLIBS += -pthread
+$(OUTPUT)/xfs_dio_admission: LDLIBS += -pthread
+$(OUTPUT)/xfs_dio_remap: LDLIBS += -pthread
+
 $(OUTPUT)/idmapped_tmpfile: LDLIBS += -lcap
 $(OUTPUT)/idmapped_tmpfile: utils.c
diff --git a/tools/testing/selftests/filesystems/xfs_dio_admission.c b/tools/testing/selftests/filesystems/xfs_dio_admission.c
new file mode 100644
index 000000000000..3162de34b2b8
--- /dev/null
+++ b/tools/testing/selftests/filesystems/xfs_dio_admission.c
@@ -0,0 +1,380 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Deterministic real-syscall admission retries; see CONFIG_XFS_DIO_TEST. */
+#define _GNU_SOURCE
+#include <errno.h>
+#include <fcntl.h>
+#include <linux/fs.h>
+#include <linux/magic.h>
+#include <linux/userfaultfd.h>
+#include <poll.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/ioctl.h>
+#include <sys/mman.h>
+#include <sys/resource.h>
+#include <sys/stat.h>
+#include <sys/statfs.h>
+#include <sys/syscall.h>
+#include <sys/uio.h>
+#include <time.h>
+#include <unistd.h>
+#include "../kselftest.h"
+
+/* Event indices in fs/xfs/xfs_dio_test.h. */
+enum { WAIT_UNLOCKED = 4, WRITE_CHECKS, WRITE_ADMIT, WRITE_REJECT, EVENTS };
+#define GATE(event) (1UL << (event))
+struct snapshot { unsigned long gates, timedout, hits[EVENTS], value[EVENTS]; };
+struct request {
+	int fd;
+	void *buffer;
+	size_t length;
+	off_t offset;
+	bool read, punch;
+	ssize_t result;
+	int error;
+};
+struct parked_read {
+	struct request req;
+	pthread_t thread;
+	int uffd;
+};
+
+static void die(const char *what)
+{
+	ksft_exit_fail_msg("%s: %s\n", what, strerror(errno));
+}
+
+static long long milliseconds(void)
+{
+	struct timespec now;
+
+	clock_gettime(CLOCK_MONOTONIC, &now);
+	return now.tv_sec * 1000LL + now.tv_nsec / 1000000;
+}
+
+static void *submit(void *data)
+{
+	struct request *r = data;
+
+	if (r->punch)
+		r->result = fallocate(r->fd, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE,
+				r->offset, r->length);
+	else if (r->read)
+		r->result = pread(r->fd, r->buffer, r->length, r->offset);
+	else
+		r->result = pwrite(r->fd, r->buffer, r->length, r->offset);
+	r->error = errno;
+	return NULL;
+}
+
+static void command(int control, const char *verb, unsigned long value)
+{
+	char text[80];
+	int len = snprintf(text, sizeof(text), "%s %lu", verb, value);
+
+	if (write(control, text, len) != len)
+		die("DIO test command");
+}
+
+static int arm(int fd, unsigned long gates)
+{
+	char text[80];
+	int control = open("/sys/kernel/debug/xfs/dio_test", O_RDWR);
+	int len;
+
+	if (control < 0)
+		return -1;
+	len = snprintf(text, sizeof(text), "arm %d %lu", fd, gates);
+	if (write(control, text, len) != len)
+		die("arm DIO test");
+	return control;
+}
+
+static struct snapshot snapshot(int control)
+{
+	struct snapshot s = { };
+	char buffer[2048], *cursor, *next;
+	ssize_t n = pread(control, buffer, sizeof(buffer) - 1, 0);
+
+	if (n <= 0)
+		die("read DIO test state");
+	buffer[n] = 0;
+	if (sscanf(buffer, "%lu %lu", &s.gates, &s.timedout) != 2)
+		die("parse DIO test state");
+	cursor = strchr(buffer, '\n');
+	for (int i = 0; i < EVENTS; i++) {
+		if (!cursor || sscanf(cursor + 1, "%lu %lu", &s.hits[i],
+				&s.value[i]) != 2)
+			die("parse DIO test event");
+		next = strchr(cursor + 1, '\n');
+		cursor = next;
+	}
+	return s;
+}
+
+static void wait_event(int control, int event, unsigned long count)
+{
+	long long deadline = milliseconds() + 5000;
+
+	do {
+		struct snapshot s = snapshot(control);
+
+		if (s.hits[event] >= count && !s.timedout)
+			return;
+		usleep(1000);
+	} while (milliseconds() < deadline);
+	ksft_exit_fail_msg("DIO event %d did not reach %lu\n", event, count);
+}
+
+static void park_read(struct parked_read *old, int fd, size_t unit)
+{
+	struct uffdio_api api = { .api = UFFD_API };
+	struct uffdio_register reg = { .mode = UFFDIO_REGISTER_MODE_MISSING };
+	struct uffd_msg message;
+	struct pollfd pollfd;
+
+	memset(old, 0, sizeof(*old));
+	old->req = (struct request){ .fd = fd, .length = unit,
+		.offset = unit, .read = true };
+	old->req.buffer = mmap(NULL, unit, PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (old->req.buffer == MAP_FAILED)
+		die("missing read destination");
+	old->uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
+	reg.range.start = (uintptr_t)old->req.buffer;
+	reg.range.len = unit;
+	if (old->uffd < 0 || ioctl(old->uffd, UFFDIO_API, &api) ||
+	    ioctl(old->uffd, UFFDIO_REGISTER, &reg))
+		die("kernel-fault userfaultfd required by admission test");
+	if (pthread_create(&old->thread, NULL, submit, &old->req))
+		die("parked read thread");
+	pollfd = (struct pollfd){ .fd = old->uffd, .events = POLLIN };
+	if (poll(&pollfd, 1, 5000) != 1 ||
+	    read(old->uffd, &message, sizeof(message)) != sizeof(message) ||
+	    message.event != UFFD_EVENT_PAGEFAULT)
+		ksft_exit_fail_msg("read did not register and reach GUP\n");
+}
+
+static void release_read(struct parked_read *old, void *buffer, size_t unit)
+{
+	struct uffdio_copy copy = { .src = (uintptr_t)buffer,
+		.dst = (uintptr_t)old->req.buffer, .len = unit };
+
+	if (ioctl(old->uffd, UFFDIO_COPY, &copy))
+		die("release registered read");
+	pthread_join(old->thread, NULL);
+	if (old->req.result != (ssize_t)unit)
+		ksft_exit_fail_msg("registered read result %zd\n", old->req.result);
+	close(old->uffd);
+	munmap(old->req.buffer, unit);
+}
+
+static void wait_exclusion(int fd, void *buffer, size_t unit, off_t offset)
+{
+	struct iovec iov = { .iov_base = buffer, .iov_len = unit };
+	long long deadline = milliseconds() + 5000;
+
+	do {
+		if (preadv2(fd, &iov, 1, offset, RWF_NOWAIT) == -1 && errno == EAGAIN)
+			return;
+		usleep(1000);
+	} while (milliseconds() < deadline);
+	ksft_exit_fail_msg("real drainer did not publish exclusion\n");
+}
+
+static void prefill(int fd, void *buffer, size_t unit)
+{
+	memset(buffer, 0x5a, 16 * unit);
+	if (ftruncate(fd, 0) || pwrite(fd, buffer, 8 * unit, 0) != 8 * (ssize_t)unit ||
+	    fsync(fd))
+		die("prefill");
+}
+
+static bool bytes_equal(const void *buffer, size_t length, unsigned char value)
+{
+	const unsigned char *bytes = buffer;
+
+	for (size_t i = 0; i < length; i++)
+		if (bytes[i] != value)
+			return false;
+	return true;
+}
+
+/* Two distinct real punch exclusions must reject the same writer twice. */
+static int repeated_admission(int fd, void *buffer, size_t unit)
+{
+	struct request writer = { .fd = fd, .buffer = buffer, .length = unit,
+		.offset = 4 * unit };
+	struct request punch = { .fd = fd, .length = unit,
+		.offset = writer.offset,
+		.punch = true };
+	struct snapshot s;
+	pthread_t writing, punching;
+	int control;
+	bool ok = true;
+
+	prefill(fd, buffer, unit);
+	control = arm(fd, GATE(WRITE_ADMIT));
+	if (control < 0)
+		return -1;
+	if (pthread_create(&writing, NULL, submit, &writer))
+		die("writer thread");
+	for (unsigned int round = 1; round <= 2; round++) {
+		struct parked_read old;
+
+		wait_event(control, WRITE_ADMIT, round);
+		park_read(&old, fd, unit);
+		if (pthread_create(&punching, NULL, submit, &punch))
+			die("punch thread");
+		wait_exclusion(fd, buffer + unit, unit, writer.offset);
+		command(control, "gate", GATE(WAIT_UNLOCKED));
+		wait_event(control, WAIT_UNLOCKED, round);
+		s = snapshot(control);
+		ok &= s.hits[WRITE_REJECT] == round;
+		release_read(&old, buffer + unit, unit);
+		pthread_join(punching, NULL);
+		ok &= punch.result == 0;
+		command(control, "gate", round == 1 ? GATE(WRITE_ADMIT) : 0);
+	}
+	pthread_join(writing, NULL);
+	s = snapshot(control);
+	ok &= writer.result == (ssize_t)unit && s.hits[WRITE_REJECT] == 2 &&
+		s.hits[WAIT_UNLOCKED] == 2 && s.hits[WRITE_ADMIT] == 3 &&
+		s.hits[WRITE_CHECKS] == 5 && !s.timedout;
+	ok &= pread(fd, buffer + unit, unit, writer.offset) == (ssize_t)unit &&
+		bytes_equal(buffer + unit, unit, 0x5a);
+	ksft_print_msg("retry evidence: checks=%lu tries=%lu rejects=%lu waits=%lu\n",
+		s.hits[WRITE_CHECKS], s.hits[WRITE_ADMIT], s.hits[WRITE_REJECT],
+		s.hits[WAIT_UNLOCKED]);
+	close(control);
+	return ok;
+}
+
+/* Move EOF and change the process limit with the writer's IOLOCK dropped. */
+static int append_revalidation(int fd, void *buffer, size_t unit, int mode)
+{
+	struct parked_read old;
+	struct request drain = { .fd = fd, .length = 2 * unit,
+		.offset = 8 * unit, .punch = true };
+	struct request writer = { .buffer = buffer, .length = 2 * unit };
+	struct rlimit saved, limit;
+	struct snapshot s;
+	struct stat st;
+	pthread_t draining, writing;
+	char path[64];
+	int control, appendfd;
+	bool ok;
+	bool expand = mode == 1, unaligned = mode == 2;
+	size_t expected = unaligned ? 1 : (expand ? 2 : 1) * unit;
+
+	prefill(fd, buffer, unit);
+	park_read(&old, fd, unit);
+	if (pthread_create(&draining, NULL, submit, &drain))
+		die("append-range punch thread");
+	wait_exclusion(fd, buffer + 2 * unit, unit, 8 * unit);
+	control = arm(fd, GATE(WAIT_UNLOCKED));
+	if (control < 0) {
+		release_read(&old, buffer + 2 * unit, unit);
+		pthread_join(draining, NULL);
+		return -1;
+	}
+	snprintf(path, sizeof(path), "/proc/self/fd/%d", fd);
+	appendfd = open(path, O_RDWR | O_DIRECT | O_APPEND);
+	if (appendfd < 0 || getrlimit(RLIMIT_FSIZE, &saved))
+		die("append descriptor/limit");
+	limit = saved;
+	limit.rlim_cur = expand ? 9 * unit : RLIM_INFINITY;
+	if (setrlimit(RLIMIT_FSIZE, &limit))
+		die("initial write limit");
+	writer.fd = appendfd;
+	if (pthread_create(&writing, NULL, submit, &writer))
+		die("append writer thread");
+	wait_event(control, WAIT_UNLOCKED, 1);
+	s = snapshot(control);
+	ok = s.value[WAIT_UNLOCKED] == 8 * unit &&
+		s.value[WRITE_CHECKS] == (expand ? unit : 2 * unit);
+	limit.rlim_cur = unaligned ? 14 * unit + 1 : (expand ? 16 : 15) * unit;
+	if (setrlimit(RLIMIT_FSIZE, &limit))
+		die("updated write limit");
+	release_read(&old, buffer + 2 * unit, unit);
+	pthread_join(draining, NULL);
+	ok &= drain.result == 0;
+	if (ftruncate(fd, 14 * unit))
+		die("move append EOF while IOLOCK dropped");
+	command(control, "gate", 0);
+	pthread_join(writing, NULL);
+	if (setrlimit(RLIMIT_FSIZE, &saved))
+		die("restore write limit");
+	s = snapshot(control);
+	ok &= (unaligned ? writer.result == -1 && writer.error == EINVAL :
+		writer.result == (ssize_t)expected) && s.hits[WRITE_CHECKS] == 2 &&
+		s.hits[WAIT_UNLOCKED] == 1 && s.value[WRITE_ADMIT] == 14 * unit &&
+		s.value[WRITE_CHECKS] == expected && !s.timedout;
+	ok &= !fstat(fd, &st) && st.st_size == (off_t)(14 * unit + (unaligned ? 0 : expected));
+	if (!unaligned)
+		ok &= pread(fd, buffer + 2 * unit, expected, 14 * unit) == (ssize_t)expected &&
+			bytes_equal(buffer + 2 * unit, expected, 0x5a);
+	ok &= pread(fd, buffer + 2 * unit, unit, 8 * unit) == (ssize_t)unit &&
+		bytes_equal(buffer + 2 * unit, unit, 0);
+	ksft_print_msg("append evidence: checks=%lu initial_count=%zu final_count=%lu "
+		"position=%lu result=%zd\n", s.hits[WRITE_CHECKS],
+		expand ? unit : 2 * unit, s.value[WRITE_CHECKS],
+		s.value[WRITE_ADMIT], writer.result);
+	close(appendfd);
+	close(control);
+	return ok;
+}
+
+static void timeout_handler(int signal)
+{
+	(void)signal;
+	_exit(KSFT_FAIL);
+}
+
+int main(int argc, char **argv)
+{
+	struct statfs fs;
+	char *path;
+	void *buffer;
+	size_t unit;
+	int fd, result, failed = 0;
+
+	ksft_print_header();
+	if (argc != 2 || statfs(argv[1], &fs) || fs.f_type != XFS_SUPER_MAGIC)
+		ksft_exit_skip("requires a writable XFS directory\n");
+	unit = sysconf(_SC_PAGESIZE);
+	if ((size_t)fs.f_bsize > unit)
+		unit = fs.f_bsize;
+	if (asprintf(&path, "%s/xfs-admission-XXXXXX", argv[1]) < 0 ||
+	    posix_memalign(&buffer, unit, 16 * unit))
+		die("test allocation");
+	fd = mkostemp(path, O_DIRECT);
+	if (fd < 0)
+		die("test file");
+	unlink(path);
+	free(path);
+	signal(SIGALRM, timeout_handler);
+	alarm(40);
+	ksft_set_plan(4);
+	result = repeated_admission(fd, buffer, unit);
+	if (result < 0)
+		ksft_exit_skip("requires CONFIG_XFS_DIO_TEST and mounted debugfs\n");
+	ksft_test_result(result, "released writer retries two successive exclusions\n");
+	failed |= !result;
+	for (int mode = 0; mode < 3; mode++) {
+		const char *names[] = { "reduced aligned", "expanded", "reduced unaligned" };
+
+		result = append_revalidation(fd, buffer, unit, mode);
+		ksft_test_result(result == 1, "dropped IOLOCK revalidates append and %s limit\n",
+				names[mode]);
+		failed |= result != 1;
+	}
+	alarm(0);
+	free(buffer);
+	close(fd);
+	ksft_exit(!failed);
+}
diff --git a/tools/testing/selftests/filesystems/xfs_dio_drain.c b/tools/testing/selftests/filesystems/xfs_dio_drain.c
new file mode 100644
index 000000000000..f2e005f6e896
--- /dev/null
+++ b/tools/testing/selftests/filesystems/xfs_dio_drain.c
@@ -0,0 +1,371 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Run against a writable XFS directory. CONFIG_USERFAULTFD and permission to
+ * handle kernel-originated faults are needed for the controlled cohort cases.
+ * The mapped-source cases need CONFIG_AIO; the COW case also needs reflink.
+ */
+#define _GNU_SOURCE
+#include <errno.h>
+#include <fcntl.h>
+#include <linux/magic.h>
+#include <linux/aio_abi.h>
+#include <linux/fs.h>
+#include <linux/userfaultfd.h>
+#include <poll.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdatomic.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/ioctl.h>
+#include <sys/mman.h>
+#include <sys/statfs.h>
+#include <sys/syscall.h>
+#include <sys/uio.h>
+#include <time.h>
+#include <unistd.h>
+#include "../kselftest.h"
+
+struct request {
+	int fd;
+	void *buffer;
+	size_t length;
+	off_t offset;
+	bool punch;
+	bool aio;
+	pthread_barrier_t *start;
+	atomic_bool done;
+	ssize_t result;
+};
+
+static ssize_t submit_aio(struct request *req)
+{
+	aio_context_t context = 0;
+	struct iocb cb = {
+		.aio_lio_opcode = IOCB_CMD_PWRITE, .aio_fildes = req->fd,
+		.aio_buf = (uintptr_t)req->buffer, .aio_nbytes = req->length,
+		.aio_offset = req->offset,
+	};
+	struct iocb *list[] = { &cb };
+	struct io_event event;
+	struct timespec timeout = { .tv_sec = 5 };
+	ssize_t result;
+
+	if (syscall(__NR_io_setup, 1, &context))
+		return -errno;
+	if (syscall(__NR_io_submit, context, 1, list) != 1)
+		result = -errno;
+	else if (syscall(__NR_io_getevents, context, 1, 1, &event, &timeout) != 1)
+		result = -ETIMEDOUT;
+	else
+		result = event.res;
+	syscall(__NR_io_destroy, context);
+	return result;
+}
+
+static void *submit(void *data)
+{
+	struct request *req = data;
+
+	if (req->start)
+		pthread_barrier_wait(req->start);
+	if (req->punch)
+		req->result = fallocate(req->fd,
+			FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE,
+			req->offset, req->length);
+	else if (req->aio)
+		req->result = submit_aio(req);
+	else
+		req->result = pwrite(req->fd, req->buffer, req->length, req->offset);
+	atomic_store(&req->done, true);
+	return NULL;
+}
+
+static long long milliseconds(void)
+{
+	struct timespec ts;
+
+	clock_gettime(CLOCK_MONOTONIC, &ts);
+	return (long long)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
+}
+
+static void timeout_handler(int sig)
+{
+	static const char message[] = "XFS DIO drain test timed out\n";
+
+	(void)sig;
+	if (write(STDERR_FILENO, message, sizeof(message) - 1) < 0)
+		_exit(KSFT_FAIL);
+	_exit(KSFT_FAIL);
+}
+
+static int faultfd(void *buffer, size_t length)
+{
+	struct uffdio_api api = { .api = UFFD_API };
+	struct uffdio_register reg = {
+		.range = { .start = (uintptr_t)buffer, .len = length },
+		.mode = UFFDIO_REGISTER_MODE_MISSING,
+	};
+	int fd;
+
+	/* Kernel-originated GUP faults must be handled, so omit USER_MODE_ONLY. */
+	fd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
+	if (fd < 0)
+		return -1;
+	if (ioctl(fd, UFFDIO_API, &api) || ioctl(fd, UFFDIO_REGISTER, &reg)) {
+		close(fd);
+		return -1;
+	}
+	return fd;
+}
+
+/*
+ * Stop a real, registered DIO in GUP. A drain must wait for this old request,
+ * although its file range is disjoint from the punch or post-EOF gap.
+ * Probe the exclusion before checking that newer disjoint I/O can complete.
+ */
+static int cohort_test(int fd, size_t unit, bool eof)
+{
+	struct request old = { .fd = fd, .length = unit, .offset = unit };
+	struct request drain = { .fd = fd, .length = unit,
+		.offset = (eof ? 10 : 4) * unit, .punch = !eof };
+	struct uffd_msg message;
+	struct uffdio_copy copy = { .len = unit };
+	struct pollfd pollfd;
+	struct iovec iov;
+	pthread_t writer, drainer;
+	void *buffer, *missing;
+	long long deadline;
+	bool excluded = false, progressed = false, rejected_write = false, ok = true;
+	int uffd;
+	ssize_t ret;
+
+	if (posix_memalign(&buffer, unit, 12 * unit))
+		ksft_exit_fail_msg("aligned allocation failed\n");
+	memset(buffer, 0x5a, 12 * unit);
+	if (ftruncate(fd, 0) || pwrite(fd, buffer, 8 * unit, 0) != 8 * (ssize_t)unit ||
+	    fsync(fd))
+		ksft_exit_fail_msg("prefill failed: %s\n", strerror(errno));
+	missing = mmap(NULL, unit, PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (missing == MAP_FAILED)
+		ksft_exit_fail_msg("mmap failed\n");
+	uffd = faultfd(missing, unit);
+	if (uffd < 0) {
+		free(buffer);
+		munmap(missing, unit);
+		return -1;
+	}
+	old.buffer = missing;
+	drain.buffer = buffer;
+	if (pthread_create(&writer, NULL, submit, &old))
+		ksft_exit_fail_msg("writer creation failed\n");
+	pollfd = (struct pollfd){ .fd = uffd, .events = POLLIN };
+	if (poll(&pollfd, 1, 5000) != 1 ||
+	    read(uffd, &message, sizeof(message)) != sizeof(message) ||
+	    message.event != UFFD_EVENT_PAGEFAULT)
+		ksft_exit_fail_msg("DIO did not reach the controlled GUP fault\n");
+	if (pthread_create(&drainer, NULL, submit, &drain))
+		ksft_exit_fail_msg("drainer creation failed\n");
+	iov = (struct iovec){ .iov_base = buffer + unit, .iov_len = unit };
+	deadline = milliseconds() + 5000;
+	do {
+		ret = preadv2(fd, &iov, 1, (eof ? 8 : 4) * unit, RWF_NOWAIT);
+		if (ret == -1 && errno == EAGAIN) {
+			excluded = true;
+			break;
+		}
+		if (atomic_load(&drain.done))
+			break;
+		usleep(1000);
+	} while (milliseconds() < deadline);
+	if (excluded) {
+		/* This must succeed before the stalled old request is released. */
+		ret = preadv2(fd, &iov, 1, 0, RWF_NOWAIT);
+		progressed = ret == (ssize_t)unit;
+		ret = pwritev2(fd, &iov, 1, (eof ? 8 : 4) * unit, RWF_NOWAIT);
+		rejected_write = ret == -1 && errno == EAGAIN;
+	}
+	ok = excluded && progressed && rejected_write && !atomic_load(&drain.done) &&
+		!atomic_load(&old.done);
+	if (!ok)
+		ksft_print_msg("%s: exclusion=%d disjoint-progress=%d drain-done=%d\n",
+			       eof ? "EOF" : "punch", excluded, progressed,
+			       atomic_load(&drain.done));
+	/* Resolve every missing source page and release both operations. */
+	memset(buffer, 0x5a, unit);
+	copy.src = (uintptr_t)buffer;
+	copy.dst = (uintptr_t)missing;
+	if (ioctl(uffd, UFFDIO_COPY, &copy))
+		ksft_exit_fail_msg("UFFDIO_COPY failed: %s\n", strerror(errno));
+	pthread_join(writer, NULL);
+	pthread_join(drainer, NULL);
+	ok &= old.result == (ssize_t)unit;
+	ok &= drain.result == (eof ? (ssize_t)unit : 0);
+	if (fsync(fd))
+		ok = false;
+	memset(buffer, 0xff, 12 * unit);
+	ret = pread(fd, buffer, (eof ? 11 : 8) * unit, 0);
+	ok &= ret == (eof ? 11 : 8) * (ssize_t)unit;
+	for (size_t i = 0; i < (eof ? 11 : 8) * unit; i++) {
+		unsigned char expected = 0x5a;
+
+		if (eof ? i >= 8 * unit && i < 10 * unit :
+			  i >= 4 * unit && i < 5 * unit)
+			expected = 0;
+		if (((unsigned char *)buffer)[i] != expected) {
+			ksft_print_msg("bad preserved data/gap at byte %zu\n", i);
+			ok = false;
+			break;
+		}
+	}
+	close(uffd);
+	munmap(missing, unit);
+	free(buffer);
+	return ok;
+}
+
+static bool bytes_equal(const void *buffer, size_t length, unsigned char value)
+{
+	const unsigned char *p = buffer;
+
+	for (size_t i = 0; i < length; i++) {
+		if (p[i] != value)
+			return false;
+	}
+	return true;
+}
+
+/* Exercise the MMAPLOCK fault dependency and real asynchronous COW end_io. */
+static int mapped_source_test(int fd, int clonefd, size_t unit, bool cow)
+{
+	void *buffer, *source;
+	bool ok = true;
+
+	if (posix_memalign(&buffer, unit, 8 * unit))
+		ksft_exit_fail_msg("aligned allocation failed\n");
+	memset(buffer, 0x5a, 8 * unit);
+	if (ftruncate(fd, 0) || pwrite(fd, buffer, 8 * unit, 0) != 8 * (ssize_t)unit ||
+	    fsync(fd))
+		ksft_exit_fail_msg("mapped source prefill failed\n");
+	source = mmap(NULL, unit, PROT_READ, MAP_PRIVATE, fd, 0);
+	if (source == MAP_FAILED)
+		ksft_exit_fail_msg("file source mmap failed\n");
+	for (unsigned int round = 0; round < 64; round++) {
+		pthread_barrier_t barrier;
+		struct request writer = { .fd = fd, .buffer = source,
+			.length = unit, .offset = (cow ? 2 : 10) * unit,
+			.aio = round & 1, .start = &barrier };
+		struct request punch = { .fd = fd, .length = unit,
+			.offset = 4 * unit, .punch = true, .start = &barrier };
+		pthread_t writing, punching;
+
+		if (ftruncate(fd, 8 * unit))
+			ksft_exit_fail_msg("mapped source truncate failed\n");
+		if (cow) {
+			memset(buffer, 0xa5, unit);
+			if (pwrite(fd, buffer, unit, 2 * unit) != (ssize_t)unit || fsync(fd))
+				ksft_exit_fail_msg("COW destination prefill failed\n");
+			if (ioctl(clonefd, FICLONE, fd)) {
+				if (errno == EOPNOTSUPP || errno == ENOTTY) {
+					munmap(source, unit);
+					free(buffer);
+					return -1;
+				}
+				ksft_exit_fail_msg("FICLONE failed: %s\n", strerror(errno));
+			}
+		}
+		/* Drop PTEs and clean cache so GUP must fault through this inode. */
+		if (madvise(source, unit, MADV_DONTNEED) ||
+		    posix_fadvise(fd, 0, unit, POSIX_FADV_DONTNEED))
+			ksft_exit_fail_msg("cannot evict mapped source\n");
+		if (pthread_barrier_init(&barrier, NULL, 3) ||
+		    pthread_create(&writing, NULL, submit, &writer) ||
+		    pthread_create(&punching, NULL, submit, &punch))
+			ksft_exit_fail_msg("mapped source thread setup failed\n");
+		pthread_barrier_wait(&barrier);
+		pthread_join(writing, NULL);
+		pthread_join(punching, NULL);
+		pthread_barrier_destroy(&barrier);
+		if (writer.result == -ENOSYS) {
+			munmap(source, unit);
+			free(buffer);
+			return -1;
+		}
+		ok &= writer.result == (ssize_t)unit && punch.result == 0;
+		ok &= pread(fd, buffer, unit, writer.offset) == (ssize_t)unit &&
+			bytes_equal(buffer, unit, 0x5a);
+		ok &= pread(fd, buffer, unit, 4 * unit) == (ssize_t)unit &&
+			bytes_equal(buffer, unit, 0);
+		if (cow)
+			ok &= pread(clonefd, buffer, unit, 2 * unit) == (ssize_t)unit &&
+				bytes_equal(buffer, unit, 0xa5);
+		if (!ok)
+			break;
+	}
+	munmap(source, unit);
+	free(buffer);
+	return ok;
+}
+
+int main(int argc, char **argv)
+{
+	struct statfs fs;
+	char *path;
+	size_t unit;
+	int fd, clonefd, result, failed = 0;
+
+	ksft_print_header();
+	if (argc != 2)
+		ksft_exit_skip("usage: %s WRITABLE_XFS_DIRECTORY\n", argv[0]);
+	if (statfs(argv[1], &fs) || fs.f_type != XFS_SUPER_MAGIC)
+		ksft_exit_skip("requires a writable XFS directory\n");
+	unit = sysconf(_SC_PAGESIZE);
+	if ((size_t)fs.f_bsize > unit)
+		unit = fs.f_bsize;
+	if (asprintf(&path, "%s/xfs-dio-drain-XXXXXX", argv[1]) < 0)
+		ksft_exit_fail_msg("path allocation failed\n");
+	fd = mkostemp(path, O_DIRECT);
+	if (fd < 0)
+		ksft_exit_fail_msg("open failed: %s\n", strerror(errno));
+	unlink(path);
+	free(path);
+	if (asprintf(&path, "%s/xfs-dio-clone-XXXXXX", argv[1]) < 0)
+		ksft_exit_fail_msg("clone path allocation failed\n");
+	clonefd = mkostemp(path, O_DIRECT);
+	if (clonefd < 0)
+		ksft_exit_fail_msg("clone file open failed\n");
+	unlink(path);
+	free(path);
+	signal(SIGALRM, timeout_handler);
+	alarm(30);
+	ksft_set_plan(3);
+	for (int eof = 0; eof < 1; eof++) {
+		result = cohort_test(fd, unit, eof);
+		if (result < 0)
+			ksft_test_result_skip("%s drain needs kernel-fault userfaultfd support\n",
+					      eof ? "EOF" : "punch");
+		else {
+			ksft_test_result(result, "%s drain permits disjoint I/O\n",
+					 eof ? "EOF" : "punch");
+			failed |= !result;
+		}
+	}
+	for (int cow = 0; cow < 2; cow++) {
+		result = mapped_source_test(fd, clonefd, unit, cow);
+		if (result < 0)
+			ksft_test_result_skip("mapped source test needs AIO%s\n",
+					      cow ? " and reflink support" : "");
+		else {
+			ksft_test_result(result, "same-inode faulting source: %s vs punch\n",
+					 cow ? "sync/AIO COW writes" : "sync/AIO extending writes");
+			failed |= !result;
+		}
+	}
+	close(clonefd);
+	alarm(0);
+	close(fd);
+	ksft_exit(!failed);
+}
diff --git a/tools/testing/selftests/filesystems/xfs_dio_remap.c b/tools/testing/selftests/filesystems/xfs_dio_remap.c
new file mode 100644
index 000000000000..8460a932d195
--- /dev/null
+++ b/tools/testing/selftests/filesystems/xfs_dio_remap.c
@@ -0,0 +1,391 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Deterministic real-syscall admission retries; see CONFIG_XFS_DIO_TEST. */
+#define _GNU_SOURCE
+#include <errno.h>
+#include <fcntl.h>
+#include <linux/fs.h>
+#include <linux/fiemap.h>
+#include <linux/magic.h>
+#include <linux/userfaultfd.h>
+#include <poll.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/ioctl.h>
+#include <sys/mman.h>
+#include <sys/resource.h>
+#include <sys/stat.h>
+#include <sys/statfs.h>
+#include <sys/syscall.h>
+#include <sys/uio.h>
+#include <time.h>
+#include <unistd.h>
+#include "../kselftest.h"
+
+/* Event indices in fs/xfs/xfs_dio_test.h. */
+enum { BATCH_COLLECT = 2, WAIT_UNLOCKED = 4, WRITE_CHECKS, WRITE_ADMIT,
+	WRITE_REJECT, UNALIGNED_RECHECK, REMAP_ACTIVE, REMAP_UPGRADE, EVENTS };
+#define GATE(event) (1UL << (event))
+struct snapshot { unsigned long gates, timedout, hits[EVENTS], value[EVENTS]; };
+struct request {
+	int fd;
+	void *buffer;
+	size_t length;
+	off_t offset;
+	bool read, punch;
+	ssize_t result;
+	int error;
+};
+struct parked_read {
+	struct request req;
+	pthread_t thread;
+	int uffd;
+};
+
+static void die(const char *what)
+{
+	ksft_exit_fail_msg("%s: %s\n", what, strerror(errno));
+}
+
+static long long milliseconds(void)
+{
+	struct timespec now;
+
+	clock_gettime(CLOCK_MONOTONIC, &now);
+	return now.tv_sec * 1000LL + now.tv_nsec / 1000000;
+}
+
+static void *submit(void *data)
+{
+	struct request *r = data;
+
+	if (r->punch)
+		r->result = fallocate(r->fd, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE,
+				r->offset, r->length);
+	else if (r->read)
+		r->result = pread(r->fd, r->buffer, r->length, r->offset);
+	else
+		r->result = pwrite(r->fd, r->buffer, r->length, r->offset);
+	r->error = errno;
+	return NULL;
+}
+
+static void command(int control, const char *verb, unsigned long value)
+{
+	char text[80];
+	int len = snprintf(text, sizeof(text), "%s %lu", verb, value);
+
+	if (write(control, text, len) != len)
+		die("DIO test command");
+}
+
+static int arm(int fd, unsigned long gates)
+{
+	char text[80];
+	int control = open("/sys/kernel/debug/xfs/dio_test", O_RDWR);
+	int len;
+
+	if (control < 0)
+		return -1;
+	len = snprintf(text, sizeof(text), "arm %d %lu", fd, gates);
+	if (write(control, text, len) != len)
+		die("arm DIO test");
+	return control;
+}
+
+static struct snapshot snapshot(int control)
+{
+	struct snapshot s = { };
+	char buffer[2048], *cursor, *next;
+	ssize_t n = pread(control, buffer, sizeof(buffer) - 1, 0);
+
+	if (n <= 0)
+		die("read DIO test state");
+	buffer[n] = 0;
+	if (sscanf(buffer, "%lu %lu", &s.gates, &s.timedout) != 2)
+		die("parse DIO test state");
+	cursor = strchr(buffer, '\n');
+	for (int i = 0; i < EVENTS; i++) {
+		if (!cursor || sscanf(cursor + 1, "%lu %lu", &s.hits[i],
+				&s.value[i]) != 2)
+			die("parse DIO test event");
+		next = strchr(cursor + 1, '\n');
+		cursor = next;
+	}
+	return s;
+}
+
+static void wait_event(int control, int event, unsigned long count)
+{
+	long long deadline = milliseconds() + 5000;
+
+	do {
+		struct snapshot s = snapshot(control);
+
+		if (s.hits[event] >= count && !s.timedout)
+			return;
+		usleep(1000);
+	} while (milliseconds() < deadline);
+	ksft_exit_fail_msg("DIO event %d did not reach %lu\n", event, count);
+}
+
+static void park_read(struct parked_read *old, int fd, size_t unit)
+{
+	struct uffdio_api api = { .api = UFFD_API };
+	struct uffdio_register reg = { .mode = UFFDIO_REGISTER_MODE_MISSING };
+	struct uffd_msg message;
+	struct pollfd pollfd;
+
+	memset(old, 0, sizeof(*old));
+	old->req = (struct request){ .fd = fd, .length = unit,
+		.offset = unit, .read = true };
+	old->req.buffer = mmap(NULL, unit, PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (old->req.buffer == MAP_FAILED)
+		die("missing read destination");
+	old->uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
+	reg.range.start = (uintptr_t)old->req.buffer;
+	reg.range.len = unit;
+	if (old->uffd < 0 || ioctl(old->uffd, UFFDIO_API, &api) ||
+	    ioctl(old->uffd, UFFDIO_REGISTER, &reg))
+		die("kernel-fault userfaultfd required by admission test");
+	if (pthread_create(&old->thread, NULL, submit, &old->req))
+		die("parked read thread");
+	pollfd = (struct pollfd){ .fd = old->uffd, .events = POLLIN };
+	if (poll(&pollfd, 1, 5000) != 1 ||
+	    read(old->uffd, &message, sizeof(message)) != sizeof(message) ||
+	    message.event != UFFD_EVENT_PAGEFAULT)
+		ksft_exit_fail_msg("read did not register and reach GUP\n");
+}
+
+static void release_read(struct parked_read *old, void *buffer, size_t unit)
+{
+	struct uffdio_copy copy = { .src = (uintptr_t)buffer,
+		.dst = (uintptr_t)old->req.buffer, .len = unit };
+
+	if (ioctl(old->uffd, UFFDIO_COPY, &copy))
+		die("release registered read");
+	pthread_join(old->thread, NULL);
+	if (old->req.result != (ssize_t)unit)
+		ksft_exit_fail_msg("registered read result %zd\n", old->req.result);
+	close(old->uffd);
+	munmap(old->req.buffer, unit);
+}
+
+static void wait_exclusion(int fd, void *buffer, size_t unit, off_t offset)
+{
+	struct iovec iov = { .iov_base = buffer, .iov_len = unit };
+	long long deadline = milliseconds() + 5000;
+
+	do {
+		if (preadv2(fd, &iov, 1, offset, RWF_NOWAIT) == -1 && errno == EAGAIN)
+			return;
+		usleep(1000);
+	} while (milliseconds() < deadline);
+	ksft_exit_fail_msg("real drainer did not publish exclusion\n");
+}
+
+static void prefill(int fd, void *buffer, size_t unit)
+{
+	memset(buffer, 0x5a, 16 * unit);
+	if (ftruncate(fd, 0) || pwrite(fd, buffer, 8 * unit, 0) != 8 * (ssize_t)unit ||
+	    fsync(fd))
+		die("prefill");
+}
+
+static bool bytes_equal(const void *buffer, size_t length, unsigned char value)
+{
+	const unsigned char *bytes = buffer;
+
+	for (size_t i = 0; i < length; i++)
+		if (bytes[i] != value)
+			return false;
+	return true;
+}
+
+/* FIEMAP proves sharing of the exact partial block before the writer resumes. */
+static bool shared_extent(int fd, off_t offset, uint64_t *physical)
+{
+	struct {
+		struct fiemap map;
+		struct fiemap_extent extent;
+	} request = { .map = { .fm_start = offset, .fm_length = 512,
+		.fm_flags = FIEMAP_FLAG_SYNC, .fm_extent_count = 1 } };
+
+	if (ioctl(fd, FS_IOC_FIEMAP, &request))
+		die("FIEMAP sharing proof");
+	if (request.map.fm_mapped_extents != 1 ||
+	    !(request.extent.fe_flags & FIEMAP_EXTENT_SHARED))
+		return false;
+	*physical = request.extent.fe_physical + offset - request.extent.fe_logical;
+	return true;
+}
+
+struct clone_request { int source, destination, result, error; };
+
+static void *clone_file(void *data)
+{
+	struct clone_request *r = data;
+
+	r->result = ioctl(r->destination, FICLONE, r->source);
+	r->error = errno;
+	return NULL;
+}
+
+/* Change COW eligibility only after observing the writer drop its IOLOCK. */
+static bool remap_revalidation(int fd, int clonefd, void *buffer, size_t unit,
+		bool active)
+{
+	struct parked_read old;
+	struct request punch = { .fd = fd, .offset = 4 * unit,
+		.length = unit, .punch = true };
+	struct request writer = { .fd = fd, .offset = 4 * unit + 512,
+		.buffer = buffer, .length = 512 };
+	struct clone_request clone = { .source = fd, .destination = clonefd };
+	struct snapshot s;
+	pthread_t punching, writing, cloning;
+	uint64_t source_physical = 0, clone_physical = 0;
+	char path[64];
+	int control, buffered;
+	bool ok;
+
+	prefill(fd, buffer, unit);
+	/* First write must really enter the unaligned direct path, before COW. */
+	if (ftruncate(clonefd, 0))
+		die("clear clone");
+	control = arm(fd, GATE(WAIT_UNLOCKED));
+	if (control < 0)
+		ksft_exit_skip("requires CONFIG_XFS_DIO_TEST and mounted debugfs\n");
+	park_read(&old, fd, unit);
+	if (pthread_create(&punching, NULL, submit, &punch))
+		die("punch thread");
+	wait_exclusion(fd, buffer + unit, unit, punch.offset);
+	memset(buffer, 0xa5, 512);
+	if (pthread_create(&writing, NULL, submit, &writer))
+		die("unaligned writer");
+	wait_event(control, WAIT_UNLOCKED, 1);
+	s = snapshot(control);
+	ok = s.value[WAIT_UNLOCKED] == (unsigned long)writer.offset &&
+		s.hits[UNALIGNED_RECHECK] == 0 && s.hits[WRITE_CHECKS] == 1;
+	release_read(&old, buffer + unit, unit);
+	pthread_join(punching, NULL);
+	ok &= punch.result == 0;
+	/* Restore a real allocated block after the punch, still with writer gated. */
+	snprintf(path, sizeof(path), "/proc/self/fd/%d", fd);
+	buffered = open(path, O_RDWR);
+	memset(buffer + unit, 0x5a, unit);
+	if (buffered < 0 || pwrite(buffered, buffer + unit, unit, 4 * unit) !=
+	    (ssize_t)unit || fsync(buffered))
+		die("restore block before clone");
+	close(buffered);
+	if (active)
+		command(control, "gate", GATE(WAIT_UNLOCKED) | GATE(REMAP_ACTIVE));
+	if (pthread_create(&cloning, NULL, clone_file, &clone))
+		die("reflink thread");
+	if (active) {
+		wait_event(control, REMAP_ACTIVE, 1);
+		/* Source IOLOCK is now shared, with IREMAPPING set. */
+		command(control, "release", GATE(WAIT_UNLOCKED));
+		wait_event(control, REMAP_UPGRADE, 1);
+		command(control, "release", GATE(REMAP_ACTIVE));
+	}
+	pthread_join(cloning, NULL);
+	if (clone.result)
+		ksft_exit_fail_msg("FICLONE failed: %s\n", strerror(clone.error));
+	/* Completed case can prove extent sharing before allowing any write. */
+	if (!active) {
+		ok &= shared_extent(fd, writer.offset, &source_physical) &&
+			shared_extent(clonefd, writer.offset, &clone_physical) &&
+			source_physical == clone_physical;
+		command(control, "gate", 0);
+	}
+	pthread_join(writing, NULL);
+	s = snapshot(control);
+	ok &= writer.result == 512 && !s.timedout && s.hits[WAIT_UNLOCKED] == 1 &&
+		s.hits[UNALIGNED_RECHECK] == 1 && s.hits[WRITE_ADMIT] == 0 &&
+		/* IOLOCK_EXCL=1, IOLOCK_SHARED=2; low bit records COW eligibility. */
+		s.value[UNALIGNED_RECHECK] == ((active ? 1UL : 2UL) << 8 | 1UL) &&
+		s.hits[REMAP_UPGRADE] == (unsigned long)active;
+	/* The last write-check pass is the buffered COW fallback. */
+	ok &= s.hits[WRITE_CHECKS] == 4 && !fsync(fd) && !fsync(clonefd);
+	ok &= pread(fd, buffer + unit, unit, 4 * unit) == (ssize_t)unit &&
+		bytes_equal(buffer + unit, 512, 0x5a) &&
+		bytes_equal(buffer + unit + 512, 512, 0xa5) &&
+		bytes_equal(buffer + unit + 1024, unit - 1024, 0x5a);
+	ok &= pread(clonefd, buffer + unit, unit, 4 * unit) == (ssize_t)unit &&
+		bytes_equal(buffer + unit, unit, 0x5a);
+	ksft_print_msg("remap evidence: active=%d waits=%lu upgrades=%lu rechecks=%lu "
+		"lock_cow=%lu checks=%lu admits=%lu shared_physical=%llu result=%zd\n", active,
+		s.hits[WAIT_UNLOCKED], s.hits[REMAP_UPGRADE], s.hits[UNALIGNED_RECHECK],
+		s.value[UNALIGNED_RECHECK], s.hits[WRITE_CHECKS], s.hits[WRITE_ADMIT],
+		(unsigned long long)source_physical, writer.result);
+	close(control);
+	return ok;
+}
+
+static void timeout_handler(int signal)
+{
+	(void)signal;
+	_exit(KSFT_FAIL);
+}
+
+int main(int argc, char **argv)
+{
+	struct statfs fs;
+	char *path;
+	void *buffer;
+	size_t unit;
+	int fd, clonefd, failed = 0;
+
+	ksft_print_header();
+	if (argc != 2 || statfs(argv[1], &fs) || fs.f_type != XFS_SUPER_MAGIC)
+		ksft_exit_skip("requires a writable XFS directory\n");
+	unit = sysconf(_SC_PAGESIZE);
+	if ((size_t)fs.f_bsize > unit)
+		unit = fs.f_bsize;
+	if (asprintf(&path, "%s/xfs-remap-XXXXXX", argv[1]) < 0 ||
+	    posix_memalign(&buffer, unit, 16 * unit))
+		die("test allocation");
+	fd = mkostemp(path, O_DIRECT);
+	if (fd < 0)
+		die("test file");
+	unlink(path);
+	free(path);
+	if (asprintf(&path, "%s/xfs-clone-XXXXXX", argv[1]) < 0)
+		die("clone path");
+	clonefd = mkostemp(path, O_DIRECT);
+	if (clonefd < 0)
+		die("clone file");
+	unlink(path);
+	free(path);
+	signal(SIGALRM, timeout_handler);
+	alarm(40);
+	ksft_set_plan(2);
+	for (int active = 0; active < 2; active++) {
+		/* Each case needs an inode that has never carried the reflink flag. */
+		bool ok = remap_revalidation(fd, clonefd, buffer, unit, active);
+
+		ksft_test_result(ok, "dropped IOLOCK rechecks %s reflink and unaligned COW\n",
+			active ? "active" : "completed");
+		failed |= !ok;
+		/* Reflink flags persist, so use a fresh writer inode for each case. */
+		if (!active) {
+			close(fd);
+			if (asprintf(&path, "%s/xfs-remap-XXXXXX", argv[1]) < 0)
+				die("second path");
+			fd = mkostemp(path, O_DIRECT);
+			if (fd < 0)
+				die("second inode");
+			unlink(path);
+			free(path);
+		}
+	}
+	alarm(0);
+	close(fd);
+	close(clonefd);
+	free(buffer);
+	ksft_exit(!failed);
+}
-- 
2.56.0

[RFC PATCH 6/6] xfs: use range-based DIO draining for post-EOF zeroing

[Vitali Lovich]

Zeroing a gap before a direct write normally drains DIO with exclusive
IOLOCK. Use range draining for block-aligned, non-atomic data-device
writes entering write checks beyond EOF with shared IOLOCK.

Drop IOLOCK before obtaining the lazy drain state and its mutex, then
reacquire through xfs_ilock_iocb_for_write() and repeat write checks.
A drainer may need IOLOCK_EXCL during these checks; a second drainer must
not hold shared IOLOCK while waiting for its mutex, or both would block.
Publish the exclusion under shared IOLOCK, demoting first if required,
and drain the pre-cutover cohort while admitting disjoint new DIO.

Round the exclusion outward to the larger of the filesystem block size
and mapping_max_folio_size() to cover every folio touched by zeroing and
invalidation. Recheck EOF after draining because completions can advance
it, then zero any remaining gap under MMAPLOCK_EXCL.

Finish the exclusion before submitting the triggering write, which can
share its final folio. Release the drain mutex on every write-check exit.
Ineligible writes retain existing handling; allocation and drain-start
errors propagate to the caller.

Extend the filesystem selftest to check the post-EOF drain consumer:
park an old registered request, observe the gap exclusion, and verify
disjoint progress, overlapping NOWAIT rejection, and preserved data.

Also hold an old extending request while preparation drains, then verify
its completed bytes survive the post-drain EOF recheck and only the
remaining gap is zeroed.

Signed-off-by: Vitali Lovich <vlovich@gmail.com>
---
 fs/xfs/xfs_file.c                             | 109 +++++++++++++++---
 .../selftests/filesystems/xfs_dio_drain.c     |  41 +++++--
 2 files changed, 125 insertions(+), 25 deletions(-)

diff --git a/fs/xfs/xfs_file.c b/fs/xfs/xfs_file.c
index ad83328b7350..53b3d6266e03 100644
--- a/fs/xfs/xfs_file.c
+++ b/fs/xfs/xfs_file.c
@@ -453,6 +453,55 @@ xfs_file_splice_read(
 	return ret;
 }
 
+/*
+ * Serialize zeroing with I/O to all cache folios that can cover the old EOF
+ * block and gap, not unrelated requests outside those folios. Clear the
+ * exclusion before submitting this write, which can share the final folio.
+ * Drain the finite cohort admitted before the generation cutover, including
+ * COW end-I/O remapping. New unrelated requests can proceed under shared IOLOCK.
+ */
+static int
+xfs_file_write_zero_eof_range(
+	struct kiocb		*iocb,
+	loff_t			isize,
+	unsigned int		*iolock,
+	struct inode_dio_drain	*drain)
+{
+	struct inode		*inode = file_inode(iocb->ki_filp);
+	struct xfs_inode	*ip = XFS_I(inode);
+	loff_t			granularity = max_t(loff_t, i_blocksize(inode),
+					mapping_max_folio_size(inode->i_mapping));
+	loff_t			start = round_down(isize, granularity);
+	loff_t			last = (iocb->ki_pos - 1) | (granularity - 1);
+	int			error = 0;
+
+	error = kiocb_modified(iocb);
+	if (error)
+		return error;
+	if (*iolock == XFS_IOLOCK_EXCL) {
+		xfs_ilock_demote(ip, XFS_IOLOCK_EXCL);
+		*iolock = XFS_IOLOCK_SHARED;
+	}
+	/* Folio invalidation is wider than zeroing, even for sub-page blocks. */
+	error = inode_dio_drain_start(inode, drain, start, last);
+	if (error)
+		return error;
+	inode_dio_drain_wait(inode, drain);
+	spin_lock(&ip->i_flags_lock);
+	isize = i_size_read(inode);
+	spin_unlock(&ip->i_flags_lock);
+	if (iocb->ki_pos > isize) {
+		trace_xfs_zero_eof(ip, isize, iocb->ki_pos - isize);
+		xfs_ilock(ip, XFS_MMAPLOCK_EXCL);
+		error = iomap_zero_range(inode, isize, iocb->ki_pos - isize,
+				NULL, &xfs_buffered_write_iomap_ops,
+				&xfs_iomap_write_ops, NULL);
+		xfs_iunlock(ip, XFS_MMAPLOCK_EXCL);
+	}
+	inode_dio_drain_finish(inode, drain);
+	return error;
+}
+
 /*
  * Take care of zeroing post-EOF blocks when they might exist.
  *
@@ -467,20 +516,17 @@ xfs_file_write_zero_eof(
 	unsigned int		*iolock,
 	size_t			count,
 	bool			*drained_dio,
-	struct xfs_zone_alloc_ctx *ac)
+	struct xfs_zone_alloc_ctx *ac,
+	struct inode_dio_drain	*drain)
 {
 	struct xfs_inode	*ip = XFS_I(iocb->ki_filp->f_mapping->host);
 	loff_t			isize;
 	int			error;
 
 	/*
-	 * We need to serialise against EOF updates that occur in IO completions
-	 * here. We want to make sure that nobody is changing the size while
-	 * we do this check until we have placed an IO barrier (i.e. hold
-	 * XFS_IOLOCK_EXCL) that prevents new IO from being dispatched.  The
-	 * spinlock effectively forms a memory barrier once we have
-	 * XFS_IOLOCK_EXCL so we are guaranteed to see the latest EOF value and
-	 * hence be able to correctly determine if we need to run zeroing.
+	 * Pair with EOF updates in DIO completion. The IOLOCK excludes truncate,
+	 * so EOF can only move forwards. Recheck after establishing either a
+	 * range exclusion or the exclusive IOLOCK and draining the relevant I/O.
 	 */
 	spin_lock(&ip->i_flags_lock);
 	isize = i_size_read(VFS_I(ip));
@@ -493,6 +539,9 @@ xfs_file_write_zero_eof(
 	if (iocb->ki_flags & IOCB_NOWAIT)
 		return -EAGAIN;
 
+	if (drain)
+		return xfs_file_write_zero_eof_range(iocb, isize, iolock, drain);
+
 	if (!*drained_dio) {
 		/*
 		 * If zeroing is needed and we are currently holding the iolock
@@ -543,23 +592,49 @@ xfs_file_write_checks(
 	struct xfs_zone_alloc_ctx *ac)
 {
 	struct inode		*inode = iocb->ki_filp->f_mapping->host;
+	struct inode_dio_drain	*drain = NULL;
 	size_t			count = iov_iter_count(from);
 	bool			drained_dio = false;
 	ssize_t			error;
 
+	/*
+	 * Another drainer may need IOLOCK_EXCL during write checks. Do not pin
+	 * shared IOLOCK while waiting for its drain mutex.
+	 */
+	if (*iolock == XFS_IOLOCK_SHARED &&
+	    (iocb->ki_flags & IOCB_DIRECT) &&
+	    !(iocb->ki_flags & IOCB_ATOMIC) &&
+	    !XFS_IS_REALTIME_INODE(XFS_I(inode)) &&
+	    IS_ALIGNED(iocb->ki_pos | count, i_blocksize(inode)) &&
+	    iocb->ki_pos > i_size_read(inode)) {
+		if (iocb->ki_flags & IOCB_NOWAIT)
+			return -EAGAIN;
+		xfs_iunlock(XFS_I(inode), *iolock);
+		drain = inode_dio_drain_get(&XFS_I(inode)->i_dio_drain);
+		if (IS_ERR(drain)) {
+			xfs_ilock(XFS_I(inode), *iolock);
+			return PTR_ERR(drain);
+		}
+		mutex_lock(&drain->lock);
+		error = xfs_ilock_iocb_for_write(iocb, iolock);
+		if (error) {
+			*iolock = 0;
+			goto out;
+		}
+	}
 restart:
 	/* Normalize append position and enforce limits before range admission. */
 	error = generic_write_checks(iocb, from);
 	xfs_dio_test_event(inode, XFS_DIO_TEST_WRITE_CHECKS,
 			iov_iter_count(from));
 	if (error <= 0)
-		return error;
+		goto out;
 
 	if (iocb->ki_flags & IOCB_DIRECT) {
 		error = xfs_dio_range_wait(iocb, iov_iter_count(from), iolock,
 				true);
 		if (error < 0)
-			return error;
+			goto out;
 		if (error) {
 			iov_iter_reexpand(from, count);
 			/* EOF and limits can change while IOLOCK is dropped. */
@@ -576,7 +651,7 @@ xfs_file_write_checks(
 	}
 
 	if (error)
-		return error;
+		goto out;
 
 	/*
 	 * For changing security info in file_remove_privs() we need i_rwsem
@@ -588,7 +663,7 @@ xfs_file_write_checks(
 		error = xfs_ilock_iocb(iocb, *iolock);
 		if (error) {
 			*iolock = 0;
-			return error;
+			goto out;
 		}
 		goto restart;
 	}
@@ -605,14 +680,18 @@ xfs_file_write_checks(
 	 */
 	if (iocb->ki_pos > i_size_read(inode)) {
 		error = xfs_file_write_zero_eof(iocb, from, iolock, count,
-				&drained_dio, ac);
+				&drained_dio, ac, drain);
 		if (error == 1)
 			goto restart;
 		if (error)
-			return error;
+			goto out;
 	}
 
-	return kiocb_modified(iocb);
+	error = kiocb_modified(iocb);
+out:
+	if (drain)
+		mutex_unlock(&drain->lock);
+	return error;
 }
 
 static ssize_t
diff --git a/tools/testing/selftests/filesystems/xfs_dio_drain.c b/tools/testing/selftests/filesystems/xfs_dio_drain.c
index f2e005f6e896..c84f298ff3e4 100644
--- a/tools/testing/selftests/filesystems/xfs_dio_drain.c
+++ b/tools/testing/selftests/filesystems/xfs_dio_drain.c
@@ -21,6 +21,7 @@
 #include <string.h>
 #include <sys/ioctl.h>
 #include <sys/mman.h>
+#include <sys/stat.h>
 #include <sys/statfs.h>
 #include <sys/syscall.h>
 #include <sys/uio.h>
@@ -122,16 +123,19 @@ static int faultfd(void *buffer, size_t length)
 }
 
 /*
- * Stop a real, registered DIO in GUP. A drain must wait for this old request,
- * although its file range is disjoint from the punch or post-EOF gap.
+ * Stop a real, registered DIO in GUP. A drain must wait for this old request.
+ * The overwrite cases use a disjoint range; the extending case starts exactly
+ * at EOF and must finish before the farther writer zeroes the remaining gap.
  * Probe the exclusion before checking that newer disjoint I/O can complete.
  */
-static int cohort_test(int fd, size_t unit, bool eof)
+static int cohort_test(int fd, size_t unit, bool eof, bool advance_eof)
 {
-	struct request old = { .fd = fd, .length = unit, .offset = unit };
+	struct request old = { .fd = fd, .length = unit,
+		.offset = (advance_eof ? 8 : 1) * unit };
 	struct request drain = { .fd = fd, .length = unit,
 		.offset = (eof ? 10 : 4) * unit, .punch = !eof };
 	struct uffd_msg message;
+	struct stat st;
 	struct uffdio_copy copy = { .len = unit };
 	struct pollfd pollfd;
 	struct iovec iov;
@@ -190,13 +194,14 @@ static int cohort_test(int fd, size_t unit, bool eof)
 	}
 	ok = excluded && progressed && rejected_write && !atomic_load(&drain.done) &&
 		!atomic_load(&old.done);
+	ok &= !fstat(fd, &st) && st.st_size == 8 * (off_t)unit;
 	if (!ok)
 		ksft_print_msg("%s: exclusion=%d disjoint-progress=%d drain-done=%d\n",
 			       eof ? "EOF" : "punch", excluded, progressed,
 			       atomic_load(&drain.done));
 	/* Resolve every missing source page and release both operations. */
-	memset(buffer, 0x5a, unit);
-	copy.src = (uintptr_t)buffer;
+	memset(buffer + 2 * unit, advance_eof ? 0xa5 : 0x5a, unit);
+	copy.src = (uintptr_t)(buffer + 2 * unit);
 	copy.dst = (uintptr_t)missing;
 	if (ioctl(uffd, UFFDIO_COPY, &copy))
 		ksft_exit_fail_msg("UFFDIO_COPY failed: %s\n", strerror(errno));
@@ -206,15 +211,19 @@ static int cohort_test(int fd, size_t unit, bool eof)
 	ok &= drain.result == (eof ? (ssize_t)unit : 0);
 	if (fsync(fd))
 		ok = false;
+	ok &= !fstat(fd, &st) &&
+		st.st_size == (eof ? 11 : 8) * (off_t)unit;
 	memset(buffer, 0xff, 12 * unit);
 	ret = pread(fd, buffer, (eof ? 11 : 8) * unit, 0);
 	ok &= ret == (eof ? 11 : 8) * (ssize_t)unit;
 	for (size_t i = 0; i < (eof ? 11 : 8) * unit; i++) {
 		unsigned char expected = 0x5a;
 
-		if (eof ? i >= 8 * unit && i < 10 * unit :
+		if (eof ? i >= (advance_eof ? 9 : 8) * unit && i < 10 * unit :
 			  i >= 4 * unit && i < 5 * unit)
 			expected = 0;
+		if (advance_eof && i >= 8 * unit && i < 9 * unit)
+			expected = 0xa5;
 		if (((unsigned char *)buffer)[i] != expected) {
 			ksft_print_msg("bad preserved data/gap at byte %zu\n", i);
 			ok = false;
@@ -341,9 +350,9 @@ int main(int argc, char **argv)
 	free(path);
 	signal(SIGALRM, timeout_handler);
 	alarm(30);
-	ksft_set_plan(3);
-	for (int eof = 0; eof < 1; eof++) {
-		result = cohort_test(fd, unit, eof);
+	ksft_set_plan(5);
+	for (int eof = 0; eof < 2; eof++) {
+		result = cohort_test(fd, unit, eof, false);
 		if (result < 0)
 			ksft_test_result_skip("%s drain needs kernel-fault userfaultfd support\n",
 					      eof ? "EOF" : "punch");
@@ -353,6 +362,18 @@ int main(int argc, char **argv)
 			failed |= !result;
 		}
 	}
+	/*
+	 * The old cohort member extends EOF from 8 to 9 units while the farther
+	 * writer is draining it. The preparer must reread i_size after draining:
+	 * zeroing from its original 8-unit snapshot would destroy the 0xa5 data.
+	 */
+	result = cohort_test(fd, unit, true, true);
+	if (result < 0)
+		ksft_test_result_skip("EOF advance needs kernel-fault userfaultfd support\n");
+	else {
+		ksft_test_result(result, "EOF advance during drain preserves completed data\n");
+		failed |= !result;
+	}
 	for (int cow = 0; cow < 2; cow++) {
 		result = mapped_source_test(fd, clonefd, unit, cow);
 		if (result < 0)
-- 
2.56.0