[PGPRO-17014] Fix the issue with early reset of AQOCacheMemCtx by ATX

The autonomous transaction can be aborted in nested levels of plPgSql.
Check this level in aqo_free_callback.

Author: Daniil Anisimov

Makefile

@@ -16,6 +16,7 @@ REGRESS =	aqo_disabled \
 			aqo_intelligent \
 			aqo_forced \
 			aqo_learn \
+			aqo_atx \
 			schema \
 			aqo_fdw \
 			aqo_CVE-2020-14350 \

aqo.c

@@ -12,6 +12,7 @@
 
 #include "access/relation.h"
 #include "access/table.h"
+#include "access/xact.h"
 #include "catalog/pg_extension.h"
 #include "commands/extension.h"
 #include "miscadmin.h"
@@ -113,7 +114,11 @@ aqo_free_callback(ResourceReleasePhase phase,
 	if (phase != RESOURCE_RELEASE_AFTER_LOCKS)
 		return;
 
-	if (isTopLevel)
+	if (isTopLevel
+#ifdef PGPRO_EE
+		&& getNestLevelATX() == 0
+#endif
+		)
 	{
 		MemoryContextReset(AQOCacheMemCtx);
 		cur_classes = NIL;

expected/aqo_atx.out

@@ -0,0 +1,79 @@
+-- Test for early reset of AQOCache memory context by ATX
+SELECT (count(*) = 0) AS skip_test FROM pg_settings
+WHERE name = 'max_autonomous_transactions' \gset
+\if :skip_test
+\quit
+\endif
+CREATE EXTENSION IF NOT EXISTS aqo;
+SELECT true AS success FROM aqo_reset();
+ success 
+---------
+ t
+(1 row)
+
+CREATE TABLE t1 (a int, b text, c int);
+CREATE unique index on t1 (a, b);
+CREATE TABLE t2 AS SELECT 0 AS a, 'cero' AS b, 1 AS c;
+CREATE FUNCTION t1_update_func() RETURNS trigger AS $$
+BEGIN
+	PERFORM * FROM t2 WHERE a >= 0;
+	RETURN new;
+END;
+$$ LANGUAGE plpgsql;
+CREATE TRIGGER t1_update
+AFTER UPDATE ON t1
+FOR EACH STATEMENT
+EXECUTE PROCEDURE t1_update_func();
+CREATE FUNCTION f2()
+RETURNS int AS $$
+BEGIN
+	BEGIN AUTONOMOUS
+		INSERT INTO t1 SELECT a,b,c FROM t2 WHERE t2.a >= 0
+			ON conflict (a,b) do UPDATE SET c = t1.c + 1;
+		SELECT 1/0;
+	EXCEPTION WHEN division_by_zero THEN
+		NULL;
+	END;
+
+	INSERT INTO t1 SELECT a,b,c FROM t2 WHERE t2.a >= 0
+		ON conflict (a,b) do UPDATE SET c = t1.c + 1;
+
+	RETURN 0;
+END;
+$$ LANGUAGE plpgsql security definer;
+SET aqo.join_threshold = 0;
+SET aqo.mode = learn;
+PREPARE prep1 AS SELECT f2();
+EXECUTE prep1;
+ f2 
+----
+  0
+(1 row)
+
+EXECUTE prep1;
+ f2 
+----
+  0
+(1 row)
+
+EXECUTE prep1;
+ f2 
+----
+  0
+(1 row)
+
+EXECUTE prep1;
+ f2 
+----
+  0
+(1 row)
+
+EXECUTE prep1;
+ f2 
+----
+  0
+(1 row)
+
+DROP TABLE t1, t2 CASCADE;
+DROP FUNCTION t1_update_func, f2;
+DROP EXTENSION aqo;

expected/aqo_atx_1.out

@@ -0,0 +1,5 @@
+-- Test for early reset of AQOCache memory context by ATX
+SELECT (count(*) = 0) AS skip_test FROM pg_settings
+WHERE name = 'max_autonomous_transactions' \gset
+\if :skip_test
+\quit

sql/aqo_atx.sql

@@ -0,0 +1,60 @@
+-- Test for early reset of AQOCache memory context by ATX
+SELECT (count(*) = 0) AS skip_test FROM pg_settings
+WHERE name = 'max_autonomous_transactions' \gset
+
+\if :skip_test
+\quit
+\endif
+
+CREATE EXTENSION IF NOT EXISTS aqo;
+SELECT true AS success FROM aqo_reset();
+
+CREATE TABLE t1 (a int, b text, c int);
+CREATE unique index on t1 (a, b);
+
+CREATE TABLE t2 AS SELECT 0 AS a, 'cero' AS b, 1 AS c;
+
+CREATE FUNCTION t1_update_func() RETURNS trigger AS $$
+BEGIN
+	PERFORM * FROM t2 WHERE a >= 0;
+	RETURN new;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE TRIGGER t1_update
+AFTER UPDATE ON t1
+FOR EACH STATEMENT
+EXECUTE PROCEDURE t1_update_func();
+
+CREATE FUNCTION f2()
+RETURNS int AS $$
+BEGIN
+	BEGIN AUTONOMOUS
+		INSERT INTO t1 SELECT a,b,c FROM t2 WHERE t2.a >= 0
+			ON conflict (a,b) do UPDATE SET c = t1.c + 1;
+		SELECT 1/0;
+	EXCEPTION WHEN division_by_zero THEN
+		NULL;
+	END;
+
+	INSERT INTO t1 SELECT a,b,c FROM t2 WHERE t2.a >= 0
+		ON conflict (a,b) do UPDATE SET c = t1.c + 1;
+
+	RETURN 0;
+END;
+$$ LANGUAGE plpgsql security definer;
+
+SET aqo.join_threshold = 0;
+SET aqo.mode = learn;
+
+PREPARE prep1 AS SELECT f2();
+EXECUTE prep1;
+EXECUTE prep1;
+EXECUTE prep1;
+EXECUTE prep1;
+EXECUTE prep1;
+
+DROP TABLE t1, t2 CASCADE;
+DROP FUNCTION t1_update_func, f2;
+
+DROP EXTENSION aqo;