v1-0013-fe_utils-Provide-cancel_set-for-fast-quit-paths.patch

From 7e897c4c8cd99cce7b475b8339e50079c984a5b1 Mon Sep 17 00:00:00 2001
From: Thomas Munro <thomas.munro@gmail.com>
Date: Mon, 6 Jul 2026 16:53:46 +1200
Subject: [PATCH v1 13/14] fe_utils: Provide cancel_set for fast quit paths.

Provide a simpler way for frontend programs managing many connections to
arrange cleanup.  Now that we can run more complicated code on eg SIGINT
via pg_signal_processor, this API provides an easy way to manage a set
of connections that need to be canceled.  We couldn't have done this
before, because it requires locks.  A follow-up patch will use this for
pg_dump/pg_restore.

XXX Experimental!

Discussion:
Reviewed-by:
---
 src/fe_utils/Makefile             |   1 +
 src/fe_utils/cancel_set.c         | 158 ++++++++++++++++++++++++++++++
 src/fe_utils/meson.build          |   1 +
 src/include/fe_utils/cancel_set.h |  20 ++++
 src/tools/pgindent/typedefs.list  |   2 +
 5 files changed, 182 insertions(+)
 create mode 100644 src/fe_utils/cancel_set.c
 create mode 100644 src/include/fe_utils/cancel_set.h

diff --git a/src/fe_utils/Makefile b/src/fe_utils/Makefile
index cbfbf93ac69..6ca39b7cf9a 100644
--- a/src/fe_utils/Makefile
+++ b/src/fe_utils/Makefile
@@ -27,6 +27,7 @@ OBJS = \
 	astreamer_tar.o \
 	astreamer_zstd.o \
 	cancel.o \
+	cancel_set.o \
 	conditional.o \
 	connect_utils.o \
 	mbprint.o \
diff --git a/src/fe_utils/cancel_set.c b/src/fe_utils/cancel_set.c
new file mode 100644
index 00000000000..cba1cb02f06
--- /dev/null
+++ b/src/fe_utils/cancel_set.c
@@ -0,0 +1,158 @@
+#include "c.h"
+
+#include "fe_utils/cancel_set.h"
+#include "lib/ilist.h"
+#include "port/pg_threads.h"
+
+typedef struct cancel_set_entry
+{
+	dlist_node	node;
+	PGconn	   *connection;
+	PGcancelConn *cancel;
+} cancel_set_entry;
+
+struct cancel_set
+{
+	pg_mtx_t	mutex;
+	dlist_head	entries;
+};
+
+cancel_set *
+cancel_set_alloc(void)
+{
+	cancel_set *cs;
+
+	if (!(cs = malloc(sizeof(*cs))))
+		return NULL;
+
+	pg_mtx_init(&cs->mutex, pg_mtx_plain);
+	dlist_init(&cs->entries);
+
+	return cs;
+}
+
+void
+cancel_set_free(cancel_set *cs)
+{
+	cancel_set_remove_all(cs);
+	pg_mtx_destroy(&cs->mutex);
+	free(cs);
+}
+
+/*
+ * Add a connection to the set.  Returns false on failure.
+ */
+bool
+cancel_set_add(cancel_set *cs, PGconn *connection)
+{
+	bool		result;
+	cancel_set_entry *entry;
+	PGcancelConn *cancel;
+
+	cancel = PQcancelCreate(connection);
+	if (cancel == NULL)
+		return false;
+
+	pg_mtx_lock(&cs->mutex);
+	entry = malloc(sizeof(*entry));
+	if (entry)
+	{
+		entry->connection = connection;
+		entry->cancel = cancel;
+		dlist_push_tail(&cs->entries, &entry->node);
+		result = true;
+	}
+	else
+	{
+		PQcancelFinish(cancel);
+		result = false;
+	}
+	pg_mtx_unlock(&cs->mutex);
+
+	return result;
+}
+
+static int
+cancel_set_scan(cancel_set *cs, PGconn *connection, bool cancel, bool remove)
+{
+	dlist_mutable_iter iter;
+	int			count = 0;
+
+	pg_mtx_lock(&cs->mutex);
+	dlist_foreach_modify(iter, &cs->entries)
+	{
+		cancel_set_entry *entry;
+
+		entry = dlist_container(cancel_set_entry, node, iter.cur);
+		if (connection == NULL || entry->connection == connection)
+		{
+			count++;
+			if (cancel)
+			{
+				/*
+				 * XXX Pin these so they can't be removed yet, and then
+				 * perform the network request without holding the lock, and
+				 * then go back and unpin them?
+				 *
+				 * XXX Would it be possible to detect connections that match,
+				 * and collect multiple cancel packets that we send together
+				 * in one network request?
+				 */
+				PQcancelBlocking(entry->cancel);
+			}
+			if (remove)
+			{
+				dlist_delete(&entry->node);
+				PQcancelFinish(entry->cancel);
+				free(entry);
+			}
+			if (connection)
+				break;
+		}
+	}
+	pg_mtx_unlock(&cs->mutex);
+
+	return count;
+}
+
+/* Cancel a connection, without removing it. */
+bool
+cancel_set_cancel(cancel_set *cs, PGconn *connection)
+{
+	return cancel_set_scan(cs, connection, true, false) != 0;
+}
+
+/* Cancel and remove a connection. */
+bool
+cancel_set_cancel_and_remove(cancel_set *cs, PGconn *connection)
+{
+	return cancel_set_scan(cs, connection, true, true) != 0;
+}
+
+/* Remove a connection from the set. */
+bool
+cancel_set_remove(cancel_set *cs, PGconn *connection)
+{
+	return cancel_set_scan(cs, connection, false, true) != 0;
+}
+
+/* Cancel all connections in the set, without removing them. */
+int
+cancel_set_cancel_all(cancel_set *cs)
+{
+	return cancel_set_scan(cs, NULL, true, false);
+}
+
+/* Cancel and remove all connections in the set. */
+int
+cancel_set_cancel_and_remove_all(cancel_set *cs)
+{
+	return cancel_set_scan(cs, NULL, true, true);
+}
+
+/* Remove all connections from the set. */
+int
+cancel_set_remove_all(cancel_set *cs)
+{
+	return cancel_set_scan(cs, NULL, false, true);
+}
diff --git a/src/fe_utils/meson.build b/src/fe_utils/meson.build
index 86befca192e..b84edde2dfd 100644
--- a/src/fe_utils/meson.build
+++ b/src/fe_utils/meson.build
@@ -8,6 +8,7 @@ fe_utils_sources = files(
   'astreamer_tar.c',
   'astreamer_zstd.c',
   'cancel.c',
+  'cancel_set.c',
   'conditional.c',
   'connect_utils.c',
   'mbprint.c',
diff --git a/src/include/fe_utils/cancel_set.h b/src/include/fe_utils/cancel_set.h
new file mode 100644
index 00000000000..b75f0b7cd04
--- /dev/null
+++ b/src/include/fe_utils/cancel_set.h
@@ -0,0 +1,20 @@
+#ifndef CANCEL_SET
+#define CANCEL_SET
+
+#include "libpq-fe.h"
+
+typedef struct cancel_set cancel_set;
+
+extern cancel_set *cancel_set_alloc(void);
+extern void cancel_set_free(cancel_set *e);
+
+extern bool cancel_set_add(cancel_set *e, PGconn *connection);
+extern bool cancel_set_remove(cancel_set *e, PGconn *connection);
+extern bool cancel_set_cancel(cancel_set *e, PGconn *connection);
+extern bool cancel_set_cancel_and_remove(cancel_set *e, PGconn *connection);
+
+extern int	cancel_set_remove_all(cancel_set *e);
+extern int	cancel_set_cancel_all(cancel_set *e);
+extern int	cancel_set_cancel_and_remove_all(cancel_set *e);
+
+#endif
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index e830092ee62..2ece9b8a92e 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -3667,6 +3667,8 @@ btree_gin_leftmost_function
 build_simple_rel_hook_type
 bytea
 cached_re_str
+cancel_set
+cancel_set_entry
 canonicalize_state
 cashKEY
 catalogid_hash
-- 
2.47.3