From 99246f4de6668d7a3f6054f672e67e66a1982edf Mon Sep 17 00:00:00 2001
From: Tomas Vondra <tomas.vondra@postgresql.org>
Date: Wed, 24 May 2023 13:12:03 +0200
Subject: [PATCH 2/3] Switch to short-lived MemoryContext

---
 src/backend/commands/trigger.c  | 39 ++++++++++++++++++++++-----------
 src/backend/executor/execMain.c | 14 ++++++------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/backend/commands/trigger.c b/src/backend/commands/trigger.c
index fb11203c473..f5872ca6da2 100644
--- a/src/backend/commands/trigger.c
+++ b/src/backend/commands/trigger.c
@@ -2867,6 +2867,7 @@ ExecBSUpdateTriggers(EState *estate, ResultRelInfo *relinfo)
 	int			i;
 	TriggerData LocTriggerData = {0};
 	Bitmapset  *updatedCols;
+	MemoryContext	oldcxt;
 
 	trigdesc = relinfo->ri_TrigDesc;
 
@@ -2883,8 +2884,12 @@ ExecBSUpdateTriggers(EState *estate, ResultRelInfo *relinfo)
 	/* statement-level triggers operate on the parent table */
 	Assert(relinfo->ri_RootResultRelInfo == NULL);
 
+	oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
+
 	updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
 
+	MemoryContextSwitchTo(oldcxt);
+
 	LocTriggerData.type = T_TriggerData;
 	LocTriggerData.tg_event = TRIGGER_EVENT_UPDATE |
 		TRIGGER_EVENT_BEFORE;
@@ -2916,8 +2921,6 @@ ExecBSUpdateTriggers(EState *estate, ResultRelInfo *relinfo)
 					(errcode(ERRCODE_E_R_I_E_TRIGGER_PROTOCOL_VIOLATED),
 					 errmsg("BEFORE STATEMENT trigger cannot return a value")));
 	}
-
-	bms_free(updatedCols);
 }
 
 void
@@ -2931,7 +2934,14 @@ ExecASUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 
 	if (trigdesc && trigdesc->trig_update_after_statement)
 	{
-		Bitmapset *updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
+		MemoryContext	oldcxt;
+		Bitmapset	   *updatedCols;
+
+		oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
+
+		updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
+
+		MemoryContextSwitchTo(oldcxt);
 
 		AfterTriggerSaveEvent(estate, relinfo, NULL, NULL,
 							  TRIGGER_EVENT_UPDATE,
@@ -2939,8 +2949,6 @@ ExecASUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 							  updatedCols,
 							  transition_capture,
 							  false);
-
-		bms_free(updatedCols);
 	}
 }
 
@@ -2962,6 +2970,7 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 	TriggerData LocTriggerData = {0};
 	int			i;
 	Bitmapset  *updatedCols;
+	MemoryContext	oldcxt;
 	LockTupleMode lockmode;
 
 	/* Determine lock mode to use */
@@ -3015,7 +3024,13 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 		TRIGGER_EVENT_ROW |
 		TRIGGER_EVENT_BEFORE;
 	LocTriggerData.tg_relation = relinfo->ri_RelationDesc;
+
+	oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
+
 	updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
+
+	MemoryContextSwitchTo(oldcxt);
+
 	LocTriggerData.tg_updatedcols = updatedCols;
 	for (i = 0; i < trigdesc->numtriggers; i++)
 	{
@@ -3051,9 +3066,6 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 				heap_freetuple(trigtuple);
 			if (should_free_new)
 				heap_freetuple(oldtuple);
-
-			bms_free(updatedCols);
-
 			return false;		/* "do nothing" */
 		}
 		else if (newtuple != oldtuple)
@@ -3079,8 +3091,6 @@ ExecBRUpdateTriggers(EState *estate, EPQState *epqstate,
 	if (should_free_trig)
 		heap_freetuple(trigtuple);
 
-	bms_free(updatedCols);
-
 	return true;
 }
 
@@ -3120,7 +3130,8 @@ ExecARUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 		 */
 		TupleTableSlot *oldslot;
 		ResultRelInfo *tupsrc;
-		Bitmapset *updatedCols;
+		MemoryContext	oldcxt;
+		Bitmapset	   *updatedCols;
 
 		Assert((src_partinfo != NULL && dst_partinfo != NULL) ||
 			   !is_crosspart_update);
@@ -3143,8 +3154,12 @@ ExecARUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 		else
 			ExecClearTuple(oldslot);
 
+		oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
+
 		updatedCols = ExecGetAllUpdatedCols(relinfo, estate);
 
+		MemoryContextSwitchTo(oldcxt);
+
 		AfterTriggerSaveEvent(estate, relinfo,
 							  src_partinfo, dst_partinfo,
 							  TRIGGER_EVENT_UPDATE,
@@ -3153,8 +3168,6 @@ ExecARUpdateTriggers(EState *estate, ResultRelInfo *relinfo,
 							  updatedCols,
 							  transition_capture,
 							  is_crosspart_update);
-
-		bms_free(updatedCols);
 	}
 }
 
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 98502b956c2..71e93486c42 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -2367,7 +2367,9 @@ ExecUpdateLockMode(EState *estate, ResultRelInfo *relinfo)
 {
 	Bitmapset  *keyCols;
 	Bitmapset  *updatedCols;
-	LockTupleMode	lockMode;
+	MemoryContext oldcxt;
+
+	oldcxt = MemoryContextSwitchTo(GetPerTupleMemoryContext(estate));
 
 	/*
 	 * Compute lock mode to use.  If columns that are part of the key have not
@@ -2378,14 +2380,12 @@ ExecUpdateLockMode(EState *estate, ResultRelInfo *relinfo)
 	keyCols = RelationGetIndexAttrBitmap(relinfo->ri_RelationDesc,
 										 INDEX_ATTR_BITMAP_KEY);
 
-	if (bms_overlap(keyCols, updatedCols))
-		lockMode = LockTupleExclusive;
-	else
-		lockMode = LockTupleNoKeyExclusive;
+	MemoryContextSwitchTo(oldcxt);
 
-	bms_free(updatedCols);
+	if (bms_overlap(keyCols, updatedCols))
+		return LockTupleExclusive;
 
-	return lockMode;
+	return LockTupleNoKeyExclusive;
 }
 
 /*
-- 
2.40.1