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.

1 files changed, 91 insertions, 16 deletions

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();
 }