From feafb27b0d572251015ecd313c6a53adfdb1a36b Mon Sep 17 00:00:00 2001
From: Tyler Wilding <xtvaser@gmail.com>
Date: Fri, 21 Jun 2024 23:48:47 -0400
Subject: [PATCH] game: Add a compressed dumping mode

---
 common/global_profiler/GlobalProfiler.cpp | 30 +++++++++++++++++------
 common/global_profiler/GlobalProfiler.h   |  8 ++++--
 common/util/compress.cpp                  | 13 ++++++++++
 common/util/compress.h                    |  3 ++-
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/common/global_profiler/GlobalProfiler.cpp b/common/global_profiler/GlobalProfiler.cpp
index 246885ab82..5c6f645123 100644
--- a/common/global_profiler/GlobalProfiler.cpp
+++ b/common/global_profiler/GlobalProfiler.cpp
@@ -26,6 +26,8 @@ u64 get_current_tid() {
 }
 #endif
 #include "common/log/log.h"
+#include "common/util/compress.h"
+#include "common/util/string_util.h"
 // clang-format on
 
 u64 get_current_ts() {
@@ -34,12 +36,11 @@ u64 get_current_ts() {
 
 GlobalProfiler::GlobalProfiler() {
   m_t0 = get_current_ts();
-  set_max_events(65536);
+  m_nodes.resize(m_max_events);
 }
 
-void GlobalProfiler::set_max_events(size_t event_count) {
-  ASSERT(!m_enabled);
-  m_nodes.resize(event_count);
+void GlobalProfiler::update_event_buffer_size() {
+  m_nodes.resize(m_max_events);
 }
 
 void GlobalProfiler::set_waiting_for_event(const std::string& event_name) {
@@ -96,8 +97,10 @@ void GlobalProfiler::set_enable(bool en) {
   m_enabled = en;
 }
 
-void GlobalProfiler::dump_to_json(const std::string& path) {
-  ASSERT(!m_enabled);
+void GlobalProfiler::dump_to_json() {
+  if (m_enabled) {
+    set_enable(false);
+  }
 
   nlohmann::json json;
   auto& trace_events = json["traceEvents"];
@@ -185,7 +188,20 @@ void GlobalProfiler::dump_to_json(const std::string& path) {
               t.second.highest_at_target);
   }
 
-  file_util::write_text_file(path, json.dump());
+  if (m_enable_compression) {
+    const auto json_str = json.dump();
+    const auto compressed_data =
+        compression::compress_zstd_no_header(json_str.data(), json_str.size());
+    auto file_path = file_util::get_jak_project_dir() / "profile_data" /
+                     fmt::format("prof-{}.json.zst", str_util::current_local_timestamp_no_colons());
+    file_util::create_dir_if_needed_for_file(file_path);
+    file_util::write_binary_file(file_path, compressed_data.data(), compressed_data.size());
+  } else {
+    auto file_path = file_util::get_jak_project_dir() / "profile_data" /
+                     fmt::format("prof-{}.json", str_util::current_local_timestamp_no_colons());
+    file_util::create_dir_if_needed_for_file(file_path);
+    file_util::write_text_file(file_path, json.dump());
+  }
 }
 
 GlobalProfiler gprof;
diff --git a/common/global_profiler/GlobalProfiler.h b/common/global_profiler/GlobalProfiler.h
index 44c9062972..eb55f421f2 100644
--- a/common/global_profiler/GlobalProfiler.h
+++ b/common/global_profiler/GlobalProfiler.h
@@ -17,7 +17,7 @@ struct ProfNode {
 class GlobalProfiler {
  public:
   GlobalProfiler();
-  void set_max_events(size_t event_count);
+  void update_event_buffer_size();
   void set_waiting_for_event(const std::string& event_name);
   void instant_event(const char* name);
   void begin_event(const char* name);
@@ -25,8 +25,12 @@ class GlobalProfiler {
   void end_event();
   void clear();
   void set_enable(bool en);
-  void dump_to_json(const std::string& path);
+  void dump_to_json();
   void root_event();
+  bool is_enabled() { return m_enabled; }
+
+  int m_max_events = 65536;
+  bool m_enable_compression = false;
 
  private:
   std::atomic_bool m_enabled = false;
diff --git a/common/util/compress.cpp b/common/util/compress.cpp
index 7c4aee82be..ec61030a94 100644
--- a/common/util/compress.cpp
+++ b/common/util/compress.cpp
@@ -46,4 +46,17 @@ std::vector<u8> decompress_zstd(const void* data, size_t size) {
   ASSERT(decomp_size == decompressed_size);
   return result;
 }
+
+std::vector<u8> compress_zstd_no_header(const void* data, size_t size) {
+  size_t max_compressed = ZSTD_compressBound(size);
+  std::vector<u8> result(max_compressed);
+
+  size_t compressed_size = ZSTD_compress(result.data(), max_compressed, data, size, 1);
+  if (ZSTD_isError(compressed_size)) {
+    ASSERT_MSG(false, fmt::format("ZSTD error: {}", ZSTD_getErrorName(compressed_size)));
+  }
+
+  result.resize(compressed_size);
+  return result;
+}
 }  // namespace compression
diff --git a/common/util/compress.h b/common/util/compress.h
index 295ccfa4f9..1f15474bbe 100644
--- a/common/util/compress.h
+++ b/common/util/compress.h
@@ -8,4 +8,5 @@ namespace compression {
 // compress and decompress data with zstd
 std::vector<u8> compress_zstd(const void* data, size_t size);
 std::vector<u8> decompress_zstd(const void* data, size_t size);
-}  // namespace compression
\ No newline at end of file
+std::vector<u8> compress_zstd_no_header(const void* data, size_t size);
+}  // namespace compression