8000 Make so subqueries do not rewrite queryId of top-level query · postgrespro/pg_wait_sampling@f4fbcce · GitHub
[go: up one dir, main page]
Skip to content

Commit f4fbcce

Browse files
author
Oleg Tselebrovskiy
committed
Make so subqueries do not rewrite queryId of top-level query
1 parent 69db35f commit f4fbcce

File tree

1 file changed

+81
-11
lines changed

1 file changed

+81
-11
lines changed

‎pg_wait_sampling.c

Lines changed: 81 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,14 @@ static bool shmem_initialized = false;
4343

4444
/* Hooks */
4545
static ExecutorStart_hook_type prev_ExecutorStart = NULL;
46+
static ExecutorRun_hook_type prev_ExecutorRun = NULL;
47+
static ExecutorFinish_hook_type prev_ExecutorFinish = NULL;
4648
static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
4749
static planner_hook_type planner_hook_next = NULL;
4850

51+
/* Current nesting depth of planner/Executor calls */
52+
static int nesting_level = 0;
53+
4954
/* Pointers to shared memory objects */
5055
shm_mq *pgws_collector_mq = NULL;
5156
uint64 *pgws_proc_queryids = NULL;
@@ -67,6 +72,10 @@ static PlannedStmt *pgws_planner_hook(Query *parse,
6772
#endif
6873
int cursorOptions, ParamListInfo boundParams);
6974
static void pgws_ExecutorStart(QueryDesc *queryDesc, int eflags);
75+
static void pgws_ExecutorRun(QueryDesc *queryDesc,
76+
ScanDirection direction,
77+
uint64 count, bool execute_once);
78+
static void pgws_ExecutorFinish(QueryDesc *queryDesc);
7079
static void pgws_ExecutorEnd(QueryDesc *queryDesc);
7180

7281
/*
@@ -395,6 +404,10 @@ _PG_init(void)
395404
planner_hook = pgws_planner_hook;
396405
prev_ExecutorStart = ExecutorStart_hook;
397406
ExecutorStart_hook = pgws_ExecutorStart;
407+
prev_ExecutorRun = ExecutorRun_hook;
408+
ExecutorRun_hook = pgws_ExecutorRun;
409+
prev_ExecutorFinish = ExecutorFinish_hook;
410+
ExecutorFinish_hook = pgws_ExecutorFinish;
398411
prev_ExecutorEnd = ExecutorEnd_hook;
399412
ExecutorEnd_hook = pgws_ExecutorEnd;
400413
}
@@ -865,23 +878,38 @@ pgws_planner_hook(Query *parse,
865878
int cursorOptions,
866879
ParamListInfo boundParams)
867880
{
881+
PlannedStmt *result;
868882
int i = MyProc - ProcGlobal->allProcs;
869-
if (!pgws_proc_queryids[i])
883+
if (nesting_level == 0)
870884
pgws_proc_queryids[i] = parse->queryId;
871885

872886
/* Invoke original hook if needed */
873-
if (planner_hook_next)
874-
return planner_hook_next(parse,
887+
nesting_level++;
888+
PG_TRY();
889+
{
890+
if (planner_hook_next)
891+
result = planner_hook_next(parse,
875892
#if PG_VERSION_NUM >= 130000
876-
query_string,
893+
query_string,
877894
#endif
878-
cursorOptions, boundParams);
879-
880-
return standard_planner(parse,
895+
cursorOptions, boundParams);
896+
else
897+
result = standard_planner(parse,
881898
#if PG_VERSION_NUM >= 130000
882-
query_string,
899+
query_string,
883900
#endif
884-
cursorOptions, boundParams);
 10000
901+
cursorOptions, boundParams);
902+
}
903+
PG_CATCH();
904+
{
905+
nesting_level--;
906+
if (nesting_level == 0)
907+
pgws_proc_queryids[i] = UINT64CONST(0);
908+
PG_RE_THROW();
909+
}
910+
PG_END_TRY();
911+
912+
return result;
885913
}
886914

887915
/*
@@ -891,7 +919,7 @@ static void
891919
pgws_ExecutorStart(QueryDesc *queryDesc, int eflags)
892920
{
893921
int i = MyProc - ProcGlobal->allProcs;
894-
if (!pgws_proc_queryids[i])
922+
if (nesting_level == 0)
895923
pgws_proc_queryids[i] = queryDesc->plannedstmt->queryId;
896924

897925
if (prev_ExecutorStart)
@@ -900,13 +928,55 @@ pgws_ExecutorStart(QueryDesc *queryDesc, int eflags)
900928
standard_ExecutorStart(queryDesc, eflags);
901929
}
902930

931+
static void
932+
pgws_ExecutorRun(QueryDesc *queryDesc,
933+
ScanDirection direction,
934+
uint64 count, bool execute_once)
935+
{
936+
nesting_level++;
937+
PG_TRY();
938+
{
939+
if (prev_ExecutorRun)
940+
prev_ExecutorRun(queryDesc, direction, count, execute_once);
941+
else
942+
standard_ExecutorRun(queryDesc, direction, count, execute_once);
943+
}
944+
PG_CATCH();
945+
{
946+
nesting_level--;
947+
PG_RE_THROW();
948+
}
949+
PG_END_TRY();
950+
}
951+
952+
static void
953+
pgws_ExecutorFinish(QueryDesc *queryDesc)
954+
{
955+
nesting_level++;
956+
PG_TRY();
957+
{
958+
if (prev_ExecutorFinish)
959+
prev_ExecutorFinish(queryDesc);
960+
else
961+
standard_ExecutorFinish(queryDesc);
962+
}
963+
PG_CATCH();
964+
{
965+
nesting_level--;
966+
PG_RE_THROW();
967+
}
968+
PG_END_TRY();
969+
}
970+
903971
/*
904972
* ExecutorEnd hook: clear queryId
905973
*/
906974
static void
907975
pgws_ExecutorEnd(QueryDesc *queryDesc)
908976
{
909-
pgws_proc_queryids[MyProc - ProcGlobal->allProcs] = UINT64CONST(0);
977+
int i = MyProc - ProcGlobal->allProcs;
978+
if (nesting_level == 0)
979+
pgws_proc_queryids[i] = UINT64CONST(0);
910980

911981
if (prev_ExecutorEnd)
912982
prev_ExecutorEnd(queryDesc);

0 commit comments

Comments
 (0)
0