This is an automated email from the ASF dual-hosted git repository.

djwang pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/cloudberry.git

commit 7027dd05bbec7826607d9c4a453a372876ec58d7
Author: NJrslv <[email protected]>
AuthorDate: Fri Jun 27 13:31:34 2025 +0300

    [yagp_hooks_collector] Fix memory leaks, add safe C++ wrappers, improve 
Makefile
    
    Fix memory leaks in C++ and PG contexts.  Add safe C++ wrappers
    around PG functions.  Improve error message logging.  Enable
    parallel make.  Fix variable expansion.
---
 Makefile                                           |   2 +
 src/Config.cpp                                     |  22 +--
 src/Config.h                                       |   2 +-
 src/EventSender.cpp                                |  20 +--
 src/EventSender.h                                  |   3 -
 src/PgUtils.cpp                                    | 106 +++------------
 src/PgUtils.h                                      |   6 +-
 src/ProcStats.cpp                                  |   8 +-
 src/ProtoUtils.cpp                                 |  73 +++++-----
 src/ProtoUtils.h                                   |   2 +
 src/UDSConnector.cpp                               |   6 +-
 src/UDSConnector.h                                 |   1 -
 src/hook_wrappers.cpp                              |   6 +-
 src/memory/gpdbwrappers.cpp                        | 148 +++++++++++++++++++++
 src/memory/gpdbwrappers.h                          | 131 ++++++++++++++++++
 .../pg_stat_statements_ya_parser.h                 |   6 +-
 16 files changed, 380 insertions(+), 162 deletions(-)

diff --git a/Makefile b/Makefile
index 91be52c4468..15c5dabb70e 100644
--- a/Makefile
+++ b/Makefile
@@ -8,8 +8,10 @@
 # to "Makefile" if it exists. PostgreSQL is shipped with a
 # "GNUmakefile". If the user hasn't run the configure script yet, the
 # GNUmakefile won't exist yet, so we catch that case as well.
+
 # AIX make defaults to building *every* target of the first rule.  Start with
 # a single-target, empty rule to make the other targets non-default.
+all:
 
 all check install installdirs installcheck installcheck-parallel uninstall 
clean distclean maintainer-clean dist distcheck world check-world install-world 
installcheck-world installcheck-resgroup installcheck-resgroup-v2:
        @if [ ! -f GNUmakefile ] ; then \
diff --git a/src/Config.cpp b/src/Config.cpp
index ac274a1e218..a1289a48891 100644
--- a/src/Config.cpp
+++ b/src/Config.cpp
@@ -1,4 +1,5 @@
 #include "Config.h"
+#include "memory/gpdbwrappers.h"
 #include <limits.h>
 #include <memory>
 #include <string>
@@ -6,7 +7,6 @@
 
 extern "C" {
 #include "postgres.h"
-#include "utils/builtins.h"
 #include "utils/guc.h"
 }
 
@@ -29,15 +29,15 @@ static void update_ignored_users(const char 
*new_guc_ignored_users) {
       std::make_unique<std::unordered_set<std::string>>();
   if (new_guc_ignored_users != nullptr && new_guc_ignored_users[0] != '\0') {
     /* Need a modifiable copy of string */
-    char *rawstring = pstrdup(new_guc_ignored_users);
+    char *rawstring = gpdb::pstrdup(new_guc_ignored_users);
     List *elemlist;
     ListCell *l;
 
     /* Parse string into list of identifiers */
-    if (!SplitIdentifierString(rawstring, ',', &elemlist)) {
+    if (!gpdb::split_identifier_string(rawstring, ',', &elemlist)) {
       /* syntax error in list */
-      pfree(rawstring);
-      list_free(elemlist);
+      gpdb::pfree(rawstring);
+      gpdb::list_free(elemlist);
       ereport(
           LOG,
           (errcode(ERRCODE_SYNTAX_ERROR),
@@ -48,8 +48,8 @@ static void update_ignored_users(const char 
*new_guc_ignored_users) {
     foreach (l, elemlist) {
       new_ignored_users_set->insert((char *)lfirst(l));
     }
-    pfree(rawstring);
-    list_free(elemlist);
+    gpdb::pfree(rawstring);
+    gpdb::list_free(elemlist);
   }
   ignored_users_set = std::move(new_ignored_users_set);
 }
@@ -119,11 +119,11 @@ size_t Config::max_text_size() { return guc_max_text_size 
* 1024; }
 size_t Config::max_plan_size() { return guc_max_plan_size * 1024; }
 int Config::min_analyze_time() { return guc_min_analyze_time; };
 
-bool Config::filter_user(const std::string *username) {
-  if (!username || !ignored_users_set) {
+bool Config::filter_user(std::string username) {
+  if (!ignored_users_set) {
     return true;
   }
-  return ignored_users_set->find(*username) != ignored_users_set->end();
+  return ignored_users_set->find(username) != ignored_users_set->end();
 }
 
 void Config::sync() {
@@ -131,4 +131,4 @@ void Config::sync() {
     update_ignored_users(guc_ignored_users);
     ignored_users_guc_dirty = false;
   }
-}
\ No newline at end of file
+}
diff --git a/src/Config.h b/src/Config.h
index dd081c41dd6..eff83f0960a 100644
--- a/src/Config.h
+++ b/src/Config.h
@@ -9,7 +9,7 @@ public:
   static bool enable_analyze();
   static bool enable_cdbstats();
   static bool enable_collector();
-  static bool filter_user(const std::string *username);
+  static bool filter_user(std::string username);
   static bool report_nested_queries();
   static size_t max_text_size();
   static size_t max_plan_size();
diff --git a/src/EventSender.cpp b/src/EventSender.cpp
index 19787fe0db0..8711c4cbd4f 100644
--- a/src/EventSender.cpp
+++ b/src/EventSender.cpp
@@ -1,15 +1,14 @@
 #include "Config.h"
 #include "UDSConnector.h"
+#include "memory/gpdbwrappers.h"
 
 #define typeid __typeid
 extern "C" {
 #include "postgres.h"
 
-#include "access/hash.h"
 #include "executor/executor.h"
 #include "utils/elog.h"
 
-#include "cdb/cdbdisp.h"
 #include "cdb/cdbexplain.h"
 #include "cdb/cdbvars.h"
 #include "cdb/ml_ipc.h"
@@ -81,7 +80,7 @@ void EventSender::executor_before_start(QueryDesc 
*query_desc, int eflags) {
         instr_time starttime;
         INSTR_TIME_SET_CURRENT(starttime);
         query_desc->showstatctx =
-            cdbexplain_showExecStatsBegin(query_desc, starttime);
+            gpdb::cdbexplain_showExecStatsBegin(query_desc, starttime);
       }
     }
   }
@@ -106,10 +105,10 @@ void EventSender::executor_after_start(QueryDesc 
*query_desc, int /* eflags*/) {
         // Make sure the space is allocated in the per-query
         // context so it will go away at executor_end.
         if (query_desc->totaltime == NULL) {
-          MemoryContext oldcxt;
-          oldcxt = MemoryContextSwitchTo(query_desc->estate->es_query_cxt);
-          query_desc->totaltime = InstrAlloc(1, INSTRUMENT_ALL);
-          MemoryContextSwitchTo(oldcxt);
+          MemoryContext oldcxt =
+              gpdb::mem_ctx_switch_to(query_desc->estate->es_query_cxt);
+          query_desc->totaltime = gpdb::instr_alloc(1, INSTRUMENT_ALL);
+          gpdb::mem_ctx_switch_to(oldcxt);
         }
       }
       yagpcc::GPMetrics stats;
@@ -240,7 +239,7 @@ void EventSender::collect_query_done(QueryDesc *query_desc,
     }
     query_msgs.erase({query_desc->gpmon_pkt->u.qexec.key.ccnt,
                       query_desc->gpmon_pkt->u.qexec.key.tmid});
-    pfree(query_desc->gpmon_pkt);
+    gpdb::pfree(query_desc->gpmon_pkt);
   }
 }
 
@@ -297,7 +296,7 @@ void EventSender::analyze_stats_collect(QueryDesc 
*query_desc) {
   }
   // Make sure stats accumulation is done.
   // (Note: it's okay if several levels of hook all do this.)
-  InstrEndLoop(query_desc->totaltime);
+  gpdb::instr_end_loop(query_desc->totaltime);
 
   double ms = query_desc->totaltime->total * 1000.0;
   if (ms >= Config::min_analyze_time()) {
@@ -364,7 +363,8 @@ EventSender::QueryItem 
*EventSender::get_query_message(QueryDesc *query_desc) {
       query_msgs.find({query_desc->gpmon_pkt->u.qexec.key.ccnt,
                        query_desc->gpmon_pkt->u.qexec.key.tmid}) ==
           query_msgs.end()) {
-    query_desc->gpmon_pkt = (gpmon_packet_t *)palloc0(sizeof(gpmon_packet_t));
+    query_desc->gpmon_pkt =
+        (gpmon_packet_t *)gpdb::palloc0(sizeof(gpmon_packet_t));
     query_desc->gpmon_pkt->u.qexec.key.ccnt = gp_command_count;
     query_desc->gpmon_pkt->u.qexec.key.tmid = nesting_level;
     query_msgs.insert({{gp_command_count, nesting_level},
diff --git a/src/EventSender.h b/src/EventSender.h
index 4d09b429fc8..f3dd1d2a528 100644
--- a/src/EventSender.h
+++ b/src/EventSender.h
@@ -1,13 +1,10 @@
 #pragma once
 
-#include <memory>
 #include <unordered_map>
-#include <string>
 
 #define typeid __typeid
 extern "C" {
 #include "utils/metrics_utils.h"
-#include "cdb/ml_ipc.h"
 #ifdef IC_TEARDOWN_HOOK
 #include "cdb/ic_udpifc.h"
 #endif
diff --git a/src/PgUtils.cpp b/src/PgUtils.cpp
index ed3e69c6d44..f36cd030a39 100644
--- a/src/PgUtils.cpp
+++ b/src/PgUtils.cpp
@@ -1,37 +1,41 @@
 #include "PgUtils.h"
 #include "Config.h"
+#include "memory/gpdbwrappers.h"
 
 extern "C" {
-#include "utils/guc.h"
-#include "commands/dbcommands.h"
 #include "commands/resgroupcmds.h"
 #include "cdb/cdbvars.h"
 }
 
-std::string *get_user_name() {
-  const char *username = GetConfigOption("session_authorization", false, 
false);
-  // username is not to be freed
-  return username ? new std::string(username) : nullptr;
+std::string get_user_name() {
+  // username is allocated on stack, we don't need to pfree it.
+  const char *username =
+      ya_gpdb::get_config_option("session_authorization", false, false);
+  return username ? std::string(username) : "";
 }
 
-std::string *get_db_name() {
-  char *dbname = get_database_name(MyDatabaseId);
-  std::string *result = nullptr;
+std::string get_db_name() {
+  char *dbname = ya_gpdb::get_database_name(MyDatabaseId);
   if (dbname) {
-    result = new std::string(dbname);
-    pfree(dbname);
+    std::string result(dbname);
+    ya_gpdb::pfree(dbname);
+    return result;
   }
-  return result;
+  return "";
 }
 
-std::string *get_rg_name() {
-  auto groupId = ResGroupGetGroupIdBySessionId(MySessionState->sessionId);
+std::string get_rg_name() {
+  auto groupId = ya_gpdb::get_rg_id_by_session_id(MySessionState->sessionId);
   if (!OidIsValid(groupId))
-    return nullptr;
-  char *rgname = GetResGroupNameForId(groupId);
+    return "";
+
+  char *rgname = ya_gpdb::get_rg_name_for_id(groupId);
   if (rgname == nullptr)
-    return nullptr;
-  return new std::string(rgname);
+    return "";
+
+  std::string result(rgname);
+  ya_gpdb::pfree(rgname);
+  return result;
 }
 
 /**
@@ -80,69 +84,3 @@ bool need_collect(QueryDesc *query_desc, int nesting_level) {
   return !filter_query(query_desc) &&
          nesting_is_valid(query_desc, nesting_level);
 }
-
-ExplainState get_explain_state(QueryDesc *query_desc, bool costs) {
-  ExplainState es;
-  ExplainInitState(&es);
-  es.costs = costs;
-  es.verbose = true;
-  es.format = EXPLAIN_FORMAT_TEXT;
-  ExplainBeginOutput(&es);
-  PG_TRY();
-  { ExplainPrintPlan(&es, query_desc); }
-  PG_CATCH();
-  {
-    // PG and GP both have known and yet unknown bugs in EXPLAIN VERBOSE
-    // implementation. We don't want any queries to fail due to those bugs, so
-    // we report the bug here for future investigatin and continue collecting
-    // metrics w/o reporting any plans
-    resetStringInfo(es.str);
-    appendStringInfo(
-        es.str,
-        "Unable to restore query plan due to PostgreSQL internal error. "
-        "See logs for more information");
-    ereport(INFO,
-            (errmsg("YAGPCC failed to reconstruct explain text for query: %s",
-                    query_desc->sourceText)));
-  }
-  PG_END_TRY();
-  ExplainEndOutput(&es);
-  return es;
-}
-
-ExplainState get_analyze_state_json(QueryDesc *query_desc, bool analyze) {
-  ExplainState es;
-  ExplainInitState(&es);
-  es.analyze = analyze;
-  es.verbose = true;
-  es.buffers = es.analyze;
-  es.timing = es.analyze;
-  es.summary = es.analyze;
-  es.format = EXPLAIN_FORMAT_JSON;
-  ExplainBeginOutput(&es);
-  if (analyze) {
-    PG_TRY();
-    {
-      ExplainPrintPlan(&es, query_desc);
-      ExplainPrintExecStatsEnd(&es, query_desc);
-    }
-    PG_CATCH();
-    {
-      // PG and GP both have known and yet unknown bugs in EXPLAIN VERBOSE
-      // implementation. We don't want any queries to fail due to those bugs, 
so
-      // we report the bug here for future investigatin and continue collecting
-      // metrics w/o reporting any plans
-      resetStringInfo(es.str);
-      appendStringInfo(
-          es.str,
-          "Unable to restore analyze plan due to PostgreSQL internal error. "
-          "See logs for more information");
-      ereport(INFO,
-              (errmsg("YAGPCC failed to reconstruct analyze text for query: 
%s",
-                      query_desc->sourceText)));
-    }
-    PG_END_TRY();
-  }
-  ExplainEndOutput(&es);
-  return es;
-}
diff --git a/src/PgUtils.h b/src/PgUtils.h
index 81282a473a8..ceb07c2e8e5 100644
--- a/src/PgUtils.h
+++ b/src/PgUtils.h
@@ -5,9 +5,9 @@ extern "C" {
 
 #include <string>
 
-std::string *get_user_name();
-std::string *get_db_name();
-std::string *get_rg_name();
+std::string get_user_name();
+std::string get_db_name();
+std::string get_rg_name();
 bool is_top_level_query(QueryDesc *query_desc, int nesting_level);
 bool nesting_is_valid(QueryDesc *query_desc, int nesting_level);
 bool need_report_nested_query();
diff --git a/src/ProcStats.cpp b/src/ProcStats.cpp
index a557a20cbb0..5c09fa0bce4 100644
--- a/src/ProcStats.cpp
+++ b/src/ProcStats.cpp
@@ -75,16 +75,16 @@ void fill_status_stats(yagpcc::SystemStat *stats) {
       stats->set_vmpeakkb(value);
       proc_stat >> measure;
       if (measure != "kB") {
-        ereport(FATAL, (errmsg("Expected memory sizes in kB, but got in %s",
-                               measure.c_str())));
+        throw std::runtime_error("Expected memory sizes in kB, but got in " +
+                                 measure);
       }
     } else if (key == "VmSize:") {
       uint64_t value;
       proc_stat >> value;
       stats->set_vmsizekb(value);
       if (measure != "kB") {
-        ereport(FATAL, (errmsg("Expected memory sizes in kB, but got in %s",
-                               measure.c_str())));
+        throw std::runtime_error("Expected memory sizes in kB, but got in " +
+                                 measure);
       }
     }
   }
diff --git a/src/ProtoUtils.cpp b/src/ProtoUtils.cpp
index 6e9fa6bd5c5..6dc39278bcd 100644
--- a/src/ProtoUtils.cpp
+++ b/src/ProtoUtils.cpp
@@ -2,6 +2,7 @@
 #include "PgUtils.h"
 #include "ProcStats.h"
 #include "Config.h"
+#include "memory/gpdbwrappers.h"
 
 #define typeid __typeid
 #define operator __operator
@@ -15,10 +16,7 @@ extern "C" {
 #ifdef IC_TEARDOWN_HOOK
 #include "cdb/ic_udpifc.h"
 #endif
-#include "gpmon/gpmon.h"
 #include "utils/workfile_mgr.h"
-
-#include "stat_statements_parser/pg_stat_statements_ya_parser.h"
 }
 #undef typeid
 #undef operator
@@ -60,18 +58,21 @@ void set_query_plan(yagpcc::SetQueryReq *req, QueryDesc 
*query_desc) {
                           ? yagpcc::PlanGenerator::PLAN_GENERATOR_OPTIMIZER
                           : yagpcc::PlanGenerator::PLAN_GENERATOR_PLANNER);
     MemoryContext oldcxt =
-        MemoryContextSwitchTo(query_desc->estate->es_query_cxt);
-    auto es = get_explain_state(query_desc, true);
-    MemoryContextSwitchTo(oldcxt);
-    *qi->mutable_plan_text() =
-        char_to_trimmed_str(es.str->data, es.str->len, 
Config::max_plan_size());
-    StringInfo norm_plan = gen_normplan(es.str->data);
-    *qi->mutable_template_plan_text() = char_to_trimmed_str(
-        norm_plan->data, norm_plan->len, Config::max_plan_size());
-    qi->set_plan_id(hash_any((unsigned char *)norm_plan->data, 
norm_plan->len));
-    qi->set_query_id(query_desc->plannedstmt->queryId);
-    pfree(es.str->data);
-    pfree(norm_plan->data);
+        gpdb::mem_ctx_switch_to(query_desc->estate->es_query_cxt);
+    ExplainState es = gpdb::get_explain_state(query_desc, true);
+    if (es.str) {
+      *qi->mutable_plan_text() = char_to_trimmed_str(es.str->data, es.str->len,
+                                                     Config::max_plan_size());
+      StringInfo norm_plan = gpdb::gen_normplan(es.str->data);
+      *qi->mutable_template_plan_text() = char_to_trimmed_str(
+          norm_plan->data, norm_plan->len, Config::max_plan_size());
+      qi->set_plan_id(
+          hash_any((unsigned char *)norm_plan->data, norm_plan->len));
+      qi->set_query_id(query_desc->plannedstmt->queryId);
+      gpdb::pfree(es.str->data);
+      gpdb::pfree(norm_plan->data);
+    }
+    gpdb::mem_ctx_switch_to(oldcxt);
   }
 }
 
@@ -81,7 +82,7 @@ void set_query_text(yagpcc::SetQueryReq *req, QueryDesc 
*query_desc) {
     *qi->mutable_query_text() = char_to_trimmed_str(
         query_desc->sourceText, strlen(query_desc->sourceText),
         Config::max_text_size());
-    char *norm_query = gen_normquery(query_desc->sourceText);
+    char *norm_query = gpdb::gen_normquery(query_desc->sourceText);
     *qi->mutable_template_query_text() = char_to_trimmed_str(
         norm_query, strlen(norm_query), Config::max_text_size());
   }
@@ -101,9 +102,9 @@ void clear_big_fields(yagpcc::SetQueryReq *req) {
 void set_query_info(yagpcc::SetQueryReq *req) {
   if (Gp_session_role == GP_ROLE_DISPATCH) {
     auto qi = req->mutable_query_info();
-    qi->set_allocated_username(get_user_name());
-    qi->set_allocated_databasename(get_db_name());
-    qi->set_allocated_rsgname(get_rg_name());
+    qi->set_username(get_user_name());
+    qi->set_databasename(get_db_name());
+    qi->set_rsgname(get_rg_name());
   }
 }
 
@@ -233,23 +234,23 @@ void set_analyze_plan_text_json(QueryDesc *query_desc,
     return;
   }
   MemoryContext oldcxt =
-      MemoryContextSwitchTo(query_desc->estate->es_query_cxt);
-
-  ExplainState es = get_analyze_state_json(
+      gpdb::mem_ctx_switch_to(query_desc->estate->es_query_cxt);
+  ExplainState es = gpdb::get_analyze_state_json(
       query_desc, query_desc->instrument_options && Config::enable_analyze());
-  // Remove last line break.
-  if (es.str->len > 0 && es.str->data[es.str->len - 1] == '\n') {
-    es.str->data[--es.str->len] = '\0';
-  }
-  // Convert JSON array to JSON object.
-  if (es.str->len > 0) {
-    es.str->data[0] = '{';
-    es.str->data[es.str->len - 1] = '}';
+  gpdb::mem_ctx_switch_to(oldcxt);
+  if (es.str) {
+    // Remove last line break.
+    if (es.str->len > 0 && es.str->data[es.str->len - 1] == '\n') {
+      es.str->data[--es.str->len] = '\0';
+    }
+    // Convert JSON array to JSON object.
+    if (es.str->len > 0) {
+      es.str->data[0] = '{';
+      es.str->data[es.str->len - 1] = '}';
+    }
+    auto trimmed_analyze =
+        char_to_trimmed_str(es.str->data, es.str->len, 
Config::max_plan_size());
+    req->mutable_query_info()->set_analyze_text(trimmed_analyze);
+    gpdb::pfree(es.str->data);
   }
-  auto trimmed_analyze =
-      char_to_trimmed_str(es.str->data, es.str->len, Config::max_plan_size());
-  req->mutable_query_info()->set_analyze_text(trimmed_analyze);
-
-  pfree(es.str->data);
-  MemoryContextSwitchTo(oldcxt);
 }
\ No newline at end of file
diff --git a/src/ProtoUtils.h b/src/ProtoUtils.h
index 6fb880c2eb8..8287b3de7ea 100644
--- a/src/ProtoUtils.h
+++ b/src/ProtoUtils.h
@@ -1,3 +1,5 @@
+#pragma once
+
 #include "protos/yagpcc_set_service.pb.h"
 
 struct QueryDesc;
diff --git a/src/UDSConnector.cpp b/src/UDSConnector.cpp
index 8a5f754f3b4..b5b70836db4 100644
--- a/src/UDSConnector.cpp
+++ b/src/UDSConnector.cpp
@@ -1,6 +1,7 @@
 #include "UDSConnector.h"
 #include "Config.h"
 #include "YagpStat.h"
+#include "memory/gpdbwrappers.h"
 
 #include <string>
 #include <unistd.h>
@@ -13,7 +14,6 @@
 
 extern "C" {
 #include "postgres.h"
-#include "cdb/cdbvars.h"
 }
 
 UDSConnector::UDSConnector() { GOOGLE_PROTOBUF_VERIFY_VERSION; }
@@ -44,7 +44,7 @@ bool UDSConnector::report_query(const yagpcc::SetQueryReq 
&req,
       if (connect(sockfd, (sockaddr *)&address, sizeof(address)) != -1) {
         auto data_size = req.ByteSize();
         auto total_size = data_size + sizeof(uint32_t);
-        uint8_t *buf = (uint8_t *)palloc(total_size);
+        uint8_t *buf = (uint8_t *)gpdb::palloc(total_size);
         uint32_t *size_payload = (uint32_t *)buf;
         *size_payload = data_size;
         req.SerializeWithCachedSizesToArray(buf + sizeof(uint32_t));
@@ -67,7 +67,7 @@ bool UDSConnector::report_query(const yagpcc::SetQueryReq 
&req,
         } else {
           YagpStat::report_send(total_size);
         }
-        pfree(buf);
+        gpdb::pfree(buf);
       } else {
         // log the error and go on
         log_tracing_failure(req, event);
diff --git a/src/UDSConnector.h b/src/UDSConnector.h
index 42e0aa20968..67504fc8529 100644
--- a/src/UDSConnector.h
+++ b/src/UDSConnector.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include "protos/yagpcc_set_service.pb.h"
-#include <queue>
 
 class UDSConnector {
 public:
diff --git a/src/hook_wrappers.cpp b/src/hook_wrappers.cpp
index 79d3ec45881..25a85f086d1 100644
--- a/src/hook_wrappers.cpp
+++ b/src/hook_wrappers.cpp
@@ -7,10 +7,10 @@ extern "C" {
 #include "utils/elog.h"
 #include "utils/builtins.h"
 #include "utils/metrics_utils.h"
-#include "cdb/cdbexplain.h"
 #include "cdb/cdbvars.h"
 #include "cdb/ml_ipc.h"
 #include "tcop/utility.h"
+#include "stat_statements_parser/pg_stat_statements_ya_parser.h"
 }
 #undef typeid
 
@@ -18,7 +18,7 @@ extern "C" {
 #include "YagpStat.h"
 #include "EventSender.h"
 #include "hook_wrappers.h"
-#include "stat_statements_parser/pg_stat_statements_ya_parser.h"
+#include "memory/gpdbwrappers.h"
 
 static ExecutorStart_hook_type previous_ExecutorStart_hook = nullptr;
 static ExecutorRun_hook_type previous_ExecutorRun_hook = nullptr;
@@ -229,7 +229,7 @@ Datum yagp_functions_get(FunctionCallInfo fcinfo) {
   values[3] = Int64GetDatum(stats.failed_connects);
   values[4] = Int64GetDatum(stats.failed_other);
   values[5] = Int32GetDatum(stats.max_message_size);
-  HeapTuple tuple = heap_form_tuple(tupdesc, values, nulls);
+  HeapTuple tuple = gpdb::heap_form_tuple(tupdesc, values, nulls);
   Datum result = HeapTupleGetDatum(tuple);
   PG_RETURN_DATUM(result);
 }
\ No newline at end of file
diff --git a/src/memory/gpdbwrappers.cpp b/src/memory/gpdbwrappers.cpp
new file mode 100644
index 00000000000..1fba702a9f5
--- /dev/null
+++ b/src/memory/gpdbwrappers.cpp
@@ -0,0 +1,148 @@
+#include "gpdbwrappers.h"
+
+extern "C" {
+#include "postgres.h"
+#include "utils/guc.h"
+#include "commands/dbcommands.h"
+#include "commands/resgroupcmds.h"
+#include "utils/builtins.h"
+#include "nodes/pg_list.h"
+#include "commands/explain.h"
+#include "executor/instrument.h"
+#include "access/tupdesc.h"
+#include "access/htup.h"
+#include "utils/elog.h"
+#include "cdb/cdbexplain.h"
+#include "stat_statements_parser/pg_stat_statements_ya_parser.h"
+}
+
+void *gpdb::palloc(Size size) { return detail::wrap_throw(::palloc, size); }
+
+void *gpdb::palloc0(Size size) { return detail::wrap_throw(::palloc0, size); }
+
+char *gpdb::pstrdup(const char *str) {
+  return detail::wrap_throw(::pstrdup, str);
+}
+
+char *gpdb::get_database_name(Oid dbid) noexcept {
+  return detail::wrap_noexcept(::get_database_name, dbid);
+}
+
+bool gpdb::split_identifier_string(char *rawstring, char separator,
+                                   List **namelist) noexcept {
+  return detail::wrap_noexcept(SplitIdentifierString, rawstring, separator,
+                               namelist);
+}
+
+ExplainState gpdb::get_explain_state(QueryDesc *query_desc,
+                                     bool costs) noexcept {
+  return detail::wrap_noexcept([&]() {
+    ExplainState es;
+    ExplainInitState(&es);
+    es.costs = costs;
+    es.verbose = true;
+    es.format = EXPLAIN_FORMAT_TEXT;
+    ExplainBeginOutput(&es);
+    ExplainPrintPlan(&es, query_desc);
+    ExplainEndOutput(&es);
+    return es;
+  });
+}
+
+ExplainState gpdb::get_analyze_state_json(QueryDesc *query_desc,
+                                          bool analyze) noexcept {
+  return detail::wrap_noexcept([&]() {
+    ExplainState es;
+    ExplainInitState(&es);
+    es.analyze = analyze;
+    es.verbose = true;
+    es.buffers = es.analyze;
+    es.timing = es.analyze;
+    es.summary = es.analyze;
+    es.format = EXPLAIN_FORMAT_JSON;
+    ExplainBeginOutput(&es);
+    if (analyze) {
+      ExplainPrintPlan(&es, query_desc);
+      ExplainPrintExecStatsEnd(&es, query_desc);
+    }
+    ExplainEndOutput(&es);
+    return es;
+  });
+}
+
+Instrumentation *gpdb::instr_alloc(size_t n, int instrument_options) {
+  return detail::wrap_throw(InstrAlloc, n, instrument_options);
+}
+
+HeapTuple gpdb::heap_form_tuple(TupleDesc tupleDescriptor, Datum *values,
+                                bool *isnull) {
+  if (!tupleDescriptor || !values || !isnull)
+    throw std::runtime_error(
+        "Invalid input parameters for heap tuple formation");
+
+  return detail::wrap_throw(::heap_form_tuple, tupleDescriptor, values, 
isnull);
+}
+
+void gpdb::pfree(void *pointer) noexcept {
+  // Note that ::pfree asserts that pointer != NULL.
+  if (!pointer)
+    return;
+
+  detail::wrap_noexcept(::pfree, pointer);
+}
+
+MemoryContext gpdb::mem_ctx_switch_to(MemoryContext context) noexcept {
+  return MemoryContextSwitchTo(context);
+}
+
+const char *gpdb::get_config_option(const char *name, bool missing_ok,
+                                    bool restrict_superuser) noexcept {
+  if (!name)
+    return nullptr;
+
+  return detail::wrap_noexcept(GetConfigOption, name, missing_ok,
+                               restrict_superuser);
+}
+
+void gpdb::list_free(List *list) noexcept {
+  if (!list)
+    return;
+
+  detail::wrap_noexcept(::list_free, list);
+}
+
+CdbExplain_ShowStatCtx *
+gpdb::cdbexplain_showExecStatsBegin(QueryDesc *query_desc,
+                                    instr_time starttime) {
+  if (!query_desc)
+    throw std::runtime_error("Invalid query descriptor");
+
+  return detail::wrap_throw(::cdbexplain_showExecStatsBegin, query_desc,
+                            starttime);
+}
+
+void gpdb::instr_end_loop(Instrumentation *instr) {
+  if (!instr)
+    throw std::runtime_error("Invalid instrumentation pointer");
+
+  detail::wrap_throw(::InstrEndLoop, instr);
+}
+
+char *gpdb::gen_normquery(const char *query) {
+  return detail::wrap_throw(::gen_normquery, query);
+}
+
+StringInfo gpdb::gen_normplan(const char *exec_plan) {
+  if (!exec_plan)
+    throw std::runtime_error("Invalid execution plan string");
+
+  return detail::wrap_throw(::gen_normplan, exec_plan);
+}
+
+char *gpdb::get_rg_name_for_id(Oid group_id) {
+  return detail::wrap_throw(GetResGroupNameForId, group_id);
+}
+
+Oid gpdb::get_rg_id_by_session_id(int session_id) {
+  return detail::wrap_throw(ResGroupGetGroupIdBySessionId, session_id);
+}
\ No newline at end of file
diff --git a/src/memory/gpdbwrappers.h b/src/memory/gpdbwrappers.h
new file mode 100644
index 00000000000..437a5dd5d29
--- /dev/null
+++ b/src/memory/gpdbwrappers.h
@@ -0,0 +1,131 @@
+#pragma once
+
+extern "C" {
+#include "postgres.h"
+#include "nodes/pg_list.h"
+#include "commands/explain.h"
+#include "executor/instrument.h"
+#include "access/htup.h"
+#include "utils/elog.h"
+#include "utils/memutils.h"
+}
+
+#include <type_traits>
+#include <stdexcept>
+#include <optional>
+#include <utility>
+#include <string>
+
+namespace gpdb {
+namespace detail {
+
+template <bool Throws, typename Func, typename... Args>
+auto wrap(Func &&func, Args &&...args) noexcept(!Throws)
+    -> decltype(func(std::forward<Args>(args)...)) {
+
+  using RetType = decltype(func(std::forward<Args>(args)...));
+
+  // Empty struct for void return type.
+  struct VoidResult {};
+  using ResultHolder = std::conditional_t<std::is_void_v<RetType>, VoidResult,
+                                          std::optional<RetType>>;
+
+  bool success;
+  ErrorData *edata;
+  ResultHolder result_holder;
+
+  PG_TRY();
+  {
+    if constexpr (!std::is_void_v<RetType>) {
+      result_holder.emplace(func(std::forward<Args>(args)...));
+    } else {
+      func(std::forward<Args>(args)...);
+    }
+    edata = NULL;
+    success = true;
+  }
+  PG_CATCH();
+  {
+    MemoryContext oldctx = MemoryContextSwitchTo(TopMemoryContext);
+    edata = CopyErrorData();
+    MemoryContextSwitchTo(oldctx);
+    FlushErrorState();
+    success = false;
+  }
+  PG_END_TRY();
+
+  if (!success) {
+    std::string err;
+    if (edata && edata->message) {
+      err = std::string(edata->message);
+    } else {
+      err = "Unknown error occurred";
+    }
+
+    if (edata) {
+      FreeErrorData(edata);
+    }
+
+    if constexpr (Throws) {
+      throw std::runtime_error(err);
+    }
+
+    if constexpr (!std::is_void_v<RetType>) {
+      return RetType{};
+    } else {
+      return;
+    }
+  }
+
+  if constexpr (!std::is_void_v<RetType>) {
+    return *std::move(result_holder);
+  } else {
+    return;
+  }
+}
+
+template <typename Func, typename... Args>
+auto wrap_throw(Func &&func, Args &&...args)
+    -> decltype(func(std::forward<Args>(args)...)) {
+  return detail::wrap<true>(std::forward<Func>(func),
+                            std::forward<Args>(args)...);
+}
+
+template <typename Func, typename... Args>
+auto wrap_noexcept(Func &&func, Args &&...args) noexcept
+    -> decltype(func(std::forward<Args>(args)...)) {
+  return detail::wrap<false>(std::forward<Func>(func),
+                             std::forward<Args>(args)...);
+}
+} // namespace detail
+
+// Functions that call palloc().
+// Make sure correct memory context is set.
+void *palloc(Size size);
+void *palloc0(Size size);
+char *pstrdup(const char *str);
+char *get_database_name(Oid dbid) noexcept;
+bool split_identifier_string(char *rawstring, char separator,
+                             List **namelist) noexcept;
+ExplainState get_explain_state(QueryDesc *query_desc, bool costs) noexcept;
+ExplainState get_analyze_state_json(QueryDesc *query_desc,
+                                    bool analyze) noexcept;
+Instrumentation *instr_alloc(size_t n, int instrument_options);
+HeapTuple heap_form_tuple(TupleDesc tupleDescriptor, Datum *values,
+                          bool *isnull);
+CdbExplain_ShowStatCtx *cdbexplain_showExecStatsBegin(QueryDesc *query_desc,
+                                                      instr_time starttime);
+void instr_end_loop(Instrumentation *instr);
+char *gen_normquery(const char *query);
+StringInfo gen_normplan(const char *executionPlan);
+char *get_rg_name_for_id(Oid group_id);
+
+// Palloc-free functions.
+void pfree(void *pointer) noexcept;
+MemoryContext mem_ctx_switch_to(MemoryContext context) noexcept;
+const char *get_config_option(const char *name, bool missing_ok,
+                              bool restrict_superuser) noexcept;
+void list_free(List *list) noexcept;
+Oid get_rg_id_by_session_id(int session_id);
+
+} // namespace gpdb
diff --git a/src/stat_statements_parser/pg_stat_statements_ya_parser.h 
b/src/stat_statements_parser/pg_stat_statements_ya_parser.h
index aa9cd217e31..b08e8533992 100644
--- a/src/stat_statements_parser/pg_stat_statements_ya_parser.h
+++ b/src/stat_statements_parser/pg_stat_statements_ya_parser.h
@@ -8,9 +8,9 @@ extern "C"
 extern void stat_statements_parser_init(void);
 extern void stat_statements_parser_deinit(void);
 
+StringInfo gen_normplan(const char *executionPlan);
+char *gen_normquery(const char *query);
+
 #ifdef __cplusplus
 }
 #endif
-
-StringInfo gen_normplan(const char *executionPlan);
-char *gen_normquery(const char *query);
\ No newline at end of file


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to