perf(libc): share one snapshot of the registered guest ranges between leases until the registry changes

core/libs/prx/libSceAgcDriver/tests/GuestAllocations.cpp

@@ -7,6 +7,7 @@
 #include <cstring>
 #include <atomic>
 #include <exception>
+#include <memory>
 #include <thread>
 #include <array>
 #include <algorithm>
@@ -28,6 +29,35 @@ void reject(TAction action) {
     throw std::runtime_error("expected guest allocation ownership rejection");
 }
 
+void leaseSnapshotTests() {
+    std::array<std::byte, 128> memory{};
+    const auto address = reinterpret_cast<std::uintptr_t>(memory.data());
+    {
+        GuestAllocations::Mutation mutation;
+        mutation.Add(memory.data(), 64, true, true, true);
+    }
+    std::weak_ptr<const GuestAllocations::Range> held;
+    {
+        const auto first = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+        const auto second = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+        Require(!first.empty() && first.size() == second.size() && &first[0] == &second[0], "two leases taken without a registry change do not share one snapshot");
+        const auto found = std::find_if(first.begin(), first.end(), [&](const auto& range) { return range->address == address; });
+        Require(found != first.end(), "the shared snapshot misses a registered range");
+        held = *found;
+        Require(held.use_count() == 2, "the shared snapshot of two leases holds a range more than once");
+        {
+            GuestAllocations::Mutation mutation;
+            mutation.Add(memory.data() + 64, 64, true, true, true);
+        }
+        const auto third = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
+        Require(&third[0] != &first[0] && std::any_of(third.begin(), third.end(), [&](const auto& range) { return range->address == address + 64; }), "a lease taken after a registry change still shares the old snapshot");
+    }
+    Require(held.use_count() == 1, "a snapshot outlived its leases");
+    GuestAllocations::Mutation mutation;
+    mutation.Remove(memory.data());
+    mutation.Remove(memory.data() + 64);
+}
+
 }
 
 void RunGuestLeaseWaitTests() {
@@ -38,8 +68,10 @@ void RunGuestLeaseWaitTests() {
         mutation.Add(memory.data(), 64, true, true, true);
         mutation.Add(memory.data() + 64, 64, true, true, true);
     }
-    auto lease = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
-    std::erase_if(lease, [&](const auto& range) { return range->address != address; });
+    GuestAllocations::Lease::Ranges lease;
+    for (const auto& range : GuestAllocations::GuestAllocationsAcquire_nid_postfix()) {
+        if (range->address == address) lease.push_back(range);
+    }
     std::atomic<bool> mutationEntered = false;
     std::atomic<bool> lookupFinished = false;
     std::exception_ptr failure;
@@ -56,8 +88,10 @@ void RunGuestLeaseWaitTests() {
         } catch (...) { failure = std::current_exception(); }
     });
     mutationEntered.wait(false);
-    auto other = GuestAllocations::GuestAllocationsAcquire_nid_postfix();
-    std::erase_if(other, [&](const auto& range) { return range->address != address + 64; });
+    GuestAllocations::Lease::Ranges other;
+    for (const auto& range : GuestAllocations::GuestAllocationsAcquire_nid_postfix()) {
+        if (range->address == address + 64) other.push_back(range);
+    }
     const bool found = other.size() == 1 && other.front()->address == address + 64;
     lookupFinished.store(true);
     lease.clear();
@@ -74,6 +108,7 @@ void RunGuestLeaseWaitTests() {
 }
 
 void RunGuestAllocationTests() {
+    leaseSnapshotTests();
     void* pointer = GuestHeap::GuestHeapAllocate_nid_postfix(32);
     Require(reinterpret_cast<std::uintptr_t>(pointer) % alignof(std::max_align_t) == 0, "guest malloc is not suitably aligned");
     std::memset(pointer, 0x55, 32);

core/libs/prx/libc/include/GuestAllocations.hpp

@@ -21,7 +21,33 @@ struct Range {
     bool gpu = false;
 };
 
-using Lease = std::vector<std::shared_ptr<const Range>>;
+class Lease {
+public:
+    using Ranges = std::vector<std::shared_ptr<const Range>>;
+    using value_type = Ranges::value_type;
+    using const_iterator = Ranges::const_iterator;
+    using iterator = const_iterator;
+
+    Lease() = default;
+    explicit Lease(std::shared_ptr<const Ranges> ranges) : ranges(std::move(ranges)) {}
+
+    const_iterator begin() const { return held().begin(); }
+    const_iterator end() const { return held().end(); }
+    std::size_t size() const { return held().size(); }
+    bool empty() const { return held().empty(); }
+    const value_type& operator[](std::size_t index) const { return held()[index]; }
+    const value_type& front() const { return held().front(); }
+    const value_type& back() const { return held().back(); }
+    void clear() { ranges.reset(); }
+
+private:
+    const Ranges& held() const {
+        static const Ranges none;
+        return ranges != nullptr ? *ranges : none;
+    }
+
+    std::shared_ptr<const Ranges> ranges;
+};
 using Mapped = std::vector<std::weak_ptr<const Range>>;
 
 extern "C" {

core/libs/prx/libc/src/GuestAllocations.cpp

@@ -36,6 +36,8 @@ struct Registry {
     std::map<std::uint64_t, std::shared_ptr<const Range>> ranges;
     bool mainImageRegistered = false;
     std::map<std::uintptr_t, std::uintptr_t> registeredImages;
+    std::weak_ptr<const Lease::Ranges> readableLease;
+    std::weak_ptr<const Lease::Ranges> allLease;
 };
 
 Registry& registry() {
@@ -109,7 +111,10 @@ bool backWritableImageAnonymously() {
 }
 
 void* GuestAllocationsBegin_nid_postfix() {
-    return new MutationState{std::unique_lock<std::mutex>(registry().mutex), {}};
+    auto* state = new MutationState{std::unique_lock<std::mutex>(registry().mutex), {}};
+    registry().readableLease.reset();
+    registry().allLease.reset();
+    return state;
 }
 
 void GuestAllocationsEnd_nid_postfix(void* mutation) noexcept {
@@ -120,7 +125,7 @@ void GuestAllocationsEnd_nid_postfix(void* mutation) noexcept {
     }
     if (!state->gpuMapped.empty()) gpuMapObserverCalls.fetch_add(1);
     const auto observer = state->gpuMapped.empty() ? nullptr : gpuMapObserver.load();
-    Lease mapped;
+    Lease::Ranges mapped;
     if (observer != nullptr) {
         const auto& ranges = registry().ranges;
         for (const auto& [address, bytes] : state->gpuMapped) {
@@ -137,6 +142,8 @@ void GuestAllocationsEnd_nid_postfix(void* mutation) noexcept {
     const Mapped weak(mapped.begin(), mapped.end());
     mapped.clear();
     const bool noted = !state->gpuMapped.empty();
+    registry().readableLease.reset();
+    registry().allLease.reset();
     delete state;
     if (!weak.empty()) observer(weak, ended);
     if (noted) gpuMapObserverCalls.fetch_sub(1);
@@ -312,6 +319,8 @@ void GuestAllocationsRequireUnpinned_nid_postfix(void* mutation, const void* poi
     const auto deadline = std::chrono::steady_clock::now() + pinWait();
     int syncedRounds = 0;
     for (;;) {
+        registry().readableLease.reset();
+        registry().allLease.reset();
         bool pinned = false;
         auto it = ranges.upper_bound(address);
         if (it != ranges.begin()) --it;
@@ -487,21 +496,25 @@ void GuestAllocationsUnmap_nid_postfix(void* mutation, const void* pointer, std:
 
 Lease GuestAllocationsAcquire_nid_postfix() {
     std::lock_guard lock(registry().mutex);
-    Lease result;
+    if (auto shared = registry().readableLease.lock()) return Lease(std::move(shared));
+    auto ranges = std::make_shared<Lease::Ranges>();
     for (const auto& [address, range] : registry().ranges) {
-        if (range->readable && range->bytes != 0) result.push_back(range);
+        if (range->readable && range->bytes != 0) ranges->push_back(range);
     }
-    return result;
+    registry().readableLease = ranges;
+    return Lease(std::move(ranges));
 }
 
 Lease GuestAllocationsAcquireAll_nid_postfix() {
     std::lock_guard lock(registry().mutex);
-    Lease result;
-    result.reserve(registry().ranges.size());
+    if (auto shared = registry().allLease.lock()) return Lease(std::move(shared));
+    auto ranges = std::make_shared<Lease::Ranges>();
+    ranges->reserve(registry().ranges.size());
     for (const auto& [address, range] : registry().ranges) {
-        if (range->bytes != 0) result.push_back(range);
+        if (range->bytes != 0) ranges->push_back(range);
     }
-    return result;
+    registry().allLease = ranges;
+    return Lease(std::move(ranges));
 }
 
 }