size_t copied, struct folio *folio)
 {
        const struct iomap *srcmap = iomap_iter_srcmap(iter);
-       loff_t old_size = iter->inode->i_size;
-       size_t ret;
-
-       if (srcmap->type == IOMAP_INLINE) {
-               ret = iomap_write_end_inline(iter, folio, pos, copied);
-       } else if (srcmap->flags & IOMAP_F_BUFFER_HEAD) {
-               ret = block_write_end(NULL, iter->inode->i_mapping, pos, len,
-                               copied, &folio->page, NULL);
-       } else {
-               ret = __iomap_write_end(iter->inode, pos, len, copied, folio);
-       }
-
-       /*
-        * Update the in-memory inode size after copying the data into the page
-        * cache.  It's up to the file system to write the updated size to disk,
-        * preferably after I/O completion so that no stale data is exposed.
-        */
-       if (pos + ret > old_size) {
-               i_size_write(iter->inode, pos + ret);
-               iter->iomap.flags |= IOMAP_F_SIZE_CHANGED;
-       }
-       __iomap_put_folio(iter, pos, ret, folio);
 
-       if (old_size < pos)
-               pagecache_isize_extended(iter->inode, old_size, pos);
-       return ret;
+       if (srcmap->type == IOMAP_INLINE)
+               return iomap_write_end_inline(iter, folio, pos, copied);
+       if (srcmap->flags & IOMAP_F_BUFFER_HEAD)
+               return block_write_end(NULL, iter->inode->i_mapping, pos, len,
+                                      copied, &folio->page, NULL);
+       return __iomap_write_end(iter->inode, pos, len, copied, folio);
 }
 
 static loff_t iomap_write_iter(struct iomap_iter *iter, struct iov_iter *i)
 
        do {
                struct folio *folio;
+               loff_t old_size;
                size_t offset;          /* Offset into folio */
                size_t bytes;           /* Bytes to write to folio */
                size_t copied;          /* Bytes copied from user */
                copied = copy_folio_from_iter_atomic(folio, offset, bytes, i);
                status = iomap_write_end(iter, pos, bytes, copied, folio);
 
+               /*
+                * Update the in-memory inode size after copying the data into
+                * the page cache.  It's up to the file system to write the
+                * updated size to disk, preferably after I/O completion so that
+                * no stale data is exposed.  Only once that's done can we
+                * unlock and release the folio.
+                */
+               old_size = iter->inode->i_size;
+               if (pos + status > old_size) {
+                       i_size_write(iter->inode, pos + status);
+                       iter->iomap.flags |= IOMAP_F_SIZE_CHANGED;
+               }
+               __iomap_put_folio(iter, pos, status, folio);
+
+               if (old_size < pos)
+                       pagecache_isize_extended(iter->inode, old_size, pos);
                if (status < bytes)
                        iomap_write_failed(iter->inode, pos + status,
                                           bytes - status);
                        bytes = folio_size(folio) - offset;
 
                bytes = iomap_write_end(iter, pos, bytes, bytes, folio);
+               __iomap_put_folio(iter, pos, bytes, folio);
                if (WARN_ON_ONCE(bytes == 0))
                        return -EIO;
 
                folio_mark_accessed(folio);
 
                bytes = iomap_write_end(iter, pos, bytes, bytes, folio);
+               __iomap_put_folio(iter, pos, bytes, folio);
                if (WARN_ON_ONCE(bytes == 0))
                        return -EIO;