Add callback to mem_guard if loaded

.unreleased/pr_7788

@@ -0,0 +1 @@
+Implements: #7788 Add callback to mem_guard for background workers

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 @@
 	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;

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);