BABEL: Support session-local OID buffer for temp tables in TSQL (babe…

…lfish-for-postgresql#295) Signed-off-by: Tim Chang <[email protected]> add full range of OID Remove ability to alter buffer size GUC, correct error messages Remove lc NULL check Remove references to TempVariableCache Revert "Remove lc NULL check" This reverts commit 07c28e5. Add proper check for toast tables Added comments, added OID macros Add temp table utilization warning Use correct backend for index_create workaround Add include Fix ALTER TABLE with TOAST issue Minor whitespace + comment Signed-off-by: Tim Chang <[email protected]> Ensure non-ENR temp works as expected Disable pgstat for temp tables to avoid concurrency issues Signed-off-by: Tim Chang <[email protected]> Moved new OID generation functions into extension hooks Signed-off-by: Tim Chang <[email protected]> Minor style fixes Signed-off-by: Tim Chang <[email protected]> Additional style fixed, removed obsolete vars Signed-off-by: Tim Chang <[email protected]> Remove 80% warning since it is misleading Signed-off-by: Tim Chang <[email protected]> Minor fixes Signed-off-by: Tim Chang <[email protected]> Move macro to transam for visibility Revert "Move macro to transam for visibility" This reverts commit ea15243. Add simple fetch for VAR_OID_PREFETCH
amazon-aurora · Oct 18, 2024 · cc2fe5e · cc2fe5e
1 parent 66a981a
commit cc2fe5e
Show file tree

Hide file tree

Showing 14 changed files with 93 additions and 12 deletions.
diff --git a/src/backend/access/transam/varsup.c b/src/backend/access/transam/varsup.c
@@ -590,6 +590,14 @@ GetNewObjectId(void)
 	return result;
 }
 
+/*
+ * Simple function to get VAR_OID_PREFETCH when needed outside of this file.
+ */
+int GetVarOidPrefetch(void)
+{
+	return VAR_OID_PREFETCH;
+}
+
 /*
  * SetNextObjectId
  *

diff --git a/src/backend/catalog/catalog.c b/src/backend/catalog/catalog.c
@@ -42,6 +42,7 @@
 #include "catalog/pg_tablespace.h"
 #include "catalog/pg_type.h"
 #include "miscadmin.h"
+#include "parser/parser.h"
 #include "storage/fd.h"
 #include "utils/fmgroids.h"
 #include "utils/fmgrprotos.h"
@@ -53,6 +54,9 @@ IsExtendedCatalogHookType IsExtendedCatalogHook;
 IsToastRelationHookType IsToastRelationHook;
 IsToastClassHookType IsToastClassHook;
 
+GetNewTempObjectId_hook_type GetNewTempObjectId_hook;
+GetNewTempOidWithIndex_hook_type GetNewTempOidWithIndex_hook;
+
 /*
  * Parameters to determine when to emit a log message in
  * GetNewOidWithIndex()
@@ -418,7 +422,7 @@ GetNewOidWithIndex(Relation relation, Oid indexId, AttrNumber oidcolumn)
 	/* Only system relations are supported */
 	Assert(IsSystemRelation(relation));
 
-	/* In bootstrap mode, we don't have any indexes to use */
+	/* In bootstrap mode, we don't have any indexes to use. No temp objects in bootstrap mode. */
 	if (IsBootstrapProcessingMode())
 		return GetNewObjectId();
 
@@ -515,12 +519,14 @@ GetNewOidWithIndex(Relation relation, Oid indexId, AttrNumber oidcolumn)
  * created by bootstrap have preassigned OIDs, so there's no need.
  */
 RelFileNumber
-GetNewRelFileNumber(Oid reltablespace, Relation pg_class, char relpersistence)
+GetNewRelFileNumber(Oid reltablespace, Relation pg_class, char relpersistence, bool override_temp)
 {
 	RelFileLocatorBackend rlocator;
 	char	   *rpath;
 	bool		collides;
 	BackendId	backend;
+	int			tries = 0;
+	bool use_bbf_oid_buffer;
 
 	/*
 	 * If we ever get here during pg_upgrade, there's something wrong; all
@@ -556,16 +562,32 @@ GetNewRelFileNumber(Oid reltablespace, Relation pg_class, char relpersistence)
 	 */
 	rlocator.backend = backend;
 
+	use_bbf_oid_buffer = (relpersistence == RELPERSISTENCE_TEMP && sql_dialect == SQL_DIALECT_TSQL 
+						&& GetNewTempOidWithIndex_hook && temp_oid_buffer_size > 0 && !override_temp);
+
 	do
 	{
 		CHECK_FOR_INTERRUPTS();
 
 		/* Generate the OID */
 		if (pg_class)
-			rlocator.locator.relNumber = GetNewOidWithIndex(pg_class, ClassOidIndexId,
+		{
+			/* Temp tables use temp OID logic */
+			if (use_bbf_oid_buffer)
+				rlocator.locator.relNumber = GetNewTempOidWithIndex_hook(pg_class, ClassOidIndexId,
+													Anum_pg_class_oid);
+			else
+				rlocator.locator.relNumber = GetNewOidWithIndex(pg_class, ClassOidIndexId,
 															Anum_pg_class_oid);
+		}	
 		else
-			rlocator.locator.relNumber = GetNewObjectId();
+		{
+			/* Temp tables use temp OID logic */
+			if (use_bbf_oid_buffer)
+				rlocator.locator.relNumber = GetNewTempObjectId_hook();
+			else
+				rlocator.locator.relNumber = GetNewObjectId();
+		}
 
 		/* Check for existing file of same name */
 		rpath = relpath(rlocator, MAIN_FORKNUM);
@@ -587,7 +609,12 @@ GetNewRelFileNumber(Oid reltablespace, Relation pg_class, char relpersistence)
 			collides = false;
 		}
 
+		if (use_bbf_oid_buffer && tries > temp_oid_buffer_size)
+			ereport(ERROR,
+				(errmsg("Unable to allocate oid for temp table. Drop some temporary tables or start a new session.")));
+
 		pfree(rpath);
+		tries++;
 	} while (collides);
 
 	return rlocator.locator.relNumber;

diff --git a/src/backend/catalog/heap.c b/src/backend/catalog/heap.c
@@ -1307,7 +1307,7 @@ heap_create_with_catalog(const char *relname,
 
 		if (!OidIsValid(relid))
 			relid = GetNewRelFileNumber(reltablespace, pg_class_desc,
-										relpersistence);
+										relpersistence, false);
 	}
 
 	/*

diff --git a/src/backend/catalog/index.c b/src/backend/catalog/index.c
@@ -961,8 +961,12 @@ index_create(Relation heapRelation,
 		}
 		else
 		{
-			indexRelationId =
-				GetNewRelFileNumber(tableSpaceId, pg_class, relpersistence);
+			/* 
+			 * Index OIDs must be kept in normal OID range due to deletion issues.
+			 * Since deletion is sorted by OID, adding indexes to temp OID range 
+			 * causes deletion order issues.
+			 */
+			indexRelationId = GetNewRelFileNumber(tableSpaceId, pg_class, relpersistence, is_enr);
 		}
 	}
 

diff --git a/src/backend/catalog/pg_depend.c b/src/backend/catalog/pg_depend.c
@@ -326,7 +326,8 @@ deleteDependencyRecordsFor(Oid classId, Oid objectId,
 			((Form_pg_depend) GETSTRUCT(tup))->deptype == DEPENDENCY_EXTENSION)
 			continue;
 
-		CatalogTupleDelete(depRel, &tup->t_self);
+		if (!ENRdropTuple(depRel, tup))
+			CatalogTupleDelete(depRel, &tup->t_self);
 		count++;
 	}
 

diff --git a/src/backend/catalog/toasting.c b/src/backend/catalog/toasting.c
@@ -200,6 +200,8 @@ create_toast_table(Relation rel, Oid toastOid, Oid toastIndexOid,
 	 */
 	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(rel))
 		pg_toast_prefix = "@pg_toast";
+	else if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && get_ENR_withoid(currentQueryEnv, rel->rd_id, ENR_TSQL_TEMP))
+		pg_toast_prefix = "#pg_toast";
 
 	snprintf(toast_relname, sizeof(toast_relname),
 			 "%s_%u", pg_toast_prefix, relOid);

diff --git a/src/backend/commands/cluster.c b/src/backend/commands/cluster.c
@@ -738,8 +738,14 @@ make_new_heap(Oid OIDOldHeap, Oid NewTableSpace, Oid NewAccessMethod,
 	 * a shared rel.  However, we do make the new heap mapped if the source is
 	 * mapped.  This simplifies swap_relation_files, and is absolutely
 	 * necessary for rebuilding pg_class, for reasons explained there.
+	 *
+	 * We also must ensure that these temp tables are properly named in TSQL
+	 * so that the metadata is properly cleaned up after in this function.
 	 */
-	snprintf(NewHeapName, sizeof(NewHeapName), "pg_temp_%u", OIDOldHeap);
+	if (sql_dialect == SQL_DIALECT_TSQL && relpersistence == RELPERSISTENCE_TEMP && get_ENR_withoid(currentQueryEnv, OIDOldHeap, ENR_TSQL_TEMP))
+		snprintf(NewHeapName, sizeof(NewHeapName), "#pg_temp_%u", OIDOldHeap);
+	else
+		snprintf(NewHeapName, sizeof(NewHeapName), "pg_temp_%u", OIDOldHeap);
 
 	OIDNewHeap = heap_create_with_catalog(NewHeapName,
 										  namespaceid,
@@ -1610,6 +1616,8 @@ finish_heap_swap(Oid OIDOldHeap, Oid OIDNewHeap,
 			/* rename the toast table ... */
 			if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(newrel))
 				pg_toast_prefix = "@pg_toast";
+			else if (sql_dialect == SQL_DIALECT_TSQL && newrel->rd_rel->relpersistence == RELPERSISTENCE_TEMP && get_ENR_withoid(currentQueryEnv, newrel->rd_id, ENR_TSQL_TEMP))
+				pg_toast_prefix = "#pg_toast";
 
 			snprintf(NewToastName, NAMEDATALEN, "%s_%u",
 					pg_toast_prefix, OIDOldHeap);

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
@@ -14603,7 +14603,7 @@ ATExecSetTableSpace(Oid tableOid, Oid newTableSpace, LOCKMODE lockmode)
 	 * need to allocate a new one in the new tablespace.
 	 */
 	newrelfilenumber = GetNewRelFileNumber(newTableSpace, NULL,
-										   rel->rd_rel->relpersistence);
+										   rel->rd_rel->relpersistence, false);
 
 	/* Open old and new relation */
 	newrlocator = rel->rd_locator;

diff --git a/src/backend/utils/activity/pgstat_relation.c b/src/backend/utils/activity/pgstat_relation.c
@@ -27,6 +27,9 @@
 #include "utils/timestamp.h"
 #include "catalog/catalog.h"
 
+#include "parser/parser.h"
+#include "utils/queryenvironment.h"
+
 
 /* Record that's written to 2PC state file when pgstat state is persisted */
 typedef struct TwoPhasePgStatRecord
@@ -169,6 +172,9 @@ pgstat_unlink_relation(Relation rel)
 void
 pgstat_create_relation(Relation rel)
 {
+	/* Skip pg_stat */
+	if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP)
+		return;
 	pgstat_create_transactional(PGSTAT_KIND_RELATION,
 								rel->rd_rel->relisshared ? InvalidOid : MyDatabaseId,
 								RelationGetRelid(rel));
@@ -183,6 +189,8 @@ pgstat_drop_relation(Relation rel)
 	int			nest_level = GetCurrentTransactionNestLevel();
 	PgStat_TableStatus *pgstat_info;
 
+	if (sql_dialect == SQL_DIALECT_TSQL && rel->rd_rel->relpersistence == RELPERSISTENCE_TEMP)
+		return;
 	pgstat_drop_transactional(PGSTAT_KIND_RELATION,
 							  rel->rd_rel->relisshared ? InvalidOid : MyDatabaseId,
 							  RelationGetRelid(rel));

diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
@@ -3721,7 +3721,7 @@ RelationSetNewRelfilenumber(Relation relation, char persistence)
 	{
 		/* Allocate a new relfilenumber */
 		newrelfilenumber = GetNewRelFileNumber(relation->rd_rel->reltablespace,
-											   NULL, persistence);
+											   NULL, persistence, false);
 	}
 	else if (relation->rd_rel->relkind == RELKIND_INDEX)
 	{

diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
@@ -87,6 +87,11 @@ char	   *GUC_check_errhint_string;
 /* Kluge: for speed, we examine this GUC variable's value directly */
 extern bool in_hot_standby_guc;
 
+/*
+ * OIDs are stored as uint32, so we will add INT_MIN to match the range.
+ */
+int			temp_oid_buffer_start;
+int			temp_oid_buffer_size;
 
 /*
  * Unit conversion tables.

diff --git a/src/include/access/transam.h b/src/include/access/transam.h
@@ -252,6 +252,12 @@ typedef struct VariableCacheData
 	 */
 	TransactionId oldestClogXid;	/* oldest it's safe to look up in clog */
 
+	/*
+	 * This field is also protected by OidGenLock. For tempOidStart, Shmem will
+	 * be the source of truth, as another process may have gotten there first and 
+	 * updated the start.
+	 */
+	Oid			tempOidStart;
 } VariableCacheData;
 
 typedef VariableCacheData *VariableCache;
@@ -295,6 +301,8 @@ extern bool ForceTransactionIdLimitUpdate(void);
 extern Oid	GetNewObjectId(void);
 extern void StopGeneratingPinnedObjectIds(void);
 
+extern int GetVarOidPrefetch(void);
+
 typedef void (*GetNewObjectId_hook_type) (VariableCache variableCache);
 extern PGDLLEXPORT GetNewObjectId_hook_type GetNewObjectId_hook;
 

diff --git a/src/include/catalog/catalog.h b/src/include/catalog/catalog.h
@@ -40,7 +40,14 @@ extern Oid	GetNewOidWithIndex(Relation relation, Oid indexId,
 							   AttrNumber oidcolumn);
 extern RelFileNumber GetNewRelFileNumber(Oid reltablespace,
 										 Relation pg_class,
-										 char relpersistence);
+										 char relpersistence,
+										 bool override_temp);
+
+typedef Oid (*GetNewTempObjectId_hook_type) (void);
+extern GetNewTempObjectId_hook_type GetNewTempObjectId_hook;
+
+typedef Oid (*GetNewTempOidWithIndex_hook_type) (Relation relation, Oid indexId, AttrNumber oidcolumn);
+extern GetNewTempOidWithIndex_hook_type GetNewTempOidWithIndex_hook;
 
 typedef bool (*IsExtendedCatalogHookType) (Oid relationId);
 extern PGDLLEXPORT IsExtendedCatalogHookType IsExtendedCatalogHook;

diff --git a/src/include/utils/guc.h b/src/include/utils/guc.h
@@ -273,6 +273,9 @@ extern PGDLLIMPORT int temp_file_limit;
 
 extern PGDLLIMPORT int num_temp_buffers;
 
+extern PGDLLIMPORT int temp_oid_buffer_start;
+extern PGDLLIMPORT int temp_oid_buffer_size;
+
 extern PGDLLIMPORT char *cluster_name;
 extern PGDLLIMPORT char *ConfigFileName;
 extern PGDLLIMPORT char *HbaFileName;