Fix jit behavior, fix Fetch stmt behavior

Author: pipichinka

src/trace_execute_nodes.c

@@ -266,10 +266,10 @@ ExecInitExprRetFunc(void *data, ExprState *resNode)
 	exprData->traceData.maxTime = 0;
 	exprData->traceData.totalCalls = 0;
 	exprData->traceData.totalTimeSum = 0;
-	exprData->traceData.evalfunc = resNode->evalfunc_private;
+	exprData->traceData.evalfunc = resNode->evalfunc;
 	exprData->traceData.tag = resNode->expr->type;
 	ListAdd(currentState->uprobeExprData, exprData);
-	resNode->evalfunc_private = TraceExprNodeHook;
+	resNode->evalfunc = TraceExprNodeHook;
 }
 
 
@@ -321,6 +321,12 @@ TraceExprNodeHook(struct ExprState *expression, struct ExprContext *econtext, bo
 	if (!currentState)
 		return result;			/* This case is for ending session trace */
 
+	if (expression->evalfunc != TraceExprNodeHook)
+	{
+		exprData->traceData.evalfunc = expression->evalfunc;
+		expression->evalfunc = TraceExprNodeHook;
+	}
+
 	clock_gettime(CLOCK_MONOTONIC, &time);
 	timeDiff = time.tv_sec * 1000000000L + time.tv_nsec - startTime;
 	exprData->traceData.totalCalls++;

src/trace_session.c

@@ -8,7 +8,7 @@
 #include "commands/explain.h"
 #include "storage/ipc.h"
 #include "storage/proc.h"
-
+#include "tcop/utility.h"
 #include "trace_parsing.h"
 #include "uprobe_internal.h"
 #include "trace_lock_on_buffers.h"
@@ -35,6 +35,7 @@ static volatile Uprobe *lastSetUprobe = NULL;
 static ExecutorRun_hook_type prev_ExecutorRun_hook = NULL;
 static ExecutorStart_hook_type prev_ExecutorStart_hook = NULL;
 static ExecutorFinish_hook_type prev_ExecutorFinish_hook = NULL;
+static ProcessUtility_hook_type prev_ProcessUtility_hook = NULL;
 
 bool		isExecuteTime = false;
 
@@ -62,6 +63,14 @@ static void TraceSessionExecutorRun(QueryDesc *queryDesc,
 									uint64 count,
 									bool execute_once);
 #endif
+static void TraceSessionExecutorProcessUtility(PlannedStmt *pstmt,
+											   const char *queryString,
+											   bool readOnlyTree,
+											   ProcessUtilityContext context,
+											   ParamListInfo params,
+											   QueryEnvironment *queryEnv,
+											   DestReceiver *dest,
+											   QueryCompletion *qc);
 static void TraceSessionExecutorStart(QueryDesc *queryDesc, int eflags);
 static void TraceSessionExecutorFinish(QueryDesc *queryDesc);
 static char *ProcessDescBeforeExec(QueryDesc *queryDesc);
@@ -289,6 +298,8 @@ SessionTraceStart(void)
 		ExecutorStart_hook = TraceSessionExecutorStart;
 		prev_ExecutorFinish_hook = ExecutorFinish_hook;
 		ExecutorFinish_hook = TraceSessionExecutorFinish;
+		prev_ProcessUtility_hook = ProcessUtility_hook;
+		ProcessUtility_hook = TraceSessionExecutorProcessUtility;
 	}
 	PG_CATCH();
 	{
@@ -464,6 +475,27 @@ ProcessDescBeforeExec(QueryDesc *queryDesc)
 	return planCopy;
 }
 
+static void
+CallOriginalExecutorFinish(PlannedStmt *pstmt,
+						   const char *queryString,
+						   bool readOnlyTree,
+						   ProcessUtilityContext context,
+						   ParamListInfo params,
+						   QueryEnvironment *queryEnv,
+						   DestReceiver *dest,
+						   QueryCompletion *qc)
+{
+	if (prev_ProcessUtility_hook)
+			(*prev_ProcessUtility_hook) (pstmt, queryString, readOnlyTree,
+										 context, params, queryEnv,
+										 dest, qc);
+		else
+			standard_ProcessUtility(pstmt, queryString, readOnlyTree,
+									context, params, queryEnv,
+									dest, qc);
+}
+
+
 #if PG_MAJORVERSION_NUM >= 18
 static void
 TraceSessionExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count)
@@ -601,6 +633,82 @@ TraceSessionExecutorFinish(QueryDesc *queryDesc)
 	PG_END_TRY();
 }
 
+static void
+TraceSessionExecutorProcessUtility(PlannedStmt *pstmt,
+								   const char *queryString,
+								   bool readOnlyTree,
+								   ProcessUtilityContext context,
+								   ParamListInfo params,
+								   QueryEnvironment *queryEnv,
+								   DestReceiver *dest,
+								   QueryCompletion *qc)
+{
+	char	   *planCopy;
+	struct timespec time;
+	uint64		executionStarted;
+	FetchStmt	*fetch;
+	Portal		portal;
+	QueryDesc	*queryDesc;
+
+	if (nodeTag(pstmt->utilityStmt) != T_FetchStmt)
+	{
+		CallOriginalExecutorFinish(pstmt, queryString, readOnlyTree,
+								   context, params, queryEnv,
+								   dest, qc);
+
+		return;
+	}
+
+	fetch = (FetchStmt*) pstmt->utilityStmt;
+	portal = GetPortalByName(fetch->portalname);
+
+	if (!PortalIsValid(portal))
+	{
+		CallOriginalExecutorFinish(pstmt, queryString, readOnlyTree,
+								   context, params, queryEnv,
+								   dest, qc);
+		return;
+	}
+
+	queryDesc = portal->queryDesc;
+
+	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();
+	{
+		CallOriginalExecutorFinish(pstmt, queryString, readOnlyTree,
+								   context, params, queryEnv,
+								   dest, qc);
+	}
+	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 *

tests/test_trace_session.py

@@ -516,11 +516,57 @@ def trace_session_correct_executor_finish(node: PostgresNode):
                         (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;")
 
+        conn.execute("drop table mlparted")
+
         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 trace_session_fetch(node: PostgresNode):
+    with node.connect("postgres", autocommit=True) as conn:
+        start_session_trace(conn)
+
+        conn.execute("CREATE TABLE INT8_TBL(q1 int8, q2 int8)")
+
+        conn.execute("INSERT INTO INT8_TBL VALUES \
+                        ('  123   ','  456'), \
+                        ('123   ','4567890123456789'), \
+                        ('4567890123456789','123'), \
+                        (+4567890123456789,'4567890123456789'), \
+                        ('+4567890123456789','-4567890123456789')")
+        conn.execute("begin")
+        conn.execute("create function nochange(int) returns int \
+                      as 'select $1 limit 1' language sql stable")
+        conn.execute("declare c cursor for select * from int8_tbl limit nochange(3)")
+        conn.execute("fetch all from c")
+        conn.execute("move backward all in c")
+        conn.execute("fetch all from c")
+        conn.execute("rollback")
+        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 trace_session_correct_with_jit(node: PostgresNode):
+    with node.connect("postgres", autocommit=True) as conn:
+        conn.execute("set jit to on")
+        conn.execute("set jit_above_cost to 0.0")
+        conn.execute("set jit_inline_above_cost to 0.0")
+        conn.execute("set jit_optimize_above_cost to 0.0")
+        start_session_trace(conn)
+
+        conn.execute("create table mlparted (a int, b int)")
+        conn.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;")
+
+        conn.execute("drop table mlparted")
+
+        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)
@@ -537,4 +583,6 @@ def run_tests(node: PostgresNode):
     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_correct_executor_finish)
+    test_wrapper(node, trace_session_correct_executor_finish)
+    test_wrapper(node, trace_session_fetch)
+    test_wrapper(node, trace_session_correct_with_jit)