From ee64684b0b1f3a1af1376ce1bd9b31cdca56a843 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?= <benjamin.antal@memgraph.io>
Date: Tue, 25 Oct 2022 13:37:57 +0200
Subject: [PATCH] Add `ScopedCustomProfile`

---
 src/query/v2/plan/profile.hpp        |  5 +++-
 src/query/v2/plan/scoped_profile.hpp | 44 ++++++++++++++++++++++++++--
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/query/v2/plan/profile.hpp b/src/query/v2/plan/profile.hpp
index 74437f82f..c18ae1e56 100644
--- a/src/query/v2/plan/profile.hpp
+++ b/src/query/v2/plan/profile.hpp
@@ -26,10 +26,13 @@ namespace plan {
  * Stores profiling statistics for a single logical operator.
  */
 struct ProfilingStats {
+  static constexpr std::string_view kNumCycles{"num_cycles"};
+
   int64_t actual_hits{0};
-  unsigned long long num_cycles{0};
+  uint64_t num_cycles{0};
   uint64_t key{0};
   const char *name{nullptr};
+  nlohmann::json custom_data;
   // TODO: This should use the allocator for query execution
   std::vector<ProfilingStats> children;
 };
diff --git a/src/query/v2/plan/scoped_profile.hpp b/src/query/v2/plan/scoped_profile.hpp
index e71cbb047..f70701c4c 100644
--- a/src/query/v2/plan/scoped_profile.hpp
+++ b/src/query/v2/plan/scoped_profile.hpp
@@ -13,6 +13,8 @@
 
 #include <cstdint>
 
+#include <json/json.hpp>
+
 #include "query/v2/context.hpp"
 #include "query/v2/plan/profile.hpp"
 #include "utils/likely.hpp"
@@ -20,6 +22,35 @@
 
 namespace memgraph::query::v2::plan {
 
+class ScopedCustomProfile {
+ public:
+  explicit ScopedCustomProfile(const std::string_view custom_data_name, ExecutionContext &context)
+      : custom_data_name_(custom_data_name), start_time_{utils::ReadTSC()}, context_{&context} {}
+
+  ScopedCustomProfile(const ScopedCustomProfile &) = delete;
+  ScopedCustomProfile(ScopedCustomProfile &&) = delete;
+  ScopedCustomProfile &operator=(const ScopedCustomProfile &) = delete;
+  ScopedCustomProfile &operator=(ScopedCustomProfile &&) = delete;
+
+  ~ScopedCustomProfile() {
+    if (nullptr != context_->stats_root) {
+      auto &custom_data = context_->stats_root->custom_data[custom_data_name_];
+      if (!custom_data.is_object()) {
+        custom_data = nlohmann::json::object();
+      }
+      const auto elapsed = utils::ReadTSC() - start_time_;
+      auto &num_cycles_json = custom_data[ProfilingStats::kNumCycles];
+      const auto num_cycles = num_cycles_json.is_null() ? 0 : num_cycles_json.get<uint64_t>();
+      num_cycles_json = num_cycles + elapsed;
+    }
+  }
+
+ private:
+  std::string_view custom_data_name_;
+  uint64_t start_time_;
+  ExecutionContext *context_;
+};
+
 /**
  * A RAII class used for profiling logical operators. Instances of this class
  * update the profiling data stored within the `ExecutionContext` object and build
@@ -29,7 +60,7 @@ namespace memgraph::query::v2::plan {
 class ScopedProfile {
  public:
   ScopedProfile(uint64_t key, const char *name, query::v2::ExecutionContext *context) noexcept : context_(context) {
-    if (UNLIKELY(context_->is_profile_query)) {
+    if (UNLIKELY(IsProfiling())) {
       root_ = context_->stats_root;
 
       // Are we the root logical operator?
@@ -60,8 +91,13 @@ class ScopedProfile {
     }
   }
 
+  ScopedProfile(const ScopedProfile &) = delete;
+  ScopedProfile(ScopedProfile &&) = delete;
+  ScopedProfile &operator=(const ScopedProfile &) = delete;
+  ScopedProfile &operator=(ScopedProfile &&) = delete;
+
   ~ScopedProfile() noexcept {
-    if (UNLIKELY(context_->is_profile_query)) {
+    if (UNLIKELY(IsProfiling())) {
       stats_->num_cycles += utils::ReadTSC() - start_time_;
 
       // Restore the old root ("pop")
@@ -70,10 +106,12 @@ class ScopedProfile {
   }
 
  private:
+  [[nodiscard]] bool IsProfiling() const { return context_->is_profile_query; }
+
   query::v2::ExecutionContext *context_;
   ProfilingStats *root_{nullptr};
   ProfilingStats *stats_{nullptr};
-  unsigned long long start_time_{0};
+  uint64_t start_time_{0};
 };
 
 }  // namespace memgraph::query::v2::plan