cancel aqo timeout action in the critical section

Author: Andrey Kazarinov

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"
@@ -622,8 +623,11 @@ 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. */
 

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();