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.

Author: mkindahl

src/bgw/job.c

@@ -56,6 +56,47 @@
 static scheduler_test_hook_type scheduler_test_hook = NULL;
 static char *job_entrypoint_function_name = "ts_bgw_job_entrypoint";
 
+/*
+ * 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.
+ */
+
+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;
+
+/*
+ * Get the mem_guard callbacks.
+ *
+ * You might get a NULL pointer back if there are no mem_guard installed, so
+ * check before using.
+ */
+static MGCallbacks *
+get_mem_guard_callbacks()
+{
+	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");
+
+	return *mem_guard_callback_ptr;
+}
+
 typedef enum JobLockLifetime
 {
 	SESSION_LOCK = 0,
@@ -1141,6 +1182,13 @@ ts_bgw_job_entrypoint(PG_FUNCTION_ARGS)
 	pqsignal(SIGTERM, die);
 	BackgroundWorkerUnblockSignals();
 
+	/*
+	 * Set up mem_guard before starting to allocate memory.
+	 */
+	MGCallbacks *callbacks = get_mem_guard_callbacks();
+	if (callbacks && 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;