├── .gitignore ├── Makefile ├── README.md ├── pg_worker_pool--0.1.sql ├── pg_worker_pool.c └── pg_worker_pool.control /.gitignore: -------------------------------------------------------------------------------- 1 | .cache/ 2 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | EXTENSION = pg_worker_pool 2 | MODULE_big = pg_worker_pool 3 | OBJS = pg_worker_pool.o 4 | PG_CONFIG = pg_config 5 | PGXS := $(shell $(PG_CONFIG) --pgxs) 6 | PG_CFLAGS = -Wall -Wextra -Wno-declaration-after-statement -ggdb3 7 | DATA = pg_worker_pool--0.1.sql 8 | 9 | include $(PGXS) 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # pg_worker_pool 2 | 3 | A Postgresql extension that creates a pool of named background workers. 4 | 5 | This uses [Postgresql background workers](https://www.postgresql.org/docs/current/bgworker.html) functionality. 6 | 7 | Keep in mind that Postgresql has limited number of available background workers. This is specified by `max_worker_processes` configuration. Changing that limit requires server restart. 8 | 9 | Currently this extension has hard-coded number of workers which is set to `8`. 10 | If you want to change it, you can set `MAX_WORKERS` variable in `pg_worker_pool.c` to any other number and recompile the extension. 11 | Increasing this number will let you start more workers. 12 | Making this larger than `max_worker_processes` will yield no effect though. This is a Postgresql limitation. 13 | You could also decrease that number, but I'm not sure why would one do that. You can control how many workers actually start by choosing appropriate worker name. 14 | 15 | Background workers are created on-demand and exit as soon as they aren't needed. 16 | 17 | ## Usage 18 | 19 | Upon creation, this extension creates `worker_pool` schema. 20 | 21 | `worker_pool.submit` can be used to add query to be executed on a specified background worker. This also starts the worker at the end of current transaction (if committed). If the current transaction is aborted, the worker is not started. 22 | If you start more than one worker in the same transaction, the order in which they are started in unspecified. It might be and likely will be different than the order of submission. Do not rely on this ordering. 23 | 24 | `worker_pool.launch` can be used to start a specified worker without adding any queries to the queue. 25 | 26 | `worker_pool.jobs` stores all queries submitted to be executed. 27 | 28 | Worker name can be any string with max length of `NAMEDATALEN` (by default 64 characters). 29 | 30 | This extension must be loaded via `shared_preload_libraries` to correctly allocate required shared memory. 31 | 32 | ## How to install 33 | 34 | Regular `git clone` + `make` + `make install` should be enough. 35 | 36 | Note that `pg_config` utility needs to be in your `$PATH`. 37 | It is commonly part of `postgresql-common` package on Debian/Ubuntu and `postgresql-libs` on Arch systems. 38 | 39 | ``` 40 | git clone https://github.com/serpent7776/pg_worker_pool.git 41 | cd pg_worker_pool 42 | make 43 | sudo make install 44 | ``` 45 | 46 | ## Example 47 | 48 | ```sql 49 | CREATE EXTENSION pg_worker_pool; 50 | CALL worker_pool.submit('foo', 'create index myindex_1 on my_big_table (id)'); 51 | CALL worker_pool.submit('foo', 'create index myindex_2 on my_big_table (name)'); 52 | CALL worker_pool.submit('bar', 'create index otherindex_1 on other_big_table (author)'); 53 | CALL worker_pool.submit('bar', 'create index otherindex_2 on other_big_table (title)'); 54 | 55 | SELECT * FROM worker_pool.jobs ORDER BY id DESC; 56 | id | worker_name | query_text | status 57 | ----+-------------+-------------------------------------------------------+-------- 58 | 4 | bar | create index otherindex_2 on other_big_table (title) | done 59 | 3 | bar | create index otherindex_1 on other_big_table (author) | done 60 | 2 | foo | create index myindex_2 on my_big_table (name) | done 61 | 1 | foo | create index myindex_1 on my_big_table (id) | done 62 | (4 rows) 63 | ``` 64 | 65 | This will start two background workers, `foo` and `bar`. 66 | `foo` will create an indices on table `my_big_table` and `bar` on table `other_big_table`. 67 | `foo` and `bar` will run independently of each other, but all indices submitted to the same worker will be created in order. 68 | 69 | ## FAQ 70 | 71 | ### How many workers should I start? 72 | 73 | It's hard to answer that in general case, but you probably don't want to start more than the number of CPUs you have on the machine, especially for long running queries. 74 | -------------------------------------------------------------------------------- /pg_worker_pool--0.1.sql: -------------------------------------------------------------------------------- 1 | CREATE SCHEMA worker_pool; 2 | 3 | CREATE TABLE worker_pool.jobs ( 4 | id SERIAL PRIMARY KEY, 5 | worker_name TEXT NOT NULL, 6 | query_text TEXT NOT NULL, 7 | status TEXT CHECK (status IN ('waiting', 'pending', 'done', 'failed')) DEFAULT 'waiting' 8 | ); 9 | 10 | CREATE PROCEDURE worker_pool.submit(worker TEXT, query TEXT) 11 | AS 'MODULE_PATHNAME', 'pg_worker_pool_submit' LANGUAGE c; 12 | 13 | CREATE PROCEDURE worker_pool.launch(worker TEXT) 14 | AS 'MODULE_PATHNAME', 'pg_worker_pool_launch' LANGUAGE c; 15 | -------------------------------------------------------------------------------- /pg_worker_pool.c: -------------------------------------------------------------------------------- 1 | #pragma GCC diagnostic push 2 | #pragma GCC diagnostic ignored "-Wunused-parameter" 3 | #include "postgres.h" 4 | #include "fmgr.h" 5 | #include "miscadmin.h" 6 | #include "storage/ipc.h" 7 | #include "storage/shmem.h" 8 | #include "storage/lwlock.h" 9 | #include "utils/builtins.h" 10 | #include "executor/spi.h" 11 | #include "postmaster/bgworker.h" 12 | #include "utils/snapmgr.h" 13 | #include "utils/memutils.h" 14 | #include "utils/backend_status.h" 15 | #pragma GCC diagnostic pop 16 | 17 | PG_MODULE_MAGIC; 18 | 19 | typedef struct WorkerInfo 20 | { 21 | char name[NAMEDATALEN]; 22 | pid_t pid; 23 | bool is_active; 24 | } WorkerInfo; 25 | 26 | typedef struct WorkerPool 27 | { 28 | LWLockId lock; 29 | WorkerInfo worker[]; 30 | } WorkerPool; 31 | 32 | typedef struct WorkerParams 33 | { 34 | Oid database; 35 | Oid user; 36 | } WorkerParams; 37 | _Static_assert (sizeof(WorkerParams) < BGW_EXTRALEN, "WorkerParams too big"); 38 | 39 | typedef struct XactArgs 40 | { 41 | char worker_name[NAMEDATALEN]; 42 | } XactArgs; 43 | 44 | #define MAX_WORKERS 8 45 | 46 | static WorkerPool* pool = NULL; 47 | static shmem_request_hook_type prev_shmem_request_hook = NULL; 48 | static shmem_startup_hook_type prev_shmem_startup_hook = NULL; 49 | static int worker_pool_size = MAX_WORKERS; 50 | 51 | PGDLLEXPORT void _PG_init(void); 52 | PGDLLEXPORT void pg_worker_main(Datum main_arg); 53 | PGDLLEXPORT void pg_worker_pool_on_xact(XactEvent event, void *arg); 54 | 55 | PG_FUNCTION_INFO_V1(pg_worker_pool_submit); 56 | Datum pg_worker_pool_submit(PG_FUNCTION_ARGS) 57 | { 58 | const char *worker_name = text_to_cstring(PG_GETARG_TEXT_PP(0)); 59 | const char *query = text_to_cstring(PG_GETARG_TEXT_PP(1)); 60 | 61 | if (strlen(worker_name) >= NAMEDATALEN) 62 | ereport(ERROR, (errmsg("Worker name too long"))); 63 | 64 | SPI_connect(); 65 | 66 | StringInfoData buf; 67 | initStringInfo(&buf); 68 | appendStringInfo(&buf, "INSERT INTO worker_pool.jobs(worker_name, query_text, status) VALUES ('%s', '%s', 'waiting')", 69 | worker_name, query); 70 | 71 | if (SPI_execute(buf.data, false, 0) != SPI_OK_INSERT) 72 | { 73 | SPI_finish(); 74 | ereport(ERROR, (errmsg("Failed to submit query to worker"))); 75 | } 76 | 77 | SPI_finish(); 78 | 79 | XactArgs* xact_args = MemoryContextAlloc(TopTransactionContext, sizeof(XactArgs)); 80 | strncpy(xact_args->worker_name, worker_name, NAMEDATALEN); 81 | RegisterXactCallback(pg_worker_pool_on_xact, xact_args); 82 | 83 | PG_RETURN_VOID(); 84 | } 85 | 86 | PG_FUNCTION_INFO_V1(pg_worker_pool_launch); 87 | Datum pg_worker_pool_launch(PG_FUNCTION_ARGS) 88 | { 89 | const char *worker_name = text_to_cstring(PG_GETARG_TEXT_PP(0)); 90 | 91 | if (strlen(worker_name) >= NAMEDATALEN) 92 | ereport(ERROR, (errmsg("Worker name too long"))); 93 | 94 | XactArgs* xact_args = MemoryContextAlloc(TopTransactionContext, sizeof(XactArgs)); 95 | strncpy(xact_args->worker_name, worker_name, NAMEDATALEN); 96 | RegisterXactCallback(pg_worker_pool_on_xact, xact_args); 97 | 98 | PG_RETURN_VOID(); 99 | } 100 | 101 | void pg_worker_pool_on_xact(XactEvent event, void *arg) 102 | { 103 | if (event == XACT_EVENT_PRE_COMMIT || event == XACT_EVENT_PARALLEL_PRE_COMMIT) return; 104 | if (event == XACT_EVENT_PREPARE || event == XACT_EVENT_PRE_PREPARE) return; 105 | UnregisterXactCallback(pg_worker_pool_on_xact, arg); 106 | if (event == XACT_EVENT_ABORT || event == XACT_EVENT_PARALLEL_ABORT) return; 107 | 108 | XactArgs* xact_args = (XactArgs*)arg; 109 | const char* worker_name = xact_args->worker_name; 110 | 111 | WorkerParams params = { 112 | .database = MyDatabaseId, 113 | .user = GetUserId(), 114 | }; 115 | 116 | BackgroundWorkerHandle* handle; 117 | BackgroundWorker worker = { 118 | .bgw_flags = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION, 119 | .bgw_start_time = BgWorkerStart_RecoveryFinished, 120 | .bgw_restart_time = BGW_NEVER_RESTART, 121 | .bgw_notify_pid = MyProcPid, 122 | }; 123 | memcpy(worker.bgw_extra, ¶ms, sizeof(WorkerParams)); 124 | 125 | volatile bool found = false; 126 | LWLockAcquire(pool->lock, LW_EXCLUSIVE); 127 | PG_TRY(); 128 | { 129 | for (int i = 0; i < worker_pool_size; i++) 130 | { 131 | if (strcmp(pool->worker[i].name, worker_name) == 0) 132 | { 133 | if (!pool->worker[i].is_active) 134 | { 135 | worker.bgw_main_arg = Int32GetDatum(i), 136 | 137 | snprintf(worker.bgw_name, BGW_MAXLEN, "pgworker: %s", worker_name); 138 | snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_worker_pool"); 139 | snprintf(worker.bgw_function_name, BGW_MAXLEN, "pg_worker_main"); 140 | 141 | if (RegisterDynamicBackgroundWorker(&worker, &handle)) 142 | { 143 | pool->worker[i].is_active = true; 144 | } 145 | else ereport(WARNING, (errmsg("Failed to start background worker"))); 146 | } 147 | found = true; 148 | break; 149 | } 150 | } 151 | if (!found) 152 | for (int i = 0; i < worker_pool_size; i++) 153 | { 154 | if (!pool->worker[i].is_active) 155 | { 156 | worker.bgw_main_arg = Int32GetDatum(i), 157 | snprintf(pool->worker[i].name, BGW_MAXLEN, "%s", worker_name); 158 | 159 | snprintf(worker.bgw_name, BGW_MAXLEN, "pgworker: %s", worker_name); 160 | snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_worker_pool"); 161 | snprintf(worker.bgw_function_name, BGW_MAXLEN, "pg_worker_main"); 162 | 163 | if (RegisterDynamicBackgroundWorker(&worker, &handle)) 164 | { 165 | pool->worker[i].is_active = true; 166 | found = true; 167 | } 168 | else ereport(WARNING, (errmsg("Failed to register background worker"))); 169 | break; 170 | } 171 | } 172 | } 173 | PG_FINALLY(); 174 | { 175 | LWLockRelease(pool->lock); 176 | } 177 | PG_END_TRY(); 178 | 179 | if (!found) 180 | ereport(WARNING, (errmsg("Worker name '%s' not found", worker_name))); 181 | } 182 | 183 | void pg_worker_main(Datum main_arg) 184 | { 185 | const int worker_index = DatumGetInt32(main_arg); 186 | WorkerParams params; 187 | memcpy(¶ms, MyBgworkerEntry->bgw_extra, sizeof(WorkerParams)); 188 | 189 | if (worker_index < 0 || worker_index >= worker_pool_size) 190 | ereport(ERROR, (errmsg("Invalid worker index"))); 191 | 192 | BackgroundWorkerInitializeConnectionByOid(params.database, params.user, 0); 193 | 194 | BackgroundWorkerUnblockSignals(); 195 | 196 | CurrentResourceOwner = ResourceOwnerCreate(NULL, "pg_worker_pool"); 197 | 198 | LWLockAcquire(pool->lock, LW_SHARED); 199 | WorkerInfo* worker = &pool->worker[worker_index]; 200 | const char* worker_name = MemoryContextStrdup(TopMemoryContext, worker->name); 201 | LWLockRelease(pool->lock); 202 | 203 | MyBackendType = B_BACKEND; 204 | pgstat_report_appname("pg_worker_pool"); 205 | pgstat_report_activity(STATE_IDLE, NULL); 206 | 207 | StringInfoData buf; 208 | initStringInfo(&buf); 209 | 210 | while (true) 211 | { 212 | StartTransactionCommand(); 213 | if (SPI_connect() != SPI_OK_CONNECT) 214 | { 215 | ereport(ERROR, 216 | (errcode(ERRCODE_INTERNAL_ERROR), 217 | errmsg("could not connect to SPI manager"), 218 | errdetail("SPI_connect failed"))); 219 | } 220 | PushActiveSnapshot(GetTransactionSnapshot()); 221 | 222 | resetStringInfo(&buf); 223 | appendStringInfo(&buf, 224 | "WITH x AS (\n" 225 | " SELECT id, worker_name, query_text FROM worker_pool.jobs\n" 226 | " WHERE worker_name = '%s' AND status = 'waiting'\n" 227 | " FOR UPDATE\n" 228 | " LIMIT 1\n" 229 | ")\n" 230 | "UPDATE worker_pool.jobs AS j SET status = 'pending'\n" 231 | "FROM x\n" 232 | "WHERE j.id = x.id\n" 233 | "RETURNING j.id, j.query_text", 234 | worker_name); 235 | SetCurrentStatementStartTimestamp(); 236 | pgstat_report_activity(STATE_RUNNING, buf.data); 237 | 238 | bool isnull; 239 | if (SPI_execute(buf.data, false, 0) == SPI_OK_UPDATE_RETURNING && SPI_processed == 1) 240 | { 241 | Datum id = SPI_getbinval(SPI_tuptable->vals[0], SPI_tuptable->tupdesc, 1, &isnull); 242 | const char *query = SPI_getvalue(SPI_tuptable->vals[0], SPI_tuptable->tupdesc, 2); 243 | SetCurrentStatementStartTimestamp(); 244 | pgstat_report_activity(STATE_RUNNING, query); 245 | 246 | volatile bool exn = false; 247 | volatile MemoryContext mycontext = CurrentMemoryContext; 248 | PG_TRY(); 249 | { 250 | int code = SPI_execute(query, false, 0); 251 | if (code > 0) 252 | { 253 | resetStringInfo(&buf); 254 | appendStringInfo(&buf, "UPDATE worker_pool.jobs SET status = 'done' WHERE id = %d", 255 | DatumGetInt32(id)); 256 | SetCurrentStatementStartTimestamp(); 257 | pgstat_report_activity(STATE_RUNNING, buf.data); 258 | 259 | if (SPI_execute(buf.data, false, 0) != SPI_OK_UPDATE) 260 | ereport(WARNING, (errmsg("Failed to update job status"))); 261 | } 262 | else 263 | { 264 | ereport(WARNING, (errmsg("Query failed with error code %d: %s", code, query))); 265 | } 266 | } 267 | PG_CATCH(); 268 | { 269 | exn = true; 270 | MemoryContextSwitchTo(mycontext); 271 | EmitErrorReport(); 272 | FlushErrorState(); 273 | } 274 | PG_END_TRY(); 275 | if (exn) 276 | { 277 | SPI_finish(); 278 | PopActiveSnapshot(); 279 | if (IsTransactionState()) 280 | AbortCurrentTransaction(); 281 | 282 | StartTransactionCommand(); 283 | if (SPI_connect() != SPI_OK_CONNECT) 284 | ereport(ERROR, (errmsg("could not connect to SPI manager"))); 285 | PushActiveSnapshot(GetTransactionSnapshot()); 286 | 287 | resetStringInfo(&buf); 288 | appendStringInfo(&buf, "UPDATE worker_pool.jobs SET status = 'failed' WHERE id = %d", 289 | DatumGetInt32(id)); 290 | SetCurrentStatementStartTimestamp(); 291 | pgstat_report_activity(STATE_RUNNING, buf.data); 292 | 293 | if (SPI_execute(buf.data, false, 0) != SPI_OK_UPDATE) 294 | ereport(WARNING, (errmsg("Failed to update job status"))); 295 | } 296 | } 297 | else 298 | { 299 | LWLockAcquire(pool->lock, LW_EXCLUSIVE); 300 | worker->is_active = false; 301 | LWLockRelease(pool->lock); 302 | SPI_finish(); 303 | PopActiveSnapshot(); 304 | CommitTransactionCommand(); 305 | break; 306 | } 307 | 308 | SPI_finish(); 309 | PopActiveSnapshot(); 310 | CommitTransactionCommand(); 311 | pgstat_report_activity(STATE_IDLE, NULL); 312 | pg_usleep(100000); 313 | } 314 | 315 | proc_exit(0); 316 | } 317 | 318 | static void worker_pool_shmem_request(void) 319 | { 320 | if (prev_shmem_request_hook) { 321 | prev_shmem_request_hook(); 322 | } 323 | 324 | RequestNamedLWLockTranche("worker_pool", 1); 325 | } 326 | 327 | static void worker_pool_shmem_startup(void) 328 | { 329 | if (prev_shmem_startup_hook) { 330 | prev_shmem_startup_hook(); 331 | } 332 | 333 | bool found; 334 | LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); 335 | InitShmemIndex(); 336 | pool = ShmemInitStruct("pg_worker_pool", sizeof(WorkerPool) + worker_pool_size * sizeof(WorkerInfo), &found); 337 | if (!found) { 338 | for (int i = 0; i < worker_pool_size; i++) 339 | { 340 | pool->worker[i].is_active = false; 341 | memset(pool->worker[i].name, 0, NAMEDATALEN); 342 | } 343 | pool->lock = &(GetNamedLWLockTranche("worker_pool"))->lock; 344 | } 345 | LWLockRelease(AddinShmemInitLock); 346 | } 347 | 348 | void _PG_init(void) 349 | { 350 | if (!process_shared_preload_libraries_in_progress) 351 | ereport(ERROR, (errmsg("pg_worker_pool must be loaded via shared_preload_libraries"))); 352 | 353 | prev_shmem_startup_hook = shmem_startup_hook; 354 | shmem_startup_hook = worker_pool_shmem_startup; 355 | 356 | prev_shmem_request_hook = shmem_request_hook; 357 | shmem_request_hook = worker_pool_shmem_request; 358 | } 359 | -------------------------------------------------------------------------------- /pg_worker_pool.control: -------------------------------------------------------------------------------- 1 | comment = 'A background worker pool for executing SQL queries' 2 | default_version = '0.1' 3 | module_pathname = '$libdir/pg_worker_pool' 4 | relocatable = true 5 | --------------------------------------------------------------------------------