Support GiST index support functions that want to cache data across calls.

pg_trgm was already doing this unofficially, but the implementation hadn't
been thought through very well and leaked memory.  Restructure the core
GiST code so that it actually works, and document it.  Ordinarily this
would have required an extra memory context creation/destruction for each
GiST index search, but I was able to avoid that in the normal case of a
non-rescanned search by finessing the handling of the RBTree.  It used to
have its own context always, but now shares a context with the
scan-lifespan data structures, unless there is more than one rescan call.
This should make the added overhead unnoticeable in typical cases.

4 files changed, 163 insertions, 62 deletions

diff --git a/src/backend/access/gist/gist.c b/src/backend/access/gist/gist.c
index 24f30099a1c..0ce56b850d3 100644
--- a/src/backend/access/gist/gist.c
+++ b/src/backend/access/gist/gist.c
@@ -94,25 +94,29 @@ gistinsert(PG_FUNCTION_ARGS)
 	IndexUniqueCheck checkUnique = (IndexUniqueCheck) PG_GETARG_INT32(5);
 #endif
 	IndexTuple	itup;
-	GISTSTATE	giststate;
-	MemoryContext oldCtx;
-	MemoryContext insertCtx;
+	GISTSTATE  *giststate;
+	MemoryContext oldCxt;
 
-	insertCtx = createTempGistContext();
-	oldCtx = MemoryContextSwitchTo(insertCtx);
+	giststate = initGISTstate(r);
 
-	initGISTstate(&giststate, r);
+	/*
+	 * We use the giststate's scan context as temp context too.  This means
+	 * that any memory leaked by the support functions is not reclaimed until
+	 * end of insert.  In most cases, we aren't going to call the support
+	 * functions very many times before finishing the insert, so this seems
+	 * cheaper than resetting a temp context for each function call.
+	 */
+	oldCxt = MemoryContextSwitchTo(giststate->tempCxt);
 
-	itup = gistFormTuple(&giststate, r,
+	itup = gistFormTuple(giststate, r,
 						 values, isnull, true /* size is currently bogus */ );
 	itup->t_tid = *ht_ctid;
 
-	gistdoinsert(r, itup, 0, &giststate);
+	gistdoinsert(r, itup, 0, giststate);
 
 	/* cleanup */
-	freeGISTstate(&giststate);
-	MemoryContextSwitchTo(oldCtx);
-	MemoryContextDelete(insertCtx);
+	MemoryContextSwitchTo(oldCxt);
+	freeGISTstate(giststate);
 
 	PG_RETURN_BOOL(false);
 }
@@ -1213,47 +1217,64 @@ gistSplit(Relation r,
 }
 
 /*
- * Fill a GISTSTATE with information about the index
+ * Create a GISTSTATE and fill it with information about the index
  */
-void
-initGISTstate(GISTSTATE *giststate, Relation index)
+GISTSTATE *
+initGISTstate(Relation index)
 {
+	GISTSTATE  *giststate;
+	MemoryContext scanCxt;
+	MemoryContext oldCxt;
 	int			i;
 
+	/* safety check to protect fixed-size arrays in GISTSTATE */
 	if (index->rd_att->natts > INDEX_MAX_KEYS)
 		elog(ERROR, "numberOfAttributes %d > %d",
 			 index->rd_att->natts, INDEX_MAX_KEYS);
 
+	/* Create the memory context that will hold the GISTSTATE */
+	scanCxt = AllocSetContextCreate(CurrentMemoryContext,
+									"GiST scan context",
+									ALLOCSET_DEFAULT_MINSIZE,
+									ALLOCSET_DEFAULT_INITSIZE,
+									ALLOCSET_DEFAULT_MAXSIZE);
+	oldCxt = MemoryContextSwitchTo(scanCxt);
+
+	/* Create and fill in the GISTSTATE */
+	giststate = (GISTSTATE *) palloc(sizeof(GISTSTATE));
+
+	giststate->scanCxt = scanCxt;
+	giststate->tempCxt = scanCxt;	/* caller must change this if needed */
 	giststate->tupdesc = index->rd_att;
 
 	for (i = 0; i < index->rd_att->natts; i++)
 	{
 		fmgr_info_copy(&(giststate->consistentFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_CONSISTENT_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->unionFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_UNION_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->compressFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_COMPRESS_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->decompressFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_DECOMPRESS_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->penaltyFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_PENALTY_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->picksplitFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_PICKSPLIT_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		fmgr_info_copy(&(giststate->equalFn[i]),
 					   index_getprocinfo(index, i + 1, GIST_EQUAL_PROC),
-					   CurrentMemoryContext);
+					   scanCxt);
 		/* opclasses are not required to provide a Distance method */
 		if (OidIsValid(index_getprocid(index, i + 1, GIST_DISTANCE_PROC)))
 			fmgr_info_copy(&(giststate->distanceFn[i]),
 						 index_getprocinfo(index, i + 1, GIST_DISTANCE_PROC),
-						   CurrentMemoryContext);
+						   scanCxt);
 		else
 			giststate->distanceFn[i].fn_oid = InvalidOid;
 
@@ -1273,10 +1294,15 @@ initGISTstate(GISTSTATE *giststate, Relation index)
 		else
 			giststate->supportCollation[i] = DEFAULT_COLLATION_OID;
 	}
+
+	MemoryContextSwitchTo(oldCxt);
+
+	return giststate;
 }
 
 void
 freeGISTstate(GISTSTATE *giststate)
 {
-	/* no work */
+	/* It's sufficient to delete the scanCxt */
+	MemoryContextDelete(giststate->scanCxt);
 }
diff --git a/src/backend/access/gist/gistbuild.c b/src/backend/access/gist/gistbuild.c
index 0046c7b3ab3..be1b202d39d 100644
--- a/src/backend/access/gist/gistbuild.c
+++ b/src/backend/access/gist/gistbuild.c
@@ -54,7 +54,7 @@ typedef enum
 typedef struct
 {
 	Relation	indexrel;
-	GISTSTATE	giststate;
+	GISTSTATE  *giststate;
 	GISTBuildBuffers *gfbb;
 
 	int64		indtuples;		/* number of tuples indexed */
@@ -63,7 +63,6 @@ typedef struct
 	Size		freespace;		/* amount of free space to leave on pages */
 
 	GistBufferingMode bufferingMode;
-	MemoryContext tmpCtx;
 } GISTBuildState;
 
 static void gistInitBuffering(GISTBuildState *buildstate);
@@ -146,7 +145,14 @@ gistbuild(PG_FUNCTION_ARGS)
 			 RelationGetRelationName(index));
 
 	/* no locking is needed */
-	initGISTstate(&buildstate.giststate, index);
+	buildstate.giststate = initGISTstate(index);
+
+	/*
+	 * Create a temporary memory context that is reset once for each tuple
+	 * processed.  (Note: we don't bother to make this a child of the
+	 * giststate's scanCxt, so we have to delete it separately at the end.)
+	 */
+	buildstate.giststate->tempCxt = createTempGistContext();
 
 	/* initialize the root page */
 	buffer = gistNewBuffer(index);
@@ -185,12 +191,6 @@ gistbuild(PG_FUNCTION_ARGS)
 	buildstate.indtuplesSize = 0;
 
 	/*
-	 * create a temporary memory context that is reset once for each tuple
-	 * processed.
-	 */
-	buildstate.tmpCtx = createTempGistContext();
-
-	/*
 	 * Do the heap scan.
 	 */
 	reltuples = IndexBuildHeapScan(heap, index, indexInfo, true,
@@ -208,9 +208,9 @@ gistbuild(PG_FUNCTION_ARGS)
 
 	/* okay, all heap tuples are indexed */
 	MemoryContextSwitchTo(oldcxt);
-	MemoryContextDelete(buildstate.tmpCtx);
+	MemoryContextDelete(buildstate.giststate->tempCxt);
 
-	freeGISTstate(&buildstate.giststate);
+	freeGISTstate(buildstate.giststate);
 
 	/*
 	 * Return statistics
@@ -440,10 +440,10 @@ gistBuildCallback(Relation index,
 	IndexTuple	itup;
 	MemoryContext oldCtx;
 
-	oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
+	oldCtx = MemoryContextSwitchTo(buildstate->giststate->tempCxt);
 
 	/* form an index tuple and point it at the heap tuple */
-	itup = gistFormTuple(&buildstate->giststate, index, values, isnull, true);
+	itup = gistFormTuple(buildstate->giststate, index, values, isnull, true);
 	itup->t_tid = htup->t_self;
 
 	if (buildstate->bufferingMode == GIST_BUFFERING_ACTIVE)
@@ -458,7 +458,7 @@ gistBuildCallback(Relation index,
 		 * locked, we call gistdoinsert directly.
 		 */
 		gistdoinsert(index, itup, buildstate->freespace,
-					 &buildstate->giststate);
+					 buildstate->giststate);
 	}
 
 	/* Update tuple count and total size. */
@@ -466,7 +466,7 @@ gistBuildCallback(Relation index,
 	buildstate->indtuplesSize += IndexTupleSize(itup);
 
 	MemoryContextSwitchTo(oldCtx);
-	MemoryContextReset(buildstate->tmpCtx);
+	MemoryContextReset(buildstate->giststate->tempCxt);
 
 	if (buildstate->bufferingMode == GIST_BUFFERING_ACTIVE &&
 		buildstate->indtuples % BUFFERING_MODE_TUPLE_SIZE_STATS_TARGET == 0)
@@ -520,7 +520,7 @@ static bool
 gistProcessItup(GISTBuildState *buildstate, IndexTuple itup,
 				GISTBufferingInsertStack *startparent)
 {
-	GISTSTATE  *giststate = &buildstate->giststate;
+	GISTSTATE  *giststate = buildstate->giststate;
 	GISTBuildBuffers *gfbb = buildstate->gfbb;
 	Relation	indexrel = buildstate->indexrel;
 	GISTBufferingInsertStack *path;
@@ -652,7 +652,7 @@ gistbufferinginserttuples(GISTBuildState *buildstate, Buffer buffer,
 
 	is_split = gistplacetopage(buildstate->indexrel,
 							   buildstate->freespace,
-							   &buildstate->giststate,
+							   buildstate->giststate,
 							   buffer,
 							   itup, ntup, oldoffnum,
 							   InvalidBuffer,
@@ -720,7 +720,7 @@ gistbufferinginserttuples(GISTBuildState *buildstate, Buffer buffer,
 		 * buffers that will eventually be inserted to them.
 		 */
 		gistRelocateBuildBuffersOnSplit(gfbb,
-										&buildstate->giststate,
+										buildstate->giststate,
 										buildstate->indexrel,
 										path, buffer, splitinfo);
 
@@ -919,7 +919,7 @@ gistProcessEmptyingQueue(GISTBuildState *buildstate)
 			}
 
 			/* Free all the memory allocated during index tuple processing */
-			MemoryContextReset(CurrentMemoryContext);
+			MemoryContextReset(buildstate->giststate->tempCxt);
 		}
 	}
 }
@@ -938,7 +938,7 @@ gistEmptyAllBuffers(GISTBuildState *buildstate)
 	MemoryContext oldCtx;
 	int			i;
 
-	oldCtx = MemoryContextSwitchTo(buildstate->tmpCtx);
+	oldCtx = MemoryContextSwitchTo(buildstate->giststate->tempCxt);
 
 	/*
 	 * Iterate through the levels from top to bottom.
@@ -970,7 +970,7 @@ gistEmptyAllBuffers(GISTBuildState *buildstate)
 					nodeBuffer->queuedForEmptying = true;
 					gfbb->bufferEmptyingQueue =
 						lcons(nodeBuffer, gfbb->bufferEmptyingQueue);
-					MemoryContextSwitchTo(buildstate->tmpCtx);
+					MemoryContextSwitchTo(buildstate->giststate->tempCxt);
 				}
 				gistProcessEmptyingQueue(buildstate);
 			}
diff --git a/src/backend/access/gist/gistget.c b/src/backend/access/gist/gistget.c
index 60116dfd46e..b565d09b388 100644
--- a/src/backend/access/gist/gistget.c
+++ b/src/backend/access/gist/gistget.c
@@ -307,12 +307,12 @@ gistScanPage(IndexScanDesc scan, GISTSearchItem *pageItem, double *myDistances,
 		 * Must call gistindex_keytest in tempCxt, and clean up any leftover
 		 * junk afterward.
 		 */
-		oldcxt = MemoryContextSwitchTo(so->tempCxt);
+		oldcxt = MemoryContextSwitchTo(so->giststate->tempCxt);
 
 		match = gistindex_keytest(scan, it, page, i, &recheck);
 
 		MemoryContextSwitchTo(oldcxt);
-		MemoryContextReset(so->tempCxt);
+		MemoryContextReset(so->giststate->tempCxt);
 
 		/* Ignore tuple if it doesn't match */
 		if (!match)
diff --git a/src/backend/access/gist/gistscan.c b/src/backend/access/gist/gistscan.c
index 5071d46150e..e6140a151bb 100644
--- a/src/backend/access/gist/gistscan.c
+++ b/src/backend/access/gist/gistscan.c
@@ -104,20 +104,28 @@ gistbeginscan(PG_FUNCTION_ARGS)
 	int			nkeys = PG_GETARG_INT32(1);
 	int			norderbys = PG_GETARG_INT32(2);
 	IndexScanDesc scan;
+	GISTSTATE  *giststate;
 	GISTScanOpaque so;
+	MemoryContext oldCxt;
 
 	scan = RelationGetIndexScan(r, nkeys, norderbys);
 
+	/* First, set up a GISTSTATE with a scan-lifespan memory context */
+	giststate = initGISTstate(scan->indexRelation);
+
+	/*
+	 * Everything made below is in the scanCxt, or is a child of the scanCxt,
+	 * so it'll all go away automatically in gistendscan.
+	 */
+	oldCxt = MemoryContextSwitchTo(giststate->scanCxt);
+
 	/* initialize opaque data */
 	so = (GISTScanOpaque) palloc0(sizeof(GISTScanOpaqueData));
-	so->queueCxt = AllocSetContextCreate(CurrentMemoryContext,
-										 "GiST queue context",
-										 ALLOCSET_DEFAULT_MINSIZE,
-										 ALLOCSET_DEFAULT_INITSIZE,
-										 ALLOCSET_DEFAULT_MAXSIZE);
-	so->tempCxt = createTempGistContext();
-	so->giststate = (GISTSTATE *) palloc(sizeof(GISTSTATE));
-	initGISTstate(so->giststate, scan->indexRelation);
+	so->giststate = giststate;
+	giststate->tempCxt = createTempGistContext();
+	so->queue = NULL;
+	so->queueCxt = giststate->scanCxt;			/* see gistrescan */
+
 	/* workspaces with size dependent on numberOfOrderBys: */
 	so->tmpTreeItem = palloc(GSTIHDRSZ + sizeof(double) * scan->numberOfOrderBys);
 	so->distances = palloc(sizeof(double) * scan->numberOfOrderBys);
@@ -125,6 +133,8 @@ gistbeginscan(PG_FUNCTION_ARGS)
 
 	scan->opaque = so;
 
+	MemoryContextSwitchTo(oldCxt);
+
 	PG_RETURN_POINTER(scan);
 }
 
@@ -137,12 +147,44 @@ gistrescan(PG_FUNCTION_ARGS)
 
 	/* nkeys and norderbys arguments are ignored */
 	GISTScanOpaque so = (GISTScanOpaque) scan->opaque;
+	bool		first_time;
 	int			i;
 	MemoryContext oldCxt;
 
 	/* rescan an existing indexscan --- reset state */
-	MemoryContextReset(so->queueCxt);
-	so->curTreeItem = NULL;
+
+	/*
+	 * The first time through, we create the search queue in the scanCxt.
+	 * Subsequent times through, we create the queue in a separate queueCxt,
+	 * which is created on the second call and reset on later calls.  Thus, in
+	 * the common case where a scan is only rescan'd once, we just put the
+	 * queue in scanCxt and don't pay the overhead of making a second memory
+	 * context.  If we do rescan more than once, the first RBTree is just left
+	 * for dead until end of scan; this small wastage seems worth the savings
+	 * in the common case.
+	 */
+	if (so->queue == NULL)
+	{
+		/* first time through */
+		Assert(so->queueCxt == so->giststate->scanCxt);
+		first_time = true;
+	}
+	else if (so->queueCxt == so->giststate->scanCxt)
+	{
+		/* second time through */
+		so->queueCxt = AllocSetContextCreate(so->giststate->scanCxt,
+											 "GiST queue context",
+											 ALLOCSET_DEFAULT_MINSIZE,
+											 ALLOCSET_DEFAULT_INITSIZE,
+											 ALLOCSET_DEFAULT_MAXSIZE);
+		first_time = false;
+	}
+	else
+	{
+		/* third or later time through */
+		MemoryContextReset(so->queueCxt);
+		first_time = false;
+	}
 
 	/* create new, empty RBTree for search queue */
 	oldCxt = MemoryContextSwitchTo(so->queueCxt);
@@ -154,11 +196,28 @@ gistrescan(PG_FUNCTION_ARGS)
 						  scan);
 	MemoryContextSwitchTo(oldCxt);
 
+	so->curTreeItem = NULL;
 	so->firstCall = true;
 
 	/* Update scan key, if a new one is given */
 	if (key && scan->numberOfKeys > 0)
 	{
+		/*
+		 * If this isn't the first time through, preserve the fn_extra
+		 * pointers, so that if the consistentFns are using them to cache
+		 * data, that data is not leaked across a rescan.
+		 */
+		if (!first_time)
+		{
+			for (i = 0; i < scan->numberOfKeys; i++)
+			{
+				ScanKey		skey = scan->keyData + i;
+
+				so->giststate->consistentFn[skey->sk_attno - 1].fn_extra =
+					skey->sk_func.fn_extra;
+			}
+		}
+
 		memmove(scan->keyData, key,
 				scan->numberOfKeys * sizeof(ScanKeyData));
 
@@ -172,6 +231,10 @@ gistrescan(PG_FUNCTION_ARGS)
 		 * Next, if any of keys is a NULL and that key is not marked with
 		 * SK_SEARCHNULL/SK_SEARCHNOTNULL then nothing can be found (ie, we
 		 * assume all indexable operators are strict).
+		 *
+		 * Note: we intentionally memcpy the FmgrInfo to sk_func rather than
+		 * using fmgr_info_copy.  This is so that the fn_extra field gets
+		 * preserved across multiple rescans.
 		 */
 		so->qual_ok = true;
 
@@ -192,6 +255,18 @@ gistrescan(PG_FUNCTION_ARGS)
 	/* Update order-by key, if a new one is given */
 	if (orderbys && scan->numberOfOrderBys > 0)
 	{
+		/* As above, preserve fn_extra if not first time through */
+		if (!first_time)
+		{
+			for (i = 0; i < scan->numberOfOrderBys; i++)
+			{
+				ScanKey		skey = scan->orderByData + i;
+
+				so->giststate->distanceFn[skey->sk_attno - 1].fn_extra =
+					skey->sk_func.fn_extra;
+			}
+		}
+
 		memmove(scan->orderByData, orderbys,
 				scan->numberOfOrderBys * sizeof(ScanKeyData));
 
@@ -201,6 +276,8 @@ gistrescan(PG_FUNCTION_ARGS)
 		 * function in the form of its strategy number, which is available
 		 * from the sk_strategy field, and its subtype from the sk_subtype
 		 * field.
+		 *
+		 * See above comment about why we don't use fmgr_info_copy here.
 		 */
 		for (i = 0; i < scan->numberOfOrderBys; i++)
 		{
@@ -239,13 +316,11 @@ gistendscan(PG_FUNCTION_ARGS)
 	IndexScanDesc scan = (IndexScanDesc) PG_GETARG_POINTER(0);
 	GISTScanOpaque so = (GISTScanOpaque) scan->opaque;
 
+	/*
+	 * freeGISTstate is enough to clean up everything made by gistbeginscan,
+	 * as well as the queueCxt if there is a separate context for it.
+	 */
 	freeGISTstate(so->giststate);
-	pfree(so->giststate);
-	MemoryContextDelete(so->queueCxt);
-	MemoryContextDelete(so->tempCxt);
-	pfree(so->tmpTreeItem);
-	pfree(so->distances);
-	pfree(so);
 
 	PG_RETURN_VOID();
 }