[PATCH v10 11/11] x86/virt/tdx: Optimize tdx_pamt_get/put()

[Rick Edgecombe]

The Dynamic PAMT (DPAMT) get/put helpers use a global spinlock to
serialize all refcount updates and SEAMCALL invocations. This gives correct
behavior for concurrent callers, but leads to contention. It is especially
bad from the KVM side, which is designed to allow faulting in EPT under a
shared lock. With the global spinlock, not only is the lock an exclusive
one, but it is for all TDs instead of just a single one.

But taking the global lock each time is actually unnecessary. Only the 0->1
and 1->0 refcount transitions actually need the lock (to pair with
SEAMCALLs that actually add and remove with the DPAMT pages). The common
case of incrementing or decrementing a non-zero refcount can be done
locklessly.

So create a fast and slow path. Check the refcount outside the lock and
only take it for the slow path (0->1 and 1->0 transitions).

On the put side make the refcount adjustment and lock taking atomic so if
a 'get' happens between them, it doesn't cause the DPAMT to be freed
incorrectly. On the get side there is no technique for doing the refcount
adjustment and lock atomically, so check the refcount again inside the
lock.

AI was used under supervision to collect/apply feedback, review code and
workshop logs. It assisted in identifying/evaluating the stale
conditionals for the races resolved from the atomic_dec_and_lock() change.
Separate from atomic_dec_and_lock() fallout, it suggested to change
atomic_inc() to atomic_set(pamt_refcount, 1) in the put error path for the
sake of being more precise, which Kiryl had also suggested in the past.
The model also suggested updated comments following the
atomic_dec_and_lock() change based on some directed prompting. The
comments were subsequently edited or further prompted for fine tuning.

Based on a patch originally by Kiryl Shutsemau.

Signed-off-by: Rick Edgecombe <rick.p.edgecombe@xxxxxxxxx>
Tested-by: Hongyu Ning <hongyu.ning@xxxxxxxxxxxxxxx>
Reviewed-by: Chao Gao <chao.gao@xxxxxxxxx>
Reviewed-by: Tony Lindgren <tony.lindgren@xxxxxxxxxxxxxxx>
Reviewed-by: Nikolay Borisov <nik.borisov@xxxxxxxx>
Reviewed-by: Dave Hansen <dave.hansen@xxxxxxxxxxxxxxx>
Reviewed-by: Vishal Annapurve <vannapurve@xxxxxxxxxx>
Acked-by: Sohil Mehta <sohil.mehta@xxxxxxxxx>
---
v10:
 - Change "Dynamic PAMT" to "DPAMT" in comments. (Dave)

v7:
 - Drop Assisted-by tag and cover AI use in log. (Dave)
 - Move to end of the series.
 - Use atomic_inc_not_zero() in this patch inside the spin_lock(), as
   suggested on the non-optimized patch by (Dave).

v6:
 - Fix "tdx_pamt_add()" typo to "tdx_pamt_get()" in lost-race comment
 - Fix error path bug: set ret = -EIO and use WARN_ON_ONCE() instead of
   pr_err() for unexpected PAMT.ADD failures (Sean)
 - Use "set the refcount 0->1" wording to match atomic_set() usage
 - Wrap comments to 80 columns
 - Switch to atomic_dec_and_lock() and remove handling of races that are
   no longer needed as a result. Adjust comments as appropriate. (Dave)
 - Adjustments from dropping error helper patches
---
 arch/x86/virt/vmx/tdx/tdx.c | 44 ++++++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/arch/x86/virt/vmx/tdx/tdx.c b/arch/x86/virt/vmx/tdx/tdx.c
index 3daa8c63f9c51..2acab6e5f5f87 100644
--- a/arch/x86/virt/vmx/tdx/tdx.c
+++ b/arch/x86/virt/vmx/tdx/tdx.c
@@ -2176,28 +2176,41 @@ int tdx_pamt_get(kvm_pfn_t pfn, struct tdx_pamt_cache *cache)
 	if (!tdx_supports_dynamic_pamt(&tdx_sysinfo))
 		return 0;
 
-	ret = alloc_pamt_array(pamt_pages, cache);
-	if (ret)
-		return ret;
-
 	dpamt_refcount = tdx_find_dpamt_refcount(pfn);
 
-	spin_lock(&dpamt_lock);
-
 	/*
 	 * If the DPAMT entry is already added (i.e. refcount >= 1),
 	 * then just increment the refcount.
 	 */
+	if (atomic_inc_not_zero(dpamt_refcount))
+		return 0;
+
+	ret = alloc_pamt_array(pamt_pages, cache);
+	if (ret)
+		return ret;
+
+	spin_lock(&dpamt_lock);
+
+	/*
+	 * Unlike tdx_pamt_put() which uses atomic_dec_and_lock() to
+	 * atomically handle the 1->0 transition, the get side has no
+	 * equivalent combined primitive for 0->1. Recheck under the
+	 * lock since another get may have already done the 0->1
+	 * transition after both saw atomic_inc_not_zero() fail.
+	 */
 	if (atomic_inc_not_zero(dpamt_refcount))
 		goto out_free;
 
-	/* Try to add the PAMT page and take the refcount 0->1. */
 	tdx_status = tdh_phymem_pamt_add(pfn, pamt_pages);
 	if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS)) {
 		ret = -EIO;
 		goto out_free;
 	}
 
+	/*
+	 * The refcount is zero, and this locked path is the
+	 * only way to increase it from 0->1.
+	 */
 	atomic_set(dpamt_refcount, 1);
 	spin_unlock(&dpamt_lock);
 	return 0;
@@ -2222,17 +2235,13 @@ void tdx_pamt_put(kvm_pfn_t pfn)
 
 	dpamt_refcount = tdx_find_dpamt_refcount(pfn);
 
-	spin_lock(&dpamt_lock);
 	/*
 	 * If there is more than 1 reference on the DPAMT entry, don't
 	 * remove it yet. Just decrement the refcount.
 	 */
-	if (atomic_read(dpamt_refcount) > 1) {
-		atomic_dec(dpamt_refcount);
-		goto out_unlock;
-	}
+	if (!atomic_dec_and_lock(dpamt_refcount, &dpamt_lock))
+		return;
 
-	/* Try to remove the pamt page and take the refcount 1->0. */
 	tdx_status = tdh_phymem_pamt_remove(pfn, pamt_pages);
 
 	/*
@@ -2242,10 +2251,15 @@ void tdx_pamt_put(kvm_pfn_t pfn)
 	 * failure indicates a kernel bug, memory is being leaked, and
 	 * the dangling DPAMT entry may cause future operations to fail.
 	 */
-	if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS))
+	if (WARN_ON_ONCE(tdx_status != TDX_SUCCESS)) {
+		/*
+		 * atomic_dec_and_lock() already decremented it to 0,
+		 * but the DPAMT entry still exists since REMOVE failed.
+		 */
+		atomic_set(dpamt_refcount, 1);
 		goto out_unlock;
+	}
 
-	atomic_set(dpamt_refcount, 0);
 	spin_unlock(&dpamt_lock);
 	free_pamt_array(pamt_pages);
 	return;
-- 
2.55.0