Add missing index_insert_cleanup calls

The optimization for inserts into BRIN indexes added by c1ec02be1d
relies on a cache that needs to be explicitly released after calling
index_insert(). The commit however failed to invoke the cleanup in
validate_index(), which calls index_insert() indirectly through
table_index_validate_scan().

After inspecting index_insert() callers, it seems unique_key_recheck()
is missing the call too.

Fixed by adding the two missing index_insert_cleanup() calls.

The commit does two additional improvements. The aminsertcleanup()
signature is modified to have the index as the first argument, to make
it more like the other AM callbacks. And the aminsertcleanup() callback
is invoked even if the ii_AmCache is NULL, so that it can decide if the
cleanup is necessary.

Author: Alvaro Herrera, Tomas Vondra
Reported-by: Alexander Lakhin
Discussion: https://postgr.es/m/202401091043.e3nrqiad6gb7@alvherre.pgsql

doc/src/sgml/indexam.sgml

@@ -367,21 +367,21 @@ aminsert (Relation indexRelation,
    within an SQL statement, it can allocate space
    in <literal>indexInfo-&gt;ii_Context</literal> and store a pointer to the
    data in <literal>indexInfo-&gt;ii_AmCache</literal> (which will be NULL
-   initially). After the index insertions complete, the memory will be freed
-   automatically. If additional cleanup is required (e.g. if the cache contains
-   buffers and tuple descriptors), the AM may define an optional function
-   <literal>aminsertcleanup</literal>, called before the memory is released.
+   initially).  If resources other than memory have to be released after
+   index insertions, <function>aminsertcleanup</function> may be provided,
+   which will be called before the memory is released.
   </para>
 
   <para>
 <programlisting>
 void
-aminsertcleanup (IndexInfo *indexInfo);
+aminsertcleanup (Relation indexRelation,
+                 IndexInfo *indexInfo);
 </programlisting>
    Clean up state that was maintained across successive inserts in
    <literal>indexInfo-&gt;ii_AmCache</literal>. This is useful if the data
-   requires additional cleanup steps, and simply releasing the memory is not
-   sufficient.
+   requires additional cleanup steps (e.g., releasing pinned buffers), and
+   simply releasing the memory is not sufficient.
   </para>
 
   <para>

src/backend/access/brin/brin.c

@@ -500,11 +500,13 @@ brininsert(Relation idxRel, Datum *values, bool *nulls,
  * Callback to clean up the BrinInsertState once all tuple inserts are done.
  */
 void
-brininsertcleanup(IndexInfo *indexInfo)
+brininsertcleanup(Relation index, IndexInfo *indexInfo)
 {
 	BrinInsertState *bistate = (BrinInsertState *) indexInfo->ii_AmCache;
 
-	Assert(bistate);
+	/* bail out if cache not initialized */
+	if (indexInfo->ii_AmCache == NULL)
+		return;
 
 	/*
 	 * Clean up the revmap. Note that the brinDesc has already been cleaned up

src/backend/access/index/indexam.c

@@ -242,10 +242,9 @@ index_insert_cleanup(Relation indexRelation,
 					 IndexInfo *indexInfo)
 {
 	RELATION_CHECKS;
-	Assert(indexInfo);
 
-	if (indexRelation->rd_indam->aminsertcleanup && indexInfo->ii_AmCache)
-		indexRelation->rd_indam->aminsertcleanup(indexInfo);
+	if (indexRelation->rd_indam->aminsertcleanup)
+		indexRelation->rd_indam->aminsertcleanup(indexRelation, indexInfo);
 }
 
 /*

src/backend/catalog/index.c

@@ -3402,6 +3402,9 @@ validate_index(Oid heapId, Oid indexId, Snapshot snapshot)
 	/* Done with tuplesort object */
 	tuplesort_end(state.tuplesort);
 
+	/* Make sure to release resources cached in indexInfo (if needed). */
+	index_insert_cleanup(indexRelation, indexInfo);
+
 	elog(DEBUG2,
 		 "validate_index found %.0f heap tuples, %.0f index tuples; inserted %.0f missing tuples",
 		 state.htups, state.itups, state.tups_inserted);

src/backend/commands/constraint.c

@@ -174,6 +174,9 @@ unique_key_recheck(PG_FUNCTION_ARGS)
 		index_insert(indexRel, values, isnull, &checktid,
 					 trigdata->tg_relation, UNIQUE_CHECK_EXISTING,
 					 false, indexInfo);
+
+		/* Cleanup cache possibly initialized by index_insert. */
+		index_insert_cleanup(indexRel, indexInfo);
 	}
 	else
 	{

src/include/access/amapi.h

@@ -114,7 +114,8 @@ typedef bool (*aminsert_function) (Relation indexRelation,
 								   struct IndexInfo *indexInfo);
 
 /* cleanup after insert */
-typedef void (*aminsertcleanup_function) (struct IndexInfo *indexInfo);
+typedef void (*aminsertcleanup_function) (Relation indexRelation,
+										  struct IndexInfo *indexInfo);
 
 /* bulk delete */
 typedef IndexBulkDeleteResult *(*ambulkdelete_function) (IndexVacuumInfo *info,

src/include/access/brin_internal.h

@@ -96,7 +96,7 @@ extern bool brininsert(Relation idxRel, Datum *values, bool *nulls,
 					   IndexUniqueCheck checkUnique,
 					   bool indexUnchanged,
 					   struct IndexInfo *indexInfo);
-extern void brininsertcleanup(struct IndexInfo *indexInfo);
+extern void brininsertcleanup(Relation index, struct IndexInfo *indexInfo);
 extern IndexScanDesc brinbeginscan(Relation r, int nkeys, int norderbys);
 extern int64 bringetbitmap(IndexScanDesc scan, TIDBitmap *tbm);
 extern void brinrescan(IndexScanDesc scan, ScanKey scankey, int nscankeys,

src/test/regress/expected/brin.out

@@ -575,6 +575,7 @@ DROP TABLE brintest_unlogged;
 -- test that the insert optimization works if no rows end up inserted
 CREATE TABLE brin_insert_optimization (a int);
 INSERT INTO brin_insert_optimization VALUES (1);
-CREATE INDEX ON brin_insert_optimization USING brin (a);
+CREATE INDEX brin_insert_optimization_idx ON brin_insert_optimization USING brin (a);
 UPDATE brin_insert_optimization SET a = a;
+REINDEX INDEX CONCURRENTLY brin_insert_optimization_idx;
 DROP TABLE brin_insert_optimization;

src/test/regress/sql/brin.sql

@@ -519,6 +519,7 @@ DROP TABLE brintest_unlogged;
 -- test that the insert optimization works if no rows end up inserted
 CREATE TABLE brin_insert_optimization (a int);
 INSERT INTO brin_insert_optimization VALUES (1);
-CREATE INDEX ON brin_insert_optimization USING brin (a);
+CREATE INDEX brin_insert_optimization_idx ON brin_insert_optimization USING brin (a);
 UPDATE brin_insert_optimization SET a = a;
+REINDEX INDEX CONCURRENTLY brin_insert_optimization_idx;
 DROP TABLE brin_insert_optimization;