From 910d8d4d3a19d37941039303d025457af87e578c Mon Sep 17 00:00:00 2001
From: Mats Kindahl <mats@timescale.com>
Date: Tue, 4 Mar 2025 13:57:06 +0100
Subject: [PATCH] Add callback to mem_guard if loaded

Add code to call mem_guard to enable memory tracking for background
workers. This is enabled for all background workers started by
timescaledb, but not to other workers started by PostgreSQL.
---
 .unreleased/pr_7788 |  1 +
 src/bgw/job.c       | 29 +++++++++++++++++++++++++++++
 src/bgw/job.h       | 27 +++++++++++++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 .unreleased/pr_7788

diff --git a/.unreleased/pr_7788 b/.unreleased/pr_7788
new file mode 100644
index 00000000000..62262cc439d
--- /dev/null
+++ b/.unreleased/pr_7788
@@ -0,0 +1 @@
+Implements: #7788 Add callback to mem_guard for background workers
diff --git a/src/bgw/job.c b/src/bgw/job.c
index c76a1d23a3c..b28f21fba53 100644
--- a/src/bgw/job.c
+++ b/src/bgw/job.c
@@ -56,6 +56,25 @@
 static scheduler_test_hook_type scheduler_test_hook = NULL;
 static char *job_entrypoint_function_name = "ts_bgw_job_entrypoint";
 
+/*
+ * Get the mem_guard callbacks.
+ *
+ * You might get a NULL pointer back if there are no mem_guard installed, so
+ * check before using.
+ */
+MGCallbacks *
+ts_get_mem_guard_callbacks(void)
+{
+	static MGCallbacks **mem_guard_callback_ptr = NULL;
+
+	if (mem_guard_callback_ptr)
+		return *mem_guard_callback_ptr;
+
+	mem_guard_callback_ptr = (MGCallbacks **) find_rendezvous_variable(MG_CALLBACKS_VAR_NAME);
+
+	return *mem_guard_callback_ptr;
+}
+
 typedef enum JobLockLifetime
 {
 	SESSION_LOCK = 0,
@@ -1141,6 +1160,16 @@ ts_bgw_job_entrypoint(PG_FUNCTION_ARGS)
 	pqsignal(SIGTERM, die);
 	BackgroundWorkerUnblockSignals();
 
+	/*
+	 * Set up mem_guard before starting to allocate (any significant amounts
+	 * of) memory but after we have unblocked signals since we have no control
+	 * over how the callback behaves.
+	 */
+	MGCallbacks *callbacks = ts_get_mem_guard_callbacks();
+	if (callbacks && callbacks->version_num == MG_CALLBACKS_VERSION &&
+		callbacks->toggle_allocation_blocking && !callbacks->enabled)
+		callbacks->toggle_allocation_blocking(/*enable=*/true);
+
 	BackgroundWorkerInitializeConnectionByOid(db_oid, params.user_oid, 0);
 
 	log_min_messages = ts_guc_bgw_log_level;
diff --git a/src/bgw/job.h b/src/bgw/job.h
index a291336fdae..974f118300f 100644
--- a/src/bgw/job.h
+++ b/src/bgw/job.h
@@ -15,6 +15,31 @@
 #define TELEMETRY_INITIAL_NUM_RUNS 12
 #define SCHEDULER_APPNAME "TimescaleDB Background Worker Scheduler"
 
+/*
+ * This is copied from mem_guard and have to be the same as the type in
+ * mem_guard.
+ *
+ * These are intended as an interim solution and will be removed once we have
+ * a stable plugin ABI for TimescaleDB.
+ */
+
+#define MG_CALLBACKS_VERSION 1
+#define MG_CALLBACKS_VAR_NAME "mg_callbacks"
+
+typedef void (*mg_toggle_allocation_blocking)(bool enable);
+typedef size_t (*mg_get_allocated_memory)();
+typedef size_t (*mg_get_total_allocated_memory)();
+typedef bool (*mg_enabled)();
+
+typedef struct MGCallbacks
+{
+	int64 version_num;
+	mg_toggle_allocation_blocking toggle_allocation_blocking;
+	mg_get_allocated_memory get_allocated_memory;
+	mg_get_total_allocated_memory get_total_allocated_memory;
+	mg_enabled enabled;
+} MGCallbacks;
+
 typedef struct BgwJobHistory
 {
 	int64 id;
@@ -85,3 +110,5 @@ ScanTupleResult ts_bgw_job_change_owner(TupleInfo *ti, void *data);
 
 extern TSDLLEXPORT Oid ts_bgw_job_get_funcid(BgwJob *job);
 extern TSDLLEXPORT const char *ts_bgw_job_function_call_string(BgwJob *job);
+
+extern TSDLLEXPORT MGCallbacks *ts_get_mem_guard_callbacks(void);