cancel aqo timeout action in the critical section

Andrey Kazarinov · Anisimov-ds · commit 9d2b27af12a9 · 2023-06-30T16:38:09.000+07:00
diff --git a/aqo.h b/aqo.h
@@ -172,6 +172,7 @@ extern bool aqo_show_details;
 extern int aqo_join_threshold;
 extern bool use_wide_search;
 extern bool aqo_learn_statement_timeout;
+extern bool aqo_learn_statement_timeout_enable;
 
 /* Parameters for current query */
 typedef struct QueryContextData
diff --git a/postprocessing.c b/postprocessing.c
@@ -22,6 +22,7 @@
 #include "optimizer/optimizer.h"
 #include "postgres_fdw.h"
 #include "utils/queryenvironment.h"
+#include "miscadmin.h"
 
 #include "aqo.h"
 #include "hash.h"
@@ -628,8 +629,12 @@ aqo_timeout_handler(void)
 	MemoryContext oldctx = MemoryContextSwitchTo(AQOLearnMemCtx);
 	aqo_obj_stat ctx = {NIL, NIL, NIL, false, false};
 
-	if (!timeoutCtl.queryDesc || !ExtractFromQueryEnv(timeoutCtl.queryDesc))
+	if (CritSectionCount > 0 || !timeoutCtl.queryDesc ||
+		!ExtractFromQueryEnv(timeoutCtl.queryDesc))
+	{
+		MemoryContextSwitchTo(oldctx);
 		return;
+	}
 
 	/* Now we can analyze execution state of the query. */
 
@@ -664,7 +669,7 @@ set_timeout_if_need(QueryDesc *queryDesc)
 {
 	int64 fintime = (int64) get_timeout_finish_time(STATEMENT_TIMEOUT)-1;
 
-	if (aqo_learn_statement_timeout && aqo_statement_timeout > 0)
+	if (aqo_learn_statement_timeout_enable && aqo_statement_timeout > 0)
 	{
 		max_timeout_value = Min(query_context.smart_timeout, (int64) aqo_statement_timeout);
 		if (max_timeout_value > fintime)
@@ -684,7 +689,7 @@ set_timeout_if_need(QueryDesc *queryDesc)
 		*/
 		return false;
 
-	if (!get_timeout_active(STATEMENT_TIMEOUT) || !aqo_learn_statement_timeout)
+	if (!get_timeout_active(STATEMENT_TIMEOUT) || !aqo_learn_statement_timeout_enable)
 		return false;
 
 	if (!ExtractFromQueryEnv(queryDesc))
@@ -829,7 +834,7 @@ aqo_ExecutorEnd(QueryDesc *queryDesc)
 
 			error = stat->est_error_aqo[stat->cur_stat_slot_aqo-1] - cardinality_sum_errors/(1 + cardinality_num_objects);
 
-			if ( aqo_learn_statement_timeout && aqo_statement_timeout > 0 && error >= 0.1)
+			if ( aqo_learn_statement_timeout_enable && aqo_statement_timeout > 0 && error >= 0.1)
 			{
 				int64 fintime = increase_smart_timeout();
 				elog(NOTICE, "[AQO] Time limit for execution of the statement was increased. Current timeout is "UINT64_FORMAT, fintime);
diff --git a/preprocessing.c b/preprocessing.c
@@ -71,7 +71,10 @@ List *cur_classes = NIL;
 
 int aqo_join_threshold = 3;
 
+bool aqo_learn_statement_timeout_enable = false;
+
 static planner_hook_type	aqo_planner_next = NULL;
+static post_parse_analyze_hook_type	aqo_post_parse_analyze_hook = NULL;
 
 static void disable_aqo_for_query(void);
 static bool isQueryUsingSystemRelation(Query *query);
@@ -478,9 +481,26 @@ isQueryUsingSystemRelation_walker(Node *node, void *context)
 								  context);
 }
 
+static void
+aqo_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
+{
+	aqo_learn_statement_timeout_enable = false;
+	/*
+	 * Enable learn_statement_timeout for
+	 * the top level SELECT statement only.
+	 */
+	if (query->commandType == CMD_SELECT)
+		aqo_learn_statement_timeout_enable = aqo_learn_statement_timeout;
+
+	if (aqo_post_parse_analyze_hook)
+		aqo_post_parse_analyze_hook(pstate, query, jstate);
+}
+
 void
 aqo_preprocessing_init(void)
 {
 	aqo_planner_next	= planner_hook ? planner_hook : standard_planner;
 	planner_hook		= aqo_planner;
-}
+	aqo_post_parse_analyze_hook = post_parse_analyze_hook;
+	post_parse_analyze_hook = aqo_post_parse_analyze;
+}
diff --git a/t/003_assertion_error.pl b/t/003_assertion_error.pl
@@ -0,0 +1,59 @@
+use strict;
+use warnings;
+
+use Config;
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+
+use Test::More tests => 1;
+
+my $node = PostgreSQL::Test::Cluster->new('aqotest');
+$node->init;
+$node->append_conf('postgresql.conf', qq{
+						shared_preload_libraries = 'aqo'
+						aqo.join_threshold = 0
+						aqo.mode = 'learn'
+						aqo.show_details = 'off'
+						aqo.learn_statement_timeout = 'on'
+					});
+
+# Test constants. Default values.
+my $TRANSACTIONS = 100;
+
+# Disable connection default settings, forced by PGOPTIONS in AQO Makefile
+# $ENV{PGOPTIONS}="";
+
+# Change pgbench parameters according to the environment variable.
+if (defined $ENV{TRANSACTIONS})
+{
+	$TRANSACTIONS = $ENV{TRANSACTIONS};
+}
+
+my $query_string = '
+CREATE TABLE IF NOT EXISTS aqo_test1(a int, b int);
+WITH RECURSIVE t(a, b)
+AS (
+   VALUES (1, 2)
+   UNION ALL
+   SELECT t.a + 1, t.b + 1 FROM t WHERE t.a < 10
+) INSERT INTO aqo_test1 (SELECT * FROM t);
+
+SET statement_timeout = 10;
+
+CREATE TABLE tmp1 AS SELECT t1.a AS a, t2.a AS b, t3.a AS c
+FROM aqo_test1 AS t1, aqo_test1 AS t2, aqo_test1 AS t3
+WHERE t1.a = t2.b AND t2.a = t3.b;
+DROP TABLE tmp1;
+';
+
+$node->start();
+
+$node->safe_psql('postgres', 'CREATE EXTENSION IF NOT EXISTS aqo;');
+
+for (1..$TRANSACTIONS) {
+	$node->psql('postgres', $query_string);
+}
+
+ok(1, "There are no segfaults");
+
+$node->stop();
