From de709e6eeac73ed3b32aaa5d5b8a2f137667cb6c Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Wed, 6 Sep 2023 17:02:37 +0900
Subject: [PATCH] Reset more aggressively hints for queries executed via
 extended query protocol

A query executed with the extended query protocol may entirely ignore
hints defined in it, as the only moment where a hint is reset is when a
query is restarted.  This commit introduces a hook for ExecutorEnd(),
which is a code path taken at the end of each individual execution
happening through a Portal when using the extended query protocol.

After this patch, it should be possible to make the reset of
current_hint_retrieved less optimistic, but the current logic does not
hurt as well as the end of executor would match with the next query to
come in for most users.

Per report and patch #12741 in yugabyte-db, also reported as #145 in
pg_hint_plan.

Reported-by: Furutani-san (kfuru)
Author: tanujnay112
Backpatch-through: 11
---
 pg_hint_plan.c | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/pg_hint_plan.c b/pg_hint_plan.c
index 559aea1e..e77fd6ec 100644
--- a/pg_hint_plan.c
+++ b/pg_hint_plan.c
@@ -586,6 +586,7 @@ static post_parse_analyze_hook_type prev_post_parse_analyze_hook = NULL;
 static planner_hook_type prev_planner = NULL;
 static join_search_hook_type prev_join_search = NULL;
 static set_rel_pathlist_hook_type prev_set_rel_pathlist = NULL;
+static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
 
 /* Hold reference to currently active hint */
 static HintState *current_hint_state = NULL;
@@ -641,6 +642,26 @@ PLpgSQL_plugin  plugin_funcs = {
 	NULL,
 };
 
+/*
+ * pg_hint_ExecutorEnd
+ *
+ * Force a hint to be retrieved when we are at the top of a PL recursion
+ * level.  This can become necessary to handle hints in queries executed
+ * in the extended protocol, where the executor can be executed multiple
+ * times in a portal, but it could be possible to fail the hint retrieval.
+ */
+static void
+pg_hint_ExecutorEnd(QueryDesc *queryDesc)
+{
+	if (plpgsql_recurse_level <= 0)
+		current_hint_retrieved = false;
+
+	if (prev_ExecutorEnd)
+		prev_ExecutorEnd(queryDesc);
+	else
+		standard_ExecutorEnd(queryDesc);
+}
+
 /*
  * Module load callbacks
  */
@@ -730,6 +751,8 @@ _PG_init(void)
 	join_search_hook = pg_hint_plan_join_search;
 	prev_set_rel_pathlist = set_rel_pathlist_hook;
 	set_rel_pathlist_hook = pg_hint_plan_set_rel_pathlist;
+	prev_ExecutorEnd = ExecutorEnd_hook;
+	ExecutorEnd_hook = pg_hint_ExecutorEnd;
 
 	/* setup PL/pgSQL plugin hook */
 	var_ptr = (PLpgSQL_plugin **) find_rendezvous_variable("PLpgSQL_plugin");
@@ -751,6 +774,7 @@ _PG_fini(void)
 	planner_hook = prev_planner;
 	join_search_hook = prev_join_search;
 	set_rel_pathlist_hook = prev_set_rel_pathlist;
+	ExecutorEnd_hook = prev_ExecutorEnd;
 
 	/* uninstall PL/pgSQL plugin hook */
 	var_ptr = (PLpgSQL_plugin **) find_rendezvous_variable("PLpgSQL_plugin");