Add separate Recycle method for allocator.

This improves debugging for pagetable-related issues.

PiperOrigin-RevId: 209827795
Change-Id: I4cfa11664b0b52f26f6bc90a14c5bb106f01e038

3 files changed, 26 insertions, 3 deletions

diff --git a/pkg/sentry/platform/kvm/address_space.go b/pkg/sentry/platform/kvm/address_space.go
index fbd11ed71..463617170 100644
--- a/pkg/sentry/platform/kvm/address_space.go
+++ b/pkg/sentry/platform/kvm/address_space.go
@@ -273,6 +273,9 @@ func (as *addressSpace) Unmap(addr usermem.Addr, length uint64) {
 			Start: addr,
 			End:   addr + usermem.Addr(length),
 		})
+
+		// Recycle any freed intermediate pages.
+		as.pageTables.Allocator.Recycle()
 	}
 }
 
diff --git a/pkg/sentry/platform/kvm/allocator.go b/pkg/sentry/platform/kvm/allocator.go
index 80066bfc5..f5cebd5b3 100644
--- a/pkg/sentry/platform/kvm/allocator.go
+++ b/pkg/sentry/platform/kvm/allocator.go
@@ -67,3 +67,10 @@ func (a allocator) LookupPTEs(physical uintptr) *pagetables.PTEs {
 func (a allocator) FreePTEs(ptes *pagetables.PTEs) {
 	a.base.FreePTEs(ptes)
 }
+
+// Recycle implements pagetables.Allocator.Recycle.
+//
+//go:nosplit
+func (a allocator) Recycle() {
+	a.base.Recycle()
+}
diff --git a/pkg/sentry/platform/ring0/pagetables/allocator.go b/pkg/sentry/platform/ring0/pagetables/allocator.go
index 1499623fb..049fd0247 100644
--- a/pkg/sentry/platform/ring0/pagetables/allocator.go
+++ b/pkg/sentry/platform/ring0/pagetables/allocator.go
@@ -27,8 +27,12 @@ type Allocator interface {
 	// LookupPTEs looks up PTEs by physical address.
 	LookupPTEs(physical uintptr) *PTEs
 
-	// FreePTEs frees a set of PTEs.
+	// FreePTEs marks a set of PTEs a freed, although they may not be available
+	// for use again until Recycle is called, below.
 	FreePTEs(ptes *PTEs)
+
+	// Recycle makes freed PTEs available for use again.
+	Recycle()
 }
 
 // RuntimeAllocator is a trivial allocator.
@@ -42,6 +46,9 @@ type RuntimeAllocator struct {
 
 	// pool is the set of free-to-use PTEs.
 	pool []*PTEs
+
+	// freed is the set of recently-freed PTEs.
+	freed []*PTEs
 }
 
 // NewRuntimeAllocator returns an allocator that uses runtime allocation.
@@ -51,8 +58,15 @@ func NewRuntimeAllocator() *RuntimeAllocator {
 	}
 }
 
+// Recycle returns freed pages to the pool.
+func (r *RuntimeAllocator) Recycle() {
+	r.pool = append(r.pool, r.freed...)
+	r.freed = r.freed[:0]
+}
+
 // Drain empties the pool.
 func (r *RuntimeAllocator) Drain() {
+	r.Recycle()
 	for i, ptes := range r.pool {
 		// Zap the entry in the underlying array to ensure that it can
 		// be properly garbage collected.
@@ -104,6 +118,5 @@ func (r *RuntimeAllocator) LookupPTEs(physical uintptr) *PTEs {
 //
 //go:nosplit
 func (r *RuntimeAllocator) FreePTEs(ptes *PTEs) {
-	// Add to the pool.
-	r.pool = append(r.pool, ptes)
+	r.freed = append(r.freed, ptes)
 }