[#24095] YSQL, ASH: Fix query id not popping from nested query stack

Summary:
D36801 / adf3c54 updated the way how
query id is being popped from the nested query id stack. There was a
bug in the previous workflow -

- Once a new query id comes, push the current query id to the stack and set the current query id to the new query id.
- When popping, the current query id is checked, if it matches the top of the stack, then we pop it.

However, this is incorrect as we never pushed the current query id to
the stack so when we are trying to pop, we will never get a match of the
query id. This would result in the wait events being tracked by incorrect
query ids if `pg_stat_statements.track="all"` was turned on and user
fired nested queries.

This diff fixes this by also pushing the current query id to the stack.

This also fixes a out of bounds error while accessing `query_id_stack.query_ids`
Jira: DB-12988

Test Plan:
./yb_build.sh --java-test TestYbAsh#testNestedQueriesWithAsh

Manually tested by logging the query id while pushing and popping and made sure they are the same

Reviewers: jason

Reviewed By: jason

Subscribers: yql

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

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

@@ -195,7 +195,9 @@ void
 YbAshInit(void)
 {
 	YbAshInstallHooks();
-	query_id_stack.top_index = -1;
+	/* Keep the default query id in the stack */
+	query_id_stack.top_index = 0;
+	query_id_stack.query_ids[0] = YBCGetQueryIdForCatalogRequests();
 	query_id_stack.num_query_ids_not_pushed = 0;
 }
 
@@ -242,7 +244,7 @@ yb_ash_cb_max_entries(void)
 static bool
 YbAshNestedQueryIdStackPush(uint64 query_id)
 {
-	if (query_id_stack.top_index < MAX_NESTED_QUERY_LEVEL)
+	if (query_id_stack.top_index < MAX_NESTED_QUERY_LEVEL - 1)
 	{
 		query_id_stack.query_ids[++query_id_stack.top_index] = query_id;
 		return true;
@@ -255,7 +257,7 @@ YbAshNestedQueryIdStackPush(uint64 query_id)
 }
 
 /*
- * Pop a query id from the stack
+ * Pop and return the top query id from the stack
  */
 static uint64
 YbAshNestedQueryIdStackPop(uint64 query_id)
@@ -270,9 +272,9 @@ YbAshNestedQueryIdStackPop(uint64 query_id)
 	 * 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 &&
+	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 query_id_stack.query_ids[--query_id_stack.top_index];
 
 	return 0;
 }
@@ -489,7 +491,7 @@ YbAshSetQueryId(uint64 query_id)
 {
 	if (set_query_id())
 	{
-		if (YbAshNestedQueryIdStackPush(MyProc->yb_ash_metadata.query_id))
+		if (YbAshNestedQueryIdStackPush(query_id))
 		{
 			LWLockAcquire(&MyProc->yb_ash_metadata_lock, LW_EXCLUSIVE);
 			MyProc->yb_ash_metadata.query_id = query_id;
@@ -516,8 +518,8 @@ YbAshResetQueryId(uint64 query_id)
 void
 YbAshSetMetadata(void)
 {
-	/* The stack must be empty at the start of a request */
-	Assert(query_id_stack.top_index == -1);
+	/* The stack should have the default query id at the start of a request */
+	Assert(query_id_stack.top_index == 0);
 
 	LWLockAcquire(&MyProc->yb_ash_metadata_lock, LW_EXCLUSIVE);
 	YBCGenerateAshRootRequestId(MyProc->yb_ash_metadata.root_request_id);
@@ -532,7 +534,7 @@ YbAshUnsetMetadata(void)
 	 * returns an error. Reset the stack here. We can remove this if we
 	 * make query_id atomic
 	 */
-	query_id_stack.top_index = -1;
+	query_id_stack.top_index = 0;
 	query_id_stack.num_query_ids_not_pushed = 0;
 
 	LWLockAcquire(&MyProc->yb_ash_metadata_lock, LW_EXCLUSIVE);