[PATCH v3 kvm/queue 05/16] KVM: Maintain ofs_tree for fast memslot lookup by file offset

[Chao Peng <chao.p.peng@linux.intel.com>]

Similar to hva_tree for hva range, maintain interval tree ofs_tree for
offset range of a fd-based memslot so the lookup by offset range can be
faster when memslot count is high.

Signed-off-by: Chao Peng <chao.p.peng@linux.intel.com>
---
 include/linux/kvm_host.h |  2 ++
 virt/kvm/kvm_main.c      | 17 +++++++++++++----
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 2cd35560c44b..3bd875f9669f 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -451,6 +451,7 @@ static inline int kvm_vcpu_exiting_guest_mode(struct kvm_vcpu *vcpu)
 struct kvm_memory_slot {
 	struct hlist_node id_node[2];
 	struct interval_tree_node hva_node[2];
+	struct interval_tree_node ofs_node[2];
 	struct rb_node gfn_node[2];
 	gfn_t base_gfn;
 	unsigned long npages;
@@ -560,6 +561,7 @@ struct kvm_memslots {
 	u64 generation;
 	atomic_long_t last_used_slot;
 	struct rb_root_cached hva_tree;
+	struct rb_root_cached ofs_tree;
 	struct rb_root gfn_tree;
 	/*
 	 * The mapping table from slot id to memslot.
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index b0f7e6eb00ff..47e96d1eb233 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -1087,6 +1087,7 @@ static struct kvm *kvm_create_vm(unsigned long type)
 
 			atomic_long_set(&slots->last_used_slot, (unsigned long)NULL);
 			slots->hva_tree = RB_ROOT_CACHED;
+			slots->ofs_tree = RB_ROOT_CACHED;
 			slots->gfn_tree = RB_ROOT;
 			hash_init(slots->id_hash);
 			slots->node_idx = j;
@@ -1363,7 +1364,7 @@ static void kvm_replace_gfn_node(struct kvm_memslots *slots,
  * With NULL @old this simply adds @new.
  * With NULL @new this simply removes @old.
  *
- * If @new is non-NULL its hva_node[slots_idx] range has to be set
+ * If @new is non-NULL its hva/ofs_node[slots_idx] range has to be set
  * appropriately.
  */
 static void kvm_replace_memslot(struct kvm *kvm,
@@ -1377,6 +1378,7 @@ static void kvm_replace_memslot(struct kvm *kvm,
 	if (old) {
 		hash_del(&old->id_node[idx]);
 		interval_tree_remove(&old->hva_node[idx], &slots->hva_tree);
+		interval_tree_remove(&old->ofs_node[idx], &slots->ofs_tree);
 
 		if ((long)old == atomic_long_read(&slots->last_used_slot))
 			atomic_long_set(&slots->last_used_slot, (long)new);
@@ -1388,20 +1390,27 @@ static void kvm_replace_memslot(struct kvm *kvm,
 	}
 
 	/*
-	 * Initialize @new's hva range.  Do this even when replacing an @old
+	 * Initialize @new's hva/ofs range.  Do this even when replacing an @old
 	 * slot, kvm_copy_memslot() deliberately does not touch node data.
 	 */
 	new->hva_node[idx].start = new->userspace_addr;
 	new->hva_node[idx].last = new->userspace_addr +
 				  (new->npages << PAGE_SHIFT) - 1;
+	if (kvm_slot_is_private(new)) {
+		new->ofs_node[idx].start = new->ofs;
+		new->ofs_node[idx].last = new->ofs +
+					  (new->npages << PAGE_SHIFT) - 1;
+	}
 
 	/*
 	 * (Re)Add the new memslot.  There is no O(1) interval_tree_replace(),
-	 * hva_node needs to be swapped with remove+insert even though hva can't
-	 * change when replacing an existing slot.
+	 * hva_node/ofs_node needs to be swapped with remove+insert even though
+	 * hva/ofs can't change when replacing an existing slot.
 	 */
 	hash_add(slots->id_hash, &new->id_node[idx], new->id);
 	interval_tree_insert(&new->hva_node[idx], &slots->hva_tree);
+	if (kvm_slot_is_private(new))
+		interval_tree_insert(&new->ofs_node[idx], &slots->ofs_tree);
 
 	/*
 	 * If the memslot gfn is unchanged, rb_replace_node() can be used to
-- 
2.17.1