diff --git a/src/yb/yql/pggate/pg_memctx.cc b/src/yb/yql/pggate/pg_memctx.cc index d396e358c956..c11eb1296696 100644 --- a/src/yb/yql/pggate/pg_memctx.cc +++ b/src/yb/yql/pggate/pg_memctx.cc @@ -91,5 +91,9 @@ void PgMemctx::GetCache(size_t hash_id, PgTableDesc **handle) { } } +void ClearGlobalPgMemctxMap() { + postgres_process_memctxs.clear(); +} + } // namespace pggate } // namespace yb diff --git a/src/yb/yql/pggate/pg_memctx.h b/src/yb/yql/pggate/pg_memctx.h index c21263b999f8..9b15777130e8 100644 --- a/src/yb/yql/pggate/pg_memctx.h +++ b/src/yb/yql/pggate/pg_memctx.h @@ -98,6 +98,8 @@ class PgMemctx { DISALLOW_COPY_AND_ASSIGN(PgMemctx); }; +void ClearGlobalPgMemctxMap(); + } // namespace pggate } // namespace yb diff --git a/src/yb/yql/pggate/test/pggate_test.cc b/src/yb/yql/pggate/test/pggate_test.cc index 62b41b4a5376..15d67577be30 100644 --- a/src/yb/yql/pggate/test/pggate_test.cc +++ b/src/yb/yql/pggate/test/pggate_test.cc @@ -32,21 +32,30 @@ extern "C" void FetchUniqueConstraintName(PgOid relation_id, char* dest, size_t CHECK(false) << "Not implemented"; } -} // namespace +YBCPgMemctx global_test_memctx = nullptr; + +YBCPgMemctx GetCurrentTestYbMemctx() { + if (!global_test_memctx) { + global_test_memctx = YBCPgCreateMemctx(); + } + return global_test_memctx; +} -YBCPgMemctx test_memctx = nullptr; -static YBCPgMemctx TestGetCurrentYbMemctx() { - if (!test_memctx) { - test_memctx = YBCPgCreateMemctx(); +void ClearCurrentTestYbMemctx() { + if (global_test_memctx != nullptr) { + CHECK_YBC_STATUS(YBCPgDestroyMemctx(global_test_memctx)); + + // We assume the memory context has actually already been deleted. + global_test_memctx = nullptr; } - return test_memctx; } +} // namespace + PggateTest::PggateTest() { } PggateTest::~PggateTest() { - CHECK_YBC_STATUS(YBCPgDestroyMemctx(test_memctx)); } //-------------------------------------------------------------------------------------------------- @@ -82,6 +91,9 @@ void PggateTest::SetUp() { } void PggateTest::TearDown() { + // It is important to destroy the memory context before destroying PgGate. + ClearCurrentTestYbMemctx(); + // Destroy the client before shutting down servers. YBCDestroyPgGate(); @@ -105,7 +117,7 @@ Status PggateTest::Init(const char *test_name, int num_tablet_servers) { YBCTestGetTypeTable(&type_table, &count); YBCPgCallbacks callbacks; callbacks.FetchUniqueConstraintName = &FetchUniqueConstraintName; - callbacks.GetCurrentYbMemctx = &TestGetCurrentYbMemctx; + callbacks.GetCurrentYbMemctx = &GetCurrentTestYbMemctx; YBCInitPgGate(type_table, count, callbacks); // Don't try to connect to tserver shared memory in pggate tests. diff --git a/src/yb/yql/pggate/test/pggate_test.h b/src/yb/yql/pggate/test/pggate_test.h index 951f691de4ac..fef09e284d6d 100644 --- a/src/yb/yql/pggate/test/pggate_test.h +++ b/src/yb/yql/pggate/test/pggate_test.h @@ -38,7 +38,8 @@ namespace pggate { //-------------------------------------------------------------------------------------------------- // Test base class. //-------------------------------------------------------------------------------------------------- -#define CHECK_YBC_STATUS(s) CheckYBCStatus((s), __FILE__, __LINE__) +#define CHECK_YBC_STATUS(s) \ + ::yb::pggate::PggateTest::CheckYBCStatus((s), __FILE__, __LINE__) class PggateTest : public YBTest { public: @@ -51,7 +52,7 @@ class PggateTest : public YBTest { virtual ~PggateTest(); //------------------------------------------------------------------------------------------------ - void CheckYBCStatus(YBCStatus status, const char* file_name, int line_number); + static void CheckYBCStatus(YBCStatus status, const char* file_name, int line_number); //------------------------------------------------------------------------------------------------ // Test start and cleanup functions. diff --git a/src/yb/yql/pggate/ybc_pggate.cc b/src/yb/yql/pggate/ybc_pggate.cc index d7c2ceb3e9e1..58fb3a35aff5 100644 --- a/src/yb/yql/pggate/ybc_pggate.cc +++ b/src/yb/yql/pggate/ybc_pggate.cc @@ -24,6 +24,7 @@ #include "yb/yql/pggate/pggate_thread_local_vars.h" #include "yb/yql/pggate/pg_txn_manager.h" #include "yb/yql/pggate/pggate_flags.h" +#include "yb/yql/pggate/pg_memctx.h" DECLARE_bool(client_suppress_created_logs); @@ -93,6 +94,7 @@ void YBCDestroyPgGate() { pggate::PgApiImpl* local_pgapi = pgapi; pgapi = nullptr; // YBCPgIsYugaByteEnabled() must return false from now on. delete local_pgapi; + ClearGlobalPgMemctxMap(); VLOG(1) << __PRETTY_FUNCTION__ << " finished"; } }