Support session-local OID buffer for temp tables in TSQL

Signed-off-by: Tim Chang <[email protected]>

src/backend/access/transam/varsup.c

@@ -24,6 +24,7 @@
 #include "postmaster/autovacuum.h"
 #include "storage/pmsignal.h"
 #include "storage/proc.h"
+#include "utils/guc.h"
 #include "utils/syscache.h"
 
 
@@ -32,6 +33,7 @@
 
 /* pointer to "variable cache" in shared memory (set up by shmem.c) */
 VariableCache ShmemVariableCache = NULL;
+VariableCache TempVariableCache = NULL;
 
 GetNewObjectId_hook_type GetNewObjectId_hook = NULL;
 
@@ -514,6 +516,68 @@ ForceTransactionIdLimitUpdate(void)
 	return false;
 }
 
+Oid
+GetNewTempObjectId(void)
+{
+	Oid			result;
+	Oid			tempOidStart;
+	static Oid 	nextTempOid = InvalidOid;
+
+	/* safety check, we should never get this far in a HS standby */
+	if (RecoveryInProgress())
+		elog(ERROR, "cannot assign OIDs during recovery");
+
+	if (!OidIsValid((Oid) temp_oid_buffer_start)) /*  If this is the first time initializing */
+	{
+		/* First check to see if another connection has already picked a start, then update. */
+		LWLockAcquire(OidGenLock, LW_EXCLUSIVE);
+		if (OidIsValid(ShmemVariableCache->tempOidStart) && (ShmemVariableCache->tempOidBufferSize == temp_oid_buffer_size))
+		{
+			/* If ShmemVariableCache->tempOidStart is valid, it could also mean the buffer size GUC has been altered, and we need to pick a new start. */
+			temp_oid_buffer_start = ShmemVariableCache->tempOidStart;
+			nextTempOid = ShmemVariableCache->tempOidStart;
+		}
+		else
+		{
+			/* The first time this is called, we need to pick a start for the buffer range. */
+			tempOidStart = ShmemVariableCache->nextOid;
+			if (ShmemVariableCache->oidCount < temp_oid_buffer_size)
+				ShmemVariableCache->oidCount = 0;
+			else
+				ShmemVariableCache->oidCount -= temp_oid_buffer_size;
+
+			/* If the OID range would wraparound, start from beginning instead. */
+			if (tempOidStart + temp_oid_buffer_size < tempOidStart) 
+			{
+				tempOidStart = FirstNormalObjectId;
+			}
+			temp_oid_buffer_start = tempOidStart;
+			ShmemVariableCache->tempOidStart = tempOidStart;
+			ShmemVariableCache->tempOidBufferSize = temp_oid_buffer_size;
+			TempVariableCache = (VariableCache) palloc0(sizeof(*TempVariableCache));
+
+			TempVariableCache->nextOid = (Oid) tempOidStart;
+			nextTempOid = (Oid) tempOidStart;
+
+			result = tempOidStart;
+		}
+		LWLockRelease(OidGenLock);
+	}
+
+	/*
+	 * Check for wraparound of the temp OID buffer.
+	 */
+	if (nextTempOid >= (Oid) (temp_oid_buffer_start + temp_oid_buffer_size) || nextTempOid < (Oid) temp_oid_buffer_start)
+	{
+		nextTempOid = (Oid) temp_oid_buffer_start;
+	}
+
+	result = nextTempOid;
+
+	nextTempOid++;
+
+	return result;
+}
 
 /*
  * GetNewObjectId -- allocate a new OID
@@ -536,6 +600,14 @@ GetNewObjectId(void)
 
 	LWLockAcquire(OidGenLock, LW_EXCLUSIVE);
 
+	/* If we're in tempOid buffer range, skip ahead. */
+	if (OidIsValid(temp_oid_buffer_start)
+	&& ShmemVariableCache->nextOid >= temp_oid_buffer_start
+	&& ShmemVariableCache->nextOid < (temp_oid_buffer_start + temp_oid_buffer_size))
+	{
+		ShmemVariableCache->nextOid = temp_oid_buffer_start + temp_oid_buffer_size;
+	}
+
 	/*
 	 * Check for wraparound of the OID counter.  We *must* not return 0
 	 * (InvalidOid), and in normal operation we mustn't return anything below

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"
@@ -416,7 +417,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();
 
@@ -433,7 +434,111 @@ GetNewOidWithIndex(Relation relation, Oid indexId, AttrNumber oidcolumn)
 	{
 		CHECK_FOR_INTERRUPTS();
 
-		newOid = GetNewObjectId();
+		if (relation->rd_rel->relpersistence == RELPERSISTENCE_TEMP && sql_dialect == SQL_DIALECT_TSQL)
+		{
+			/* Temp tables use temp OID logic */
+			newOid = GetNewTempObjectId();
+		}
+		else
+		{
+			newOid = GetNewObjectId();
+		}
+
+		ScanKeyInit(&key,
+					oidcolumn,
+					BTEqualStrategyNumber, F_OIDEQ,
+					ObjectIdGetDatum(newOid));
+
+		/* see notes above about using SnapshotAny */
+		scan = systable_beginscan(relation, indexId, true,
+								SnapshotAny, 1, &key);
+
+		collides = HeapTupleIsValid(systable_getnext(scan));
+
+		systable_endscan(scan);
+
+
+
+		/*
+		 * Log that we iterate more than GETNEWOID_LOG_THRESHOLD but have not
+		 * yet found OID unused in the relation. Then repeat logging with
+		 * exponentially increasing intervals until we iterate more than
+		 * GETNEWOID_LOG_MAX_INTERVAL. Finally repeat logging every
+		 * GETNEWOID_LOG_MAX_INTERVAL unless an unused OID is found. This
+		 * logic is necessary not to fill up the server log with the similar
+		 * messages.
+		 */
+		if (retries >= retries_before_log)
+		{
+			ereport(LOG,
+					(errmsg("still searching for an unused OID in relation \"%s\"",
+							RelationGetRelationName(relation)),
+					 errdetail_plural("OID candidates have been checked %llu time, but no unused OID has been found yet.",
+									  "OID candidates have been checked %llu times, but no unused OID has been found yet.",
+									  retries,
+									  (unsigned long long) retries)));
+
+			/*
+			 * Double the number of retries to do before logging next until it
+			 * reaches GETNEWOID_LOG_MAX_INTERVAL.
+			 */
+			if (retries_before_log * 2 <= GETNEWOID_LOG_MAX_INTERVAL)
+				retries_before_log *= 2;
+			else
+				retries_before_log += GETNEWOID_LOG_MAX_INTERVAL;
+		}
+
+		retries++;
+	} while (collides);
+
+	/*
+	 * If at least one log message is emitted, also log the completion of OID
+	 * assignment.
+	 */
+	if (retries > GETNEWOID_LOG_THRESHOLD)
+	{
+		ereport(LOG,
+				(errmsg_plural("new OID has been assigned in relation \"%s\" after %llu retry",
+							   "new OID has been assigned in relation \"%s\" after %llu retries",
+							   retries,
+							   RelationGetRelationName(relation), (unsigned long long) retries)));
+	}
+
+	return newOid;
+}
+
+Oid
+GetNewTempOidWithIndex(Relation relation, Oid indexId, AttrNumber oidcolumn)
+{
+	Oid			newOid;
+	SysScanDesc scan;
+	ScanKeyData key;
+	bool		collides;
+	uint64		retries = 0;
+	uint64		retries_before_log = GETNEWOID_LOG_THRESHOLD;
+
+	/* Only system relations are supported */
+	Assert(IsSystemRelation(relation));
+
+	/* In bootstrap mode, we don't have any indexes to use. No temp objects in bootstrap mode. */
+	if (IsBootstrapProcessingMode())
+		return GetNewObjectId();
+
+	/*
+	 * We should never be asked to generate a new pg_type OID during
+	 * pg_upgrade; doing so would risk collisions with the OIDs it wants to
+	 * assign.  Hitting this assert means there's some path where we failed to
+	 * ensure that a type OID is determined by commands in the dump script.
+	 */
+	Assert(!IsBinaryUpgrade || RelationGetRelid(relation) != TypeRelationId);
+
+	/* Generate new OIDs until we find one not in the table */
+	do
+	{
+		CHECK_FOR_INTERRUPTS();
+
+		newOid = GetNewTempObjectId();
+		// elog(WARNING, "GetNewTempOidWithIndex got a newOid of %d for indexId %d", newOid, indexId);
 
 		ScanKeyInit(&key,
 					oidcolumn,
@@ -442,12 +547,14 @@ GetNewOidWithIndex(Relation relation, Oid indexId, AttrNumber oidcolumn)
 
 		/* see notes above about using SnapshotAny */
 		scan = systable_beginscan(relation, indexId, true,
-								  SnapshotAny, 1, &key);
+								SnapshotAny, 1, &key);
 
 		collides = HeapTupleIsValid(systable_getnext(scan));
 
 		systable_endscan(scan);
 
+
+
 		/*
 		 * Log that we iterate more than GETNEWOID_LOG_THRESHOLD but have not
 		 * yet found OID unused in the relation. Then repeat logging with
@@ -558,10 +665,31 @@ GetNewRelFileNode(Oid reltablespace, Relation pg_class, char relpersistence)
 
 		/* Generate the OID */
 		if (pg_class)
-			rnode.node.relNode = GetNewOidWithIndex(pg_class, ClassOidIndexId,
+		{
+			if (relpersistence == RELPERSISTENCE_TEMP && sql_dialect == SQL_DIALECT_TSQL)
+			{
+				/* Temp tables use temp OID logic */
+				rnode.node.relNode = GetNewTempOidWithIndex(pg_class, ClassOidIndexId,
+													Anum_pg_class_oid);
+			}
+			else
+			{
+				rnode.node.relNode = GetNewOidWithIndex(pg_class, ClassOidIndexId,
 													Anum_pg_class_oid);
+			}
+		}	
 		else
-			rnode.node.relNode = GetNewObjectId();
+		{
+			if (relpersistence == RELPERSISTENCE_TEMP && sql_dialect == SQL_DIALECT_TSQL)
+			{
+				/* Temp tables use temp OID logic */
+				rnode.node.relNode = GetNewTempObjectId();
+			}
+			else
+			{
+				rnode.node.relNode = GetNewObjectId();
+			}
+		}
 
 		/* Check for existing file of same name */
 		rpath = relpath(rnode, MAIN_FORKNUM);

src/backend/catalog/index.c

@@ -950,8 +950,16 @@ index_create(Relation heapRelation,
 		}
 		else
 		{
-			indexRelationId =
-				GetNewRelFileNode(tableSpaceId, pg_class, relpersistence);
+			if (is_enr)
+			{
+				/* Index OIDs must be kept in normal OID range. */
+				indexRelationId = GetNewRelFileNode(tableSpaceId, pg_class, RELPERSISTENCE_PERMANENT);
+			}
+			else
+			{
+				indexRelationId =
+					GetNewRelFileNode(tableSpaceId, pg_class, relpersistence);
+			}
 		}
 	}
 

src/backend/catalog/toasting.c

@@ -198,7 +198,7 @@ create_toast_table(Relation rel, Oid toastOid, Oid toastIndexOid,
 	/*
 	 * Create the toast table and its index
 	 */
-	if (sql_dialect == SQL_DIALECT_TSQL && RelationIsBBFTableVariable(rel))
+	if (sql_dialect == SQL_DIALECT_TSQL && (RelationIsBBFTableVariable(rel) || RelationIsBBFTempTable(rel)))
 		pg_toast_prefix = "@pg_toast";
 
 	snprintf(toast_relname, sizeof(toast_relname),

src/backend/utils/init/postinit.c

@@ -1141,6 +1141,8 @@ InitPostgres(const char *in_dbname, Oid dboid,
 	/* Initialize this backend's session state. */
 	InitializeSession();
 
+	TempVariableCache = (VariableCache) palloc0(sizeof(*TempVariableCache));
+
 	/*
 	 * If this is an interactive session, load any libraries that should be
 	 * preloaded at backend start.  Since those are determined by GUCs, this

src/backend/utils/misc/guc.c

@@ -694,6 +694,12 @@ int			ssl_renegotiation_limit;
 int			huge_pages;
 int			huge_page_size;
 
+/*
+ * OIDs are stored as uint32, so we have to store it in a double.
+ */
+int 		temp_oid_buffer_start;
+int			temp_oid_buffer_size;
+
 /*
  * These variables are all dummies that don't do anything, except in some
  * cases provide the value for SHOW to display.  The real state is elsewhere

src/backend/utils/misc/queryenvironment.c

@@ -878,15 +878,21 @@ static bool _ENR_tuple_operation(Relation catalog_rel, HeapTuple tup, ENRTupleOp
 					CacheInvalidateHeapTuple(catalog_rel, newtup, NULL);
 					break;
 				case ENR_OP_UPDATE:
-					oldtup = lfirst(lc);
-					lfirst(lc) = heap_copytuple(tup);
-					CacheInvalidateHeapTuple(catalog_rel, oldtup, tup);
+					if (lc)
+					{
+						oldtup = lfirst(lc);
+						lfirst(lc) = heap_copytuple(tup);
+						CacheInvalidateHeapTuple(catalog_rel, oldtup, tup);
+					}
 					break;
 				case ENR_OP_DROP:
-					tmp = lfirst(lc);
-					*list_ptr = list_delete_ptr(*list_ptr, tmp);
-					CacheInvalidateHeapTuple(catalog_rel, tup, NULL);
-					heap_freetuple(tmp);
+					if (lc)
+					{
+						tmp = lfirst(lc);
+						*list_ptr = list_delete_ptr(*list_ptr, tmp);
+						CacheInvalidateHeapTuple(catalog_rel, tup, NULL);
+						heap_freetuple(tmp);
+					}
 					break;
 				default:
 					break;

src/include/access/transam.h

@@ -252,6 +252,13 @@ typedef struct VariableCacheData
 	 */
 	TransactionId oldestClogXid;	/* oldest it's safe to look up in clog */
 
+	/*
+	 * These fields are 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. For tempOidBufferSize, the GUC will be the source of truth.
+	 */
+	Oid			tempOidStart;
+	uint32		tempOidBufferSize;
 } VariableCacheData;
 
 typedef VariableCacheData *VariableCache;
@@ -267,6 +274,7 @@ extern bool TransactionStartedDuringRecovery(void);
 
 /* in transam/varsup.c */
 extern PGDLLIMPORT VariableCache ShmemVariableCache;
+extern PGDLLIMPORT VariableCache TempVariableCache;
 
 /*
  * prototypes for functions in transam/transam.c
@@ -292,6 +300,7 @@ extern void SetTransactionIdLimit(TransactionId oldest_datfrozenxid,
 								  Oid oldest_datoid);
 extern void AdvanceOldestClogXid(TransactionId oldest_datfrozenxid);
 extern bool ForceTransactionIdLimitUpdate(void);
+extern Oid	GetNewTempObjectId(void);
 extern Oid	GetNewObjectId(void);
 extern void StopGeneratingPinnedObjectIds(void);
 

src/include/catalog/catalog.h

@@ -36,6 +36,8 @@ extern bool IsSharedRelation(Oid relationId);
 
 extern bool IsPinnedObject(Oid classId, Oid objectId);
 
+extern Oid	GetNewTempOidWithIndex(Relation relation, Oid indexId,
+							   AttrNumber oidcolumn);
 extern Oid	GetNewOidWithIndex(Relation relation, Oid indexId,
 							   AttrNumber oidcolumn);
 extern Oid	GetNewRelFileNode(Oid reltablespace, Relation pg_class,