Fixed pg re2 cache cleanup (#6978)

5 files changed, 61 insertions, 5 deletions

diff --git a/ydb/library/yql/parser/pg_wrapper/arena_ctx.cpp b/ydb/library/yql/parser/pg_wrapper/arena_ctx.cpp
index e424dcd306c..e91cd6ad7a2 100644
--- a/ydb/library/yql/parser/pg_wrapper/arena_ctx.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/arena_ctx.cpp
@@ -118,6 +118,7 @@ TArenaMemoryContext::TArenaMemoryContext() {
 }
 
 TArenaMemoryContext::~TArenaMemoryContext() {
+    MemoryContextDeleteChildren(MyContext);
     Release();
     free(MyContext);
 }
diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index cab7f9f6049..373e3abb16a 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -4948,7 +4948,10 @@ void* PgInitializeMainContext() {
 }
 
 void PgDestroyMainContext(void* ctx) {
-    delete (TMainContext*)ctx;
+    auto typedCtx = (TMainContext*)ctx;
+    MemoryContextDeleteChildren((MemoryContext)&typedCtx->Data);
+    MemoryContextDeleteChildren((MemoryContext)&typedCtx->ErrorData);
+    delete typedCtx;
 }
 
 void PgAcquireThreadContext(void* ctx) {
diff --git a/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/regexp.c b/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/regexp.c
index 830767f35f8..23c295fbc82 100644
--- a/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/regexp.c
+++ b/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/regexp.c
@@ -114,10 +114,9 @@ static __thread int	num_res = 0;		/* # of cached re's */
 static __thread cached_re_str re_array[MAX_CACHED_RES];	/* cached re's */
 
 void RE_cleanup_cache(void) {
-    int i;
-    for (i = 0; i < num_res; ++i) {
-        pg_regfree(&re_array[i].cre_re);
-        free(re_array[i].cre_pat);
+    if (RegexpCacheMemoryContext) {
+        MemoryContextDelete(RegexpCacheMemoryContext);
+        RegexpCacheMemoryContext = NULL;
     }
 
     num_res = 0;
diff --git a/ydb/library/yql/parser/pg_wrapper/ut/memory_ut.cpp b/ydb/library/yql/parser/pg_wrapper/ut/memory_ut.cpp
new file mode 100644
index 00000000000..2a14ea0d436
--- /dev/null
+++ b/ydb/library/yql/parser/pg_wrapper/ut/memory_ut.cpp
@@ -0,0 +1,52 @@
+#include "../pg_compat.h"
+
+extern "C" {
+#include <ydb/library/yql/parser/pg_wrapper/postgresql/src/include/postgres.h>
+#include <ydb/library/yql/parser/pg_wrapper/postgresql/src/include/utils/memutils.h>
+}
+
+#include "arena_ctx.h"
+
+#include <ydb/library/yql/minikql/mkql_alloc.h>
+
+#include <library/cpp/testing/unittest/registar.h>
+
+namespace NYql {
+
+Y_UNIT_TEST_SUITE(TPGMemoryTests) {
+    Y_UNIT_TEST(TestArenaContextBasic) {
+        TArenaMemoryContext arena;
+        auto p1 = palloc(123);
+        auto p2 = palloc(456);
+        Y_UNUSED(p2);
+        pfree(p1);
+    }
+
+    Y_UNIT_TEST(TestMkqlContextBasic) {
+        NKikimr::NMiniKQL::TScopedAlloc alloc(__LOCATION__);
+        auto p1 = palloc(123);
+        auto p2 = palloc(456);
+        Y_UNUSED(p2);
+        pfree(p1);
+    }
+
+    Y_UNIT_TEST(TestArenaContextCleanupChild) {
+        TArenaMemoryContext arena;
+        auto tmpContext = AllocSetContextCreate(CurrentMemoryContext, "Tmp", ALLOCSET_SMALL_SIZES);
+        auto oldcontext = MemoryContextSwitchTo(tmpContext);
+        auto p1 = palloc(123);
+        Y_UNUSED(p1);
+        MemoryContextSwitchTo(oldcontext);
+    }
+
+    Y_UNIT_TEST(TestMkqlContextCleanupChild) {
+        NKikimr::NMiniKQL::TScopedAlloc alloc(__LOCATION__);
+        auto tmpContext = AllocSetContextCreate(CurrentMemoryContext, "Tmp", ALLOCSET_SMALL_SIZES);
+        auto oldcontext = MemoryContextSwitchTo(tmpContext);
+        auto p1 = palloc(123);
+        Y_UNUSED(p1);
+        MemoryContextSwitchTo(oldcontext);
+    }
+}
+
+}
diff --git a/ydb/library/yql/parser/pg_wrapper/ut/ya.make b/ydb/library/yql/parser/pg_wrapper/ut/ya.make
index e95e2209803..96234587342 100644
--- a/ydb/library/yql/parser/pg_wrapper/ut/ya.make
+++ b/ydb/library/yql/parser/pg_wrapper/ut/ya.make
@@ -9,6 +9,7 @@ SRCS(
     arrow_ut.cpp
     codegen_ut.cpp
     error_ut.cpp
+    memory_ut.cpp
     parser_ut.cpp
     sort_ut.cpp
     type_cache_ut.cpp