Add logic for ExecutorFinish + add test to check this logic. Some fixes for tests.

Author: pipichinka

src/trace_execute_nodes.c

@@ -188,7 +188,6 @@ TraceExecProcNodeHook(PlanState *node)
 
 	clock_gettime(CLOCK_MONOTONIC, &time);
 	startTime = time.tv_sec * 1000000000L + time.tv_nsec;
-	nodeData->traceData.totalCalls++;
 
 	/* call real node */
 	result = nodeData->traceData.execProcNode(node);
@@ -197,6 +196,7 @@ TraceExecProcNodeHook(PlanState *node)
 
 	clock_gettime(CLOCK_MONOTONIC, &time);
 	timeDiff = time.tv_sec * 1000000000L + time.tv_nsec - startTime;
+	nodeData->traceData.totalCalls++;
 
 	if (timeDiff > nodeData->traceData.maxTime)
 		nodeData->traceData.maxTime = timeDiff;
@@ -315,7 +315,6 @@ TraceExprNodeHook(struct ExprState *expression, struct ExprContext *econtext, bo
 
 	clock_gettime(CLOCK_MONOTONIC, &time);
 	startTime = time.tv_sec * 1000000000L + time.tv_nsec;
-	exprData->traceData.totalCalls++;
 
 	/* call real node */
 	result = exprData->traceData.evalfunc(expression, econtext, isNull);
@@ -324,6 +323,7 @@ TraceExprNodeHook(struct ExprState *expression, struct ExprContext *econtext, bo
 
 	clock_gettime(CLOCK_MONOTONIC, &time);
 	timeDiff = time.tv_sec * 1000000000L + time.tv_nsec - startTime;
+	exprData->traceData.totalCalls++;
 
 	if (timeDiff > exprData->traceData.maxTime)
 		exprData->traceData.maxTime = timeDiff;

src/trace_session.c

@@ -6,10 +6,6 @@
 #include "utils/jsonb.h"
 #include "utils/guc.h"
 #include "commands/explain.h"
-#if PG_MAJORVERSION_NUM >= 18
-#include "commands/explain_state.h"
-#include "commands/explain_format.h"
-#endif
 #include "storage/ipc.h"
 #include "storage/proc.h"
 
@@ -33,13 +29,12 @@ static UprobeList *uprobesList = NULL;
 
 static MemoryContext traceMemoryContext = NULL;
 
+/* We need it in case that ListAdd fails and we need to delete this uprobe */
 static volatile Uprobe *lastSetUprobe = NULL;
 
-/* we need it in case that ListAdd fails and we need to delete this uprobe */
-
 static ExecutorRun_hook_type prev_ExecutorRun_hook = NULL;
-
 static ExecutorStart_hook_type prev_ExecutorStart_hook = NULL;
+static ExecutorFinish_hook_type prev_ExecutorFinish_hook = NULL;
 
 bool		isExecuteTime = false;
 
@@ -68,6 +63,7 @@ static void TraceSessionExecutorRun(QueryDesc *queryDesc,
 									bool execute_once);
 #endif
 static void TraceSessionExecutorStart(QueryDesc *queryDesc, int eflags);
+static void TraceSessionExecutorFinish(QueryDesc *queryDesc);
 static char *ProcessDescBeforeExec(QueryDesc *queryDesc);
 
 static UprobeList *MakeNodesPlanStringsText(char *startPlanExplain, char *endPlanExplain);
@@ -291,6 +287,8 @@ SessionTraceStart(void)
 		ExecutorRun_hook = TraceSessionExecutorRun;
 		prev_ExecutorStart_hook = ExecutorStart_hook;
 		ExecutorStart_hook = TraceSessionExecutorStart;
+		prev_ExecutorFinish_hook = ExecutorFinish_hook;
+		ExecutorFinish_hook = TraceSessionExecutorFinish;
 	}
 	PG_CATCH();
 	{
@@ -548,6 +546,62 @@ TraceSessionExecutorStart(QueryDesc *queryDesc, int eflags)
 }
 
 
+static void
+TraceSessionExecutorFinish(QueryDesc *queryDesc)
+{
+	char	   *planCopy;
+	struct timespec time;
+	uint64		executionStarted;
+
+	/*No Executor nodes will be called, so no need for additional set up*/
+	if (queryDesc->estate->es_auxmodifytables == NULL)
+	{
+		if (prev_ExecutorFinish_hook)
+			(*prev_ExecutorFinish_hook) (queryDesc);
+		else
+			standard_ExecutorFinish(queryDesc);
+		return;
+	}
+
+	if (writeMode == JSON_WRITE_MODE && !isFirstNodeCall)
+	{
+		TracePrintf(",\n");
+	}
+
+	planCopy = BeforeExecution(queryDesc);
+
+	clock_gettime(CLOCKTYPE, &time);
+	executionStarted = time.tv_sec * 1000000000L + time.tv_nsec;
+
+	ExecutorRunNestLevel++;
+
+	PG_TRY();
+	{
+		if (prev_ExecutorFinish_hook)
+			(*prev_ExecutorFinish_hook) (queryDesc);
+		else
+			standard_ExecutorFinish(queryDesc);
+	}
+	PG_FINALLY();
+	{
+		uint64		timeDiff;
+
+		ExecutorRunNestLevel--;
+
+		clock_gettime(CLOCKTYPE, &time);
+		timeDiff = time.tv_sec * 1000000000L + time.tv_nsec - executionStarted;
+		if (writeMode == TEXT_WRITE_MODE)
+			TracePrintf("TRACE. Execution finished for %lu nanosec\n", timeDiff);
+		else
+			TracePrintf("\n],\n\"executionTime\": \"%lu nanosec\"\n", timeDiff);
+
+		AfterExecution(queryDesc, planCopy);
+		isFirstNodeCall = false;
+	}
+	PG_END_TRY();
+}
+
+
 /* find the plan field in explain jsonb */
 static JsonbValue *
 FindField(JsonbValue *json, char *field, size_t len)

tests/test_trace_session.py

@@ -496,7 +496,7 @@ def trace_session_plpgsql_functions_exceptions(node: PostgresNode):
         start_session_trace(conn)
 
         try:
-            node.execute("SELECT process_complete_order(1, ARRAY[5,8], ARRAY[10, 5])")
+            conn.execute("SELECT process_complete_order(1, ARRAY[5,8], ARRAY[10, 5])")
         except:
             pass
 
@@ -506,6 +506,22 @@ def trace_session_plpgsql_functions_exceptions(node: PostgresNode):
         validate_each_session_trace_result(json.loads(result), conn.pid)
 
 
+def trace_session_correct_executor_finish(node: PostgresNode):
+    with node.connect("postgres", autocommit=True) as conn:
+        start_session_trace(conn)
+
+        node.execute("create table mlparted (a int, b int)")
+
+        node.execute("with ins (a, b, c) as \
+                        (insert into mlparted (b, a) select s.a, 1 from generate_series(2, 39) s(a) returning tableoid::regclass, *) \
+                        select a, b, min(c), max(c) from ins group by a, b order by 1;")
+
+        stop_session_trace(conn)
+
+        result = node_read_file_one_line(node, f"/pg_uprobe/trace_file.txt_{conn.pid}")
+        validate_each_session_trace_result(json.loads(result), conn.pid)
+
+
 def run_tests(node: PostgresNode):
     test_wrapper(node, trace_current_session_trace)
     test_wrapper(node, trace_current_session_trace_non_sleep_buffer_locks)
@@ -520,4 +536,5 @@ def run_tests(node: PostgresNode):
     test_wrapper(node, trace_current_session_large)
     test_wrapper(node, trace_session_pid)
     test_wrapper(node, trace_session_plpgsql_functions)
-    test_wrapper(node, trace_session_plpgsql_functions_exceptions)
+    test_wrapper(node, trace_session_plpgsql_functions_exceptions)
+    test_wrapper(node, trace_session_correct_executor_finish)