[#23272] YSQL, ASH: Fix incorrect popping of query id from nested que…

…ry ids stack

Summary:
When ASH was enabled by default to run the tests, some tests failed with this assertion
```
TRAP: FailedAssertion("!(query_id_stack.top_index >= 0)", File: "../../../../../../../src/postgres/src/backend/utils/misc/yb_ash.c", Line: 248)

@        0x102aa1960  YbAshNestedQueryIdStackPop
@        0x102aa16e4  YbAshResetQueryId
@        0x102aa0e28  yb_ash_ExecutorEnd
@        0x10251770c  ExecutorEnd
@        0x1024526b8  PortalCleanup
@        0x102ab7b6c  PortalDrop
@        0x102ab89f8  PreCommit_Portals
@        0x10226d2d8  CommitTransaction
@        0x10226caa4  CommitTransactionCommand
@        0x102802d44  finish_xact_command
@        0x102807b2c  exec_simple_query
@        0x102804fa8  yb_exec_simple_query_impl
@        0x10280510c  yb_exec_query_wrapper_one_attempt
@        0x102804f70  yb_exec_query_wrapper
@        0x1027ffb04  yb_exec_simple_query
@        0x1027fe42c  PostgresMain
@        0x1026fd7b0  BackendRun
@        0x1026fc6bc  BackendStartup
@        0x1026fad30  ServerLoop
@        0x1026f6cac  PostmasterMain
@        0x1025cb848  PostgresServerProcessMain
@        0x1025cbd28  main
@        0x1a389bf28  start
```

`finish_xact_command` is called, which calls `yb_ash_ExecutorEnd`. In `yb_ash_ExecutorEnd`,
we assumed that `yb_ash_ExecutorStart` must have pushed a query id into the nested
query id stack. But in this case, it was a `COMMIT` statement, so it goes through the
`yb_ash_ProcessUtility` path, which pushes and pops a query id from the stack.
`yb_ash_ExecutorStart` is never called in this case. So we try to pop from an empty stack.

This diff fixes the assertion failure by checking that the stack is non-empty before popping.
And in case of nested queries, where the stack size can be greater than 1, we check for the
query id before popping so that only the correct query id is popped.
Jira: DB-12200

Test Plan:
Jenkins

The assertion failure was fixed here after diff id 213903 with the same change from this diff
Jenkins results: https://detective.dev.yugabyte.com/D36430?show_all_diffs=1#tests-failing-in-213903

Not triggering jenkins with ASH enabled by default because there are different kinds of tests which fails.

Reviewers: jason

Reviewed By: jason

Subscribers: amitanand, hbhanawat, yql

Differential Revision: https://phorge.dev.yugabyte.com/D36801

src/postgres/src/backend/utils/misc/yb_ash.c

@@ -103,13 +103,13 @@ static int nested_level = 0;
 static void YbAshInstallHooks(void);
 static int yb_ash_cb_max_entries(void);
 static void YbAshSetQueryId(uint64 query_id);
-static void YbAshResetQueryId(void);
+static void YbAshResetQueryId(uint64 query_id);
 static uint64 yb_ash_utility_query_id(const char *query, int query_len,
 									  int query_location);
 static void YbAshAcquireBufferLock(bool exclusive);
 static void YbAshReleaseBufferLock();
 static bool YbAshNestedQueryIdStackPush(uint64 query_id);
-static uint64 YbAshNestedQueryIdStackPop(void);
+static uint64 YbAshNestedQueryIdStackPop(uint64 query_id);
 
 static void yb_ash_ExecutorStart(QueryDesc *queryDesc, int eflags);
 static void yb_ash_ExecutorRun(QueryDesc *queryDesc,
@@ -237,16 +237,23 @@ YbAshNestedQueryIdStackPush(uint64 query_id)
  * Pop a query id from the stack
  */
 static uint64
-YbAshNestedQueryIdStackPop(void)
+YbAshNestedQueryIdStackPop(uint64 query_id)
 {
 	if (query_id_stack.num_query_ids_not_pushed > 0)
 	{
 		--query_id_stack.num_query_ids_not_pushed;
 		return 0;
 	}
 
-	Assert(query_id_stack.top_index >= 0);
-	return query_id_stack.query_ids[query_id_stack.top_index--];
+	/*
+	 * When an extra ExecutorEnd is called during PortalCleanup,
+	 * we shouldn't pop the incorrect query_id from the stack.
+	 */
+	if (query_id_stack.top_index >= 0 &&
+		query_id_stack.query_ids[query_id_stack.top_index] == query_id)
+		return query_id_stack.query_ids[query_id_stack.top_index--];
+
+	return 0;
 }
 
 /*
@@ -292,14 +299,16 @@ YbAshShmemInit(void)
 static void
 yb_ash_ExecutorStart(QueryDesc *queryDesc, int eflags)
 {
+	uint64 query_id;
+
 	if (yb_enable_ash)
 	{
 		/* Query id can be zero here only if pg_stat_statements is disabled */
-		uint64 query_id = queryDesc->plannedstmt->queryId != 0
-						  ? queryDesc->plannedstmt->queryId
-						  : yb_ash_utility_query_id(queryDesc->sourceText,
-					   								queryDesc->plannedstmt->stmt_len,
-													queryDesc->plannedstmt->stmt_location);
+		query_id = queryDesc->plannedstmt->queryId != 0
+				   ? queryDesc->plannedstmt->queryId
+				   : yb_ash_utility_query_id(queryDesc->sourceText,
+					   						 queryDesc->plannedstmt->stmt_len,
+											 queryDesc->plannedstmt->stmt_location);
 		YbAshSetQueryId(query_id);
 	}
 
@@ -313,7 +322,7 @@ yb_ash_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	PG_CATCH();
 	{
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(query_id);
 
 		PG_RE_THROW();
 	}
@@ -338,7 +347,7 @@ yb_ash_ExecutorRun(QueryDesc *queryDesc, ScanDirection direction, uint64 count,
 		--nested_level;
 
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(queryDesc->plannedstmt->queryId);
 
 		PG_RE_THROW();
 	}
@@ -362,7 +371,7 @@ yb_ash_ExecutorFinish(QueryDesc *queryDesc)
 		--nested_level;
 
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(queryDesc->plannedstmt->queryId);
 
 		PG_RE_THROW();
 	}
@@ -380,12 +389,12 @@ yb_ash_ExecutorEnd(QueryDesc *queryDesc)
 			standard_ExecutorEnd(queryDesc);
 
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(queryDesc->plannedstmt->queryId);
 	}
 	PG_CATCH();
 	{
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(queryDesc->plannedstmt->queryId);
 
 		PG_RE_THROW();
 	}
@@ -398,13 +407,15 @@ yb_ash_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 					  QueryEnvironment *queryEnv, DestReceiver *dest,
 					  char *completionTag)
 {
+	uint64 query_id;
+
 	if (yb_enable_ash)
 	{
-		uint64 query_id = pstmt->queryId != 0
-						  ? pstmt->queryId
-						  : yb_ash_utility_query_id(queryString,
-					   								pstmt->stmt_len,
-													pstmt->stmt_location);
+		query_id = pstmt->queryId != 0
+				   ? pstmt->queryId
+				   : yb_ash_utility_query_id(queryString,
+					   						 pstmt->stmt_len,
+											 pstmt->stmt_location);
 		YbAshSetQueryId(query_id);
 	}
 
@@ -422,14 +433,14 @@ yb_ash_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 		--nested_level;
 
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(query_id);
 	}
 	PG_CATCH();
 	{
 		--nested_level;
 
 		if (yb_enable_ash)
-			YbAshResetQueryId();
+			YbAshResetQueryId(query_id);
 
 		PG_RE_THROW();
 	}
@@ -451,11 +462,11 @@ YbAshSetQueryId(uint64 query_id)
 }
 
 static void
-YbAshResetQueryId(void)
+YbAshResetQueryId(uint64 query_id)
 {
 	if (set_query_id())
 	{
-		uint64 prev_query_id = YbAshNestedQueryIdStackPop();
+		uint64 prev_query_id = YbAshNestedQueryIdStackPop(query_id);
 		if (prev_query_id != 0)
 		{
 			LWLockAcquire(&MyProc->yb_ash_metadata_lock, LW_EXCLUSIVE);