Skip to content

Commit

Permalink
Add callback to mem_guard if loaded
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mkindahl committed Mar 4, 2025
1 parent 9115e12 commit 330539e
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 0 deletions.
1 change: 1 addition & 0 deletions .unreleased/pr_7788
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Implements: #7788 Add callback to mem_guard for background workers
52 changes: 52 additions & 0 deletions src/bgw/job.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Check warning on line 96 in src/bgw/job.c

View check run for this annotation

Codecov / codecov/patch

src/bgw/job.c#L96

Added line #L96 was not covered by tests

mem_guard_callback_ptr = (MGCallbacks **) find_rendezvous_variable(MG_CALLBACKS_VAR_NAME);

return *mem_guard_callback_ptr;
}

typedef enum JobLockLifetime
{
SESSION_LOCK = 0,
Expand Down Expand Up @@ -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);

Check warning on line 1194 in src/bgw/job.c

View check run for this annotation

Codecov / codecov/patch

src/bgw/job.c#L1194

Added line #L1194 was not covered by tests

BackgroundWorkerInitializeConnectionByOid(db_oid, params.user_oid, 0);

log_min_messages = ts_guc_bgw_log_level;
Expand Down

0 comments on commit 330539e

Please sign in to comment.