fix(agc): keep Windows host imports watched while late driver marks settle

core/libs/prx/libSceAgcDriver/Execution/include/GuestMemory.hpp

@@ -59,7 +59,7 @@ void WriteChanged(std::uint64_t address, std::span<const std::byte> current, std
 // UnchangedSince is always false, so callers fall back to comparing bytes.
 bool WriteWatched();
 void Unwatch(std::uint64_t address, std::size_t bytes);
-bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function<bool()>& import);
+bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function<bool()>& import, bool settle = false);
 bool Watched(std::uint64_t address, std::size_t bytes);
 std::uint64_t CollectWrites(std::uint64_t address, std::size_t bytes);
 bool UnchangedSince(std::uint64_t address, std::size_t bytes, std::uint64_t generation);

core/libs/prx/libSceAgcDriver/Execution/src/GuestMemory.cpp

@@ -768,6 +768,34 @@ struct WriteTracker {
     std::vector<std::uint32_t> cpuBlocks;
     std::vector<std::uint32_t> writtenBlocks;
     WriteWatchCoverage coverage;
+    struct SettlingImport {
+        std::uint64_t first;
+        std::uint64_t stop;
+        std::chrono::steady_clock::time_point until;
+        std::vector<std::byte> bytes;
+        bool closing = false;
+    };
+    std::vector<SettlingImport> settling;
+    std::uint64_t settlingBytes = 0;
+    bool finishingSettled = false;
+
+    void settlingStored(std::uint64_t first, std::uint64_t stop) {
+        for (auto& entry : settling) {
+            const auto begin = std::max(first, entry.first) & ~std::uint64_t{4095};
+            const auto end = std::min(stop, entry.stop);
+            for (auto page = begin; page < end; page += 4096) {
+                if (Accessible(reinterpret_cast<const void*>(page), 4096)) std::memcpy(entry.bytes.data() + (page - entry.first), reinterpret_cast<const void*>(page), 4096);
+            }
+        }
+    }
+
+    void dropSettling(std::uint64_t first, std::uint64_t stop) {
+        std::erase_if(settling, [&](const SettlingImport& entry) {
+            if (entry.closing || entry.stop <= first || stop <= entry.first) return false;
+            settlingBytes -= entry.bytes.size();
+            return true;
+        });
+    }
 #else
     static constexpr std::size_t LeafBlocks = std::size_t{1} << 16;
     static constexpr std::size_t LeafCount = std::size_t{1} << 15;
@@ -961,6 +989,7 @@ void watchPrivateMapping(std::uintptr_t address, std::size_t bytes, std::uint64_
     const auto lock = lockTracker(tracker);
     if (!tracker.watched || bytes == 0 || bytes > tracker.size || address < tracker.base || address - tracker.base > tracker.size - bytes) return;
     tracker.coverage.Restore(address, bytes, generation);
+    tracker.dropSettling(address, address + bytes);
     ++tracker.generation;
     for (auto block = tracker.blockOf(address); block <= tracker.blockOf(address + bytes - 1); ++block) {
         tracker.driverPieces.erase(block);
@@ -1019,7 +1048,64 @@ void stampWrittenRun(void* context, std::uintptr_t begin, std::uintptr_t end) {
 }
 #endif
 
+#ifdef _WIN32
+std::chrono::milliseconds importSettle() {
+    static const std::chrono::milliseconds settle{[] {
+        const char* text = std::getenv("APS5_IMPORT_SETTLE_MS");
+        return text != nullptr ? std::strtoull(text, nullptr, 10) : 1000ull;
+    }()};
+    return settle;
+}
+
+constexpr std::uint64_t MaxSettleSnapshotBytes = std::uint64_t{256} << 20;
+constexpr std::uint64_t MaxSettlingBytes = std::uint64_t{512} << 20;
+
+StampKind settlingKind(WriteTracker& tracker, std::uintptr_t page, StampKind kind) {
+    if (kind != StampKind::Cpu) return kind;
+    for (auto& entry : tracker.settling) {
+        if (page < entry.first || page >= entry.stop) continue;
+        auto* saved = entry.bytes.data() + (page - entry.first);
+        const auto* current = reinterpret_cast<const std::byte*>(page);
+        if (!Accessible(current, WritePageBytes)) return kind;
+        if (std::memcmp(saved, current, WritePageBytes) == 0) return StampKind::ImportWindow;
+        std::memcpy(saved, current, WritePageBytes);
+        return kind;
+    }
+    return kind;
+}
+#endif
+
+bool walkWrites(WriteTracker& tracker, std::uint64_t first, std::uint64_t stop, StampKind kind);
+
+#ifdef _WIN32
+void finishSettledImports(WriteTracker& tracker) {
+    const auto now = std::chrono::steady_clock::now();
+    bool expired = false;
+    for (auto& entry : tracker.settling) {
+        if (entry.until > now) continue;
+        entry.closing = true;
+        expired = true;
+    }
+    if (!expired) return;
+    tracker.finishingSettled = true;
+    for (std::size_t i = 0; i < tracker.settling.size(); ++i) {
+        const auto first = tracker.settling[i].first;
+        const auto stop = tracker.settling[i].stop;
+        if (tracker.settling[i].closing) walkWrites(tracker, first, stop, StampKind::Cpu);
+    }
+    tracker.finishingSettled = false;
+    std::erase_if(tracker.settling, [&](const WriteTracker::SettlingImport& entry) {
+        if (!entry.closing) return false;
+        tracker.settlingBytes -= entry.bytes.size();
+        return true;
+    });
+}
+#endif
+
 bool walkWrites(WriteTracker& tracker, std::uint64_t first, std::uint64_t stop, StampKind kind) {
+#ifdef _WIN32
+    if (kind == StampKind::Cpu && !tracker.settling.empty() && !tracker.finishingSettled) finishSettledImports(tracker);
+#endif
     ++tracker.generation;
 #ifdef _WIN32
     constexpr std::uint64_t page = 4096;
@@ -1040,8 +1126,9 @@ bool walkWrites(WriteTracker& tracker, std::uint64_t first, std::uint64_t stop,
         if (count != 0) dirty = true;
         for (ULONG_PTR i = 0; i < count; ++i) {
             const auto page = reinterpret_cast<std::uintptr_t>(tracker.pages[i]);
-            tracker.stamp(tracker.blockOf(page), tracker.generation, kind);
-            if (kind != StampKind::Driver) tracker.noteCpuStore(page, page + WritePageBytes, tracker.generation);
+            const auto pageKind = settlingKind(tracker, page, kind);
+            tracker.stamp(tracker.blockOf(page), tracker.generation, pageKind);
+            if (pageKind != StampKind::Driver) tracker.noteCpuStore(page, page + WritePageBytes, tracker.generation);
         }
         if (count < tracker.pages.size()) break;
         cursor = reinterpret_cast<std::uintptr_t>(tracker.pages[count - 1]) + (granularity != 0 ? granularity : page);
@@ -1056,8 +1143,9 @@ bool walkWrites(WriteTracker& tracker, std::uint64_t first, std::uint64_t stop,
             if (!GuestArena::GuestArenaCollectWrites_nid_postfix(cursor, static_cast<std::size_t>(stop - cursor), tracker.pages.data(), &count, true)) return false;
             for (ULONG_PTR i = 0; i < count; ++i) {
                 const auto page = reinterpret_cast<std::uintptr_t>(tracker.pages[i]);
-                tracker.stamp(tracker.blockOf(page), tracker.generation, kind);
-                if (kind != StampKind::Driver) tracker.noteCpuStore(page, page + WritePageBytes, tracker.generation);
+                const auto pageKind = settlingKind(tracker, page, kind);
+                tracker.stamp(tracker.blockOf(page), tracker.generation, pageKind);
+                if (pageKind != StampKind::Driver) tracker.noteCpuStore(page, page + WritePageBytes, tracker.generation);
             }
             if (count < tracker.pages.size()) break;
             cursor = reinterpret_cast<std::uintptr_t>(tracker.pages[count - 1]) + (granularity != 0 ? granularity : page);
@@ -1144,6 +1232,7 @@ namespace {
 void unwatchLocked(WriteTracker& tracker, std::uint64_t address, std::size_t bytes) {
     if (!tracker.watched) return;
 #ifdef _WIN32
+    tracker.dropSettling(address, address + bytes);
     if (tracker.coverage.Exclude(address, bytes, GuestArena::GuestArenaCommitGeneration_nid_postfix())) unwatchSerial.fetch_add(1, std::memory_order_release);
 #else
     if (GuestWriteWatch::GuestWriteWatchUnregister_nid_postfix(reinterpret_cast<const void*>(address), bytes)) unwatchSerial.fetch_add(1, std::memory_order_release);
@@ -1160,7 +1249,7 @@ void Unwatch(std::uint64_t address, std::size_t bytes) {
     unwatchLocked(tracker, address, bytes);
 }
 
-bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function<bool()>& import) {
+bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function<bool()>& import, bool settle) {
     if (bytes == 0 || bytes > std::numeric_limits<std::uint64_t>::max() - address) return import();
     auto& tracker = Tracker();
     auto lock = lockTracker(tracker);
@@ -1174,8 +1263,23 @@ bool ImportWatched(std::uint64_t address, std::size_t bytes, const std::function
     const auto stop = (address + bytes + page - 1) & ~(page - 1);
     const bool covered = tracker.covers(first, static_cast<std::size_t>(stop - first));
     const bool before = covered && walkWrites(tracker, first, stop, StampKind::Cpu);
+#ifdef _WIN32
+    const bool settles = settle && before && importSettle().count() != 0 && stop - first <= MaxSettleSnapshotBytes && tracker.settlingBytes + (stop - first) <= MaxSettlingBytes;
+    std::vector<std::byte> snapshot;
+    if (settles) snapshot.assign(reinterpret_cast<const std::byte*>(first), reinterpret_cast<const std::byte*>(stop));
+#else
+    static_cast<void>(settle);
+#endif
     if (!import()) return false;
     if (!before || !walkWrites(tracker, first, stop, StampKind::ImportWindow)) unwatchLocked(tracker, first, static_cast<std::size_t>(stop - first));
+#ifdef _WIN32
+    else if (settles) {
+        tracker.settlingBytes += stop - first;
+        tracker.settling.push_back({first, stop, std::chrono::steady_clock::now() + importSettle(), std::move(snapshot)});
+    } else if (settle) {
+        unwatchLocked(tracker, first, static_cast<std::size_t>(stop - first));
+    }
+#endif
     return true;
 }
 
@@ -1226,6 +1330,9 @@ std::uint64_t storeOwn(std::uint64_t address, std::size_t bytes, const std::func
     if (bytes == 0) return 0;
     const auto stampStored = [](WriteTracker& tracker, std::pair<std::uint64_t, std::uint64_t> stored) -> std::uint64_t {
         if (stored.second <= stored.first || !tracker.watched || !tracker.covers(stored.first, static_cast<std::size_t>(stored.second - stored.first))) return 0;
+#ifdef _WIN32
+        tracker.settlingStored(stored.first, stored.second);
+#endif
         ++tracker.generation;
         for (auto block = tracker.blockOf(stored.first); block <= tracker.blockOf(stored.second - 1); ++block) {
             tracker.stamp(block, tracker.generation, StampKind::Driver);

core/libs/prx/libSceAgcDriver/Graphics/src/GuestBufferMemory.cpp

@@ -136,6 +136,7 @@ struct HostImports {
     VkDevice watchDevice = VK_NULL_HANDLE;
     bool unwatchImports = false;
     bool unwatchDmaBufImports = false;
+    bool settleImports = false;
 };
 
 HostImports& Imports() {
@@ -268,7 +269,6 @@ const char* createImport(const Context& context, HostImport& entry, VkResult& fa
     return step;
 }
 
-#ifndef _WIN32
 enum class ImportWatchRequest : std::uint8_t { Probe, Watch, Unwatch };
 
 ImportWatchRequest importWatchRequest() {
@@ -281,15 +281,19 @@ ImportWatchRequest importWatchRequest() {
     }();
     return request;
 }
-#endif
 
 void decideImportWatch(const Context& context, HostImports& state) {
     if (state.watchDevice == context.device) return;
 #ifdef _WIN32
     state.watchDevice = context.device;
-    state.unwatchImports = true;
+    const auto request = importWatchRequest();
+    state.unwatchImports = request == ImportWatchRequest::Unwatch;
     state.unwatchDmaBufImports = false;
-    if (context.hostImportAlignment != 0 && GuestMemory::WriteWatched()) std::fprintf(stderr, "[write-watch] host imports are compared on Windows because driver writes can arrive after the import window\n");
+    state.settleImports = request == ImportWatchRequest::Probe;
+    if (context.hostImportAlignment == 0 || !GuestMemory::WriteWatched()) return;
+    if (state.unwatchImports) std::fprintf(stderr, "[write-watch] host imports are compared, not watched (APS5_WRITE_WATCH_IMPORTS=unwatch)\n");
+    else if (state.settleImports) std::fprintf(stderr, "[write-watch] host imports stay watched; pages marked written while an import settles are compared, not trusted\n");
+    else std::fprintf(stderr, "[write-watch] host imports stay watched without settling (APS5_WRITE_WATCH_IMPORTS=watch)\n");
 #else
     const auto request = importWatchRequest();
     state.watchDevice = context.device;
@@ -407,11 +411,12 @@ const HostImport* importAllocation(const Context& context, HostImports& state, s
     decideImportWatch(context, state);
     VkResult result = VK_SUCCESS;
     const char* step = nullptr;
+    const bool settle = state.settleImports && !state.unwatchImports;
     GuestMemory::ImportWatched(base, bytes, [&] {
         step = createImport(context, entry, result);
         return step == nullptr;
-    });
-    entry.unwatched = step == nullptr && (entry.dmaBuf ? state.unwatchDmaBufImports : state.unwatchImports);
+    }, settle);
+    entry.unwatched = step == nullptr && (entry.dmaBuf ? state.unwatchDmaBufImports : state.unwatchImports || (settle && !GuestMemory::Watched(base, bytes)));
     if (entry.unwatched) GuestMemory::Unwatch(base, bytes);
     if (step != nullptr) {
 #ifdef _WIN32

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

@@ -1692,8 +1692,8 @@ void importWatchTests(const Device& device) {
     using namespace AgcDriver::GuestMemory;
     const auto& context = device.GetContext();
 #ifdef _WIN32
-    Require(PrepareImportWatch(context) == ImportWatch::Unwatch, "Windows host imports stayed watched by default");
-    std::cout << "import watch decisions: Windows host imports use comparisons\n";
+    Require(PrepareImportWatch(context) == ImportWatch::Watch, "Windows host imports were compared by default");
+    std::cout << "import watch decisions: Windows host imports stay watched and settle\n";
 #else
     if (context.hostImportAlignment == 0 || !WriteWatched()) {
         std::cout << "host imports or write watching unavailable: import watch decisions not tested\n";

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

@@ -11,6 +11,7 @@
 #endif
 #include <array>
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
 #include <iostream>
 #include <stdexcept>
@@ -177,7 +178,70 @@ void CheckPrivateMappingReuse() {
 #endif
 }
 
+#ifdef _WIN32
+void CheckImportSettle() {
+    void* memory = AllocateWatched(3 * Block);
+    const auto base = reinterpret_cast<std::uint64_t>(memory);
+    auto* bytes = static_cast<volatile std::uint8_t*>(memory);
+    std::memset(memory, 0x11, 3 * Block);
+    CollectWritesUncached(base, 3 * Block);
+    Require(ImportWatched(base, Block, [] { return true; }, true) && ImportWatched(base + Block, Block, [] { return true; }, true), "the import callbacks failed");
+    std::array<std::uint64_t, 1> generations{TrackerGeneration()};
+    std::array<std::uint8_t, 1> changed{};
+    const auto state = [&](std::uint64_t address) {
+        Require(ChangedBlocks(address, Block, generations, changed), "a settling import is not tracked");
+        return changed[0];
+    };
+
+    bytes[8] = bytes[8];
+    bytes[Block + 8] = 0x22;
+    bytes[2 * Block + 8] = 0x22;
+    CollectWritesUncached(base, 3 * Block);
+    Require(state(base) == BlockMaybeWritten, "a page rewritten with its own bytes while the import settles was trusted as a CPU store");
+    Require(state(base + Block) == BlockWritten, "a CPU store that changed a settling import was not trusted");
+    Require(state(base + 2 * Block) == BlockWritten, "a CPU store next to a settling import was not trusted");
+
+    generations[0] = TrackerGeneration();
+    bytes[16] = bytes[16];
+    bytes[Block + 16] = 0x33;
+    Sleep(150);
+    CollectWritesUncached(base + 2 * Block, Block);
+    Require(state(base) == BlockMaybeWritten, "the closing walk trusted a page rewritten with its own bytes");
+    Require(state(base + Block) == BlockWritten, "the closing walk missed a CPU store that changed the import");
+
+    generations[0] = TrackerGeneration();
+    bytes[24] = bytes[24];
+    CollectWritesUncached(base, Block);
+    Require(state(base) == BlockWritten, "a page marked after the import settled was not trusted");
+
+    Require(ImportWatched(base + 2 * Block, Block, [] { return true; }), "the plain import callback failed");
+    generations[0] = TrackerGeneration();
+    bytes[2 * Block + 32] = bytes[2 * Block + 32];
+    CollectWritesUncached(base + 2 * Block, Block);
+    Require(state(base + 2 * Block) == BlockWritten, "an import watched without settling did not trust its marks");
+
+    Require(ImportWatched(base, Block, [] { return true; }, true), "the second import of the first block failed");
+    StoreOwnBytes(base + 8, 1, [&] { bytes[8] = 0x44; });
+    generations[0] = TrackerGeneration();
+    bytes[8] = bytes[8];
+    CollectWritesUncached(base, Block);
+    Require(state(base) == BlockMaybeWritten, "a late mark over the driver's own store read as a CPU store");
+    Require(MarkWritten(base + 2 * 4096, 4096) != 0, "a GPU write into a settling import was not stamped");
+    generations[0] = TrackerGeneration();
+    bytes[2 * 4096 + 8] = 0x55;
+    CollectWritesUncached(base, Block);
+    Require(state(base) == BlockWritten, "a CPU store after a GPU write into a settling import was not trusted");
+    Unwatch(base, Block);
+    Require(!Watched(base, Block), "an unwatched settling import stayed watched");
+    GuestArena::GuestArenaReset_nid_postfix(memory, 3 * Block);
+    GuestArena::GuestArenaRelease_nid_postfix(memory, 3 * Block);
+}
+#endif
+
 int main() {
+#ifdef _WIN32
+    _putenv_s("APS5_IMPORT_SETTLE_MS", "100");
+#endif
     try {
         if (!WriteWatched()) {
             std::cout << "no write watching: skipped\n";
@@ -188,6 +252,7 @@ int main() {
         CheckUnwatch();
 #ifdef _WIN32
         CheckPrivateMappingReuse();
+        CheckImportSettle();
 #endif
     } catch (const std::exception& error) {
         std::cerr << "write tracking test failed: " << error.what() << "\n";

core/libs/prx/libSceAgcDriver/tests/execution/DccKeysFirstWrite.cpp

@@ -175,6 +175,13 @@ void RequireKeys(const std::uint8_t* keys, std::uint8_t expected, const std::str
     Require(found == KeyBytes, message);
 }
 
+void RepeatLateImportMarks(std::uint8_t* begin, std::size_t bytes) {
+    for (std::size_t offset = 0; offset < bytes; offset += 4096) {
+        auto* page = reinterpret_cast<volatile std::uint8_t*>(begin + offset);
+        *page = *page;
+    }
+}
+
 void Run(AgcDriver::VulkanDevice& device, std::uint8_t* block) {
     bool watched = AgcDriver::GuestMemory::Watched(AddressOf(block), BlockBytes);
     auto* texels = block;
@@ -195,6 +202,9 @@ void Run(AgcDriver::VulkanDevice& device, std::uint8_t* block) {
     std::lock_guard gpu(AgcDriver::GuestMemory::GpuMutex());
     Dispatch(device, WriteCode, writeData);
     Dispatch(device, CopyKeysCode, copyData);
+#ifdef _WIN32
+    if (watched) RepeatLateImportMarks(block, BlockBytes);
+#endif
     RequireKeys(copied, 0xff, "a kernel reading the keys after the first write of a 0000-cleared surface");
     Require(AgcDriver::Graphics::CurrentDccKeys(AddressOf(copied), SurfaceBytes) == DccKeys::Uncompressed, "the copied keys do not read uncompressed");
     const auto image = StorageTexture::FindPending(surface, SurfaceBytes);
@@ -332,6 +342,9 @@ void Run(AgcDriver::VulkanDevice& device, std::uint8_t* block) {
 }
 
 int main() {
+#ifdef _WIN32
+    _putenv_s("APS5_IMPORT_SETTLE_MS", "60000");
+#endif
     try {
         std::optional<GuestBlock> block;
         std::optional<GuestBlock> watched;