DAX: __dax_fault race question

index fc2e314..89bb1f8 100644
@@ -349,7 +349,7 @@ static int copy_user_bh(struct page *to, struct inode *inode,
 #define NO_SECTOR -1
 #define DAX_PMD_INDEX(page_index) (page_index & (PMD_MASK >> PAGE_CACHE_SHIFT))
 
-static int dax_radix_entry(struct address_space *mapping, pgoff_t index,
+static int __dax_radix_entry_insert(struct address_space *mapping, pgoff_t index,
 		sector_t sector, bool pmd_entry, bool dirty)
 {
 	struct radix_tree_root *page_tree = &mapping->page_tree;

@@ -358,10 +358,6 @@ static int dax_radix_entry(struct address_space *mapping, pgoff_t index,
 	void *entry;
 
 	WARN_ON_ONCE(pmd_entry && !dirty);
-	if (dirty)
-		__mark_inode_dirty(mapping->host, I_DIRTY_PAGES);
-
-	spin_lock_irq(&mapping->tree_lock);
 
 	entry = radix_tree_lookup(page_tree, pmd_index);
 	if (entry && RADIX_DAX_TYPE(entry) == RADIX_DAX_PMD) {

@@ -374,8 +370,7 @@ static int dax_radix_entry(struct address_space *mapping, pgoff_t index,
 		type = RADIX_DAX_TYPE(entry);
 		if (WARN_ON_ONCE(type != RADIX_DAX_PTE &&
 					type != RADIX_DAX_PMD)) {
-			error = -EIO;
-			goto unlock;
+		        return -EIO;
 		}
 
 		if (!pmd_entry || type == RADIX_DAX_PMD)

@@ -402,19 +397,31 @@ static int dax_radix_entry(struct address_space *mapping, pgoff_t index,
 		 * pte_same() check will fail, eventually causing page fault
 		 * to be retried by the CPU.
 		 */
-		goto unlock;
+		return 0;
 	}
 
 	error = radix_tree_insert(page_tree, index,
 			RADIX_DAX_ENTRY(sector, pmd_entry));
 	if (error)
-		goto unlock;
+		return error;
 
 	mapping->nrexceptional++;
- dirty:
+dirty:
 	if (dirty)
 		radix_tree_tag_set(page_tree, index, PAGECACHE_TAG_DIRTY);
- unlock:
+	return error;
+}
+
+static int dax_radix_entry_insert(struct address_space *mapping, pgoff_t index,
+		sector_t sector, bool pmd_entry, bool dirty)
+{
+	int error;
+
+	if (dirty)
+		__mark_inode_dirty(mapping->host, I_DIRTY_PAGES);
+
+	spin_lock_irq(&mapping->tree_lock);
+	error =__dax_radix_entry_insert(mapping, index, sector, pmd_entry, dirty);
 	spin_unlock_irq(&mapping->tree_lock);
 	return error;
 }

@@ -579,8 +586,8 @@ static int dax_insert_mapping(struct inode *inode, struct buffer_head *bh,
 	}
 	dax_unmap_atomic(bdev, &dax);
 
-	error = dax_radix_entry(mapping, vmf->pgoff, dax.sector, false,
-			vmf->flags & FAULT_FLAG_WRITE);
+	error = dax_radix_entry_insert(mapping, vmf->pgoff, dax.sector, false,
+				vmf->flags & FAULT_FLAG_WRITE, vmf->page);
 	if (error)
 		goto out;
 

@@ -984,7 +991,7 @@ int __dax_pmd_fault(struct vm_area_struct *vma, unsigned long address,
 		 * the write to insert a dirty entry.
 		 */
 		if (write) {
-			error = dax_radix_entry(mapping, pgoff, dax.sector,
+			error = dax_radix_entry_insert(mapping, pgoff, dax.sector,
 					true, true);
 			if (error) {
 				dax_pmd_dbg(&bh, address,

@@ -1057,14 +1064,14 @@ int dax_pfn_mkwrite(struct vm_area_struct *vma, struct vm_fault *vmf)
 	struct file *file = vma->vm_file;
 
 	/*
-	 * We pass NO_SECTOR to dax_radix_entry() because we expect that a
+	 * We pass NO_SECTOR to dax_radix_entry_insert() because we expect that a
 	 * RADIX_DAX_PTE entry already exists in the radix tree from a
 	 * previous call to __dax_fault().  We just want to look up that PTE
 	 * entry using vmf->pgoff and make sure the dirty tag is set.  This
 	 * saves us from having to make a call to get_block() here to look
 	 * up the sector.
 	 */
-	dax_radix_entry(file->f_mapping, vmf->pgoff, NO_SECTOR, false, true);
+	dax_radix_entry_insert(file->f_mapping, vmf->pgoff, NO_SECTOR, false, true);
 	return VM_FAULT_NOPAGE;
 }
 EXPORT_SYMBOL_GPL(dax_pfn_mkwrite);

index 89bb1f8..0294fc9 100644
@@ -424,6 +424,31 @@ static int dax_radix_entry_insert(struct address_space *mapping, pgoff_t index,
 	error =__dax_radix_entry_insert(mapping, index, sector, pmd_entry, dirty);
 	spin_unlock_irq(&mapping->tree_lock);
 	return error;
+
+}
+
+static int dax_radix_entry_replace(struct address_space *mapping, pgoff_t index,
+				   sector_t sector, bool pmd_entry, bool dirty,
+				   struct page* old_page)
+{
+	int error;
+
+	BUG_ON(old_page && !PageLocked(old_page));
+	if (dirty)
+		__mark_inode_dirty(mapping->host, I_DIRTY_PAGES);
+
+	spin_lock_irq(&mapping->tree_lock);
+	if (old_page)
+		__delete_from_page_cache(old_page, NULL);
+	error =__dax_radix_entry_insert(mapping, index, sector, pmd_entry, dirty);
+	spin_unlock_irq(&mapping->tree_lock);
+	if (old_page) {
+		if (mapping->a_ops->freepage)
+			mapping->a_ops->freepage(old_page);
+		page_cache_release(old_page);
+	}
+	return error;
+
 }
 
 static int dax_writeback_one(struct block_device *bdev,

@@ -586,7 +611,7 @@ static int dax_insert_mapping(struct inode *inode, struct buffer_head *bh,
 	}
 	dax_unmap_atomic(bdev, &dax);
 
-	error = dax_radix_entry_insert(mapping, vmf->pgoff, dax.sector, false,
+	error = dax_radix_entry_replace(mapping, vmf->pgoff, dax.sector, false,
 				vmf->flags & FAULT_FLAG_WRITE, vmf->page);
 	if (error)
 		goto out;

@@ -711,14 +736,16 @@ int __dax_fault(struct vm_area_struct *vma, struct vm_fault *vmf,
 		page = find_lock_page(mapping, vmf->pgoff);
 
 	if (page) {
+		vmf->page = page;
 		unmap_mapping_range(mapping, vmf->pgoff << PAGE_SHIFT,
 							PAGE_CACHE_SIZE, 0);
-		delete_from_page_cache(page);
+	}
+	error = dax_insert_mapping(inode, &bh, vma, vmf);
+	if (page) {
 		unlock_page(page);
 		page_cache_release(page);
-		page = NULL;
+		vmf->page = page = NULL;
 	}
-
 	/*
 	 * If we successfully insert the new mapping over an unwritten extent,
 	 * we need to ensure we convert the unwritten extent. If there is an

@@ -729,14 +756,12 @@ int __dax_fault(struct vm_area_struct *vma, struct vm_fault *vmf,
 	 * indicate what the callback should do via the uptodate variable, same
 	 * as for normal BH based IO completions.
 	 */
-	error = dax_insert_mapping(inode, &bh, vma, vmf);
 	if (buffer_unwritten(&bh)) {
 		if (complete_unwritten)
 			complete_unwritten(&bh, !error);
 		else
 			WARN_ON_ONCE(!(vmf->flags & FAULT_FLAG_WRITE));
 	}
-
  out:
 	if (error == -ENOMEM)
 		return VM_FAULT_OOM | major;