From 8cee8430327ab1922d9b3b82134e3eae7b1129b2 Mon Sep 17 00:00:00 2001
From: Daniel Gustafsson <dgustafsson@postgresql.org>
Date: Fri, 11 Apr 2025 11:50:47 +0200
Subject: [PATCH v4] Don't try to enlarge resourceowner when releasing

If DSA or DSM segment is attached to, or created, while a
transaction's resource owner is being released, it throws
an error when trying to enlarge the resource owner.

 Session 1 (PID=1000):
 =# begin;
 =# select 1/0;
 ERROR:  division by zero

 Session 2:
 =# select * from pg_get_process_memory_contexts(1000,false,1);

 Session 1 terminates with:
 ERROR:  ResourceOwnerEnlarge called after release started
 FATAL:  terminating connection because protocol synchronization was lost

Fix by checking whether the resource owner is being released
before adding the segment to it.  In order to test this an
injection point is added.  While the bug exists further back,
backpatching is currently limited to v17 where the injection
point framework was introduced as it cannot be easily tested
without it.

Author: Rahila Syed <rahilasyed90@gmail.com>
Reported-by: Fujii Masao <masao.fujii@oss.nttdata.com>
Discussion: https://postgr.es/m/a1a7e2b7-8f33-4313-baff-42e92ec14fd3%40oss.nttdata.com
Discussion: https://postgr.es/m/CAH2L28shr0j3JE5V3CXDFmDH-agTSnh2V8pR23X0UhRMbDQD9Q@mail.gmail.com
Backpatch-through: 17
---
 src/backend/storage/ipc/dsm.c                 | 13 ++++--
 src/backend/utils/mmgr/dsa.c                  | 10 ++++-
 src/backend/utils/resowner/resowner.c         |  8 ++++
 src/include/utils/resowner.h                  |  1 +
 .../modules/test_dsa/expected/test_dsa.out    | 10 +++++
 src/test/modules/test_dsa/sql/test_dsa.sql    |  5 +++
 src/test/modules/test_dsa/test_dsa--1.0.sql   |  4 ++
 src/test/modules/test_dsa/test_dsa.c          | 44 +++++++++++++++++--
 8 files changed, 86 insertions(+), 9 deletions(-)

diff --git a/src/backend/storage/ipc/dsm.c b/src/backend/storage/ipc/dsm.c
index f92a52a00e6..479983882bb 100644
--- a/src/backend/storage/ipc/dsm.c
+++ b/src/backend/storage/ipc/dsm.c
@@ -934,6 +934,10 @@ void
 dsm_unpin_mapping(dsm_segment *seg)
 {
 	Assert(seg->resowner == NULL);
+
+	if (CurrentResourceOwner && IsResourceOwnerReleasing(CurrentResourceOwner))
+		return;
+
 	ResourceOwnerEnlarge(CurrentResourceOwner);
 	seg->resowner = CurrentResourceOwner;
 	ResourceOwnerRememberDSM(seg->resowner, seg);
@@ -1202,7 +1206,7 @@ dsm_create_descriptor(void)
 {
 	dsm_segment *seg;
 
-	if (CurrentResourceOwner)
+	if (CurrentResourceOwner && !IsResourceOwnerReleasing(CurrentResourceOwner))
 		ResourceOwnerEnlarge(CurrentResourceOwner);
 
 	seg = MemoryContextAlloc(TopMemoryContext, sizeof(dsm_segment));
@@ -1213,10 +1217,13 @@ dsm_create_descriptor(void)
 	seg->impl_private = NULL;
 	seg->mapped_address = NULL;
 	seg->mapped_size = 0;
+	seg->resowner = NULL;
 
-	seg->resowner = CurrentResourceOwner;
-	if (CurrentResourceOwner)
+	if (CurrentResourceOwner && !IsResourceOwnerReleasing(CurrentResourceOwner))
+	{
 		ResourceOwnerRememberDSM(CurrentResourceOwner, seg);
+		seg->resowner = CurrentResourceOwner;
+	}
 
 	slist_init(&seg->on_detach);
 
diff --git a/src/backend/utils/mmgr/dsa.c b/src/backend/utils/mmgr/dsa.c
index 17d4f7a7a06..7b2c5a40157 100644
--- a/src/backend/utils/mmgr/dsa.c
+++ b/src/backend/utils/mmgr/dsa.c
@@ -1282,7 +1282,10 @@ create_internal(void *place, size_t size,
 	 */
 	area = palloc(sizeof(dsa_area));
 	area->control = control;
-	area->resowner = CurrentResourceOwner;
+	if (CurrentResourceOwner && !IsResourceOwnerReleasing(CurrentResourceOwner))
+		area->resowner = CurrentResourceOwner;
+	else
+		area->resowner = NULL;
 	memset(area->segment_maps, 0, sizeof(dsa_segment_map) * DSA_MAX_SEGMENTS);
 	area->high_segment_index = 0;
 	area->freed_segment_counter = 0;
@@ -1338,7 +1341,10 @@ attach_internal(void *place, dsm_segment *segment, dsa_handle handle)
 	/* Build the backend-local area object. */
 	area = palloc(sizeof(dsa_area));
 	area->control = control;
-	area->resowner = CurrentResourceOwner;
+	if (CurrentResourceOwner && !IsResourceOwnerReleasing(CurrentResourceOwner))
+		area->resowner = CurrentResourceOwner;
+	else
+		area->resowner = NULL;
 	memset(&area->segment_maps[0], 0,
 		   sizeof(dsa_segment_map) * DSA_MAX_SEGMENTS);
 	area->high_segment_index = 0;
diff --git a/src/backend/utils/resowner/resowner.c b/src/backend/utils/resowner/resowner.c
index d39f3e1b655..aec0ade61ec 100644
--- a/src/backend/utils/resowner/resowner.c
+++ b/src/backend/utils/resowner/resowner.c
@@ -52,6 +52,7 @@
 #include "storage/ipc.h"
 #include "storage/predicate.h"
 #include "storage/proc.h"
+#include "utils/injection_point.h"
 #include "utils/memutils.h"
 #include "utils/resowner.h"
 
@@ -702,6 +703,7 @@ ResourceOwnerReleaseInternal(ResourceOwner owner,
 		Assert(phase == RESOURCE_RELEASE_BEFORE_LOCKS);
 		Assert(!owner->sorted);
 		owner->releasing = true;
+		INJECTION_POINT("resowner-release");
 	}
 	else
 	{
@@ -1111,3 +1113,9 @@ ResourceOwnerForgetAioHandle(ResourceOwner owner, struct dlist_node *ioh_node)
 {
 	dlist_delete_from(&owner->aio_handles, ioh_node);
 }
+
+bool
+IsResourceOwnerReleasing(ResourceOwner owner)
+{
+	return owner->releasing;
+}
diff --git a/src/include/utils/resowner.h b/src/include/utils/resowner.h
index aede4bfc820..7b4c3765a1f 100644
--- a/src/include/utils/resowner.h
+++ b/src/include/utils/resowner.h
@@ -163,6 +163,7 @@ extern void ReleaseAuxProcessResources(bool isCommit);
 struct LOCALLOCK;
 extern void ResourceOwnerRememberLock(ResourceOwner owner, struct LOCALLOCK *locallock);
 extern void ResourceOwnerForgetLock(ResourceOwner owner, struct LOCALLOCK *locallock);
+extern bool IsResourceOwnerReleasing(ResourceOwner owner);
 
 /* special support for AIO */
 struct dlist_node;
diff --git a/src/test/modules/test_dsa/expected/test_dsa.out b/src/test/modules/test_dsa/expected/test_dsa.out
index 266010e77fe..e6bd8184143 100644
--- a/src/test/modules/test_dsa/expected/test_dsa.out
+++ b/src/test/modules/test_dsa/expected/test_dsa.out
@@ -11,3 +11,13 @@ SELECT test_dsa_resowners();
  
 (1 row)
 
+SELECT inject_dsa_in_resowner_release();
+ inject_dsa_in_resowner_release 
+--------------------------------
+ 
+(1 row)
+
+begin;
+select 1/0;
+ERROR:  division by zero
+commit;
diff --git a/src/test/modules/test_dsa/sql/test_dsa.sql b/src/test/modules/test_dsa/sql/test_dsa.sql
index c3d8db94372..4d0c90bd5ec 100644
--- a/src/test/modules/test_dsa/sql/test_dsa.sql
+++ b/src/test/modules/test_dsa/sql/test_dsa.sql
@@ -2,3 +2,8 @@ CREATE EXTENSION test_dsa;
 
 SELECT test_dsa_basic();
 SELECT test_dsa_resowners();
+SELECT inject_dsa_in_resowner_release();
+
+begin;
+select 1/0;
+commit;
diff --git a/src/test/modules/test_dsa/test_dsa--1.0.sql b/src/test/modules/test_dsa/test_dsa--1.0.sql
index 2904cb23525..2629b469197 100644
--- a/src/test/modules/test_dsa/test_dsa--1.0.sql
+++ b/src/test/modules/test_dsa/test_dsa--1.0.sql
@@ -10,3 +10,7 @@ CREATE FUNCTION test_dsa_basic()
 CREATE FUNCTION test_dsa_resowners()
 	RETURNS pg_catalog.void
 	AS 'MODULE_PATHNAME' LANGUAGE C;
+
+CREATE FUNCTION inject_dsa_in_resowner_release()
+	RETURNS pg_catalog.void
+	AS 'MODULE_PATHNAME' LANGUAGE C;
diff --git a/src/test/modules/test_dsa/test_dsa.c b/src/test/modules/test_dsa/test_dsa.c
index cd24d0f4873..c9931c63d8b 100644
--- a/src/test/modules/test_dsa/test_dsa.c
+++ b/src/test/modules/test_dsa/test_dsa.c
@@ -14,15 +14,19 @@
 
 #include "fmgr.h"
 #include "storage/lwlock.h"
+#include "utils/builtins.h"
 #include "utils/dsa.h"
+#include "utils/injection_point.h"
 #include "utils/resowner.h"
 
 PG_MODULE_MAGIC;
 
-/* Test basic DSA functionality */
-PG_FUNCTION_INFO_V1(test_dsa_basic);
-Datum
-test_dsa_basic(PG_FUNCTION_ARGS)
+static void test_dsa_common(void);
+
+extern PGDLLEXPORT void inj_create_dsa(const char *name, const void *private_data);
+
+static void
+test_dsa_common(void)
 {
 	int			tranche_id;
 	dsa_area   *a;
@@ -54,7 +58,14 @@ test_dsa_basic(PG_FUNCTION_ARGS)
 	}
 
 	dsa_detach(a);
+}
 
+/* Test basic DSA functionality */
+PG_FUNCTION_INFO_V1(test_dsa_basic);
+Datum
+test_dsa_basic(PG_FUNCTION_ARGS)
+{
+	test_dsa_common();
 	PG_RETURN_VOID();
 }
 
@@ -111,3 +122,28 @@ test_dsa_resowners(PG_FUNCTION_ARGS)
 
 	PG_RETURN_VOID();
 }
+
+/*
+ * Test to create a DSA in a resource owner which is currently releasing.
+ */
+PG_FUNCTION_INFO_V1(inject_dsa_in_resowner_release);
+Datum
+inject_dsa_in_resowner_release(PG_FUNCTION_ARGS)
+{
+#ifdef USE_INJECTION_POINTS
+	InjectionPointAttach("resowner-release",
+						 "test_dsa",
+						 "inj_create_dsa",
+						 NULL,
+						 0);
+#endif
+	PG_RETURN_VOID();
+}
+void
+inj_create_dsa(const char *name, const void *private_data)
+{
+	(void) name;
+	(void) private_data;
+
+	test_dsa_common();
+}
-- 
2.39.3 (Apple Git-146)

