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..0926eb1c96a 100644 --- a/src/bgw/job.c +++ b/src/bgw/job.c @@ -56,6 +56,50 @@ 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. + */ + +#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; + +/* + * 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_VAR_NAME); + + return *mem_guard_callback_ptr; +} + typedef enum JobLockLifetime { SESSION_LOCK = 0, @@ -1141,6 +1185,14 @@ 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->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;