Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 0b5c61b

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

File tree

1 file changed

+87
-12
lines changed

1 file changed

+87
-12
lines changed

‎pg_wait_sampling.c

+87-12
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,41 @@ 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

872-
/* Invoke original hook if needed */
873-
if (planner_hook_next)
874-
return planner_hook_next(parse,
886+
nesting_level++;
887+
PG_TRY();
888+
{
889+
/* Invoke original hook if needed */
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);
901+
cursorOptions, boundParams);
902+
nesting_level--;
903+
if (nesting_level == 0)
904+
pgws_proc_queryids[i] = UINT64CONST(0);
905+
}
906+
PG_CATCH();
907+
{
908+
nesting_level--;
909+
if (nesting_level == 0)
910+
pgws_proc_queryids[i] = UINT64CONST(0);
911+
PG_RE_THROW();
912+
}
913+
PG_END_TRY();
914+
915+
return result;
885916
}
886917

887918
/*
@@ -891,7 +922,7 @@ static void
891922
pgws_ExecutorStart(QueryDesc *queryDesc, int eflags)
892923
{
893924
int i = MyProc - ProcGlobal->allProcs;
894-
if (!pgws_proc_queryids[i])
925+
if (nesting_level == 0)
895926
pgws_proc_queryids[i] = queryDesc->plannedstmt->queryId;
896927

897928
if (prev_ExecutorStart)
@@ -900,13 +931,57 @@ pgws_ExecutorStart(QueryDesc *queryDesc, int eflags)
900931
standard_ExecutorStart(queryDesc, eflags);
901932
}
902933

934+
static void
935+
pgws_ExecutorRun(QueryDesc *queryDesc,
936+
ScanDirection direction,
937+
uint64 count, bool execute_once)
938+
{
939+
nesting_level++;
940+
PG_TRY();
941+
{
942+
if (prev_ExecutorRun)
943+
prev_ExecutorRun(queryDesc, direction, count, execute_once);
944+
else
945+
standard_ExecutorRun(queryDesc, direction, count, execute_once);
946+
nesting_level--;
947+
}
948+
PG_CATCH();
949+
{
950+
nesting_level--;
951+
PG_RE_THROW();
952+
}
953+
PG_END_TRY();
954+
}
955+
956+
static void
957+
pgws_ExecutorFinish(QueryDesc *queryDesc)
958+
{
959+
nesting_level++;
960+
PG_TRY();
961+
{
962+
if (prev_ExecutorFinish)
963+
prev_ExecutorFinish(queryDesc);
964+
else
965+
standard_ExecutorFinish(queryDesc);
966+
nesting_level--;
967+
}
968+
PG_CATCH();
969+
{
970+
nesting_level--;
971+
PG_RE_THROW();
972+
}
973+
PG_END_TRY();
974+
}
975+
903976
/*
904977
* ExecutorEnd hook: clear queryId
905978
*/
906979
static void
907980
pgws_ExecutorEnd(QueryDesc *queryDesc)
908981
{
909-
pgws_proc_queryids[MyProc - ProcGlobal->allProcs] = UINT64CONST(0);
982+
int i = MyProc - ProcGlobal->allProcs;
983+
if (nesting_level == 0)
984+
pgws_proc_queryids[i] = UINT64CONST(0);
910985

911986
if (prev_ExecutorEnd)
912987
prev_ExecutorEnd(queryDesc);

0 commit comments

Comments
 (0)