Enable ORCA to be tracked by Mem Accounting (#3378)

Before this commit all memory allocations made by ORCA/GPOS were a
blackbox to GPDB. However the ground work had been in place to allow
GPDB's Memory Accounting Framework to track memory consumption by ORCA.
This commit introduces two new functions
Ext_OptimizerAlloc and Ext_OptimizerFree which
pass through their parameters to gp_malloc and gp_free and do some bookeeping
against the Optimizer Memory Account. This introduces very little
overhead to the GPOS memory management framework.
Signed-off-by: NMelanie Plageman <mplageman@pivotal.io>
Signed-off-by: NSambitesh Dash <sdash@pivotal.io>

src/backend/gpopt/CGPOptimizer.cpp

@@ -20,8 +20,10 @@
 #include "naucrates/init.h"
 #include "gpopt/init.h"
 #include "gpos/_api.h"
+#include "gpopt/gpdbwrappers.h"
 
 #include "naucrates/exception.h"
+#include "utils/guc.h"
 
 extern MemoryContext MessageContext;
 
@@ -126,7 +128,14 @@ void
 CGPOptimizer::InitGPOPT ()
 {
   // Use GPORCA's default allocators
-  struct gpos_init_params params = { NULL, NULL };
+  void *(*gpos_alloc)(size_t) = NULL;
+  void (*gpos_free)(void *) = NULL;
+  if (optimizer_use_gpdb_allocators)
+  {
+	gpos_alloc = gpdb::OptimizerAlloc;
+	gpos_free = gpdb::OptimizerFree;
+  }
+  struct gpos_init_params params = {gpos_alloc, gpos_free};
   gpos_init(&params);
   gpdxl_init();
   gpopt_init();

src/backend/gpopt/gpdbwrappers.cpp

@@ -30,6 +30,8 @@
 
 #include "gpopt/gpdbwrappers.h"
 
+#include "utils/ext_alloc.h"
+
 #define GP_WRAP_START	\
 	sigjmp_buf local_sigjmp_buf;	\
 	{	\
@@ -3058,4 +3060,33 @@ gpdb::FMDCacheNeedsReset
 	return true;
 }
 
+// Functions for ORCA's memory consumption to be tracked by GPDB
+void *
+gpdb::OptimizerAlloc
+		(
+			size_t size
+		)
+{
+	GP_WRAP_START;
+	{
+		return Ext_OptimizerAlloc(size);
+	}
+	GP_WRAP_END;
+
+	return NULL;
+}
+
+void
+gpdb::OptimizerFree
+		(
+			void *ptr
+		)
+{
+	GP_WRAP_START;
+	{
+		Ext_OptimizerFree(ptr);
+	}
+	GP_WRAP_END;
+}
+
 // EOF

src/backend/utils/init/postinit.c

@@ -598,7 +598,11 @@ InitPostgres(const char *in_dbname, Oid dboid, const char *username,
 
 #ifdef USE_ORCA
 	/* Initialize GPOPT */
-	InitGPOPT();
+	START_MEMORY_ACCOUNT(MemoryAccounting_CreateAccount(0, MEMORY_OWNER_TYPE_Optimizer));
+	{
+		InitGPOPT();
+	}
+	END_MEMORY_ACCOUNT();
 #endif
 
 	/*
@@ -1100,7 +1104,11 @@ ShutdownPostgres(int code, Datum arg)
 	ReportOOMConsumption();
 
 #ifdef USE_ORCA
-  TerminateGPOPT();
+	START_MEMORY_ACCOUNT(MemoryAccounting_CreateAccount(0, MEMORY_OWNER_TYPE_Optimizer));
+	{
+		TerminateGPOPT();
+	}
+	END_MEMORY_ACCOUNT();
 #endif
 
 	/* Disable memory protection */

src/backend/utils/misc/guc_gp.c

@@ -476,6 +476,7 @@ bool		optimizer_minidump;
 int			optimizer_cost_model;
 bool		optimizer_metadata_caching;
 int			optimizer_mdcache_size;
+bool		optimizer_use_gpdb_allocators;
 
 /* Optimizer debugging GUCs */
 bool		optimizer_print_query;
@@ -3199,6 +3200,16 @@ struct config_bool ConfigureNamesBool_gp[] =
 		false, NULL, NULL
 	},
 
+	{
+		{"optimizer_use_gpdb_allocators", PGC_POSTMASTER, RESOURCES_MEM,
+			gettext_noop("Enable ORCA to use GPDB Memory Accounting"),
+			NULL,
+			GUC_NO_SHOW_ALL | GUC_NOT_IN_SAMPLE
+		},
+		&optimizer_use_gpdb_allocators,
+		false, NULL, NULL
+	},
+
 	{
 		{"init_codegen", PGC_POSTMASTER, DEVELOPER_OPTIONS,
 			gettext_noop("Enable just-in-time code generation."),

src/backend/utils/mmgr/Makefile

@@ -12,7 +12,7 @@ subdir = src/backend/utils/mmgr
 top_builddir = ../../../..
 include $(top_builddir)/src/Makefile.global
 
-OBJS =  aset.o mcxt.o memaccounting.o mpool.o portalmem.o memprot.o vmem_tracker.o redzone_handler.o runaway_cleaner.o idle_tracker.o event_version.o
+OBJS =  aset.o mcxt.o memaccounting.o mpool.o portalmem.o memprot.o vmem_tracker.o redzone_handler.o runaway_cleaner.o idle_tracker.o event_version.o ext_alloc.o
 
 # In PostgreSQL, this is under src/common. It has been backported, but because
 # we haven't merged the changes that introduced the src/common directory, it

src/backend/utils/mmgr/ext_alloc.c

+/*-------------------------------------------------------------------------
+ *
+ * ext_alloc.c
+ *
+ * Portions Copyright (c) 2017-Present Pivotal Software, Inc.
+ *
+ *
+ * IDENTIFICATION
+ *	    src/backend/utils/mmgr/ext_alloc.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+#include "utils/vmem_tracker.h"
+#include "utils/gp_alloc.h"
+#include "utils/memaccounting.h"
+#include "utils/memaccounting_private.h"
+
+/*
+ * This variable is used to track memory that is not freed by Orca during a
+ * single short living Optimizer account.
+ *
+ * Note: Although theoretically this value could overflow, before that happens
+ * the underlying system is likely to hit some kind of Vmem limit
+ */
+uint64 OptimizerOutstandingMemoryBalance = 0;
+
+/*
+ * Allocation & Deallocation functions for GPOS
+ *
+ *   These functions provide GPOS the facility to have all its memory managed by gp_malloc/free
+ */
+
+void*
+Ext_OptimizerAlloc(size_t size)
+{
+#ifdef USE_ASSERT_CHECKING
+	MemoryAccount *account = MemoryAccounting_ConvertIdToAccount(ActiveMemoryAccountId);
+	Assert(account->ownerType == MEMORY_OWNER_TYPE_Optimizer);
+#endif
+
+	MemoryAccounting_Allocate(ActiveMemoryAccountId, size);
+	OptimizerOutstandingMemoryBalance += size;
+	return gp_malloc(size);
+}
+
+void
+Ext_OptimizerFree(void *ptr)
+{
+	MemoryAccount *account = MemoryAccounting_ConvertIdToAccount(ActiveMemoryAccountId);
+	void *malloc_pointer = UserPtr_GetVmemPtr(ptr);
+	size_t freed_size = VmemPtr_GetUserPtrSize((VmemHeader*) malloc_pointer);
+	MemoryAccounting_Free(ActiveMemoryAccountId, freed_size);
+	OptimizerOutstandingMemoryBalance -= freed_size;
+	gp_free(ptr);
+}
+
+uint64
+GetOptimizerOutstandingMemoryBalance()
+{
+	return OptimizerOutstandingMemoryBalance;
+}
+

src/backend/utils/mmgr/memaccounting.c

@@ -25,6 +25,8 @@
 #include "miscadmin.h"
 #include "utils/vmem_tracker.h"
 #include "utils/memaccounting_private.h"
+#include "utils/gp_alloc.h"
+#include "utils/ext_alloc.h"
 
 #define MEMORY_REPORT_FILE_NAME_LENGTH 255
 #define SHORT_LIVING_MEMORY_ACCOUNT_ARRAY_INIT_LEN 64
@@ -627,6 +629,23 @@ InitializeMemoryAccount(MemoryAccount *newAccount, long maxLimit, MemoryOwnerTyp
 	newAccount->maxLimit = maxLimit;
 
 	newAccount->allocated = 0;
+
+	/*
+	 * Every call to ORCA to optimize a query maps to a new short living memory
+	 * account. However, the nature of Orca's memory usage is that it holds data
+	 * in a cache. Thus, GetOptimizerOutstandingMemoryBalance() returns the
+	 * current amount of memory that Orca has not yet freed according to the
+	 * Memory Accounting framework. Each new Orca memory account will start off
+	 * its 'allocated' amount from the outstanding amount. This approach ensures
+	 * that when Orca does release memory it allocated during an earlier
+	 * generation that the accounting math does not lead to an underflow but
+	 * properly accounts for the outstanding amount.
+	 */
+	if (ownerType == MEMORY_OWNER_TYPE_Optimizer)
+	{
+		elog(DEBUG2, "Rolling over previous outstanding Optimizer allocated memory %lu", GetOptimizerOutstandingMemoryBalance());
+		newAccount->allocated = GetOptimizerOutstandingMemoryBalance();
+	}
 	newAccount->freed = 0;
 	newAccount->peak = 0;
 	newAccount->parentId = parentAccountId;
@@ -669,7 +688,8 @@ CreateMemoryAccountImpl(long maxLimit, MemoryOwnerType ownerType, MemoryAccountI
 	 * TopMemoryContext, and not under MemoryAccountMemoryContext
 	 */
 	Assert(ownerType == MEMORY_OWNER_TYPE_LogicalRoot || ownerType == MEMORY_OWNER_TYPE_SharedChunkHeader ||
-			ownerType == MEMORY_OWNER_TYPE_Rollover || ownerType == MEMORY_OWNER_TYPE_MemAccount ||
+			ownerType == MEMORY_OWNER_TYPE_Rollover ||
+			ownerType == MEMORY_OWNER_TYPE_MemAccount ||
 			ownerType == MEMORY_OWNER_TYPE_Exec_AlienShared ||
 			(MemoryAccountMemoryContext != NULL && MemoryAccountMemoryAccount != NULL));
 
@@ -1289,6 +1309,7 @@ AdvanceMemoryAccountingGeneration()
 	 * reset.
 	 */
 	MemoryAccounting_SwitchAccount((MemoryAccountIdType)MEMORY_OWNER_TYPE_Rollover);
+
 	MemoryContextReset(MemoryAccountMemoryContext);
 	Assert(MemoryAccountMemoryAccount->allocated == MemoryAccountMemoryAccount->freed);
 	shortLivingMemoryAccountArray = NULL;

src/backend/utils/mmgr/test/memaccounting_test.c

@@ -1245,7 +1245,8 @@ test__MemoryAccounting_SaveToFile__GeneratesCorrectString(void **state)
 
 	int memoryOwnerTypes[] = {MEMORY_STAT_TYPE_VMEM_RESERVED, MEMORY_STAT_TYPE_MEMORY_ACCOUNTING_PEAK,
 			MEMORY_OWNER_TYPE_LogicalRoot, MEMORY_OWNER_TYPE_Top, MEMORY_OWNER_TYPE_Exec_Hash ,
-			MEMORY_OWNER_TYPE_Exec_AlienShared, MEMORY_OWNER_TYPE_MemAccount, MEMORY_OWNER_TYPE_Rollover,
+			MEMORY_OWNER_TYPE_Exec_AlienShared, MEMORY_OWNER_TYPE_MemAccount,
+			MEMORY_OWNER_TYPE_Rollover,
 			MEMORY_OWNER_TYPE_SharedChunkHeader};
 
 	char runId[80];

src/include/gpopt/gpdbwrappers.h

@@ -629,6 +629,11 @@ namespace gpdb {
 	// table has been changed?)
 	bool FMDCacheNeedsReset(void);
 
+	// functions for tracking ORCA memory consumption
+	void *OptimizerAlloc(size_t size);
+
+	void OptimizerFree(void *ptr);
+
 } //namespace gpdb
 
 #define ForEach(cell, l)	\

src/include/utils/ext_alloc.h

+/*-------------------------------------------------------------------------
+ *
+ * ext_alloc.h
+ *	  This file contains declarations for external memory management
+ *	  functions.
+ *
+ * Portions Copyright (c) 2017-Present Pivotal Software, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+#ifndef EXT_ALLOC_H
+#define EXT_ALLOC_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+extern uint64 OptimizerOustandingMemoryBalance;
+
+extern void
+Ext_OptimizerFree(void *ptr);
+
+extern void*
+Ext_OptimizerAlloc(size_t size);
+
+extern uint64
+GetOptimizerOutstandingMemoryBalance();
+
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

src/include/utils/guc.h

@@ -493,6 +493,8 @@ extern bool optimizer_enable_space_pruning;
 extern bool optimizer_analyze_root_partition;
 extern bool optimizer_analyze_midlevel_partition;
 
+extern bool optimizer_use_gpdb_allocators;
+
 
 /**
  * GUCs related to code generation.

src/include/utils/memaccounting.h

@@ -233,5 +233,4 @@ MemoryAccounting_SaveToLog(void);
 extern void
 MemoryAccounting_PrettyPrint(void);
 
-
 #endif   /* MEMACCOUNTING_H */