Import batch-callback API from upstream (partial)

src/backend/commands/trigger.c

@@ -3928,13 +3928,26 @@ typedef struct AfterTriggersData
 	/* per-subtransaction-level data: */
 	AfterTriggersTransData *trans_stack;	/* array of structs shown below */
 	int			maxtransdepth;	/* allocated len of above array */
+
+	List	   *batch_callbacks;	/* List of AfterTriggerCallbackItem; for
+									 * deferred constraints */
+	bool		firing_batch_callbacks; /* true when in
+										 * FireAfterTriggerBatchCallbacks() */
+
+	/*
+	 * Incremented around the trigger-firing loops in AfterTriggerEndQuery,
+	 * AfterTriggerFireDeferred, and AfterTriggerSetState.  Used by
+	 * AfterTriggerIsActive() to signal that after-trigger firing is active.
+	 */
+	int			firing_depth;
 } AfterTriggersData;
 
 struct AfterTriggersQueryData
 {
 	AfterTriggerEventList events;	/* events pending from this query */
 	Tuplestorestate *fdw_tuplestore;	/* foreign tuples for said events */
 	List	   *tables;			/* list of AfterTriggersTableData, see below */
+	List	   *batch_callbacks;	/* List of AfterTriggerCallbackItem */
 };
 
 struct AfterTriggersTransData
@@ -3975,6 +3988,13 @@ struct AfterTriggersTableData
 	TupleTableSlot *storeslot;	/* for converting to tuplestore's format */
 };
 
+/* Entry in afterTriggers.batch_callbacks */
+typedef struct AfterTriggerCallbackItem
+{
+	AfterTriggerBatchCallback callback;
+	void	   *arg;
+} AfterTriggerCallbackItem;
+
 static AfterTriggersData afterTriggers;
 
 static void AfterTriggerExecute(EState *estate,
@@ -4009,6 +4029,7 @@ static SetConstraintState SetConstraintStateAddItem(SetConstraintState state,
 													Oid tgoid, bool tgisdeferred);
 static void cancel_prior_stmt_triggers(Oid relid, CmdType cmdType, int tgevent);
 
+static void FireAfterTriggerBatchCallbacks(List *callbacks);
 
 /*
  * Get the FDW tuplestore for the current trigger query level, creating it
@@ -5090,6 +5111,9 @@ AfterTriggerBeginXact(void)
 	 */
 	afterTriggers.firing_counter = (CommandId) 1;	/* mustn't be 0 */
 	afterTriggers.query_depth = -1;
+	afterTriggers.firing_depth = 0;
+	afterTriggers.batch_callbacks = NIL;
+	afterTriggers.firing_batch_callbacks = false;
 
 	/*
 	 * Verify that there is no leftover state remaining.  If these assertions
@@ -5174,6 +5198,7 @@ AfterTriggerEndQuery(EState *estate)
 	 */
 	qs = &afterTriggers.query_stack[afterTriggers.query_depth];
 
+	afterTriggers.firing_depth++;
 	for (;;)
 	{
 		if (afterTriggerMarkEvents(&qs->events, &afterTriggers.events, true))
@@ -5211,10 +5236,18 @@ AfterTriggerEndQuery(EState *estate)
 			break;
 	}
 
+	/*
+	 * Fire batch callbacks before releasing query-level storage and before
+	 * decrementing query_depth.  Callbacks may do real work (index probes,
+	 * error reporting).
+	 */
+	FireAfterTriggerBatchCallbacks(qs->batch_callbacks);
+
 	/* Release query-level-local storage, including tuplestores if any */
 	AfterTriggerFreeQuery(&afterTriggers.query_stack[afterTriggers.query_depth]);
 
 	afterTriggers.query_depth--;
+	afterTriggers.firing_depth--;
 }
 
 
@@ -5279,6 +5312,9 @@ AfterTriggerFreeQuery(AfterTriggersQueryData *qs)
 	 */
 	qs->tables = NIL;
 	list_free_deep(tables);
+
+	list_free_deep(qs->batch_callbacks);
+	qs->batch_callbacks = NIL;
 }
 
 
@@ -5318,6 +5354,7 @@ AfterTriggerFireDeferred(void)
 	 * Run all the remaining triggers.  Loop until they are all gone, in case
 	 * some trigger queues more for us to do.
 	 */
+	afterTriggers.firing_depth++;
 	while (afterTriggerMarkEvents(events, NULL, false))
 	{
 		CommandId	firing_id = afterTriggers.firing_counter++;
@@ -5326,9 +5363,15 @@ AfterTriggerFireDeferred(void)
 			break;				/* all fired */
 	}
 
+	/* Flush any fast-path batches accumulated by the triggers just fired. */
+	FireAfterTriggerBatchCallbacks(afterTriggers.batch_callbacks);
+
+	afterTriggers.firing_depth--;
+
 	/*
-	 * We don't bother freeing the event list, since it will go away anyway
-	 * (and more efficiently than via pfree) in AfterTriggerEndXact.
+	 * We don't bother freeing the event list or batch_callbacks, since they
+	 * will go away anyway (and more efficiently than via pfree) in
+	 * AfterTriggerEndXact.
 	 */
 
 	if (snap_pushed)
@@ -5390,6 +5433,12 @@ AfterTriggerEndXact(bool isCommit)
 
 	/* No more afterTriggers manipulation until next transaction starts. */
 	afterTriggers.query_depth = -1;
+
+	afterTriggers.firing_depth = 0;
+
+	list_free_deep(afterTriggers.batch_callbacks);
+	afterTriggers.batch_callbacks = NIL;
+	afterTriggers.firing_batch_callbacks = false;
 }
 
 /*
@@ -5536,6 +5585,9 @@ AfterTriggerEndSubXact(bool isCommit)
 			}
 		}
 	}
+
+	/* Reset in case a callback threw an error while firing. */
+	afterTriggers.firing_batch_callbacks = false;
 }
 
 /*
@@ -5672,6 +5724,7 @@ AfterTriggerEnlargeQueryState(void)
 		qs->events.tailfree = NULL;
 		qs->fdw_tuplestore = NULL;
 		qs->tables = NIL;
+		qs->batch_callbacks = NIL;
 
 		++init_depth;
 	}
@@ -6021,6 +6074,7 @@ AfterTriggerSetState(ConstraintsSetStmt *stmt)
 		AfterTriggerEventList *events = &afterTriggers.events;
 		bool		snapshot_set = false;
 
+		afterTriggers.firing_depth++;
 		while (afterTriggerMarkEvents(events, NULL, true))
 		{
 			CommandId	firing_id = afterTriggers.firing_counter++;
@@ -6050,6 +6104,14 @@ AfterTriggerSetState(ConstraintsSetStmt *stmt)
 				break;			/* all fired */
 		}
 
+		/*
+		 * Flush any fast-path batches accumulated by the triggers just fired.
+		 */
+		FireAfterTriggerBatchCallbacks(afterTriggers.batch_callbacks);
+		afterTriggers.firing_depth--;
+		list_free_deep(afterTriggers.batch_callbacks);
+		afterTriggers.batch_callbacks = NIL;
+
 		if (snapshot_set)
 			PopActiveSnapshot();
 	}
@@ -6680,3 +6742,85 @@ pg_trigger_depth(PG_FUNCTION_ARGS)
 {
 	PG_RETURN_INT32(MyTriggerDepth);
 }
+
+/*
+ * RegisterAfterTriggerBatchCallback
+ *		Register a function to be called when the current trigger-firing
+ *		batch completes.
+ *
+ * Must be called from within a trigger function's execution context
+ * (i.e., while afterTriggers state is active).
+ *
+ * The callback list is cleared after invocation, so the caller must
+ * re-register for each new batch if needed.
+ */
+void
+RegisterAfterTriggerBatchCallback(AfterTriggerBatchCallback callback,
+								  void *arg)
+{
+	AfterTriggerCallbackItem *item;
+	MemoryContext oldcxt;
+
+	/*
+	 * Allocate in TopTransactionContext so the item survives for the duration
+	 * of the batch, which may span multiple trigger invocations.
+	 *
+	 * Must be called while afterTriggers is active; callbacks registered
+	 * outside a trigger-firing context would never fire.
+	 */
+	Assert(afterTriggers.firing_depth > 0);
+	Assert(!afterTriggers.firing_batch_callbacks);
+	oldcxt = MemoryContextSwitchTo(TopTransactionContext);
+	item = palloc(sizeof(AfterTriggerCallbackItem));
+	item->callback = callback;
+	item->arg = arg;
+	if (afterTriggers.query_depth >= 0)
+	{
+		AfterTriggersQueryData *qs =
+			&afterTriggers.query_stack[afterTriggers.query_depth];
+
+		qs->batch_callbacks = lappend(qs->batch_callbacks, item);
+	}
+	else
+		afterTriggers.batch_callbacks =
+			lappend(afterTriggers.batch_callbacks, item);
+	MemoryContextSwitchTo(oldcxt);
+}
+
+/*
+ * FireAfterTriggerBatchCallbacks
+ *		Invoke all callbacks in the given list.
+ *
+ * Memory cleanup of the list and its items is handled by the caller
+ * (AfterTriggerFreeQuery for query-level callbacks, AfterTriggerEndXact
+ * for top-level deferred callbacks).
+ */
+static void
+FireAfterTriggerBatchCallbacks(List *callbacks)
+{
+	ListCell   *lc;
+
+	Assert(afterTriggers.firing_depth > 0);
+	afterTriggers.firing_batch_callbacks = true;
+	foreach(lc, callbacks)
+	{
+		AfterTriggerCallbackItem *item = lfirst(lc);
+
+		item->callback(item->arg);
+	}
+	afterTriggers.firing_batch_callbacks = false;
+}
+
+/*
+ * AfterTriggerIsActive
+ *		Returns true if we're inside the after-trigger framework where
+ *		registered batch callbacks will actually be invoked.
+ *
+ * This is false during validateForeignKeyConstraint(), which calls
+ * RI trigger functions directly outside the after-trigger framework.
+ */
+bool
+AfterTriggerIsActive(void)
+{
+	return afterTriggers.firing_depth > 0;
+}

src/include/commands/trigger.h

@@ -303,4 +303,25 @@ extern void RI_PartitionRemove_Check(Trigger *trigger, Relation fk_rel,
 
 extern int	RI_FKey_trigger_type(Oid tgfoid);
 
+/*
+ * Callback type for end-of-trigger-batch callbacks.
+ *
+ * Currently used by ri_triggers.c to flush fast-path FK batches and
+ * clean up associated resources.
+ *
+ * Registered via RegisterAfterTriggerBatchCallback().  Invoked when
+ * the current trigger-firing batch completes:
+ *	- AfterTriggerEndQuery()      (immediate constraints)
+ *	- AfterTriggerFireDeferred()  (deferred constraints at COMMIT)
+ *	- AfterTriggerSetState()      (SET CONSTRAINTS IMMEDIATE)
+ *
+ * The callback list is cleared after each batch.  Callers must
+ * re-register if they need to be called again in a subsequent batch.
+ */
+typedef void (*AfterTriggerBatchCallback) (void *arg);
+
+extern void RegisterAfterTriggerBatchCallback(AfterTriggerBatchCallback callback,
+											  void *arg);
+extern bool AfterTriggerIsActive(void);
+
 #endif							/* TRIGGER_H */

src/tools/pgindent/typedefs.list

@@ -28,6 +28,8 @@ ActiveSnapshotElt
 AddForeignUpdateTargets_function
 AffixNode
 AffixNodeData
+AfterTriggerBatchCallback
+AfterTriggerCallbackItem
 AfterTriggerEvent
 AfterTriggerEventChunk
 AfterTriggerEventData