28687656621

Committed 03 Jul 2026 11:24PM UTC coverage: 93.135% (-0.3%) from 93.47%

Build # 28687656621

Build Type

Pull #264

github

Committed by

utkarash2991

Commit Message

log at LOG level when the "net" schema exists without the extension tables

Pull Request Pull Request #264: fix: worker crash loop when a "net" schema exists without the extension

Coverage Stats

1 of 3 new or added lines in 1 file covered. (33.33%)

502 of 539 relevant lines covered (93.14%)

231.31 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.26

/src/worker.c

#include <errno.h>
#include <inttypes.h>
#include <string.h>
#include <unistd.h>

#define PG_PRELUDE_IMPL
#include "pg_prelude.h"

#include "curl_prelude.h"

#include "core.h"
#include "errors.h"
#include "event.h"
#include "util.h"

#define MIN_LIBCURL_VERSION_NUM                                                                    \
  0x075300 // This is the 7.83.0 version in hex as defined in curl/curlver.h
#define REQUIRED_LIBCURL_ERR_MSG                                                                   \
  "libcurl >= 7.83.0 is required, we use the curl_easy_nextheader() function added in this "       \
  "version"
_Static_assert(LIBCURL_VERSION_NUM,
               REQUIRED_LIBCURL_ERR_MSG); // test for older libcurl versions that don't even have
                                          // LIBCURL_VERSION_NUM defined (e.g. libcurl 6.5).
_Static_assert(LIBCURL_VERSION_NUM >= MIN_LIBCURL_VERSION_NUM, REQUIRED_LIBCURL_ERR_MSG);

PG_MODULE_MAGIC;

typedef enum {
  WORKER_WAIT_NO_TIMEOUT,
  WORKER_WAIT_ONE_SECOND,
} WorkerWait;

static WorkerState *worker_state = NULL;

static const int    curl_handle_event_timeout_ms = 1000;
static const int    net_worker_restart_time_sec  = 1;
static const long   no_timeout                   = -1L;
static bool         wake_commit_cb_active        = false;
static bool         worker_should_restart        = false;
static const size_t total_extension_tables       = 2;

static char *guc_ttl;
static int   guc_batch_size;
static char *guc_database_name;
static char *guc_username;

#if PG15_GTE
static shmem_request_hook_type prev_shmem_request_hook = NULL;
#endif

static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
static volatile sig_atomic_t   got_sighup              = false;

void _PG_init(void);

#if PG_VERSION_NUM >= 180000
PGDLLEXPORT pg_noreturn void pg_net_worker(Datum main_arg);
#else
PGDLLEXPORT void pg_net_worker(Datum main_arg) pg_attribute_noreturn();
#endif

PG_FUNCTION_INFO_V1(worker_restart);
Datum worker_restart(__attribute__((unused)) PG_FUNCTION_ARGS) {
  bool result = DatumGetBool(DirectFunctionCall1(pg_reload_conf, (Datum)NULL)); // reload the config
  pg_atomic_write_u32(&worker_state->got_restart, 1);
  pg_write_barrier();
  if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);
  PG_RETURN_BOOL(result); // TODO is not necessary to return a bool here, but we do it to maintain
                          // backward compatibility
}

static void wait_until_state(WorkerState *ws, WorkerStatus expected_status) {
  if (pg_atomic_read_u32(&ws->status) ==
      expected_status) // fast return without sleeping, in case condition is fulfilled
    return;

  ConditionVariablePrepareToSleep(&ws->cv);
  while (pg_atomic_read_u32(&ws->status) != expected_status) {
    ConditionVariableSleep(&ws->cv, PG_WAIT_EXTENSION);
  }
  ConditionVariableCancelSleep();
}

PG_FUNCTION_INFO_V1(wait_until_running);
Datum wait_until_running(__attribute__((unused)) PG_FUNCTION_ARGS) {
  wait_until_state(worker_state, WS_RUNNING);

  PG_RETURN_VOID();
}

// only wake at commit time to prevent excessive and unnecessary wakes.
// e.g only one wake when doing `select
// net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100000);`
static void wake_at_commit(XactEvent event, __attribute__((unused)) void *arg) {
  elog(DEBUG2, "pg_net xact callback received: %s", xact_event_name(event));

  switch (event) {
  case XACT_EVENT_COMMIT:
  case XACT_EVENT_PARALLEL_COMMIT:
    if (wake_commit_cb_active) {
      uint32 expected = 0;
      bool   success  = pg_atomic_compare_exchange_u32(&worker_state->should_wake, &expected, 1);
      pg_write_barrier();

      if (success) // only wake the worker on first put, so if many concurrent wakes come we only
                   // wake once
        SetLatch(worker_state->shared_latch);

      wake_commit_cb_active = false;
    }
    break;
  // TODO: `PREPARE TRANSACTION 'xx';` and `COMMIT PREPARED TRANSACTION 'xx';` do not wake the
  // worker automatically, they require a manual `net.wake()` These are disabled by default and
  // rarely used, see `max_prepared_transactions`
  // https://www.postgresql.org/docs/17/runtime-config-resource.html#GUC-MAX-PREPARED-TRANSACTIONS
  case XACT_EVENT_PREPARE:
  // abort the callback on rollback
  case XACT_EVENT_ABORT:
  case XACT_EVENT_PARALLEL_ABORT: wake_commit_cb_active = false; break;
  default                       : break;
  }
}

PG_FUNCTION_INFO_V1(wake);
Datum wake(__attribute__((unused)) PG_FUNCTION_ARGS) {
  if (!wake_commit_cb_active) { // register only one callback per transaction
    RegisterXactCallback(wake_at_commit, NULL);
    wake_commit_cb_active = true;
  }

  PG_RETURN_VOID();
}

static void handle_sigterm(__attribute__((unused)) SIGNAL_ARGS) {
  int save_errno = errno;
  pg_atomic_write_u32(&worker_state->got_restart, 1);
  pg_write_barrier();
  if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);
  errno = save_errno;
}

static void handle_sighup(__attribute__((unused)) SIGNAL_ARGS) {
  int save_errno = errno;
  got_sighup     = true;
  if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);
  errno = save_errno;
}

/*
 *We have to handle sigusr1 explicitly because the default
 *procsignal_sigusr1_handler doesn't `SetLatch`, this would prevent
 *DROP DATATABASE from finishing since our worker would be sleeping and not reach
 *CHECK_FOR_INTERRUPTS()
 */
static void handle_sigusr1(SIGNAL_ARGS) {
  int save_errno = errno;
  if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);
  errno = save_errno;
  procsignal_sigusr1_handler(postgres_signal_arg);
}

static void publish_state(WorkerStatus s) {
  pg_atomic_write_u32(&worker_state->status, (uint32)s);
  pg_write_barrier();
  ConditionVariableBroadcast(&worker_state->cv);
}

static void net_on_exit(__attribute__((unused)) int code, __attribute__((unused)) Datum arg) {
  worker_should_restart = false;
  pg_atomic_write_u32(&worker_state->should_wake,
                      1); // ensure the remaining work will continue since we'll restart

  worker_state->shared_latch = NULL;

  ev_monitor_close(worker_state);

  curl_multi_cleanup(worker_state->curl_mhandle);
  curl_global_cleanup();
}

// wait according to the wait type while ensuring interrupts are processed while waiting
static void wait_while_processing_interrupts(WorkerWait ww, bool *should_restart) {
  switch (ww) {
  case WORKER_WAIT_NO_TIMEOUT:
    WaitLatch(worker_state->shared_latch, WL_LATCH_SET | WL_EXIT_ON_PM_DEATH, no_timeout,
              PG_WAIT_EXTENSION);
    ResetLatch(worker_state->shared_latch);
    break;
  case WORKER_WAIT_ONE_SECOND:
    WaitLatch(worker_state->shared_latch, WL_LATCH_SET | WL_TIMEOUT | WL_EXIT_ON_PM_DEATH, 1000,
              PG_WAIT_EXTENSION);
    ResetLatch(worker_state->shared_latch);
    break;
  }

  CHECK_FOR_INTERRUPTS();

  if (got_sighup) {
    got_sighup = false;
    ProcessConfigFile(PGC_SIGHUP);
  }

  if (pg_atomic_exchange_u32(&worker_state->got_restart, 0)) {
    *should_restart = true;
  }
}

static bool is_extension_locked(Oid ext_table_oids[static total_extension_tables]) {
  Oid net_oid = get_namespace_oid("net", true);

  if (!OidIsValid(net_oid)) {
    return false;
  }

  Oid queue_oid = get_relname_relid("http_request_queue", net_oid);
  Oid resp_oid  = get_relname_relid("_http_response", net_oid);

  /*
   * The "net" schema can exist without the extension tables, e.g. when another
   * extension is installed into a schema named "net". ConditionalLockRelationOid
   * doesn't validate the oid, so locking InvalidOid would succeed and the worker
   * would crash loop on the queries that follow.
   */
  if (!OidIsValid(queue_oid) || !OidIsValid(resp_oid)) {
    ereport(LOG,
            errmsg("schema \"net\" exists but the pg_net extension tables are missing, skipping "
                   "request processing"),
            errhint("The schema \"net\" might be used by another extension or be left over from a "
                    "partially dropped pg_net installation."));
    return false;
  }

  bool is_locked = ConditionalLockRelationOid(queue_oid, AccessShareLock) &&
                   ConditionalLockRelationOid(resp_oid, AccessShareLock);

  if (is_locked) {
    ext_table_oids[0] = queue_oid;
    ext_table_oids[1] = resp_oid;
  }

  return is_locked;
}

static void unlock_extension(Oid ext_table_oids[static total_extension_tables]) {
  UnlockRelationOid(ext_table_oids[0], AccessShareLock);
  UnlockRelationOid(ext_table_oids[1], AccessShareLock);
}

void pg_net_worker(__attribute__((unused)) Datum main_arg) {
  worker_state->shared_latch = &MyProc->procLatch;
  on_proc_exit(net_on_exit, 0);

  BackgroundWorkerUnblockSignals();
  pqsignal(SIGTERM, handle_sigterm);
  pqsignal(SIGHUP, handle_sighup);
  pqsignal(SIGUSR1, handle_sigusr1);

  BackgroundWorkerInitializeConnection(guc_database_name, guc_username, 0);
  pgstat_report_appname("pg_net " EXTVERSION); // set appname for pg_stat_activity

  elog(INFO,
       "pg_net worker started with a config of: pg_net.ttl=%s, pg_net.batch_size=%d, "
       "pg_net.username=%s, pg_net.database_name=%s",
       guc_ttl, guc_batch_size, guc_username, guc_database_name);

  int curl_ret = curl_global_init(CURL_GLOBAL_ALL);
  if (curl_ret != CURLE_OK)
    ereport(ERROR, errmsg("curl_global_init() returned %s\n", curl_easy_strerror(curl_ret)));

  worker_state->epfd = event_monitor();

  if (worker_state->epfd < 0) {
    ereport(ERROR, errmsg("Failed to create event monitor file descriptor"));
  }

  worker_state->curl_mhandle = curl_multi_init();
  if (!worker_state->curl_mhandle) ereport(ERROR, errmsg("curl_multi_init()"));

  set_curl_mhandle(worker_state);

  publish_state(WS_RUNNING);

  // Initial state: we go straight into the outer loop and wait for a wake.
  pgstat_report_activity(STATE_IDLE, NULL);

  do {

    uint32 expected = 1;
    if (!pg_atomic_compare_exchange_u32(&worker_state->should_wake, &expected, 0)) {
      elog(DEBUG1, "pg_net worker waiting for wake");
      wait_while_processing_interrupts(WORKER_WAIT_NO_TIMEOUT, &worker_should_restart);
      continue;
    }

    pgstat_report_activity(STATE_RUNNING, NULL);

    uint64 requests_consumed = 0;
    uint64 expired_responses = 0;

    do {
      SetCurrentStatementStartTimestamp();
      StartTransactionCommand();
      PushActiveSnapshot(GetTransactionSnapshot());

      Oid ext_table_oids[total_extension_tables];

      if (!is_extension_locked(ext_table_oids)) {
        elog(DEBUG1, "pg_net extension not loaded");
        PopActiveSnapshot();
        AbortCurrentTransaction();
        break;
      }

      SPI_connect();

      expired_responses = delete_expired_responses(guc_ttl, guc_batch_size);

      elog(DEBUG1, "Deleted " UINT64_FORMAT " expired rows", expired_responses);

      requests_consumed = consume_request_queue(guc_batch_size);

      elog(DEBUG1, "Consumed " UINT64_FORMAT " request rows", requests_consumed);

      if (requests_consumed > 0) {
        CurlHandle *handles = palloc(mul_size(sizeof(CurlHandle), requests_consumed));

        // initialize curl handles
        for (size_t j = 0; j < requests_consumed; j++) {
          init_curl_handle(&handles[j],
                           get_request_queue_row(SPI_tuptable->vals[j], SPI_tuptable->tupdesc));

          EREPORT_MULTI(curl_multi_add_handle(worker_state->curl_mhandle, handles[j].ez_handle));
        }

        // start curl event loop
        int   running_handles = 0;
        int   maxevents       = requests_consumed + 1; // 1 extra for the timer
        event events[maxevents];

        do {
          int nfds =
              wait_event(worker_state->epfd, events, maxevents, curl_handle_event_timeout_ms);

          if (nfds < 0) {
            int save_errno = errno;
            if (save_errno == EINTR) { // can happen when the wait is interrupted, for example when
                                       // running under GDB. Just continue in this case.
              elog(DEBUG1, "wait_event() got %s, continuing", strerror(save_errno));
              continue;
            } else {
              ereport(ERROR, errmsg("wait_event() failed: %s", strerror(save_errno)));
              break;
            }
          }

          for (int i = 0; i < nfds; i++) {
            if (is_timer(events[i])) {
              EREPORT_MULTI(curl_multi_socket_action(worker_state->curl_mhandle,
                                                     CURL_SOCKET_TIMEOUT, 0, &running_handles));
            } else {
              int curl_event = get_curl_event(events[i]);
              int sockfd     = get_socket_fd(events[i]);

              EREPORT_MULTI(curl_multi_socket_action(worker_state->curl_mhandle, sockfd, curl_event,
                                                     &running_handles));
            }
          }

          // insert finished responses
          CURLMsg *msg       = NULL;
          int      msgs_left = 0;
          while ((msg = curl_multi_info_read(worker_state->curl_mhandle, &msgs_left))) {
            if (msg->msg == CURLMSG_DONE) {
              CurlHandle *handle = NULL;
              EREPORT_CURL_GETINFO(msg->easy_handle, CURLINFO_PRIVATE, &handle);
              insert_response(handle, msg->data.result);
            } else {
              ereport(ERROR, errmsg("curl_multi_info_read(), CURLMsg=%d\n", msg->msg));
            }
          }

          elog(DEBUG1, "Pending curl running_handles: %d", running_handles);
          // run while there are curl handles, some won't finish in a single iteration since they
          // could be slow and waiting for a timeout
        } while (running_handles > 0);

        // cleanup
        for (uint64 i = 0; i < requests_consumed; i++) {
          EREPORT_MULTI(curl_multi_remove_handle(worker_state->curl_mhandle, handles[i].ez_handle));

          curl_easy_cleanup(handles[i].ez_handle);

          pfree_handle(&handles[i]);
        }

        pfree(handles);
      }

      SPI_finish();

      unlock_extension(ext_table_oids);

      PopActiveSnapshot();
      CommitTransactionCommand();

      // Background workers that modify tables must flush their pending
      // pgstat counters themselves. Regular user backends do this
      // automatically after each query via the main loop in
      // tcop/postgres.c; background workers have no equivalent. Without
      // this call, per-write counters (n_tup_ins, n_tup_del,
      // n_mod_since_analyze) for the worker's writes never reach shared
      // stats.
      pgstat_report_stat(false);

      // slow down queue processing to avoid using too much CPU
      wait_while_processing_interrupts(WORKER_WAIT_ONE_SECOND, &worker_should_restart);

    } while (!worker_should_restart && (requests_consumed > 0 || expired_responses > 0));

    // Inner loop drained; back to waiting for the next wake.
    pgstat_report_activity(STATE_IDLE, NULL);

  } while (!worker_should_restart);

  publish_state(WS_EXITED);

  // causing a failure on exit will make the postmaster process restart the bg worker
  proc_exit(EXIT_FAILURE);
}

static Size net_memsize(void) {
  return MAXALIGN(sizeof(WorkerState));
}

#if PG15_GTE
static void net_shmem_request(void) {
  if (prev_shmem_request_hook) prev_shmem_request_hook();

  RequestAddinShmemSpace(net_memsize());
}
#endif

static void net_shmem_startup(void) {
  if (prev_shmem_startup_hook) prev_shmem_startup_hook();

  bool found;

  LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);

  worker_state = ShmemInitStruct("pg_net worker state", sizeof(WorkerState), &found);

  if (!found) {
    pg_atomic_init_u32(&worker_state->got_restart, 0);
    pg_atomic_init_u32(&worker_state->status, WS_NOT_YET);
    pg_atomic_init_u32(&worker_state->should_wake, 1);
    worker_state->shared_latch = NULL;

    ConditionVariableInit(&worker_state->cv);
    worker_state->epfd         = 0;
    worker_state->curl_mhandle = NULL;
  }

  LWLockRelease(AddinShmemInitLock);
}

void _PG_init(void) {
  if (IsBinaryUpgrade) {
    return;
  }

  if (!process_shared_preload_libraries_in_progress) {
    ereport(ERROR, errmsg("pg_net is not in shared_preload_libraries"),
            errhint("Add pg_net to the shared_preload_libraries "
                    "configuration variable in postgresql.conf."));
  }

  RegisterBackgroundWorker(&(BackgroundWorker){
    .bgw_flags         = BGWORKER_SHMEM_ACCESS | BGWORKER_BACKEND_DATABASE_CONNECTION,
    .bgw_start_time    = BgWorkerStart_RecoveryFinished,
    .bgw_library_name  = "pg_net",
    .bgw_function_name = "pg_net_worker",
    .bgw_name          = "pg_net " EXTVERSION " worker",
    .bgw_restart_time  = net_worker_restart_time_sec,
  });

#if PG15_GTE
  prev_shmem_request_hook = shmem_request_hook;
  shmem_request_hook      = net_shmem_request;
#else
  RequestAddinShmemSpace(net_memsize());
#endif

  prev_shmem_startup_hook = shmem_startup_hook;
  shmem_startup_hook      = net_shmem_startup;

  DefineCustomStringVariable("pg_net.ttl", "time to live for request/response rows",
                             "should be a valid interval type", &guc_ttl, "6 hours", PGC_SIGHUP, 0,
                             NULL, NULL, NULL);

  DefineCustomIntVariable(
      "pg_net.batch_size", "number of requests executed in one iteration of the background worker",
      NULL, &guc_batch_size, 200, 0, PG_INT16_MAX, PGC_SIGHUP, 0, NULL, NULL, NULL);

  DefineCustomStringVariable("pg_net.database_name", "Database where the worker will connect to",
                             NULL, &guc_database_name, "postgres", PGC_SU_BACKEND, 0, NULL, NULL,
                             NULL);

  DefineCustomStringVariable("pg_net.username", "Connection user for the worker", NULL,
                             &guc_username, NULL, PGC_SU_BACKEND, 0, NULL, NULL, NULL);
}

1	#include <errno.h>
2	#include <inttypes.h>
3	#include <string.h>
4	#include <unistd.h>
5
6	#define PG_PRELUDE_IMPL
7	#include "pg_prelude.h"
8
9	#include "curl_prelude.h"
10
11	#include "core.h"
12	#include "errors.h"
13	#include "event.h"
14	#include "util.h"
15
16	#define MIN_LIBCURL_VERSION_NUM \
17	0x075300 // This is the 7.83.0 version in hex as defined in curl/curlver.h
18	#define REQUIRED_LIBCURL_ERR_MSG \
19	"libcurl >= 7.83.0 is required, we use the curl_easy_nextheader() function added in this " \
20	"version"
21	_Static_assert(LIBCURL_VERSION_NUM,
22	REQUIRED_LIBCURL_ERR_MSG); // test for older libcurl versions that don't even have
23	// LIBCURL_VERSION_NUM defined (e.g. libcurl 6.5).
24	_Static_assert(LIBCURL_VERSION_NUM >= MIN_LIBCURL_VERSION_NUM, REQUIRED_LIBCURL_ERR_MSG);
25
26	PG_MODULE_MAGIC;	111✔
27
28	typedef enum {
29	WORKER_WAIT_NO_TIMEOUT,
30	WORKER_WAIT_ONE_SECOND,
31	} WorkerWait;
32
33	static WorkerState *worker_state = NULL;
34
35	static const int curl_handle_event_timeout_ms = 1000;
36	static const int net_worker_restart_time_sec = 1;
37	static const long no_timeout = -1L;
38	static bool wake_commit_cb_active = false;
39	static bool worker_should_restart = false;
40	static const size_t total_extension_tables = 2;
41
42	static char *guc_ttl;
43	static int guc_batch_size;
44	static char *guc_database_name;
45	static char *guc_username;
46
47	#if PG15_GTE
48	static shmem_request_hook_type prev_shmem_request_hook = NULL;
49	#endif
50
51	static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
52	static volatile sig_atomic_t got_sighup = false;
53
54	void _PG_init(void);
55
56	#if PG_VERSION_NUM >= 180000
57	PGDLLEXPORT pg_noreturn void pg_net_worker(Datum main_arg);
58	#else
59	PGDLLEXPORT void pg_net_worker(Datum main_arg) pg_attribute_noreturn();
60	#endif
61
62	PG_FUNCTION_INFO_V1(worker_restart);	71✔
63	Datum worker_restart(__attribute__((unused)) PG_FUNCTION_ARGS) {	17✔
64	bool result = DatumGetBool(DirectFunctionCall1(pg_reload_conf, (Datum)NULL)); // reload the config	17✔
65	pg_atomic_write_u32(&worker_state->got_restart, 1);	17✔
66	pg_write_barrier();	17✔
67	if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);	17✔
68	PG_RETURN_BOOL(result); // TODO is not necessary to return a bool here, but we do it to maintain	17✔
69	// backward compatibility
70	}
71
72	static void wait_until_state(WorkerState *ws, WorkerStatus expected_status) {	21✔
73	if (pg_atomic_read_u32(&ws->status) ==	21✔
74	expected_status) // fast return without sleeping, in case condition is fulfilled
75	return;
76
77	ConditionVariablePrepareToSleep(&ws->cv);	16✔
78	while (pg_atomic_read_u32(&ws->status) != expected_status) {	32✔
79	ConditionVariableSleep(&ws->cv, PG_WAIT_EXTENSION);	16✔
80	}
81	ConditionVariableCancelSleep();	16✔
82	}
83
84	PG_FUNCTION_INFO_V1(wait_until_running);	75✔
85	Datum wait_until_running(__attribute__((unused)) PG_FUNCTION_ARGS) {	21✔
86	wait_until_state(worker_state, WS_RUNNING);	21✔
87
88	PG_RETURN_VOID();	21✔
89	}
90
91	// only wake at commit time to prevent excessive and unnecessary wakes.
92	// e.g only one wake when doing `select
93	// net.http_get('http://localhost:8080/pathological?status=200') from generate_series(1,100000);`
94	static void wake_at_commit(XactEvent event, __attribute__((unused)) void *arg) {	308✔
95	elog(DEBUG2, "pg_net xact callback received: %s", xact_event_name(event));	308✔
96
97	switch (event) {	308✔
98	case XACT_EVENT_COMMIT:	127✔
99	case XACT_EVENT_PARALLEL_COMMIT:
100	if (wake_commit_cb_active) {	127✔
101	uint32 expected = 0;	51✔
102	bool success = pg_atomic_compare_exchange_u32(&worker_state->should_wake, &expected, 1);	51✔
103	pg_write_barrier();	51✔
104
105	if (success) // only wake the worker on first put, so if many concurrent wakes come we only	51✔
106	// wake once
107	SetLatch(worker_state->shared_latch);	25✔
108
109	wake_commit_cb_active = false;	51✔
110	}
111	break;
112	// TODO: `PREPARE TRANSACTION 'xx';` and `COMMIT PREPARED TRANSACTION 'xx';` do not wake the
113	// worker automatically, they require a manual `net.wake()` These are disabled by default and
114	// rarely used, see `max_prepared_transactions`
115	// https://www.postgresql.org/docs/17/runtime-config-resource.html#GUC-MAX-PREPARED-TRANSACTIONS
116	case XACT_EVENT_PREPARE:	54✔
117	// abort the callback on rollback
118	case XACT_EVENT_ABORT:
119	case XACT_EVENT_PARALLEL_ABORT: wake_commit_cb_active = false; break;	54✔
120	default : break;
121	}
122	}	308✔
123
124	PG_FUNCTION_INFO_V1(wake);	109✔
125	Datum wake(__attribute__((unused)) PG_FUNCTION_ARGS) {	495✔
126	if (!wake_commit_cb_active) { // register only one callback per transaction	495✔
127	RegisterXactCallback(wake_at_commit, NULL);	58✔
128	wake_commit_cb_active = true;	58✔
129	}
130
131	PG_RETURN_VOID();	495✔
132	}
133
134	static void handle_sigterm(__attribute__((unused)) SIGNAL_ARGS) {	3✔
135	int save_errno = errno;	3✔
136	pg_atomic_write_u32(&worker_state->got_restart, 1);	3✔
137	pg_write_barrier();	3✔
138	if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);	3✔
139	errno = save_errno;	3✔
140	}	3✔
141
142	static void handle_sighup(__attribute__((unused)) SIGNAL_ARGS) {	20✔
143	int save_errno = errno;	20✔
144	got_sighup = true;	20✔
145	if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);	20✔
146	errno = save_errno;	20✔
147	}	20✔
148
149	/*
150	*We have to handle sigusr1 explicitly because the default
151	*procsignal_sigusr1_handler doesn't `SetLatch`, this would prevent
152	*DROP DATATABASE from finishing since our worker would be sleeping and not reach
153	*CHECK_FOR_INTERRUPTS()
154	*/
155	static void handle_sigusr1(SIGNAL_ARGS) {	4✔
156	int save_errno = errno;	4✔
157	if (worker_state->shared_latch) SetLatch(worker_state->shared_latch);	4✔
158	errno = save_errno;	4✔
159	procsignal_sigusr1_handler(postgres_signal_arg);	4✔
160	}	4✔
161
162	static void publish_state(WorkerStatus s) {	40✔
163	pg_atomic_write_u32(&worker_state->status, (uint32)s);	40✔
164	pg_write_barrier();	40✔
165	ConditionVariableBroadcast(&worker_state->cv);	40✔
166	}	40✔
167
168	static void net_on_exit(__attribute__((unused)) int code, __attribute__((unused)) Datum arg) {	20✔
169	worker_should_restart = false;	20✔
170	pg_atomic_write_u32(&worker_state->should_wake,	20✔
171	1); // ensure the remaining work will continue since we'll restart
172
173	worker_state->shared_latch = NULL;	20✔
174
175	ev_monitor_close(worker_state);	20✔
176
177	curl_multi_cleanup(worker_state->curl_mhandle);	20✔
178	curl_global_cleanup();	20✔
179	}	20✔
180
181	// wait according to the wait type while ensuring interrupts are processed while waiting
182	static void wait_while_processing_interrupts(WorkerWait ww, bool *should_restart) {	117✔
183	switch (ww) {	117✔
184	case WORKER_WAIT_NO_TIMEOUT:	31✔
185	WaitLatch(worker_state->shared_latch, WL_LATCH_SET \| WL_EXIT_ON_PM_DEATH, no_timeout,	31✔
186	PG_WAIT_EXTENSION);
187	ResetLatch(worker_state->shared_latch);	31✔
188	break;	31✔
189	case WORKER_WAIT_ONE_SECOND:	86✔
190	WaitLatch(worker_state->shared_latch, WL_LATCH_SET \| WL_TIMEOUT \| WL_EXIT_ON_PM_DEATH, 1000,	86✔
191	PG_WAIT_EXTENSION);
192	ResetLatch(worker_state->shared_latch);	86✔
193	break;	86✔
194	}
195
196	CHECK_FOR_INTERRUPTS();	117✔
197
198	if (got_sighup) {	117✔
199	got_sighup = false;	4✔
200	ProcessConfigFile(PGC_SIGHUP);	4✔
201	}
202
203	if (pg_atomic_exchange_u32(&worker_state->got_restart, 0)) {	117✔
204	*should_restart = true;	20✔
205	}
206	}	117✔
207
208	static bool is_extension_locked(Oid ext_table_oids[static total_extension_tables]) {	90✔
209	Oid net_oid = get_namespace_oid("net", true);	90✔
210
211	if (!OidIsValid(net_oid)) {	90✔
212	return false;
213	}
214
215	Oid queue_oid = get_relname_relid("http_request_queue", net_oid);	87✔
216	Oid resp_oid = get_relname_relid("_http_response", net_oid);	87✔
217
218	/*
219	* The "net" schema can exist without the extension tables, e.g. when another
220	* extension is installed into a schema named "net". ConditionalLockRelationOid
221	* doesn't validate the oid, so locking InvalidOid would succeed and the worker
222	* would crash loop on the queries that follow.
223	*/
224	if (!OidIsValid(queue_oid) \|\| !OidIsValid(resp_oid)) {	87✔
NEW 225	ereport(LOG,	×
226	errmsg("schema \"net\" exists but the pg_net extension tables are missing, skipping "
227	"request processing"),
228	errhint("The schema \"net\" might be used by another extension or be left over from a "
229	"partially dropped pg_net installation."));
NEW 230	return false;	×
231	}
232
233	bool is_locked = ConditionalLockRelationOid(queue_oid, AccessShareLock) &&	173✔
234	ConditionalLockRelationOid(resp_oid, AccessShareLock);	86✔
235
236	if (is_locked) {	86✔
237	ext_table_oids[0] = queue_oid;	86✔
238	ext_table_oids[1] = resp_oid;	86✔
239	}
240
241	return is_locked;
242	}
243
244	static void unlock_extension(Oid ext_table_oids[static total_extension_tables]) {	86✔
245	UnlockRelationOid(ext_table_oids[0], AccessShareLock);	86✔
246	UnlockRelationOid(ext_table_oids[1], AccessShareLock);	86✔
247	}	86✔
248
249	void pg_net_worker(__attribute__((unused)) Datum main_arg) {	20✔
250	worker_state->shared_latch = &MyProc->procLatch;	20✔
251	on_proc_exit(net_on_exit, 0);	20✔
252
253	BackgroundWorkerUnblockSignals();	20✔
254	pqsignal(SIGTERM, handle_sigterm);	20✔
255	pqsignal(SIGHUP, handle_sighup);	20✔
256	pqsignal(SIGUSR1, handle_sigusr1);	20✔
257
258	BackgroundWorkerInitializeConnection(guc_database_name, guc_username, 0);	20✔
259	pgstat_report_appname("pg_net " EXTVERSION); // set appname for pg_stat_activity	20✔
260
261	elog(INFO,	20✔
262	"pg_net worker started with a config of: pg_net.ttl=%s, pg_net.batch_size=%d, "
263	"pg_net.username=%s, pg_net.database_name=%s",
264	guc_ttl, guc_batch_size, guc_username, guc_database_name);
265
266	int curl_ret = curl_global_init(CURL_GLOBAL_ALL);	20✔
267	if (curl_ret != CURLE_OK)	20✔
268	ereport(ERROR, errmsg("curl_global_init() returned %s\n", curl_easy_strerror(curl_ret)));	×
269
270	worker_state->epfd = event_monitor();	20✔
271
272	if (worker_state->epfd < 0) {	20✔
273	ereport(ERROR, errmsg("Failed to create event monitor file descriptor"));	×
274	}
275
276	worker_state->curl_mhandle = curl_multi_init();	20✔
277	if (!worker_state->curl_mhandle) ereport(ERROR, errmsg("curl_multi_init()"));	20✔
278
279	set_curl_mhandle(worker_state);	20✔
280
281	publish_state(WS_RUNNING);	20✔
282
283	// Initial state: we go straight into the outer loop and wait for a wake.
284	pgstat_report_activity(STATE_IDLE, NULL);	20✔
285
286	do {	72✔
287
288	uint32 expected = 1;	72✔
289	if (!pg_atomic_compare_exchange_u32(&worker_state->should_wake, &expected, 0)) {	72✔
290	elog(DEBUG1, "pg_net worker waiting for wake");	31✔
291	wait_while_processing_interrupts(WORKER_WAIT_NO_TIMEOUT, &worker_should_restart);	31✔
292	continue;	31✔
293	}
294
295	pgstat_report_activity(STATE_RUNNING, NULL);	41✔
296
297	uint64 requests_consumed = 0;	41✔
298	uint64 expired_responses = 0;	41✔
299
300	do {	90✔
301	SetCurrentStatementStartTimestamp();	90✔
302	StartTransactionCommand();	90✔
303	PushActiveSnapshot(GetTransactionSnapshot());	90✔
304
305	Oid ext_table_oids[total_extension_tables];	90✔
306
307	if (!is_extension_locked(ext_table_oids)) {	90✔
308	elog(DEBUG1, "pg_net extension not loaded");	4✔
309	PopActiveSnapshot();	4✔
310	AbortCurrentTransaction();	4✔
311	break;	4✔
312	}
313
314	SPI_connect();	86✔
315
316	expired_responses = delete_expired_responses(guc_ttl, guc_batch_size);	86✔
317
318	elog(DEBUG1, "Deleted " UINT64_FORMAT " expired rows", expired_responses);	86✔
319
320	requests_consumed = consume_request_queue(guc_batch_size);	86✔
321
322	elog(DEBUG1, "Consumed " UINT64_FORMAT " request rows", requests_consumed);	86✔
323
324	if (requests_consumed > 0) {	86✔
325	CurlHandle *handles = palloc(mul_size(sizeof(CurlHandle), requests_consumed));	54✔
326
327	// initialize curl handles
328	for (size_t j = 0; j < requests_consumed; j++) {	423✔
329	init_curl_handle(&handles[j],	369✔
330	get_request_queue_row(SPI_tuptable->vals[j], SPI_tuptable->tupdesc));	369✔
331
332	EREPORT_MULTI(curl_multi_add_handle(worker_state->curl_mhandle, handles[j].ez_handle));	369✔
333	}
334
335	// start curl event loop
336	int running_handles = 0;	54✔
337	int maxevents = requests_consumed + 1; // 1 extra for the timer	54✔
338	event events[maxevents];	54✔
339
340	do {	179✔
341	int nfds =	179✔
342	wait_event(worker_state->epfd, events, maxevents, curl_handle_event_timeout_ms);	179✔
343
344	if (nfds < 0) {	179✔
345	int save_errno = errno;	×
346	if (save_errno == EINTR) { // can happen when the wait is interrupted, for example when	×
347	// running under GDB. Just continue in this case.
348	elog(DEBUG1, "wait_event() got %s, continuing", strerror(save_errno));	×
349	continue;	×
350	} else {
351	ereport(ERROR, errmsg("wait_event() failed: %s", strerror(save_errno)));	×
352	break;
353	}
354	}
355
356	for (int i = 0; i < nfds; i++) {	1,169✔
357	if (is_timer(events[i])) {	990✔
358	EREPORT_MULTI(curl_multi_socket_action(worker_state->curl_mhandle,	65✔
359	CURL_SOCKET_TIMEOUT, 0, &running_handles));
360	} else {
361	int curl_event = get_curl_event(events[i]);	925✔
362	int sockfd = get_socket_fd(events[i]);	925✔
363
364	EREPORT_MULTI(curl_multi_socket_action(worker_state->curl_mhandle, sockfd, curl_event,	990✔
365	&running_handles));
366	}
367	}
368
369	// insert finished responses
370	CURLMsg *msg = NULL;	179✔
371	int msgs_left = 0;	179✔
372	while ((msg = curl_multi_info_read(worker_state->curl_mhandle, &msgs_left))) {	727✔
373	if (msg->msg == CURLMSG_DONE) {	369✔
374	CurlHandle *handle = NULL;	369✔
375	EREPORT_CURL_GETINFO(msg->easy_handle, CURLINFO_PRIVATE, &handle);	369✔
376	insert_response(handle, msg->data.result);	369✔
377	} else {
378	ereport(ERROR, errmsg("curl_multi_info_read(), CURLMsg=%d\n", msg->msg));	548✔
379	}
380	}
381
382	elog(DEBUG1, "Pending curl running_handles: %d", running_handles);	179✔
383	// run while there are curl handles, some won't finish in a single iteration since they
384	// could be slow and waiting for a timeout
385	} while (running_handles > 0);	179✔
386
387	// cleanup
388	for (uint64 i = 0; i < requests_consumed; i++) {	423✔
389	EREPORT_MULTI(curl_multi_remove_handle(worker_state->curl_mhandle, handles[i].ez_handle));	369✔
390
391	curl_easy_cleanup(handles[i].ez_handle);	369✔
392
393	pfree_handle(&handles[i]);	369✔
394	}
395
396	pfree(handles);	54✔
397	}
398
399	SPI_finish();	86✔
400
401	unlock_extension(ext_table_oids);	86✔
402
403	PopActiveSnapshot();	86✔
404	CommitTransactionCommand();	86✔
405
406	// Background workers that modify tables must flush their pending
407	// pgstat counters themselves. Regular user backends do this
408	// automatically after each query via the main loop in
409	// tcop/postgres.c; background workers have no equivalent. Without
410	// this call, per-write counters (n_tup_ins, n_tup_del,
411	// n_mod_since_analyze) for the worker's writes never reach shared
412	// stats.
413	pgstat_report_stat(false);	86✔
414
415	// slow down queue processing to avoid using too much CPU
416	wait_while_processing_interrupts(WORKER_WAIT_ONE_SECOND, &worker_should_restart);	86✔
417
418	} while (!worker_should_restart && (requests_consumed > 0 \|\| expired_responses > 0));	86✔
419
420	// Inner loop drained; back to waiting for the next wake.
421	pgstat_report_activity(STATE_IDLE, NULL);	41✔
422
423	} while (!worker_should_restart);	72✔
424
425	publish_state(WS_EXITED);	20✔
426
427	// causing a failure on exit will make the postmaster process restart the bg worker
428	proc_exit(EXIT_FAILURE);	20✔
429	}
430
431	static Size net_memsize(void) {
432	return MAXALIGN(sizeof(WorkerState));
433	}
434
435	#if PG15_GTE
436	static void net_shmem_request(void) {	111✔
437	if (prev_shmem_request_hook) prev_shmem_request_hook();	111✔
438
439	RequestAddinShmemSpace(net_memsize());	111✔
440	}	111✔
441	#endif
442
443	static void net_shmem_startup(void) {	111✔
444	if (prev_shmem_startup_hook) prev_shmem_startup_hook();	111✔
445
446	bool found;	111✔
447
448	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);	111✔
449
450	worker_state = ShmemInitStruct("pg_net worker state", sizeof(WorkerState), &found);	111✔
451
452	if (!found) {	111✔
453	pg_atomic_init_u32(&worker_state->got_restart, 0);	111✔
454	pg_atomic_init_u32(&worker_state->status, WS_NOT_YET);	111✔
455	pg_atomic_init_u32(&worker_state->should_wake, 1);	111✔
456	worker_state->shared_latch = NULL;	111✔
457
458	ConditionVariableInit(&worker_state->cv);	111✔
459	worker_state->epfd = 0;	111✔
460	worker_state->curl_mhandle = NULL;	111✔
461	}
462
463	LWLockRelease(AddinShmemInitLock);	111✔
464	}	111✔
465
466	void _PG_init(void) {	111✔
467	if (IsBinaryUpgrade) {	111✔
468	return;	×
469	}
470
471	if (!process_shared_preload_libraries_in_progress) {	111✔
472	ereport(ERROR, errmsg("pg_net is not in shared_preload_libraries"),	×
473	errhint("Add pg_net to the shared_preload_libraries "
474	"configuration variable in postgresql.conf."));
475	}
476
477	RegisterBackgroundWorker(&(BackgroundWorker){	111✔
478	.bgw_flags = BGWORKER_SHMEM_ACCESS \| BGWORKER_BACKEND_DATABASE_CONNECTION,
479	.bgw_start_time = BgWorkerStart_RecoveryFinished,
480	.bgw_library_name = "pg_net",
481	.bgw_function_name = "pg_net_worker",
482	.bgw_name = "pg_net " EXTVERSION " worker",
483	.bgw_restart_time = net_worker_restart_time_sec,
484	});
485
486	#if PG15_GTE
487	prev_shmem_request_hook = shmem_request_hook;	111✔
488	shmem_request_hook = net_shmem_request;	111✔
489	#else
490	RequestAddinShmemSpace(net_memsize());
491	#endif
492
493	prev_shmem_startup_hook = shmem_startup_hook;	111✔
494	shmem_startup_hook = net_shmem_startup;	111✔
495
496	DefineCustomStringVariable("pg_net.ttl", "time to live for request/response rows",	111✔
497	"should be a valid interval type", &guc_ttl, "6 hours", PGC_SIGHUP, 0,
498	NULL, NULL, NULL);
499
500	DefineCustomIntVariable(	111✔
501	"pg_net.batch_size", "number of requests executed in one iteration of the background worker",
502	NULL, &guc_batch_size, 200, 0, PG_INT16_MAX, PGC_SIGHUP, 0, NULL, NULL, NULL);
503
504	DefineCustomStringVariable("pg_net.database_name", "Database where the worker will connect to",	111✔
505	NULL, &guc_database_name, "postgres", PGC_SU_BACKEND, 0, NULL, NULL,
506	NULL);
507
508	DefineCustomStringVariable("pg_net.username", "Connection user for the worker", NULL,	111✔
509	&guc_username, NULL, PGC_SU_BACKEND, 0, NULL, NULL, NULL);
510	}

supabase / pg_net / 28687656621

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous