FEXServer: Implement automatic aggregation of pending code map data

Source/Tools/FEXServer/ProcessPipe.cpp

@@ -17,8 +17,10 @@
 #include <cassert>
 #include <fcntl.h>
 #include <filesystem>
+#include <fstream>
 #include <poll.h>
 #include <string>
+#include <sys/file.h>
 #include <sys/resource.h>
 #include <sys/socket.h>
 #include <sys/un.h>
@@ -26,6 +28,19 @@
 
 #include <xxhash.h>
 
+namespace FEXCore {
+inline bool operator<(const FEXCore::ExecutableFileInfo& a, const FEXCore::ExecutableFileInfo& b) noexcept {
+  return a.FileId < b.FileId;
+}
+} // namespace FEXCore
+
+template<>
+struct std::hash<FEXCore::ExecutableFileInfo> {
+  std::size_t operator()(const FEXCore::ExecutableFileInfo& Val) const noexcept {
+    return Val.FileId;
+  }
+};
+
 namespace ProcessPipe {
 constexpr int USER_PERMS = S_IRWXU | S_IRWXG | S_IRWXO;
 int ServerLockFD {-1};
@@ -275,6 +290,147 @@ void SendFDSuccessPacket(fasio::tcp_socket& Socket, int FD) {
   write(Socket, Data, ec);
 }
 
+// Discovers any pending code maps, parses their contents into a runtime data structure, and deletes them
+static std::map<FEXCore::ExecutableFileInfo, fextl::set<uintptr_t>>
+ImportPendingCodeMaps(const FEXCore::ExecutableFileInfo& MainFileId, bool HasMultiblock) {
+  // Detect code maps by checking file name suffixes by counting up an index.
+  // Code maps that are ready for reading must be non-empty and flock(FLOCK_EX) must succeed:
+  // - If empty, we tried generating the cache before the client could even lock it
+  // - If exclusively lockable, we know the client either closed or crashed
+  std::vector<std::string> CodeMaps;
+  for (int Index = 0; true; ++Index) {
+    auto CodeMap = fmt::format("{}/new/{}.{}.bin", CodeMapDirectory, FEXCore::CodeMap::GetBaseFilename(MainFileId, !HasMultiblock), Index);
+    auto FD = open(CodeMap.c_str(), O_RDONLY);
+    if (FD == -1) {
+      break;
+    }
+
+    // Acquire exclusive lock to ensure the client process is done writing data.
+    // Also ensure the file is non-empty, otherwise we're racing the client in acquiring the initial lock.
+    struct stat FileStats;
+    fstat(FD, &FileStats);
+    if (FileStats.st_size == 0 || flock(FD, LOCK_EX | LOCK_NB) != 0) {
+      fmt::print("Code map {} is still in use, skipping\n", CodeMap);
+      // Still being written to by a client process, so skip this file
+      // TODO: Rename from X.n.bin to X.0.bin (once the latter has been removed!) to ensure we'll catch it on next run
+      close(FD);
+      continue;
+    } else {
+      fmt::print("Found code map {}, queuing for merge\n", CodeMap);
+    }
+    close(FD);
+    CodeMaps.push_back(CodeMap);
+  }
+
+  // Update merged code map
+  std::map<FEXCore::ExecutableFileInfo, fextl::set<uintptr_t>> ImportedCodeMaps;
+  if (!CodeMaps.empty()) {
+    fmt::print("Found {} new code maps, updating reference code map\n", CodeMaps.size());
+
+    for (auto& CodeMap : CodeMaps) {
+      std::ifstream Incoming(CodeMap, std::ios_base::binary);
+      auto NewBlocks = FEXCore::CodeMap::ParseCodeMap(Incoming);
+      for (auto& [FileId, Contents] : NewBlocks) {
+        ImportedCodeMaps.emplace(std::piecewise_construct, std::forward_as_tuple(nullptr, FileId, std::move(Contents.Filename)),
+                                 std::forward_as_tuple(std::move(Contents.Blocks)));
+      }
+    }
+  }
+
+  // Delete all imported code maps
+  for (auto& CodeMapFile : CodeMaps) {
+    std::filesystem::remove(CodeMapFile);
+    // TODO: Rename any pending (not finalized) code maps to PROGRAMNAME.0.bin so it will be found on the next run
+  }
+
+  return ImportedCodeMaps;
+}
+
+/**
+ * Writes aggregated code map data into a single code map file that is ready to be used for cache generation
+ */
+static void WriteNewCodeMap(const FEXCore::ExecutableFileInfo& File, const std::string& OutputName, const fextl::set<uintptr_t>& Blocks,
+                            bool IsMainFile, const auto& Dependencies) {
+  fmt::print("Writing {} blocks to {}\n", Blocks.size(), OutputName);
+
+  struct CodeMapOpener : FEXCore::CodeMapOpener {
+    CodeMapOpener(const std::string& Filename) {
+      FD = creat(Filename.c_str(), 0644);
+    }
+
+    int OpenCodeMapFile() override {
+      return FD;
+    }
+
+    int FD;
+  };
+
+  CodeMapOpener CodeMapOpener(OutputName);
+  FEXCore::CodeMapWriter OutputCodeMap(CodeMapOpener, true);
+  if (IsMainFile) {
+    // List the main executable and all used libraries
+    OutputCodeMap.AppendSetMainExecutable(File);
+
+    for (auto& Dependency : Dependencies) {
+      OutputCodeMap.AppendLibraryLoad(Dependency);
+    }
+  } else {
+    // List only the library itself
+    OutputCodeMap.AppendLibraryLoad(File);
+  }
+
+  for (auto& Block : Blocks) {
+    OutputCodeMap.AppendBlock(FEXCore::ExecutableFileSectionInfo {const_cast<FEXCore::ExecutableFileInfo&>(File), 0}, Block);
+  }
+}
+
+/**
+ * Checks and processes new code maps generated by FEX for the given application.
+ *
+ * Processed code maps are merged into the reference code map and deleted afterwards.
+ */
+static void AggregateCodeMaps(const FEXCore::ExecutableFileInfo& MainFileId, bool HasMultiblock) {
+  std::set<FEXCore::ExecutableFileInfo> Result;
+
+  // Read all dependencies discovered in previous runs
+  {
+    auto MainFileCodeMapPath = fmt::format("{}/ready/{}", CodeMapDirectory, FEXCore::CodeMap::GetBaseFilename(MainFileId, !HasMultiblock));
+    std::ifstream MainFileCodeMap(MainFileCodeMapPath, std::ios_base::binary);
+    for (auto& [FileId, Contents] : FEXCore::CodeMap::ParseCodeMap(MainFileCodeMap)) {
+      Result.emplace(nullptr, FileId, Contents.Filename);
+    }
+  }
+
+  // Accumulate information from new code maps
+  auto IncomingCodeMap = ImportPendingCodeMaps(MainFileId, HasMultiblock);
+  for (auto& [File, _] : IncomingCodeMap) {
+    Result.emplace(nullptr, File.FileId, File.Filename);
+  }
+
+  // For each referenced library, add referenced offsets to that library's reference code map
+  for (auto& [File, Blocks] : IncomingCodeMap) {
+    const auto BinaryName = std::string {FEXCore::CodeMap::GetBaseFilename(File, !HasMultiblock)};
+    auto OutputName = fmt::format("{}/ready/{}", CodeMapDirectory, BinaryName);
+
+    // Check if the new code maps add any new information to the previous code map
+    if (auto ReferenceCodeMap = std::ifstream(OutputName, std::ios_base::binary)) {
+      auto PreviousBlocks = FEXCore::CodeMap::ParseCodeMap(ReferenceCodeMap).at(File.FileId).Blocks;
+      auto NumPreviousBlocks = PreviousBlocks.size();
+      Blocks.merge(std::move(PreviousBlocks));
+      if (Blocks.size() == NumPreviousBlocks) {
+        // No new blocks => no need to regenerate the corresponding cache
+        continue;
+      } else {
+        fmt::println("  Found {} new blocks ({} total) in code map {} for {}", Blocks.size(), NumPreviousBlocks, BinaryName, File.Filename);
+      }
+    }
+
+    // Update code map
+    std::set<FEXCore::ExecutableFileInfo> Empty;
+    WriteNewCodeMap(File, OutputName, Blocks, true, File.FileId == MainFileId.FileId ? Result : Empty);
+  }
+}
+
 void HandleSocketData(fasio::tcp_socket& Socket) {
   std::vector<uint8_t> Data(1500);
 
@@ -398,6 +554,9 @@ void HandleSocketData(fasio::tcp_socket& Socket) {
       FEXCore::ExecutableFileInfo MainFileId = {nullptr, filename_hash, fextl::string(Tmp, TmpLen)};
       fmt::print("Requested {}cache generation for {}\n", HasMultiblock ? "" : "nomb-", MainFileId.Filename);
 
+      // Update code maps
+      AggregateCodeMaps(MainFileId, HasMultiblock);
+
       // TODO: Implement
 
       FEXServerClient::FEXServerResultPacket Res {