diff --git a/.gitignore b/.gitignore
index 566f026..e066fb5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,6 @@
-.deps
 *.o
 *.so
-/results
-*pg_wait_sampling--1.1.sql
-.log
+/.deps/
+/log/
+/results/
+/tmp_check/
diff --git a/.travis.yml b/.travis.yml
index 4038ea8..2c63eff 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,25 +1,35 @@
-os:
-    - linux
-
-sudo: required
-dist: trusty
-
+dist: jammy
 language: c
-
-services:
-    - docker
-
-install:
-    - sed -e 's/${CHECK_CODE}/'${CHECK_CODE}/g -e 's/${PG_VERSION}/'${PG_VERSION}/g Dockerfile.tmpl > Dockerfile
-    - docker-compose build
-
-script:
-    - docker-compose run tests
-
 env:
-    - PG_VERSION=9.6 CHECK_CODE=clang
-    - PG_VERSION=9.6 CHECK_CODE=cppcheck
-    - PG_VERSION=9.6 CHECK_CODE=false
-    - PG_VERSION=10  CHECK_CODE=clang
-    - PG_VERSION=10  CHECK_CODE=cppcheck
-    - PG_VERSION=10  CHECK_CODE=false
+- PG_MAJOR=18 SNAPSHOT=1
+- PG_MAJOR=17
+- PG_MAJOR=16
+- PG_MAJOR=15
+- PG_MAJOR=14
+- PG_MAJOR=13
+before_script:
+- curl https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add -
+- |
+  if [ -n "${SNAPSHOT}" ]; then
+    echo "deb https://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg-snapshot main ${PG_MAJOR}"
+  elif [ -n "${BETA}" ]; then
+    echo "deb https://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg main ${PG_MAJOR}"
+  else
+    echo "deb https://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg main"
+  fi | sudo tee /etc/apt/sources.list.d/postgresql.list
+- |
+  if [ -n "${SNAPSHOT}" ]; then
+    {
+      echo "Package: *"
+      echo "Pin: origin apt.postgresql.org"
+      echo "Pin-Priority: 600"
+    } | sudo tee /etc/apt/preferences.d/pgdg.pref
+  fi
+- sudo apt-get update
+- sudo systemctl stop postgresql
+- sudo apt-get install -y --no-install-recommends postgresql-client-${PG_MAJOR} postgresql-${PG_MAJOR} postgresql-server-dev-${PG_MAJOR}
+- sudo systemctl stop postgresql
+script: ./run-tests.sh
+after_script:
+- cat regression.diffs
+- cat logfile
diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
deleted file mode 100644
index 5e95748..0000000
--- a/Dockerfile.tmpl
+++ /dev/null
@@ -1,29 +0,0 @@
-FROM postgres:${PG_VERSION}-alpine
-
-ENV LANG=C.UTF-8 PGDATA=/pg/data
-
-RUN if [ "${CHECK_CODE}" = "clang" ] ; then \
-	echo 'http://dl-3.alpinelinux.org/alpine/edge/main' > /etc/apk/repositories; \
-	apk --no-cache add clang-analyzer make musl-dev gcc; \
-	fi
-
-RUN if [ "${CHECK_CODE}" = "cppcheck" ] ; then \
-	apk --no-cache add cppcheck --repository http://dl-cdn.alpinelinux.org/alpine/v3.6/community; \
-	fi
-
-RUN if [ "${CHECK_CODE}" = "false" ] ; then \
-	echo 'http://dl-3.alpinelinux.org/alpine/edge/main' > /etc/apk/repositories; \
-	apk --no-cache add curl python3 gcc make musl-dev;\
-	fi
-
-RUN mkdir -p /pg/data && \
-	mkdir /pg/src && \
-	chown postgres:postgres ${PGDATA} && \
-	chmod a+rwx /usr/local/lib/postgresql && \
-	chmod a+rwx /usr/local/share/postgresql/extension
-
-ADD . /pg/src
-WORKDIR /pg/src
-RUN chmod -R go+rwX /pg/src
-USER postgres
-ENTRYPOINT PGDATA=${PGDATA} CHECK_CODE=${CHECK_CODE} bash run_tests.sh
diff --git a/LICENSE b/LICENSE
index f4b38be..46c4b8f 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,7 +1,7 @@
 pg_wait_sampling is released under the PostgreSQL License, a liberal Open Source license, similar to the BSD or MIT licenses.
 
-Copyright (c) 2015-2017, Postgres Professional
-Portions Copyright (c) 1996-2017, PostgreSQL Global Development Group
+Copyright (c) 2015-2025, Postgres Professional
+Portions Copyright (c) 1996-2025, PostgreSQL Global Development Group
 Portions Copyright (c) 1994, The Regents of the University of California
 
 Permission to use, copy, modify, and distribute this software and its documentation for any purpose, without fee, and without a written agreement is hereby granted, provided that the above copyright notice and this paragraph and the following two paragraphs appear in all copies.
diff --git a/META.json b/META.json
deleted file mode 100644
index c5a2a51..0000000
--- a/META.json
+++ /dev/null
@@ -1,51 +0,0 @@
-{
-   "name": "pg_wait_sampling",
-   "abstract": "Sampling based statistics of wait events",
-   "description": "pg_wait_sampling provides functions for detailed per backend and per query statistics about PostgreSQL wait events",
-   "version": "1.1.0",
-   "maintainer": [
-      "Alexander Korotkov <a.korotkov@postgrespro.ru>",
-      "Ildus Kurbangaliev <i.kurbangaliev@gmail.com>"
-   ],
-   "license": {
-      "PostgreSQL": "http://www.postgresql.org/about/licence"
-   },
-   "prereqs": {
-      "runtime": {
-         "requires": {
-            "PostgreSQL": "9.6.0"
-         }
-      }
-   },
-   "provides": {
-     "pg_wait_sampling": {
-       "file": "pg_wait_sampling--1.1.sql",
-       "docfile": "README.md",
-       "version": "1.1.0",
-       "abstract": "Sampling based statistics of wait events"
-     }
-   },
-   "resources": {
-      "bugtracker": {
-         "web": "https://github.com/postgrespro/pg_wait_sampling/issues"
-      },
-      "repository": {
-        "url":  "https://github.com/postgrespro/pg_wait_sampling.git",
-        "web":  "https://github.com/postgrespro/pg_wait_sampling",
-        "type": "git"
-      }
-   },
-   "generated_by": "Ildus Kurbangaliev",
-   "meta-spec": {
-      "version": "1.0.0",
-      "url": "http://pgxn.org/meta/spec.txt"
-   },
-   "tags": [
-      "waits",
-      "sampling",
-      "background worker",
-      "wait events",
-      "waits history",
-      "waits profile"
-   ]
-}
diff --git a/Makefile b/Makefile
index c698365..32711a3 100644
--- a/Makefile
+++ b/Makefile
@@ -4,14 +4,11 @@ MODULE_big = pg_wait_sampling
 OBJS = pg_wait_sampling.o collector.o
 
 EXTENSION = pg_wait_sampling
-EXTVERSION = 1.1
-DATA_built = pg_wait_sampling--$(EXTVERSION).sql
-DATA = pg_wait_sampling--1.0--1.1.sql
+DATA = pg_wait_sampling--1.1.sql pg_wait_sampling--1.0--1.1.sql
 
 REGRESS = load queries
 
 EXTRA_REGRESS_OPTS=--temp-config=$(top_srcdir)/$(subdir)/conf.add
-EXTRA_CLEAN = pg_wait_sampling--$(EXTVERSION).sql
 
 ifdef USE_PGXS
 PG_CONFIG = pg_config
@@ -23,6 +20,3 @@ top_builddir = ../..
 include $(top_builddir)/src/Makefile.global
 include $(top_srcdir)/contrib/contrib-global.mk
 endif
-
-$(EXTENSION)--$(EXTVERSION).sql: setup.sql
-	cat $^ > $@
diff --git a/README.md b/README.md
index ee36bdf..bbdbd20 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,4 @@
-[![Build Status](https://travis-ci.org/postgrespro/pg_wait_sampling.svg?branch=master)](https://travis-ci.org/postgrespro/pg_wait_sampling)
-[![PGXN version](https://badge.fury.io/pg/pg_wait_sampling.svg)](https://badge.fury.io/pg/pg_wait_sampling)
+[![Build Status](https://app.travis-ci.com/postgrespro/pg_wait_sampling.svg?branch=master)](https://app.travis-ci.com/postgrespro/pg_wait_sampling)
 [![GitHub license](https://img.shields.io/badge/license-PostgreSQL-blue.svg)](https://raw.githubusercontent.com/postgrespro/pg_wait_sampling/master/LICENSE)
 
 `pg_wait_sampling` – sampling based statistics of wait events
@@ -8,7 +7,7 @@
 Introduction
 ------------
 
-PostgreSQL 9.6+ provides an information about current wait event of particular
+PostgreSQL provides information about current wait event of particular
 process.  However, in order to gather descriptive statistics of server
 behavior user have to sample current wait event multiple times.
 `pg_wait_sampling` is an extension for collecting sampling statistics of wait
@@ -19,6 +18,10 @@ The module must be loaded by adding `pg_wait_sampling` to
 shared memory and launches background worker.  This means that a server restart
 is needed to add or remove the module.
 
+When used with `pg_stat_statements` it is recommended to put `pg_stat_statements`
+before `pg_wait_sampling` in `shared_preload_libraries` so queryIds of
+utility statements are not rewritten by the former.
+
 When `pg_wait_sampling` is enabled, it collects two kinds of statistics.
 
  * History of waits events.  It's implemented as in-memory ring buffer where
@@ -48,15 +51,21 @@ PostgreSQL installation. It is available from
 [github](https://github.com/postgrespro/pg_wait_sampling)
 under the same license as
 [PostgreSQL](http://www.postgresql.org/about/licence/)
-and supports PostgreSQL 9.6+.
+and supports PostgreSQL 13+.
 
 Installation
 ------------
 
-`pg_wait_sampling` is PostgreSQL extension which requires PostgreSQL 9.6 or
+Pre-built `pg_wait_sampling` packages are provided in official PostgreSQL
+repository: https://download.postgresql.org/pub/repos/
+
+Manual build
+------------
+
+`pg_wait_sampling` is PostgreSQL extension which requires PostgreSQL 13 or
 higher. Before build and install you should ensure following:
 
- * PostgreSQL version is 9.6 or higher.
+ * PostgreSQL version is 13 or higher.
  * You have development package of PostgreSQL installed or you built
    PostgreSQL from source.
  * Your PATH variable is configured so that `pg_config` command available, or
@@ -68,8 +77,17 @@ Typical installation procedure may look like this:
     $ cd pg_wait_sampling
     $ make USE_PGXS=1
     $ sudo make USE_PGXS=1 install
+
+Then add `shared_preload_libraries = pg_wait_sampling` to `postgresql.conf` and
+restart the server.
+
+To test your installation:
+
     $ make USE_PGXS=1 installcheck
-    $ psql DB -c "CREATE EXTENSION pg_wait_sampling;"
+
+To create the extension in the target database:
+
+    CREATE EXTENSION pg_wait_sampling;
 
 Compilation on Windows is not supported, since the extension uses symbols from PostgreSQL
 that are not exported.
@@ -119,24 +137,30 @@ in-memory hash table.
 The work of wait event statistics collector worker is controlled by following
 GUCs.
 
-|         Parameter name              | Data type |                  Description                | Default value |
-| ----------------------------------- | --------- | ------------------------------------------- | ------------: |
-| pg_wait_sampling.history_size       | int4      | Size of history in-memory ring buffer       |          5000 |
-| pg_wait_sampling.history_period     | int4      | Period for history sampling in milliseconds |            10 |
-| pg_wait_sampling.profile_period     | int4      | Period for profile sampling in milliseconds |            10 |
-| pg_wait_sampling.profile_pid        | bool      | Whether profile should be per pid           |          true |
-| pg_wait_sampling.profile_queries    | bool      | Whether profile should be per query			|         false |
+| Parameter name                   | Data type | Description                                 | Default value |
+|----------------------------------| --------- |---------------------------------------------|--------------:|
+| pg_wait_sampling.history_size    | int4      | Size of history in-memory ring buffer       |          5000 |
+| pg_wait_sampling.history_period  | int4      | Period for history sampling in milliseconds |            10 |
+| pg_wait_sampling.profile_period  | int4      | Period for profile sampling in milliseconds |            10 |
+| pg_wait_sampling.profile_pid     | bool      | Whether profile should be per pid           |          true |
+| pg_wait_sampling.profile_queries | enum      | Whether profile should be per query         |           top |
+| pg_wait_sampling.sample_cpu      | bool      | Whether on CPU backends should be sampled   |          true |
 
 If `pg_wait_sampling.profile_pid` is set to false, sampling profile wouldn't be
 collected in per-process manner.  In this case the value of pid could would
 be always zero and corresponding row contain samples among all the processes.
 
-While `pg_wait_sampling.profile_queries` is set to false `queryid` field in
-views will be zero.
+If `pg_wait_sampling.profile_queries` is set to `none`, `queryid` field in
+views will be zero. If it is set to `top`, queryIds only of top level statements
+are recorded. If it is set to `all`, queryIds of nested statements are recorded.
+
+If `pg_wait_sampling.sample_cpu` is set to true then processes that are not
+waiting on anything are also sampled. The wait event columns for such processes
+will be NULL.
 
-These GUCs are allowed to be changed by superuser.  Also, they are placed into
-shared memory.  Thus, they could be changed from any backend and affects worker
-runtime.
+Values of these GUC variables can be changed only in config file or with ALTER SYSTEM.
+Then you need to reload server's configuration (such as with pg_reload_conf function)
+for changes to take effect.
 
 See
 [PostgreSQL documentation](http://www.postgresql.org/docs/devel/static/monitoring-stats.html#WAIT-EVENT-TABLE)
diff --git a/collector.c b/collector.c
index b4672a9..721299f 100644
--- a/collector.c
+++ b/collector.c
@@ -2,28 +2,33 @@
  * collector.c
  *		Collector of wait event history and profile.
  *
- * Copyright (c) 2015-2016, Postgres Professional
+ * Copyright (c) 2015-2025, Postgres Professional
  *
  * IDENTIFICATION
  *	  contrib/pg_wait_sampling/pg_wait_sampling.c
  */
 #include "postgres.h"
 
-#include "catalog/pg_type.h"
-#include "funcapi.h"
+#include <signal.h>
+
+#include "compat.h"
 #include "miscadmin.h"
+#include "pg_wait_sampling.h"
+#include "pgstat.h"
 #include "postmaster/bgworker.h"
+#include "postmaster/interrupt.h"
 #include "storage/ipc.h"
-#include "storage/procarray.h"
+#include "storage/latch.h"
+#include "storage/lock.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
 #include "storage/procsignal.h"
 #include "storage/shm_mq.h"
-#include "storage/shm_toc.h"
-#include "storage/spin.h"
+#include "utils/guc.h"
+#include "utils/hsearch.h"
 #include "utils/memutils.h"
 #include "utils/resowner.h"
-#include "pgstat.h"
-
-#include "pg_wait_sampling.h"
+#include "utils/timestamp.h"
 
 static volatile sig_atomic_t shutdown_requested = false;
 
@@ -33,7 +38,7 @@ static void handle_sigterm(SIGNAL_ARGS);
  * Register background worker for collecting waits history.
  */
 void
-register_wait_collector(void)
+pgws_register_wait_collector(void)
 {
 	BackgroundWorker worker;
 
@@ -41,10 +46,10 @@ register_wait_collector(void)
 	memset(&worker, 0, sizeof(worker));
 	worker.bgw_flags = BGWORKER_SHMEM_ACCESS;
 	worker.bgw_start_time = BgWorkerStart_ConsistentState;
-	worker.bgw_restart_time = 0;
+	worker.bgw_restart_time = 1;
 	worker.bgw_notify_pid = 0;
 	snprintf(worker.bgw_library_name, BGW_MAXLEN, "pg_wait_sampling");
-	snprintf(worker.bgw_function_name, BGW_MAXLEN, CppAsString(collector_main));
+	snprintf(worker.bgw_function_name, BGW_MAXLEN, CppAsString(pgws_collector_main));
 	snprintf(worker.bgw_name, BGW_MAXLEN, "pg_wait_sampling collector");
 	worker.bgw_main_arg = (Datum) 0;
 	RegisterBackgroundWorker(&worker);
@@ -53,7 +58,7 @@ register_wait_collector(void)
 /*
  * Allocate memory for waits history.
  */
-void
+static void
 alloc_history(History *observations, int count)
 {
 	observations->items = (HistoryItem *) palloc0(sizeof(HistoryItem) * count);
@@ -68,10 +73,10 @@ alloc_history(History *observations, int count)
 static void
 realloc_history(History *observations, int count)
 {
-	HistoryItem	   *newitems;
-	int				copyCount,
-					i,
-					j;
+	HistoryItem *newitems;
+	int			copyCount,
+				i,
+				j;
 
 	/* Allocate new array for history */
 	newitems = (HistoryItem *) palloc0(sizeof(HistoryItem) * count);
@@ -109,7 +114,8 @@ realloc_history(History *observations, int count)
 static void
 handle_sigterm(SIGNAL_ARGS)
 {
-	int save_errno = errno;
+	int			save_errno = errno;
+
 	shutdown_requested = true;
 	if (MyProc)
 		SetLatch(&MyProc->procLatch);
@@ -124,6 +130,7 @@ get_next_observation(History *observations)
 {
 	HistoryItem *result;
 
+	/* Check for wraparound */
 	if (observations->index >= observations->count)
 	{
 		observations->index = 0;
@@ -144,10 +151,10 @@ probe_waits(History *observations, HTAB *profile_hash,
 {
 	int			i,
 				newSize;
-	TimestampTz	ts = GetCurrentTimestamp();
+	TimestampTz ts = GetCurrentTimestamp();
 
 	/* Realloc waits history if needed */
-	newSize = collector_hdr->historySize;
+	newSize = pgws_historySize;
 	if (observations->count != newSize)
 		realloc_history(observations, newSize);
 
@@ -155,22 +162,15 @@ probe_waits(History *observations, HTAB *profile_hash,
 	LWLockAcquire(ProcArrayLock, LW_SHARED);
 	for (i = 0; i < ProcGlobal->allProcCount; i++)
 	{
-		HistoryItem		item,
-					   *observation;
-		PGPROC		   *proc = &ProcGlobal->allProcs[i];
+		HistoryItem item,
+				   *observation;
+		PGPROC	   *proc = &ProcGlobal->allProcs[i];
 
-		if (proc->pid == 0)
+		if (!pgws_should_sample_proc(proc, &item.pid, &item.wait_event_info))
 			continue;
 
-		if (proc->wait_event_info == 0)
-			continue;
-
-		/* Collect next wait event sample */
-		item.pid = proc->pid;
-		item.wait_event_info = proc->wait_event_info;
-
-		if (collector_hdr->profileQueries)
-			item.queryId = proc_queryids[i];
+		if (pgws_profileQueries)
+			item.queryId = pgws_proc_queryids[i];
 		else
 			item.queryId = 0;
 
@@ -186,8 +186,8 @@ probe_waits(History *observations, HTAB *profile_hash,
 		/* Write to the profile if needed */
 		if (write_profile)
 		{
-			ProfileItem	   *profileItem;
-			bool			found;
+			ProfileItem *profileItem;
+			bool		found;
 
 			if (!profile_pid)
 				item.pid = 0;
@@ -208,17 +208,39 @@ probe_waits(History *observations, HTAB *profile_hash,
 static void
 send_history(History *observations, shm_mq_handle *mqh)
 {
-	Size	count,
-			i;
+	Size		count,
+				i;
+	shm_mq_result mq_result;
 
 	if (observations->wraparound)
 		count = observations->count;
 	else
 		count = observations->index;
 
-	shm_mq_send(mqh, sizeof(count), &count, false);
+	/* Send array size first since receive_array expects this */
+	mq_result = shm_mq_send_compat(mqh, sizeof(count), &count, false, true);
+	if (mq_result == SHM_MQ_DETACHED)
+	{
+		ereport(WARNING,
+				(errmsg("pg_wait_sampling collector: "
+						"receiver of message queue has been detached")));
+		return;
+	}
 	for (i = 0; i < count; i++)
-		shm_mq_send(mqh, sizeof(HistoryItem), &observations->items[i], false);
+	{
+		mq_result = shm_mq_send_compat(mqh,
+									   sizeof(HistoryItem),
+									   &observations->items[i],
+									   false,
+									   true);
+		if (mq_result == SHM_MQ_DETACHED)
+		{
+			ereport(WARNING,
+					(errmsg("pg_wait_sampling collector: "
+							"receiver of message queue has been detached")));
+			return;
+		}
+	}
 }
 
 /*
@@ -227,15 +249,33 @@ send_history(History *observations, shm_mq_handle *mqh)
 static void
 send_profile(HTAB *profile_hash, shm_mq_handle *mqh)
 {
-	HASH_SEQ_STATUS	scan_status;
-	ProfileItem	   *item;
-	Size			count = hash_get_num_entries(profile_hash);
-
-	shm_mq_send(mqh, sizeof(count), &count, false);
+	HASH_SEQ_STATUS scan_status;
+	ProfileItem *item;
+	Size		count = hash_get_num_entries(profile_hash);
+	shm_mq_result mq_result;
+
+	/* Send array size first since receive_array expects this */
+	mq_result = shm_mq_send_compat(mqh, sizeof(count), &count, false, true);
+	if (mq_result == SHM_MQ_DETACHED)
+	{
+		ereport(WARNING,
+				(errmsg("pg_wait_sampling collector: "
+						"receiver of message queue has been detached")));
+		return;
+	}
 	hash_seq_init(&scan_status, profile_hash);
 	while ((item = (ProfileItem *) hash_seq_search(&scan_status)) != NULL)
 	{
-		shm_mq_send(mqh, sizeof(ProfileItem), item, false);
+		mq_result = shm_mq_send_compat(mqh, sizeof(ProfileItem), item, false,
+									   true);
+		if (mq_result == SHM_MQ_DETACHED)
+		{
+			hash_seq_term(&scan_status);
+			ereport(WARNING,
+					(errmsg("pg_wait_sampling collector: "
+							"receiver of message queue has been detached")));
+			return;
+		}
 	}
 }
 
@@ -245,19 +285,16 @@ send_profile(HTAB *profile_hash, shm_mq_handle *mqh)
 static HTAB *
 make_profile_hash()
 {
-	HASHCTL hash_ctl;
+	HASHCTL		hash_ctl;
 
-	hash_ctl.hash = tag_hash;
-	hash_ctl.hcxt = TopMemoryContext;
-
-	if (collector_hdr->profileQueries)
+	if (pgws_profileQueries)
 		hash_ctl.keysize = offsetof(ProfileItem, count);
 	else
 		hash_ctl.keysize = offsetof(ProfileItem, queryId);
 
 	hash_ctl.entrysize = sizeof(ProfileItem);
 	return hash_create("Waits profile hash", 1024, &hash_ctl,
-					   HASH_FUNCTION | HASH_ELEM);
+					   HASH_ELEM | HASH_BLOBS);
 }
 
 /*
@@ -266,12 +303,12 @@ make_profile_hash()
 static int64
 millisecs_diff(TimestampTz tz1, TimestampTz tz2)
 {
-	long	secs;
-	int		millisecs;
+	long		secs;
+	int			microsecs;
 
-	TimestampDifference(tz1, tz2, &secs, &millisecs);
+	TimestampDifference(tz1, tz2, &secs, &microsecs);
 
-	return secs * 1000 + millisecs;
+	return secs * 1000 + microsecs / 1000;
 
 }
 
@@ -279,68 +316,82 @@ millisecs_diff(TimestampTz tz1, TimestampTz tz2)
  * Main routine of wait history collector.
  */
 void
-collector_main(Datum main_arg)
+pgws_collector_main(Datum main_arg)
 {
- 	HTAB		   *profile_hash = NULL;
-	History			observations;
-	MemoryContext	old_context,
-					collector_context;
-	TimestampTz		current_ts,
-					history_ts,
-					profile_ts;
+	HTAB	   *profile_hash = NULL;
+	History		observations;
+	MemoryContext old_context,
+				collector_context;
+	TimestampTz current_ts,
+				history_ts,
+				profile_ts;
 
 	/*
 	 * Establish signal handlers.
 	 *
-	 * We want CHECK_FOR_INTERRUPTS() to kill off this worker process just as
-	 * it would a normal user backend.  To make that happen, we establish a
-	 * signal handler that is a stripped-down version of die().  We don't have
-	 * any equivalent of the backend's command-read loop, where interrupts can
-	 * be processed immediately, so make sure ImmediateInterruptOK is turned
-	 * off.
+	 * We want to respond to the ProcSignal notifications.  This is done in
+	 * the upstream provided procsignal_sigusr1_handler, which is
+	 * automatically used if a bgworker connects to a database.  But since our
+	 * worker doesn't connect to any database even though it calls
+	 * InitPostgres, which will still initializze a new backend and thus
+	 * partitipate to the ProcSignal infrastructure.
 	 */
 	pqsignal(SIGTERM, handle_sigterm);
+	pqsignal(SIGHUP, SignalHandlerForConfigReload);
+	pqsignal(SIGUSR1, procsignal_sigusr1_handler);
 	BackgroundWorkerUnblockSignals();
+	InitPostgresCompat(NULL, InvalidOid, NULL, InvalidOid, 0, NULL);
+	SetProcessingMode(NormalProcessing);
+
+	/* Make pg_wait_sampling recognisable in pg_stat_activity */
+	pgstat_report_appname("pg_wait_sampling collector");
 
 	profile_hash = make_profile_hash();
-	collector_hdr->latch = &MyProc->procLatch;
+	pgws_collector_hdr->latch = &MyProc->procLatch;
 
 	CurrentResourceOwner = ResourceOwnerCreate(NULL, "pg_wait_sampling collector");
 	collector_context = AllocSetContextCreate(TopMemoryContext,
-			"pg_wait_sampling context", ALLOCSET_DEFAULT_SIZES);
+											  "pg_wait_sampling context", ALLOCSET_DEFAULT_SIZES);
 	old_context = MemoryContextSwitchTo(collector_context);
-	alloc_history(&observations, collector_hdr->historySize);
+	alloc_history(&observations, pgws_historySize);
 	MemoryContextSwitchTo(old_context);
 
+	ereport(LOG, (errmsg("pg_wait_sampling collector started")));
+
 	/* Start counting time for history and profile samples */
 	profile_ts = history_ts = GetCurrentTimestamp();
 
 	while (1)
 	{
-		int				rc;
-		shm_mq_handle  *mqh;
-		int64			history_diff,
-						profile_diff;
-		int				history_period,
-						profile_period;
-		bool			write_history,
-						write_profile;
-
-		/* Wait calculate time to next sample for history or profile */
+		int			rc;
+		shm_mq_handle *mqh;
+		int64		history_diff,
+					profile_diff;
+		bool		write_history,
+					write_profile;
+
+		/* We need an explicit call for at least ProcSignal notifications. */
+		CHECK_FOR_INTERRUPTS();
+
+		if (ConfigReloadPending)
+		{
+			ConfigReloadPending = false;
+			ProcessConfigFile(PGC_SIGHUP);
+		}
+
+		/* Calculate time to next sample for history or profile */
 		current_ts = GetCurrentTimestamp();
 
 		history_diff = millisecs_diff(history_ts, current_ts);
 		profile_diff = millisecs_diff(profile_ts, current_ts);
-		history_period = collector_hdr->historyPeriod;
-		profile_period = collector_hdr->profilePeriod;
 
-		write_history = (history_diff >= (int64)history_period);
-		write_profile = (profile_diff >= (int64)profile_period);
+		write_history = (history_diff >= (int64) pgws_historyPeriod);
+		write_profile = (profile_diff >= (int64) pgws_profilePeriod);
 
 		if (write_history || write_profile)
 		{
 			probe_waits(&observations, profile_hash,
-						write_history, write_profile, collector_hdr->profilePid);
+						write_history, write_profile, pgws_profilePid);
 
 			if (write_history)
 			{
@@ -363,15 +414,9 @@ collector_main(Datum main_arg)
 		 * Wait until next sample time or request to do something through
 		 * shared memory.
 		 */
-#if PG_VERSION_NUM >= 100000
 		rc = WaitLatch(&MyProc->procLatch, WL_LATCH_SET | WL_TIMEOUT | WL_POSTMASTER_DEATH,
-				Min(history_period - (int)history_diff,
-					profile_period - (int)profile_diff), PG_WAIT_EXTENSION);
-#else
-		rc = WaitLatch(&MyProc->procLatch, WL_LATCH_SET | WL_TIMEOUT | WL_POSTMASTER_DEATH,
-				Min(history_period - (int)history_diff,
-					profile_period - (int)profile_diff));
-#endif
+					   Min(pgws_historyPeriod - (int) history_diff,
+						   pgws_historyPeriod - (int) profile_diff), PG_WAIT_EXTENSION);
 
 		if (rc & WL_POSTMASTER_DEATH)
 			proc_exit(1);
@@ -379,38 +424,50 @@ collector_main(Datum main_arg)
 		ResetLatch(&MyProc->procLatch);
 
 		/* Handle request if any */
-		if (collector_hdr->request != NO_REQUEST)
+		if (pgws_collector_hdr->request != NO_REQUEST)
 		{
 			LOCKTAG		tag;
-			SHMRequest	request = collector_hdr->request;
+			SHMRequest	request;
 
-			init_lock_tag(&tag, PGWS_COLLECTOR_LOCK);
+			pgws_init_lock_tag(&tag, PGWS_COLLECTOR_LOCK);
 
 			LockAcquire(&tag, ExclusiveLock, false, false);
-			collector_hdr->request = NO_REQUEST;
+			request = pgws_collector_hdr->request;
+			pgws_collector_hdr->request = NO_REQUEST;
 
 			if (request == HISTORY_REQUEST || request == PROFILE_REQUEST)
 			{
+				shm_mq_result mq_result;
+
 				/* Send history or profile */
-				shm_mq_set_sender(collector_mq, MyProc);
-				mqh = shm_mq_attach(collector_mq, NULL, NULL);
-				shm_mq_wait_for_attach(mqh);
-				if (shm_mq_get_receiver(collector_mq) != NULL)
+				shm_mq_set_sender(pgws_collector_mq, MyProc);
+				mqh = shm_mq_attach(pgws_collector_mq, NULL, NULL);
+				mq_result = shm_mq_wait_for_attach(mqh);
+				switch (mq_result)
 				{
-					if (request == HISTORY_REQUEST)
-					{
-						send_history(&observations, mqh);
-					}
-					else if (request == PROFILE_REQUEST)
-					{
-						send_profile(profile_hash, mqh);
-					}
+					case SHM_MQ_SUCCESS:
+						switch (request)
+						{
+							case HISTORY_REQUEST:
+								send_history(&observations, mqh);
+								break;
+							case PROFILE_REQUEST:
+								send_profile(profile_hash, mqh);
+								break;
+							default:
+								Assert(false);
+						}
+						break;
+					case SHM_MQ_DETACHED:
+						ereport(WARNING,
+								(errmsg("pg_wait_sampling collector: "
+										"receiver of message queue have been "
+										"detached")));
+						break;
+					default:
+						Assert(false);
 				}
-#if PG_VERSION_NUM >= 100000
 				shm_mq_detach(mqh);
-#else
-				shm_mq_detach(collector_mq);
-#endif
 			}
 			else if (request == PROFILE_RESET)
 			{
@@ -424,11 +481,6 @@ collector_main(Datum main_arg)
 
 	MemoryContextReset(collector_context);
 
-	/*
-	 * We're done.  Explicitly detach the shared memory segment so that we
-	 * don't get a resource leak warning at commit time.  This will fire any
-	 * on_dsm_detach callbacks we've registered, as well.  Once that's done,
-	 * we can go ahead and exit.
-	 */
+	ereport(LOG, (errmsg("pg_wait_sampling collector shutting down")));
 	proc_exit(0);
 }
diff --git a/compat.h b/compat.h
new file mode 100644
index 0000000..32aef49
--- /dev/null
+++ b/compat.h
@@ -0,0 +1,50 @@
+/*
+ * compat.h
+ *		Definitions for function wrappers compatible between PG versions.
+ *
+ * Copyright (c) 2015-2025, Postgres Professional
+ *
+ * IDENTIFICATION
+ *	  contrib/pg_wait_sampling/compat.h
+ */
+#ifndef __COMPAT_H__
+#define __COMPAT_H__
+
+#include "miscadmin.h"
+#include "storage/shm_mq.h"
+
+static inline shm_mq_result
+shm_mq_send_compat(shm_mq_handle *mqh, Size nbytes, const void *data,
+				   bool nowait, bool force_flush)
+{
+#if PG_VERSION_NUM >= 150000
+	return shm_mq_send(mqh, nbytes, data, nowait, force_flush);
+#else
+	return shm_mq_send(mqh, nbytes, data, nowait);
+#endif
+}
+
+#if PG_VERSION_NUM < 170000
+#define INIT_PG_LOAD_SESSION_LIBS		0x0001
+#define INIT_PG_OVERRIDE_ALLOW_CONNS	0x0002
+#endif
+
+static inline void
+InitPostgresCompat(const char *in_dbname, Oid dboid,
+				   const char *username, Oid useroid,
+				   bits32 flags,
+				   char *out_dbname)
+{
+#if PG_VERSION_NUM >= 170000
+	InitPostgres(in_dbname, dboid, username, useroid, flags, out_dbname);
+#elif PG_VERSION_NUM >= 150000
+	InitPostgres(in_dbname, dboid, username, useroid,
+				 flags & INIT_PG_LOAD_SESSION_LIBS,
+				 flags & INIT_PG_OVERRIDE_ALLOW_CONNS, out_dbname);
+#else
+	InitPostgres(in_dbname, dboid, username, useroid, out_dbname,
+				 flags & INIT_PG_OVERRIDE_ALLOW_CONNS);
+#endif
+}
+
+#endif
diff --git a/docker-compose.yml b/docker-compose.yml
deleted file mode 100644
index 471ab77..0000000
--- a/docker-compose.yml
+++ /dev/null
@@ -1,2 +0,0 @@
-tests:
-    build: .
diff --git a/expected/queries.out b/expected/queries.out
index 885967b..722df5f 100644
--- a/expected/queries.out
+++ b/expected/queries.out
@@ -20,6 +20,25 @@ WITH t as (SELECT sum(0) FROM pg_wait_sampling_profile)
    0
 (1 row)
 
+-- Some dummy checks just to be sure that all our functions work and return something.
+SELECT count(*) = 1 as test FROM pg_wait_sampling_get_current(pg_backend_pid());
+ test 
+------
+ t
+(1 row)
+
+SELECT count(*) >= 0 as test FROM pg_wait_sampling_get_profile();
+ test 
+------
+ t
+(1 row)
+
+SELECT count(*) >= 0 as test FROM pg_wait_sampling_get_history();
+ test 
+------
+ t
+(1 row)
+
 SELECT pg_wait_sampling_reset_profile();
  pg_wait_sampling_reset_profile 
 --------------------------------
diff --git a/meson.build b/meson.build
new file mode 100644
index 0000000..c3c3dc9
--- /dev/null
+++ b/meson.build
@@ -0,0 +1,41 @@
+# Copyright (c) 2025, Postgres Professional
+
+# Does not support the PGXS infrastructure at this time. Please, compile as part
+# of the contrib source tree.
+
+pg_wait_sampling_sources = files(
+  'collector.c',
+  'pg_wait_sampling.c',
+)
+
+if host_system == 'windows'
+  pg_wait_sampling_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'pg_wait_sampling',
+    '--FILEDESC', 'pg_wait_sampling - provides information about the current wait events for particular processes.',])
+endif
+
+pg_wait_sampling = shared_module('pg_wait_sampling',
+  pg_wait_sampling_sources,
+  kwargs: contrib_mod_args,
+)
+contrib_targets += pg_wait_sampling
+
+install_data(
+  'pg_wait_sampling.control',
+  'pg_wait_sampling--1.0--1.1.sql',
+  'pg_wait_sampling--1.1.sql',
+  kwargs: contrib_data_args,
+)
+
+tests += {
+  'name': 'pg_wait_sampling',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'regress': {
+    'sql': [
+      'load',
+      'queries',
+    ],
+    'regress_args': ['--temp-config', files('conf.add')],
+  },
+}
diff --git a/setup.sql b/pg_wait_sampling--1.1.sql
similarity index 100%
rename from setup.sql
rename to pg_wait_sampling--1.1.sql
diff --git a/pg_wait_sampling.c b/pg_wait_sampling.c
index 64d6d41..f5bd6e0 100644
--- a/pg_wait_sampling.c
+++ b/pg_wait_sampling.c
@@ -2,73 +2,197 @@
  * pg_wait_sampling.c
  *		Track information about wait events.
  *
- * Copyright (c) 2015-2017, Postgres Professional
+ * Copyright (c) 2015-2025, Postgres Professional
  *
  * IDENTIFICATION
  *	  contrib/pg_wait_sampling/pg_wait_sampling.c
  */
 #include "postgres.h"
-#include "fmgr.h"
-#include "funcapi.h"
+
+#include <limits.h>
+
 #include "access/htup_details.h"
-#include "access/twophase.h"
-#include "catalog/pg_type.h"
+#include "catalog/pg_type_d.h"
+#include "executor/executor.h"
+#include "funcapi.h"
 #include "miscadmin.h"
 #include "optimizer/planner.h"
+#include "pg_wait_sampling.h"
 #include "pgstat.h"
-#include "storage/spin.h"
 #include "storage/ipc.h"
-#include "storage/pg_shmem.h"
-#include "storage/procarray.h"
+#include "storage/latch.h"
+#include "storage/lock.h"
+#include "storage/lwlock.h"
+#include "storage/proc.h"
 #include "storage/shm_mq.h"
 #include "storage/shm_toc.h"
+#include "storage/shmem.h"
+#include "tcop/utility.h"
 #include "utils/builtins.h"
-#include "utils/datetime.h"
 #include "utils/guc.h"
-#include "utils/guc_tables.h"
+#include "utils/memutils.h"
+#include "utils/timestamp.h"
 
-#include "pg_wait_sampling.h"
+#if PG_VERSION_NUM < 150000
+#include "postmaster/autovacuum.h"
+#include "replication/walsender.h"
+#endif
 
 PG_MODULE_MAGIC;
 
 void		_PG_init(void);
-void		_PG_fini(void);
 
-/* Global variables */
-bool					shmem_initialized = false;
+static bool shmem_initialized = false;
 
-/* Hooks */
-static ExecutorEnd_hook_type	prev_ExecutorEnd = NULL;
-static planner_hook_type		planner_hook_next = NULL;
+/* Hooks variables */
+static ExecutorStart_hook_type prev_ExecutorStart = NULL;
+static ExecutorRun_hook_type prev_ExecutorRun = NULL;
+static ExecutorFinish_hook_type prev_ExecutorFinish = NULL;
+static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
+static planner_hook_type planner_hook_next = NULL;
+static ProcessUtility_hook_type prev_ProcessUtility = NULL;
 
-/* Shared memory variables */
-shm_toc				   *toc = NULL;
-shm_mq				   *collector_mq = NULL;
-uint64				   *proc_queryids = NULL;
-CollectorShmqHeader	   *collector_hdr = NULL;
+/* Current nesting depth of planner/Executor calls */
+static int	nesting_level = 0;
 
+/* Pointers to shared memory objects */
+shm_mq	   *pgws_collector_mq = NULL;
+uint64	   *pgws_proc_queryids = NULL;
+CollectorShmqHeader *pgws_collector_hdr = NULL;
+
+/* Receiver (backend) local shm_mq pointers and lock */
+static shm_mq *recv_mq = NULL;
+static shm_mq_handle *recv_mqh = NULL;
+static LOCKTAG queueTag;
+
+/* Hook functions */
+#if PG_VERSION_NUM >= 150000
+static shmem_request_hook_type prev_shmem_request_hook = NULL;
+#endif
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
-static PGPROC * search_proc(int backendPid);
+static PGPROC *search_proc(int backendPid);
 static PlannedStmt *pgws_planner_hook(Query *parse,
-		int cursorOptions, ParamListInfo boundParams);
+#if PG_VERSION_NUM >= 130000
+									  const char *query_string,
+#endif
+									  int cursorOptions, ParamListInfo boundParams);
+static void pgws_ExecutorStart(QueryDesc *queryDesc, int eflags);
+static void pgws_ExecutorRun(QueryDesc *queryDesc,
+							 ScanDirection direction,
+							 uint64 count
+#if PG_VERSION_NUM >= 100000 && PG_VERSION_NUM < 180000
+							 ,bool execute_once
+#endif
+);
+static void pgws_ExecutorFinish(QueryDesc *queryDesc);
 static void pgws_ExecutorEnd(QueryDesc *queryDesc);
+static void pgws_ProcessUtility(PlannedStmt *pstmt,
+								const char *queryString,
+#if PG_VERSION_NUM >= 140000
+								bool readOnlyTree,
+#endif
+								ProcessUtilityContext context,
+								ParamListInfo params,
+								QueryEnvironment *queryEnv,
+								DestReceiver *dest,
+#if PG_VERSION_NUM >= 130000
+								QueryCompletion *qc
+#else
+								char *completionTag
+#endif
+);
+
+/*---- GUC variables ----*/
+
+typedef enum
+{
+	PGWS_PROFILE_QUERIES_NONE,	/* profile no statements */
+	PGWS_PROFILE_QUERIES_TOP,	/* only top level statements */
+	PGWS_PROFILE_QUERIES_ALL	/* all statements, including nested ones */
+} PGWSTrackLevel;
+
+static const struct config_enum_entry pgws_profile_queries_options[] =
+{
+	{"none", PGWS_PROFILE_QUERIES_NONE, false},
+	{"off", PGWS_PROFILE_QUERIES_NONE, false},
+	{"no", PGWS_PROFILE_QUERIES_NONE, false},
+	{"false", PGWS_PROFILE_QUERIES_NONE, false},
+	{"0", PGWS_PROFILE_QUERIES_NONE, false},
+	{"top", PGWS_PROFILE_QUERIES_TOP, false},
+	{"on", PGWS_PROFILE_QUERIES_TOP, false},
+	{"yes", PGWS_PROFILE_QUERIES_TOP, false},
+	{"true", PGWS_PROFILE_QUERIES_TOP, false},
+	{"1", PGWS_PROFILE_QUERIES_TOP, false},
+	{"all", PGWS_PROFILE_QUERIES_ALL, false},
+	{NULL, 0, false}
+};
+
+int			pgws_historySize = 5000;
+int			pgws_historyPeriod = 10;
+int			pgws_profilePeriod = 10;
+bool		pgws_profilePid = true;
+int			pgws_profileQueries = PGWS_PROFILE_QUERIES_TOP;
+bool		pgws_sampleCpu = true;
+
+#define pgws_enabled(level) \
+	((pgws_profileQueries == PGWS_PROFILE_QUERIES_ALL) || \
+	 (pgws_profileQueries == PGWS_PROFILE_QUERIES_TOP && (level) == 0))
 
 /*
  * Calculate max processes count.
- * Look at InitProcGlobal (proc.c) and TotalProcs variable in it
- * if something wrong here.
+ *
+ * The value has to be in sync with ProcGlobal->allProcCount, initialized in
+ * InitProcGlobal() (proc.c).
+ *
  */
 static int
 get_max_procs_count(void)
 {
-	int count = 0;
+	int			count = 0;
 
-	/* MyProcs, including autovacuum workers and launcher */
+	/* First, add the maximum number of backends (MaxBackends). */
+#if PG_VERSION_NUM >= 150000
+
+	/*
+	 * On pg15+, we can directly access the MaxBackends variable, as it will
+	 * have already been initialized in shmem_request_hook.
+	 */
+	Assert(MaxBackends > 0);
 	count += MaxBackends;
-	/* AuxiliaryProcs */
+#else
+
+	/*
+	 * On older versions, we need to compute MaxBackends: bgworkers,
+	 * autovacuum workers and launcher. This has to be in sync with the value
+	 * computed in InitializeMaxBackends() (postinit.c)
+	 *
+	 * Note that we need to calculate the value as it won't initialized when
+	 * we need it during _PG_init().
+	 *
+	 * Note also that the value returned during _PG_init() might be different
+	 * from the value returned later if some third-party modules change one of
+	 * the underlying GUC.  This isn't ideal but can't lead to a crash, as the
+	 * value returned during _PG_init() is only used to ask for additional
+	 * shmem with RequestAddinShmemSpace(), and postgres has an extra 100kB of
+	 * shmem to compensate some small unaccounted usage.  So if the value
+	 * later changes, we will allocate and initialize the new (and correct)
+	 * memory size, which will either work thanks for the extra 100kB of
+	 * shmem, of fail (and prevent postgres startup) due to an out of shared
+	 * memory error.
+	 */
+	count += MaxConnections + autovacuum_max_workers + 1
+		+ max_worker_processes;
+
+	/*
+	 * Starting with pg12, wal senders aren't part of MaxConnections anymore
+	 * and have to be accounted for.
+	 */
+	count += max_wal_senders;
+#endif							/* pg 15- */
+	/* End of MaxBackends calculation. */
+
+	/* Add AuxiliaryProcs */
 	count += NUM_AUXILIARY_PROCS;
-	/* Prepared xacts */
-	count += max_prepared_xacts;
 
 	return count;
 }
@@ -79,9 +203,9 @@ get_max_procs_count(void)
 static Size
 pgws_shmem_size(void)
 {
-	shm_toc_estimator	e;
-	Size				size;
-	int					nkeys;
+	shm_toc_estimator e;
+	Size		size;
+	int			nkeys;
 
 	shm_toc_initialize_estimator(&e);
 
@@ -97,133 +221,22 @@ pgws_shmem_size(void)
 	return size;
 }
 
-static bool
-shmem_int_guc_check_hook(int *newval, void **extra, GucSource source)
-{
-	if (UsedShmemSegAddr == NULL)
-		return false;
-	return true;
-}
-
-static bool
-shmem_bool_guc_check_hook(bool *newval, void **extra, GucSource source)
-{
-	if (UsedShmemSegAddr == NULL)
-		return false;
-	return true;
-}
-
+#if PG_VERSION_NUM >= 150000
 /*
- * This union allows us to mix the numerous different types of structs
- * that we are organizing.
- */
-typedef union
-{
-	struct config_generic generic;
-	struct config_bool _bool;
-	struct config_real real;
-	struct config_int integer;
-	struct config_string string;
-	struct config_enum _enum;
-} mixedStruct;
-
-/*
- * Setup new GUCs or modify existsing.
+ * shmem_request hook: request additional shared memory resources.
+ *
+ * If you change code here, don't forget to also report the modifications in
+ * _PG_init() for pg14 and below.
  */
 static void
-setup_gucs()
+pgws_shmem_request(void)
 {
-	struct config_generic **guc_vars;
-	int			numOpts,
-				i;
-	bool		history_size_found = false,
-				history_period_found = false,
-				profile_period_found = false,
-				profile_pid_found = false,
-				profile_queries_found = false;
+	if (prev_shmem_request_hook)
+		prev_shmem_request_hook();
 
-	guc_vars = get_guc_variables();
-	numOpts = GetNumConfigOptions();
-
-	for (i = 0; i < numOpts; i++)
-	{
-		mixedStruct *var = (mixedStruct *) guc_vars[i];
-		const char *name = var->generic.name;
-
-		if (var->generic.flags & GUC_CUSTOM_PLACEHOLDER)
-			continue;
-
-		if (!strcmp(name, "pg_wait_sampling.history_size"))
-		{
-			history_size_found = true;
-			var->integer.variable = &collector_hdr->historySize;
-			collector_hdr->historySize = 5000;
-		}
-		else if (!strcmp(name, "pg_wait_sampling.history_period"))
-		{
-			history_period_found = true;
-			var->integer.variable = &collector_hdr->historyPeriod;
-			collector_hdr->historyPeriod = 10;
-		}
-		else if (!strcmp(name, "pg_wait_sampling.profile_period"))
-		{
-			profile_period_found = true;
-			var->integer.variable = &collector_hdr->profilePeriod;
-			collector_hdr->profilePeriod = 10;
-		}
-		else if (!strcmp(name, "pg_wait_sampling.profile_pid"))
-		{
-			profile_pid_found = true;
-			var->_bool.variable = &collector_hdr->profilePid;
-			collector_hdr->profilePid = true;
-		}
-		else if (!strcmp(name, "pg_wait_sampling.profile_queries"))
-		{
-			profile_queries_found = true;
-			var->_bool.variable = &collector_hdr->profileQueries;
-			collector_hdr->profileQueries = true;
-		}
-	}
-
-	if (!history_size_found)
-		DefineCustomIntVariable("pg_wait_sampling.history_size",
-				"Sets size of waits history.", NULL,
-				&collector_hdr->historySize, 5000, 100, INT_MAX,
-				PGC_SUSET, 0, shmem_int_guc_check_hook, NULL, NULL);
-
-	if (!history_period_found)
-		DefineCustomIntVariable("pg_wait_sampling.history_period",
-				"Sets period of waits history sampling.", NULL,
-				&collector_hdr->historyPeriod, 10, 1, INT_MAX,
-				PGC_SUSET, 0, shmem_int_guc_check_hook, NULL, NULL);
-
-	if (!profile_period_found)
-		DefineCustomIntVariable("pg_wait_sampling.profile_period",
-				"Sets period of waits profile sampling.", NULL,
-				&collector_hdr->profilePeriod, 10, 1, INT_MAX,
-				PGC_SUSET, 0, shmem_int_guc_check_hook, NULL, NULL);
-
-	if (!profile_pid_found)
-		DefineCustomBoolVariable("pg_wait_sampling.profile_pid",
-				"Sets whether profile should be collected per pid.", NULL,
-				&collector_hdr->profilePid, true,
-				PGC_SUSET, 0, shmem_bool_guc_check_hook, NULL, NULL);
-
-	if (!profile_queries_found)
-		DefineCustomBoolVariable("pg_wait_sampling.profile_queries",
-				"Sets whether profile should be collected per query.", NULL,
-				&collector_hdr->profileQueries, true,
-				PGC_SUSET, 0, shmem_bool_guc_check_hook, NULL, NULL);
-
-	if (history_size_found
-		|| history_period_found
-		|| profile_period_found
-		|| profile_pid_found
-		|| profile_queries_found)
-	{
-		ProcessConfigFile(PGC_SIGHUP);
-	}
+	RequestAddinShmemSpace(pgws_shmem_size());
 }
+#endif
 
 /*
  * Distribute shared memory.
@@ -231,41 +244,34 @@ setup_gucs()
 static void
 pgws_shmem_startup(void)
 {
-	bool	found;
-	Size	segsize = pgws_shmem_size();
-	void   *pgws;
+	bool		found;
+	Size		segsize = pgws_shmem_size();
+	void	   *pgws;
+	shm_toc    *toc;
 
 	pgws = ShmemInitStruct("pg_wait_sampling", segsize, &found);
 
 	if (!found)
 	{
+		/* Create shared objects */
 		toc = shm_toc_create(PG_WAIT_SAMPLING_MAGIC, pgws, segsize);
 
-		collector_hdr = shm_toc_allocate(toc, sizeof(CollectorShmqHeader));
-		shm_toc_insert(toc, 0, collector_hdr);
-		collector_mq = shm_toc_allocate(toc, COLLECTOR_QUEUE_SIZE);
-		shm_toc_insert(toc, 1, collector_mq);
-		proc_queryids = shm_toc_allocate(toc,
-									sizeof(uint64) * get_max_procs_count());
-		shm_toc_insert(toc, 2, proc_queryids);
-		MemSet(proc_queryids, 0, sizeof(uint64) * get_max_procs_count());
-
-		/* Initialize GUC variables in shared memory */
-		setup_gucs();
+		pgws_collector_hdr = shm_toc_allocate(toc, sizeof(CollectorShmqHeader));
+		shm_toc_insert(toc, 0, pgws_collector_hdr);
+		pgws_collector_mq = shm_toc_allocate(toc, COLLECTOR_QUEUE_SIZE);
+		shm_toc_insert(toc, 1, pgws_collector_mq);
+		pgws_proc_queryids = shm_toc_allocate(toc,
+											  sizeof(uint64) * get_max_procs_count());
+		shm_toc_insert(toc, 2, pgws_proc_queryids);
+		MemSet(pgws_proc_queryids, 0, sizeof(uint64) * get_max_procs_count());
 	}
 	else
 	{
+		/* Attach to existing shared objects */
 		toc = shm_toc_attach(PG_WAIT_SAMPLING_MAGIC, pgws);
-
-#if PG_VERSION_NUM >= 100000
-		collector_hdr = shm_toc_lookup(toc, 0, false);
-		collector_mq = shm_toc_lookup(toc, 1, false);
-		proc_queryids = shm_toc_lookup(toc, 2, false);
-#else
-		collector_hdr = shm_toc_lookup(toc, 0);
-		collector_mq = shm_toc_lookup(toc, 1);
-		proc_queryids = shm_toc_lookup(toc, 2);
-#endif
+		pgws_collector_hdr = shm_toc_lookup(toc, 0, false);
+		pgws_collector_mq = shm_toc_lookup(toc, 1, false);
+		pgws_proc_queryids = shm_toc_lookup(toc, 2, false);
 	}
 
 	shmem_initialized = true;
@@ -277,7 +283,7 @@ pgws_shmem_startup(void)
 /*
  * Check shared memory is initialized. Report an error otherwise.
  */
-void
+static void
 check_shmem(void)
 {
 	if (!shmem_initialized)
@@ -287,6 +293,14 @@ check_shmem(void)
 	}
 }
 
+static void
+pgws_cleanup_callback(int code, Datum arg)
+{
+	elog(DEBUG3, "pg_wait_sampling cleanup: detaching shm_mq and releasing queue lock");
+	shm_mq_detach(recv_mqh);
+	LockRelease(&queueTag, ExclusiveLock, false);
+}
+
 /*
  * Module load callback
  */
@@ -296,34 +310,120 @@ _PG_init(void)
 	if (!process_shared_preload_libraries_in_progress)
 		return;
 
+#if PG_VERSION_NUM < 150000
+
 	/*
 	 * Request additional shared resources.  (These are no-ops if we're not in
 	 * the postmaster process.)  We'll allocate or attach to the shared
 	 * resources in pgws_shmem_startup().
+	 *
+	 * If you change code here, don't forget to also report the modifications
+	 * in pgsp_shmem_request() for pg15 and later.
 	 */
 	RequestAddinShmemSpace(pgws_shmem_size());
+#endif
 
-	register_wait_collector();
+	pgws_register_wait_collector();
 
 	/*
 	 * Install hooks.
 	 */
+#if PG_VERSION_NUM >= 150000
+	prev_shmem_request_hook = shmem_request_hook;
+	shmem_request_hook = pgws_shmem_request;
+#endif
 	prev_shmem_startup_hook = shmem_startup_hook;
-	shmem_startup_hook		= pgws_shmem_startup;
-	planner_hook_next		= planner_hook;
-	planner_hook			= pgws_planner_hook;
-	prev_ExecutorEnd		= ExecutorEnd_hook;
-	ExecutorEnd_hook		= pgws_ExecutorEnd;
-}
-
-/*
- * Module unload callback
- */
-void
-_PG_fini(void)
-{
-	/* Uninstall hooks. */
-	shmem_startup_hook = prev_shmem_startup_hook;
+	shmem_startup_hook = pgws_shmem_startup;
+	planner_hook_next = planner_hook;
+	planner_hook = pgws_planner_hook;
+	prev_ExecutorStart = ExecutorStart_hook;
+	ExecutorStart_hook = pgws_ExecutorStart;
+	prev_ExecutorRun = ExecutorRun_hook;
+	ExecutorRun_hook = pgws_ExecutorRun;
+	prev_ExecutorFinish = ExecutorFinish_hook;
+	ExecutorFinish_hook = pgws_ExecutorFinish;
+	prev_ExecutorEnd = ExecutorEnd_hook;
+	ExecutorEnd_hook = pgws_ExecutorEnd;
+	prev_ProcessUtility = ProcessUtility_hook;
+	ProcessUtility_hook = pgws_ProcessUtility;
+
+	/* Define GUC variables */
+	DefineCustomIntVariable("pg_wait_sampling.history_size",
+							"Sets size of waits history.",
+							NULL,
+							&pgws_historySize,
+							5000,
+							100,
+							INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
+	DefineCustomIntVariable("pg_wait_sampling.history_period",
+							"Sets period of waits history sampling.",
+							NULL,
+							&pgws_historyPeriod,
+							10,
+							1,
+							INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
+	DefineCustomIntVariable("pg_wait_sampling.profile_period",
+							"Sets period of waits profile sampling.",
+							NULL,
+							&pgws_profilePeriod,
+							10,
+							1,
+							INT_MAX,
+							PGC_SIGHUP,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
+	DefineCustomBoolVariable("pg_wait_sampling.profile_pid",
+							 "Sets whether profile should be collected per pid.",
+							 NULL,
+							 &pgws_profilePid,
+							 true,
+							 PGC_SIGHUP,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
+	DefineCustomEnumVariable("pg_wait_sampling.profile_queries",
+							 "Sets whether profile should be collected per query.",
+							 NULL,
+							 &pgws_profileQueries,
+							 PGWS_PROFILE_QUERIES_TOP,
+							 pgws_profile_queries_options,
+							 PGC_SIGHUP,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
+	DefineCustomBoolVariable("pg_wait_sampling.sample_cpu",
+							 "Sets whether not waiting backends should be sampled.",
+							 NULL,
+							 &pgws_sampleCpu,
+							 true,
+							 PGC_SIGHUP,
+							 0,
+							 NULL,
+							 NULL,
+							 NULL);
+
+#if PG_VERSION_NUM >= 150000
+	MarkGUCPrefixReserved("pg_wait_sampling");
+#endif
 }
 
 /*
@@ -333,14 +433,15 @@ _PG_fini(void)
 static PGPROC *
 search_proc(int pid)
 {
-	int i;
+	int			i;
 
 	if (pid == 0)
 		return MyProc;
 
 	for (i = 0; i < ProcGlobal->allProcCount; i++)
 	{
-		PGPROC	*proc = &ProcGlobal->allProcs[i];
+		PGPROC	   *proc = &ProcGlobal->allProcs[i];
+
 		if (proc->pid && proc->pid == pid)
 		{
 			return proc;
@@ -352,35 +453,62 @@ search_proc(int pid)
 	return NULL;
 }
 
+/*
+ * Decide whether this PGPROC entry should be included in profiles and output
+ * views.
+ */
+bool
+pgws_should_sample_proc(PGPROC *proc, int *pid_p, uint32 *wait_event_info_p)
+{
+	int			pid = proc->pid;
+	uint32		wait_event_info = proc->wait_event_info;
+
+	*pid_p = pid;
+	*wait_event_info_p = wait_event_info;
+
+	if (wait_event_info == 0 && !pgws_sampleCpu)
+		return false;
+
+	/*
+	 * On PostgreSQL versions < 17 the PGPROC->pid field is not reset on
+	 * process exit. This would lead to such processes getting counted for
+	 * null wait events. So instead we make use of DisownLatch() resetting
+	 * owner_pid during ProcKill().
+	 */
+	if (pid == 0 || proc->procLatch.owner_pid == 0 || pid == MyProcPid)
+		return false;
+
+	return true;
+}
+
 typedef struct
 {
-	HistoryItem	   *items;
-	TimestampTz		ts;
+	HistoryItem *items;
+	TimestampTz ts;
 } WaitCurrentContext;
 
 PG_FUNCTION_INFO_V1(pg_wait_sampling_get_current);
 Datum
 pg_wait_sampling_get_current(PG_FUNCTION_ARGS)
 {
-	FuncCallContext 	*funcctx;
-	WaitCurrentContext 	*params;
+	FuncCallContext *funcctx;
+	WaitCurrentContext *params;
 
 	check_shmem();
 
 	if (SRF_IS_FIRSTCALL())
 	{
-		MemoryContext		oldcontext;
-		TupleDesc			tupdesc;
-		WaitCurrentContext 	*params;
+		MemoryContext oldcontext;
+		TupleDesc	tupdesc;
 
 		funcctx = SRF_FIRSTCALL_INIT();
 
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
-		params = (WaitCurrentContext *)palloc0(sizeof(WaitCurrentContext));
+		params = (WaitCurrentContext *) palloc0(sizeof(WaitCurrentContext));
 		params->ts = GetCurrentTimestamp();
 
 		funcctx->user_fctx = params;
-		tupdesc = CreateTemplateTupleDesc(4, false);
+		tupdesc = CreateTemplateTupleDesc(4);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "pid",
 						   INT4OID, -1, 0);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "type",
@@ -396,35 +524,39 @@ pg_wait_sampling_get_current(PG_FUNCTION_ARGS)
 
 		if (!PG_ARGISNULL(0))
 		{
-			HistoryItem	   *item;
-			PGPROC		   *proc;
+			/* pg_wait_sampling_get_current(pid int4) function */
+			HistoryItem *item;
+			PGPROC	   *proc;
 
 			proc = search_proc(PG_GETARG_UINT32(0));
 			params->items = (HistoryItem *) palloc0(sizeof(HistoryItem));
 			item = &params->items[0];
 			item->pid = proc->pid;
 			item->wait_event_info = proc->wait_event_info;
-			item->queryId = proc_queryids[proc - ProcGlobal->allProcs];
+			item->queryId = pgws_proc_queryids[proc - ProcGlobal->allProcs];
 			funcctx->max_calls = 1;
 		}
 		else
 		{
-			int		procCount = ProcGlobal->allProcCount,
-					i,
-					j = 0;
+			/* pg_wait_sampling_current view */
+			int			procCount = ProcGlobal->allProcCount,
+						i,
+						j = 0;
 
 			params->items = (HistoryItem *) palloc0(sizeof(HistoryItem) * procCount);
 			for (i = 0; i < procCount; i++)
 			{
-				PGPROC *proc = &ProcGlobal->allProcs[i];
-
-				if (proc != NULL && proc->pid != 0 && proc->wait_event_info)
-				{
-					params->items[j].pid = proc->pid;
-					params->items[j].wait_event_info = proc->wait_event_info;
-					params->items[j].queryId = proc_queryids[i];
-					j++;
-				}
+				PGPROC	   *proc = &ProcGlobal->allProcs[i];
+
+				if (!pgws_should_sample_proc(proc,
+											 &params->items[j].pid,
+											 &params->items[j].wait_event_info))
+					continue;
+
+				params->items[j].pid = proc->pid;
+				params->items[j].wait_event_info = proc->wait_event_info;
+				params->items[j].queryId = pgws_proc_queryids[i];
+				j++;
 			}
 			funcctx->max_calls = j;
 		}
@@ -465,25 +597,26 @@ pg_wait_sampling_get_current(PG_FUNCTION_ARGS)
 		else
 			nulls[2] = true;
 
-		values[3] = Int64GetDatumFast(item->queryId);
+		values[3] = UInt64GetDatum(item->queryId);
 		tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
 
 		SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple));
 	}
 	else
 	{
+		/* nothing left */
 		SRF_RETURN_DONE(funcctx);
 	}
 }
 
 typedef struct
 {
-	Size			count;
-	ProfileItem	   *items;
+	Size		count;
+	ProfileItem *items;
 } Profile;
 
 void
-init_lock_tag(LOCKTAG *tag, uint32 lock)
+pgws_init_lock_tag(LOCKTAG *tag, uint32 lock)
 {
 	tag->locktag_field1 = PG_WAIT_SAMPLING_MAGIC;
 	tag->locktag_field2 = lock;
@@ -493,64 +626,83 @@ init_lock_tag(LOCKTAG *tag, uint32 lock)
 	tag->locktag_lockmethodid = USER_LOCKMETHOD;
 }
 
+/* Get array (history or profile data) from shared memory */
 static void *
 receive_array(SHMRequest request, Size item_size, Size *count)
 {
-	LOCKTAG			queueTag;
-	LOCKTAG			collectorTag;
-	shm_mq		   *mq;
-	shm_mq_handle  *mqh;
-	shm_mq_result	res;
-	Size			len,
-					i;
-	void		   *data;
-	Pointer			result,
-					ptr;
+	LOCKTAG		collectorTag;
+	shm_mq_result res;
+	Size		len,
+				i;
+	void	   *data;
+	Pointer		result,
+				ptr;
+	MemoryContext oldctx;
 
 	/* Ensure nobody else trying to send request to queue */
-	init_lock_tag(&queueTag, PGWS_QUEUE_LOCK);
+	pgws_init_lock_tag(&queueTag, PGWS_QUEUE_LOCK);
 	LockAcquire(&queueTag, ExclusiveLock, false, false);
 
-	/* Ensure collector has processed previous request */
-	init_lock_tag(&collectorTag, PGWS_COLLECTOR_LOCK);
+	pgws_init_lock_tag(&collectorTag, PGWS_COLLECTOR_LOCK);
 	LockAcquire(&collectorTag, ExclusiveLock, false, false);
+	recv_mq = shm_mq_create(pgws_collector_mq, COLLECTOR_QUEUE_SIZE);
+	pgws_collector_hdr->request = request;
 	LockRelease(&collectorTag, ExclusiveLock, false);
 
-	mq = shm_mq_create(collector_mq, COLLECTOR_QUEUE_SIZE);
-	collector_hdr->request = request;
-
-	if (!collector_hdr->latch)
+	if (!pgws_collector_hdr->latch)
 		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR),
 						errmsg("pg_wait_sampling collector wasn't started")));
 
-	SetLatch(collector_hdr->latch);
-
-	shm_mq_set_receiver(mq, MyProc);
-	mqh = shm_mq_attach(mq, NULL, NULL);
+	SetLatch(pgws_collector_hdr->latch);
 
-	res = shm_mq_receive(mqh, &len, &data, false);
-	if (res != SHM_MQ_SUCCESS || len != sizeof(*count))
-		elog(ERROR, "Error reading mq.");
-	memcpy(count, data, sizeof(*count));
+	shm_mq_set_receiver(recv_mq, MyProc);
 
-	result = palloc(item_size * (*count));
-	ptr = result;
+	/*
+	 * We switch to TopMemoryContext, so that recv_mqh is allocated there and
+	 * is guaranteed to survive until before_shmem_exit callbacks are fired.
+	 * Anyway, shm_mq_detach() will free handler on its own.
+	 *
+	 * NB: we do not pass `seg` to shm_mq_attach(), so it won't set its own
+	 * callback, i.e. we do not interfere here with shm_mq_detach_callback().
+	 */
+	oldctx = MemoryContextSwitchTo(TopMemoryContext);
+	recv_mqh = shm_mq_attach(recv_mq, NULL, NULL);
+	MemoryContextSwitchTo(oldctx);
 
-	for (i = 0; i < *count; i++)
+	/*
+	 * Now we surely attached to the shm_mq and got collector's attention. If
+	 * anything went wrong (e.g. Ctrl+C received from the client) we have to
+	 * cleanup some things, i.e. detach from the shm_mq, so collector was able
+	 * to continue responding to other requests.
+	 *
+	 * PG_ENSURE_ERROR_CLEANUP() guaranties that cleanup callback will be
+	 * fired for both ERROR and FATAL.
+	 */
+	PG_ENSURE_ERROR_CLEANUP(pgws_cleanup_callback, 0);
 	{
-		res = shm_mq_receive(mqh, &len, &data, false);
-		if (res != SHM_MQ_SUCCESS || len != item_size)
-			elog(ERROR, "Error reading mq.");
-		memcpy(ptr, data, item_size);
-		ptr += item_size;
-	}
+		res = shm_mq_receive(recv_mqh, &len, &data, false);
+		if (res != SHM_MQ_SUCCESS || len != sizeof(*count))
+			elog(ERROR, "error reading mq");
 
-#if PG_VERSION_NUM >= 100000
-	shm_mq_detach(mqh);
-#else
-	shm_mq_detach(mq);
-#endif
+		memcpy(count, data, sizeof(*count));
+
+		result = palloc(item_size * (*count));
+		ptr = result;
+
+		for (i = 0; i < *count; i++)
+		{
+			res = shm_mq_receive(recv_mqh, &len, &data, false);
+			if (res != SHM_MQ_SUCCESS || len != item_size)
+				elog(ERROR, "error reading mq");
+
+			memcpy(ptr, data, item_size);
+			ptr += item_size;
+		}
+	}
+	PG_END_ENSURE_ERROR_CLEANUP(pgws_cleanup_callback, 0);
 
+	/* We still have to detach and release lock during normal operation. */
+	shm_mq_detach(recv_mqh);
 	LockRelease(&queueTag, ExclusiveLock, false);
 
 	return result;
@@ -561,15 +713,15 @@ PG_FUNCTION_INFO_V1(pg_wait_sampling_get_profile);
 Datum
 pg_wait_sampling_get_profile(PG_FUNCTION_ARGS)
 {
-	Profile			   *profile;
-	FuncCallContext	   *funcctx;
+	Profile    *profile;
+	FuncCallContext *funcctx;
 
 	check_shmem();
 
 	if (SRF_IS_FIRSTCALL())
 	{
-		MemoryContext		oldcontext;
-		TupleDesc			tupdesc;
+		MemoryContext oldcontext;
+		TupleDesc	tupdesc;
 
 		funcctx = SRF_FIRSTCALL_INIT();
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
@@ -577,13 +729,13 @@ pg_wait_sampling_get_profile(PG_FUNCTION_ARGS)
 		/* Receive profile from shmq */
 		profile = (Profile *) palloc0(sizeof(Profile));
 		profile->items = (ProfileItem *) receive_array(PROFILE_REQUEST,
-										sizeof(ProfileItem), &profile->count);
+													   sizeof(ProfileItem), &profile->count);
 
 		funcctx->user_fctx = profile;
 		funcctx->max_calls = profile->count;
 
 		/* Make tuple descriptor */
-		tupdesc = CreateTemplateTupleDesc(5, false);
+		tupdesc = CreateTemplateTupleDesc(5);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "pid",
 						   INT4OID, -1, 0);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "type",
@@ -632,12 +784,12 @@ pg_wait_sampling_get_profile(PG_FUNCTION_ARGS)
 		else
 			nulls[2] = true;
 
-		if (collector_hdr->profileQueries)
-			values[3] = Int64GetDatumFast(item->queryId);
+		if (pgws_profileQueries)
+			values[3] = UInt64GetDatum(item->queryId);
 		else
 			values[3] = (Datum) 0;
 
-		values[4] = Int64GetDatumFast(item->count);
+		values[4] = UInt64GetDatum(item->count);
 
 		tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
 
@@ -654,23 +806,22 @@ PG_FUNCTION_INFO_V1(pg_wait_sampling_reset_profile);
 Datum
 pg_wait_sampling_reset_profile(PG_FUNCTION_ARGS)
 {
-	LOCKTAG		tag;
-	LOCKTAG		tagCollector;
+	LOCKTAG		collectorTag;
 
 	check_shmem();
 
-	init_lock_tag(&tag, false);
+	pgws_init_lock_tag(&queueTag, PGWS_QUEUE_LOCK);
 
-	LockAcquire(&tag, ExclusiveLock, false, false);
+	LockAcquire(&queueTag, ExclusiveLock, false, false);
 
-	init_lock_tag(&tagCollector, true);
-	LockAcquire(&tagCollector, ExclusiveLock, false, false);
-	LockRelease(&tagCollector, ExclusiveLock, false);
+	pgws_init_lock_tag(&collectorTag, PGWS_COLLECTOR_LOCK);
+	LockAcquire(&collectorTag, ExclusiveLock, false, false);
+	pgws_collector_hdr->request = PROFILE_RESET;
+	LockRelease(&collectorTag, ExclusiveLock, false);
 
-	collector_hdr->request = PROFILE_RESET;
-	SetLatch(collector_hdr->latch);
+	SetLatch(pgws_collector_hdr->latch);
 
-	LockRelease(&tag, ExclusiveLock, false);
+	LockRelease(&queueTag, ExclusiveLock, false);
 
 	PG_RETURN_VOID();
 }
@@ -679,15 +830,15 @@ PG_FUNCTION_INFO_V1(pg_wait_sampling_get_history);
 Datum
 pg_wait_sampling_get_history(PG_FUNCTION_ARGS)
 {
-	History				*history;
-	FuncCallContext		*funcctx;
+	History    *history;
+	FuncCallContext *funcctx;
 
 	check_shmem();
 
 	if (SRF_IS_FIRSTCALL())
 	{
-		MemoryContext	oldcontext;
-		TupleDesc		tupdesc;
+		MemoryContext oldcontext;
+		TupleDesc	tupdesc;
 
 		funcctx = SRF_FIRSTCALL_INIT();
 		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
@@ -695,13 +846,13 @@ pg_wait_sampling_get_history(PG_FUNCTION_ARGS)
 		/* Receive history from shmq */
 		history = (History *) palloc0(sizeof(History));
 		history->items = (HistoryItem *) receive_array(HISTORY_REQUEST,
-										sizeof(HistoryItem), &history->count);
+													   sizeof(HistoryItem), &history->count);
 
 		funcctx->user_fctx = history;
 		funcctx->max_calls = history->count;
 
 		/* Make tuple descriptor */
-		tupdesc = CreateTemplateTupleDesc(5, false);
+		tupdesc = CreateTemplateTupleDesc(5);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "pid",
 						   INT4OID, -1, 0);
 		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "sample_ts",
@@ -750,7 +901,7 @@ pg_wait_sampling_get_history(PG_FUNCTION_ARGS)
 		else
 			nulls[3] = true;
 
-		values[4] = Int64GetDatumFast(item->queryId);
+		values[4] = UInt64GetDatum(item->queryId);
 		tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
 
 		history->index++;
@@ -769,33 +920,146 @@ pg_wait_sampling_get_history(PG_FUNCTION_ARGS)
  * planner_hook hook, save queryId for collector
  */
 static PlannedStmt *
-pgws_planner_hook(Query *parse, int cursorOptions,
+pgws_planner_hook(Query *parse,
+#if PG_VERSION_NUM >= 130000
+				  const char *query_string,
+#endif
+				  int cursorOptions,
 				  ParamListInfo boundParams)
 {
-	if (MyProc)
+	PlannedStmt *result;
+	int			i = MyProc - ProcGlobal->allProcs;
+	uint64		save_queryId = 0;
+
+	if (pgws_enabled(nesting_level))
 	{
-		int i = MyProc - ProcGlobal->allProcs;
-#if PG_VERSION_NUM >= 110000
-		/*
-		 * since we depend on queryId we need to check that its size
-		 * is uint64 as we coded in pg_wait_sampling
-		 */
-		StaticAssertExpr(sizeof(parse->queryId) == sizeof(uint64),
-				"queryId size is not uint64");
-#else
-		StaticAssertExpr(sizeof(parse->queryId) == sizeof(uint32),
-				"queryId size is not uint32");
+		save_queryId = pgws_proc_queryids[i];
+		pgws_proc_queryids[i] = parse->queryId;
+	}
+
+	nesting_level++;
+	PG_TRY();
+	{
+		/* Invoke original hook if needed */
+		if (planner_hook_next)
+			result = planner_hook_next(parse,
+#if PG_VERSION_NUM >= 130000
+									   query_string,
+#endif
+									   cursorOptions, boundParams);
+		else
+			result = standard_planner(parse,
+#if PG_VERSION_NUM >= 130000
+									  query_string,
 #endif
-		if (!proc_queryids[i])
-			proc_queryids[i] = parse->queryId;
+									  cursorOptions, boundParams);
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else if (pgws_enabled(nesting_level))
+			pgws_proc_queryids[i] = save_queryId;
+	}
+	PG_CATCH();
+	{
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else if (pgws_enabled(nesting_level))
+			pgws_proc_queryids[i] = save_queryId;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
+
+	return result;
+}
 
+/*
+ * ExecutorStart hook: save queryId for collector
+ */
+static void
+pgws_ExecutorStart(QueryDesc *queryDesc, int eflags)
+{
+	int			i = MyProc - ProcGlobal->allProcs;
+
+	if (pgws_enabled(nesting_level))
+		pgws_proc_queryids[i] = queryDesc->plannedstmt->queryId;
+	if (prev_ExecutorStart)
+		prev_ExecutorStart(queryDesc, eflags);
+	else
+		standard_ExecutorStart(queryDesc, eflags);
+}
+
+static void
+pgws_ExecutorRun(QueryDesc *queryDesc,
+				 ScanDirection direction,
+				 uint64 count
+#if PG_VERSION_NUM >= 100000 && PG_VERSION_NUM < 180000
+				 ,bool execute_once
+#endif
+)
+{
+	int			i = MyProc - ProcGlobal->allProcs;
+	uint64		save_queryId = pgws_proc_queryids[i];
+
+	nesting_level++;
+	PG_TRY();
+	{
+		if (prev_ExecutorRun)
+#if PG_VERSION_NUM >= 100000 && PG_VERSION_NUM < 180000
+			prev_ExecutorRun(queryDesc, direction, count, execute_once);
+#else
+			prev_ExecutorRun(queryDesc, direction, count);
+#endif
+		else
+#if PG_VERSION_NUM >= 100000 && PG_VERSION_NUM < 180000
+			standard_ExecutorRun(queryDesc, direction, count, execute_once);
+#else
+			standard_ExecutorRun(queryDesc, direction, count);
+#endif
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else
+			pgws_proc_queryids[i] = save_queryId;
+	}
+	PG_CATCH();
+	{
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else
+			pgws_proc_queryids[i] = save_queryId;
+		PG_RE_THROW();
 	}
+	PG_END_TRY();
+}
 
-	/* Invoke original hook if needed */
-	if (planner_hook_next)
-		return planner_hook_next(parse, cursorOptions, boundParams);
+static void
+pgws_ExecutorFinish(QueryDesc *queryDesc)
+{
+	int			i = MyProc - ProcGlobal->allProcs;
+	uint64		save_queryId = pgws_proc_queryids[i];
 
-	return standard_planner(parse, cursorOptions, boundParams);
+	nesting_level++;
+	PG_TRY();
+	{
+		if (prev_ExecutorFinish)
+			prev_ExecutorFinish(queryDesc);
+		else
+			standard_ExecutorFinish(queryDesc);
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else
+			pgws_proc_queryids[i] = save_queryId;
+	}
+	PG_CATCH();
+	{
+		nesting_level--;
+		pgws_proc_queryids[i] = save_queryId;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
 }
 
 /*
@@ -804,11 +1068,86 @@ pgws_planner_hook(Query *parse, int cursorOptions,
 static void
 pgws_ExecutorEnd(QueryDesc *queryDesc)
 {
-	if (MyProc)
-		proc_queryids[MyProc - ProcGlobal->allProcs] = UINT64CONST(0);
+	int			i = MyProc - ProcGlobal->allProcs;
+
+	if (nesting_level == 0)
+		pgws_proc_queryids[i] = UINT64CONST(0);
 
 	if (prev_ExecutorEnd)
 		prev_ExecutorEnd(queryDesc);
 	else
 		standard_ExecutorEnd(queryDesc);
 }
+
+static void
+pgws_ProcessUtility(PlannedStmt *pstmt,
+					const char *queryString,
+#if PG_VERSION_NUM >= 140000
+					bool readOnlyTree,
+#endif
+					ProcessUtilityContext context,
+					ParamListInfo params,
+					QueryEnvironment *queryEnv,
+					DestReceiver *dest,
+#if PG_VERSION_NUM >= 130000
+					QueryCompletion *qc
+#else
+					char *completionTag
+#endif
+)
+{
+	int			i = MyProc - ProcGlobal->allProcs;
+	uint64		save_queryId = 0;
+
+	if (pgws_enabled(nesting_level))
+	{
+		save_queryId = pgws_proc_queryids[i];
+		pgws_proc_queryids[i] = pstmt->queryId;
+	}
+
+	nesting_level++;
+	PG_TRY();
+	{
+		if (prev_ProcessUtility)
+			prev_ProcessUtility(pstmt, queryString,
+#if PG_VERSION_NUM >= 140000
+								readOnlyTree,
+#endif
+								context, params, queryEnv,
+								dest,
+#if PG_VERSION_NUM >= 130000
+								qc
+#else
+								completionTag
+#endif
+				);
+		else
+			standard_ProcessUtility(pstmt, queryString,
+#if PG_VERSION_NUM >= 140000
+									readOnlyTree,
+#endif
+									context, params, queryEnv,
+									dest,
+#if PG_VERSION_NUM >= 130000
+									qc
+#else
+									completionTag
+#endif
+				);
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else if (pgws_enabled(nesting_level))
+			pgws_proc_queryids[i] = save_queryId;
+	}
+	PG_CATCH();
+	{
+		nesting_level--;
+		if (nesting_level == 0)
+			pgws_proc_queryids[i] = UINT64CONST(0);
+		else if (pgws_enabled(nesting_level))
+			pgws_proc_queryids[i] = save_queryId;
+		PG_RE_THROW();
+	}
+	PG_END_TRY();
+}
diff --git a/pg_wait_sampling.h b/pg_wait_sampling.h
index 209114b..dab773c 100644
--- a/pg_wait_sampling.h
+++ b/pg_wait_sampling.h
@@ -2,7 +2,7 @@
  * pg_wait_sampling.h
  *		Headers for pg_wait_sampling extension.
  *
- * Copyright (c) 2015-2016, Postgres Professional
+ * Copyright (c) 2015-2025, Postgres Professional
  *
  * IDENTIFICATION
  *	  contrib/pg_wait_sampling/pg_wait_sampling.h
@@ -10,16 +10,10 @@
 #ifndef __PG_WAIT_SAMPLING_H__
 #define __PG_WAIT_SAMPLING_H__
 
-#include <postgres.h>
-
-/* Check PostgreSQL version */
-#if PG_VERSION_NUM < 90600
-	#error "You are trying to build pg_wait_sampling with PostgreSQL version lower than 9.6.  Please, check you environment."
-#endif
-
-#include "storage/proc.h"
+#include "datatype/timestamp.h"
+#include "storage/latch.h"
+#include "storage/lock.h"
 #include "storage/shm_mq.h"
-#include "utils/timestamp.h"
 
 #define	PG_WAIT_SAMPLING_MAGIC		0xCA94B107
 #define COLLECTOR_QUEUE_SIZE		(16 * 1024)
@@ -29,26 +23,26 @@
 
 typedef struct
 {
-	uint32			pid;
-	uint32			wait_event_info;
-	uint64			queryId;
-	uint64			count;
+	int			pid;
+	uint32		wait_event_info;
+	uint64		queryId;
+	uint64		count;
 } ProfileItem;
 
 typedef struct
 {
-	uint32			pid;
-	uint32			wait_event_info;
-	uint64			queryId;
-	TimestampTz		ts;
+	int			pid;
+	uint32		wait_event_info;
+	uint64		queryId;
+	TimestampTz ts;
 } HistoryItem;
 
 typedef struct
 {
-	bool			wraparound;
-	Size			index;
-	Size			count;
-	HistoryItem	   *items;
+	bool		wraparound;
+	Size		index;
+	Size		count;
+	HistoryItem *items;
 } History;
 
 typedef enum
@@ -61,26 +55,27 @@ typedef enum
 
 typedef struct
 {
-	Latch		   *latch;
-	SHMRequest		request;
-	int				historySize;
-	int				historyPeriod;
-	int				profilePeriod;
-	bool			profilePid;
-	bool			profileQueries;
+	Latch	   *latch;
+	SHMRequest	request;
 } CollectorShmqHeader;
 
+/* GUC variables */
+extern int	pgws_historySize;
+extern int	pgws_historyPeriod;
+extern int	pgws_profilePeriod;
+extern bool pgws_profilePid;
+extern int	pgws_profileQueries;
+extern bool pgws_sampleCpu;
+
 /* pg_wait_sampling.c */
-extern void check_shmem(void);
-extern CollectorShmqHeader *collector_hdr;
-extern shm_mq			   *collector_mq;
-extern uint64			   *proc_queryids;
-extern void read_current_wait(PGPROC *proc, HistoryItem *item);
-extern void init_lock_tag(LOCKTAG *tag, uint32 lock);
+extern CollectorShmqHeader *pgws_collector_hdr;
+extern shm_mq *pgws_collector_mq;
+extern uint64 *pgws_proc_queryids;
+extern void pgws_init_lock_tag(LOCKTAG *tag, uint32 lock);
+extern bool pgws_should_sample_proc(PGPROC *proc, int *pid_p, uint32 *wait_event_info_p);
 
 /* collector.c */
-extern void register_wait_collector(void);
-extern void alloc_history(History *, int);
-extern void collector_main(Datum main_arg);
+extern void pgws_register_wait_collector(void);
+extern PGDLLEXPORT void pgws_collector_main(Datum main_arg);
 
 #endif
diff --git a/run-tests.sh b/run-tests.sh
new file mode 100755
index 0000000..f3f1bba
--- /dev/null
+++ b/run-tests.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+set -ev
+
+PATH=/usr/lib/postgresql/$PG_MAJOR/bin:$PATH
+export PGDATA=/var/lib/postgresql/$PG_MAJOR/test
+export COPT=-Werror
+export USE_PGXS=1
+
+sudo chmod 1777 /var/lib/postgresql/$PG_MAJOR
+sudo chmod 1777 /var/run/postgresql
+
+make clean
+make
+
+sudo -E env PATH=$PATH make install
+
+initdb
+echo "shared_preload_libraries = pg_wait_sampling" >> $PGDATA/postgresql.conf
+
+pg_ctl -l logfile start
+make installcheck
+pg_ctl stop
diff --git a/run_tests.sh b/run_tests.sh
deleted file mode 100755
index 5e66bc9..0000000
--- a/run_tests.sh
+++ /dev/null
@@ -1,67 +0,0 @@
-#!/bin/bash
-
-# This is a main testing script for:
-#	* regression tests
-#	* testgres-based tests
-#	* cmocka-based tests
-# Copyright (c) 2017, Postgres Professional
-
-set -eux
-
-echo CHECK_CODE=$CHECK_CODE
-
-status=0
-
-# perform code analysis if necessary
-if [ "$CHECK_CODE" = "clang" ]; then
-    scan-build --status-bugs make USE_PGXS=1 || status=$?
-    exit $status
-
-elif [ "$CHECK_CODE" = "cppcheck" ]; then
-    cppcheck \
-		--template "{file} ({line}): {severity} ({id}): {message}" \
-        --enable=warning,portability,performance \
-        --suppress=redundantAssignment \
-        --suppress=uselessAssignmentPtrArg \
-		--suppress=literalWithCharPtrCompare \
-        --suppress=incorrectStringBooleanError \
-        --std=c89 *.c *.h 2> cppcheck.log
-
-    if [ -s cppcheck.log ]; then
-        cat cppcheck.log
-        status=1 # error
-    fi
-
-    exit $status
-fi
-
-# don't forget to "make clean"
-make USE_PGXS=1 clean
-
-# initialize database
-initdb
-
-# build extension
-make USE_PGXS=1 install
-
-# check build
-status=$?
-if [ $status -ne 0 ]; then exit $status; fi
-
-# add pg_pathman to shared_preload_libraries and restart cluster 'test'
-echo "shared_preload_libraries = 'pg_wait_sampling'" >> $PGDATA/postgresql.conf
-echo "port = 55435" >> $PGDATA/postgresql.conf
-pg_ctl start -l /tmp/postgres.log -w
-
-# check startup
-status=$?
-if [ $status -ne 0 ]; then cat /tmp/postgres.log; fi
-
-# run regression tests
-export PG_REGRESS_DIFF_OPTS="-w -U3" # for alpine's diff (BusyBox)
-PGPORT=55435 make USE_PGXS=1 installcheck || status=$?
-
-# show diff if it exists
-if test -f regression.diffs; then cat regression.diffs; fi
-
-exit $status
diff --git a/sql/queries.sql b/sql/queries.sql
index 65602ee..de44c6d 100644
--- a/sql/queries.sql
+++ b/sql/queries.sql
@@ -9,6 +9,10 @@ WITH t as (SELECT sum(0) FROM pg_wait_sampling_history)
 WITH t as (SELECT sum(0) FROM pg_wait_sampling_profile)
 	SELECT sum(0) FROM generate_series(1, 2), t;
 
+-- Some dummy checks just to be sure that all our functions work and return something.
+SELECT count(*) = 1 as test FROM pg_wait_sampling_get_current(pg_backend_pid());
+SELECT count(*) >= 0 as test FROM pg_wait_sampling_get_profile();
+SELECT count(*) >= 0 as test FROM pg_wait_sampling_get_history();
 SELECT pg_wait_sampling_reset_profile();
 
 DROP EXTENSION pg_wait_sampling;