From 0d6cef45218f175125db8451890125908e4c564f Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@rivosinc.com>
Date: Thu, 21 Nov 2024 15:41:22 +0000
Subject: [PATCH] pci: vfio: Release memory slots upon unmap

This prevents starvation of the limited set of memory slots in the
kernel.

Signed-off-by: Rob Bradford <rbradford@rivosinc.com>
---
 pci/src/vfio.rs      | 3 +++
 pci/src/vfio_user.rs | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/pci/src/vfio.rs b/pci/src/vfio.rs
index 85b3d75dc..cd468e8dc 100644
--- a/pci/src/vfio.rs
+++ b/pci/src/vfio.rs
@@ -1699,6 +1699,9 @@ impl VfioPciDevice {
                     error!("Could not remove the userspace memory region: {}", e);
                 }
 
+                self.memory_slot_allocator
+                    .free_memory_slot(user_memory_region.slot);
+
                 // SAFETY: FFI call with correct arguments
                 let ret = unsafe {
                     libc::munmap(
diff --git a/pci/src/vfio_user.rs b/pci/src/vfio_user.rs
index ce3e70f9e..1cd481421 100644
--- a/pci/src/vfio_user.rs
+++ b/pci/src/vfio_user.rs
@@ -222,6 +222,9 @@ impl VfioUserPciDevice {
                     error!("Could not remove the userspace memory region: {}", e);
                 }
 
+                self.memory_slot_allocator
+                    .free_memory_slot(user_memory_region.slot);
+
                 // Remove mmaps
                 // SAFETY: FFI call with correct arguments
                 let ret = unsafe {