From 29bf6acba8c8ad41dad7406862437ef4fc84cede Mon Sep 17 00:00:00 2001
From: Tomas Vondra <tomas.vondra@postgresql.org>
Date: Mon, 4 Oct 2021 15:10:45 +0200
Subject: [PATCH 1/2] v2

---
 src/backend/access/brin/brin.c                |  1 +
 src/backend/access/gin/ginutil.c              |  1 +
 src/backend/access/gist/gist.c                |  1 +
 src/backend/access/hash/hash.c                |  1 +
 src/backend/access/heap/heapam.c              |  2 +-
 src/backend/access/nbtree/nbtree.c            |  1 +
 src/backend/access/spgist/spgutils.c          |  1 +
 src/backend/utils/cache/relcache.c            | 19 +++++++
 src/include/access/amapi.h                    |  2 +
 src/include/utils/rel.h                       |  1 +
 src/include/utils/relcache.h                  |  3 +-
 .../modules/dummy_index_am/dummy_index_am.c   |  1 +
 src/test/regress/expected/brin.out            | 49 +++++++++++++++++++
 src/test/regress/sql/brin.sql                 | 46 +++++++++++++++++
 14 files changed, 127 insertions(+), 2 deletions(-)

diff --git a/src/backend/access/brin/brin.c b/src/backend/access/brin/brin.c
index ccc9fa0959..f521bb9635 100644
--- a/src/backend/access/brin/brin.c
+++ b/src/backend/access/brin/brin.c
@@ -108,6 +108,7 @@ brinhandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = false;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = false;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_CLEANUP;
 	amroutine->amkeytype = InvalidOid;
diff --git a/src/backend/access/gin/ginutil.c b/src/backend/access/gin/ginutil.c
index 6d2d71be32..066cf3e11a 100644
--- a/src/backend/access/gin/ginutil.c
+++ b/src/backend/access/gin/ginutil.c
@@ -56,6 +56,7 @@ ginhandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = false;
 	amroutine->amusemaintenanceworkmem = true;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_BULKDEL | VACUUM_OPTION_PARALLEL_CLEANUP;
 	amroutine->amkeytype = InvalidOid;
diff --git a/src/backend/access/gist/gist.c b/src/backend/access/gist/gist.c
index 0683f42c25..d96ce1c0a9 100644
--- a/src/backend/access/gist/gist.c
+++ b/src/backend/access/gist/gist.c
@@ -77,6 +77,7 @@ gisthandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = true;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_BULKDEL | VACUUM_OPTION_PARALLEL_COND_CLEANUP;
 	amroutine->amkeytype = InvalidOid;
diff --git a/src/backend/access/hash/hash.c b/src/backend/access/hash/hash.c
index eb3810494f..81c7da7ec6 100644
--- a/src/backend/access/hash/hash.c
+++ b/src/backend/access/hash/hash.c
@@ -74,6 +74,7 @@ hashhandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = false;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_BULKDEL;
 	amroutine->amkeytype = INT4OID;
diff --git a/src/backend/access/heap/heapam.c b/src/backend/access/heap/heapam.c
index 2da2be1696..c7f332d673 100644
--- a/src/backend/access/heap/heapam.c
+++ b/src/backend/access/heap/heapam.c
@@ -3224,7 +3224,7 @@ heap_update(Relation relation, ItemPointer otid, HeapTuple newtup,
 	 * Note that we get copies of each bitmap, so we need not worry about
 	 * relcache flush happening midway through.
 	 */
-	hot_attrs = RelationGetIndexAttrBitmap(relation, INDEX_ATTR_BITMAP_ALL);
+	hot_attrs = RelationGetIndexAttrBitmap(relation, INDEX_ATTR_BITMAP_HOT_BLOCKING);
 	key_attrs = RelationGetIndexAttrBitmap(relation, INDEX_ATTR_BITMAP_KEY);
 	id_attrs = RelationGetIndexAttrBitmap(relation,
 										  INDEX_ATTR_BITMAP_IDENTITY_KEY);
diff --git a/src/backend/access/nbtree/nbtree.c b/src/backend/access/nbtree/nbtree.c
index 40ad0956e0..bd6d6b1cc9 100644
--- a/src/backend/access/nbtree/nbtree.c
+++ b/src/backend/access/nbtree/nbtree.c
@@ -113,6 +113,7 @@ bthandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = true;
 	amroutine->amcaninclude = true;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_BULKDEL | VACUUM_OPTION_PARALLEL_COND_CLEANUP;
 	amroutine->amkeytype = InvalidOid;
diff --git a/src/backend/access/spgist/spgutils.c b/src/backend/access/spgist/spgutils.c
index 03a9cd36e6..a2cabb9d81 100644
--- a/src/backend/access/spgist/spgutils.c
+++ b/src/backend/access/spgist/spgutils.c
@@ -61,6 +61,7 @@ spghandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = true;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions =
 		VACUUM_OPTION_PARALLEL_BULKDEL | VACUUM_OPTION_PARALLEL_COND_CLEANUP;
 	amroutine->amkeytype = InvalidOid;
diff --git a/src/backend/utils/cache/relcache.c b/src/backend/utils/cache/relcache.c
index 13d9994af3..946e8d95f3 100644
--- a/src/backend/utils/cache/relcache.c
+++ b/src/backend/utils/cache/relcache.c
@@ -5006,6 +5006,7 @@ RelationGetIndexAttrBitmap(Relation relation, IndexAttrBitmapKind attrKind)
 	Bitmapset  *uindexattrs;	/* columns in unique indexes */
 	Bitmapset  *pkindexattrs;	/* columns in the primary index */
 	Bitmapset  *idindexattrs;	/* columns in the replica identity */
+	Bitmapset  *hotblockingattrs;   /* columns with HOT blocking indexes */
 	List	   *indexoidlist;
 	List	   *newindexoidlist;
 	Oid			relpkindex;
@@ -5026,6 +5027,8 @@ RelationGetIndexAttrBitmap(Relation relation, IndexAttrBitmapKind attrKind)
 				return bms_copy(relation->rd_pkattr);
 			case INDEX_ATTR_BITMAP_IDENTITY_KEY:
 				return bms_copy(relation->rd_idattr);
+			case INDEX_ATTR_BITMAP_HOT_BLOCKING:
+				return bms_copy(relation->rd_hotblockingattr);
 			default:
 				elog(ERROR, "unknown attrKind %u", attrKind);
 		}
@@ -5069,6 +5072,7 @@ restart:
 	uindexattrs = NULL;
 	pkindexattrs = NULL;
 	idindexattrs = NULL;
+	hotblockingattrs = NULL;
 	foreach(l, indexoidlist)
 	{
 		Oid			indexOid = lfirst_oid(l);
@@ -5136,6 +5140,10 @@ restart:
 				indexattrs = bms_add_member(indexattrs,
 											attrnum - FirstLowInvalidHeapAttributeNumber);
 
+				if (indexDesc->rd_indam->amhotblocking)
+					hotblockingattrs = bms_add_member(hotblockingattrs,
+												 attrnum - FirstLowInvalidHeapAttributeNumber);
+
 				if (isKey && i < indexDesc->rd_index->indnkeyatts)
 					uindexattrs = bms_add_member(uindexattrs,
 												 attrnum - FirstLowInvalidHeapAttributeNumber);
@@ -5152,9 +5160,14 @@ restart:
 
 		/* Collect all attributes used in expressions, too */
 		pull_varattnos(indexExpressions, 1, &indexattrs);
+		if (indexDesc->rd_indam->amhotblocking)
+			pull_varattnos(indexExpressions, 1, &hotblockingattrs);
 
 		/* Collect all attributes in the index predicate, too */
 		pull_varattnos(indexPredicate, 1, &indexattrs);
+		if (indexDesc->rd_indam->amhotblocking)
+			pull_varattnos(indexPredicate, 1, &hotblockingattrs);
+
 
 		index_close(indexDesc, AccessShareLock);
 	}
@@ -5182,6 +5195,7 @@ restart:
 		bms_free(uindexattrs);
 		bms_free(pkindexattrs);
 		bms_free(idindexattrs);
+		bms_free(hotblockingattrs);
 		bms_free(indexattrs);
 
 		goto restart;
@@ -5196,6 +5210,8 @@ restart:
 	relation->rd_pkattr = NULL;
 	bms_free(relation->rd_idattr);
 	relation->rd_idattr = NULL;
+	bms_free(relation->rd_hotblockingattr);
+	relation->rd_hotblockingattr = NULL;
 
 	/*
 	 * Now save copies of the bitmaps in the relcache entry.  We intentionally
@@ -5208,6 +5224,7 @@ restart:
 	relation->rd_keyattr = bms_copy(uindexattrs);
 	relation->rd_pkattr = bms_copy(pkindexattrs);
 	relation->rd_idattr = bms_copy(idindexattrs);
+	relation->rd_hotblockingattr = bms_copy(hotblockingattrs);
 	relation->rd_indexattr = bms_copy(indexattrs);
 	MemoryContextSwitchTo(oldcxt);
 
@@ -5222,6 +5239,8 @@ restart:
 			return pkindexattrs;
 		case INDEX_ATTR_BITMAP_IDENTITY_KEY:
 			return idindexattrs;
+		case INDEX_ATTR_BITMAP_HOT_BLOCKING:
+			return hotblockingattrs;
 		default:
 			elog(ERROR, "unknown attrKind %u", attrKind);
 			return NULL;
diff --git a/src/include/access/amapi.h b/src/include/access/amapi.h
index d357ebb559..a0ab70df89 100644
--- a/src/include/access/amapi.h
+++ b/src/include/access/amapi.h
@@ -244,6 +244,8 @@ typedef struct IndexAmRoutine
 	bool		amcaninclude;
 	/* does AM use maintenance_work_mem? */
 	bool		amusemaintenanceworkmem;
+	/* does AM block HOT update? */
+	bool        amhotblocking;
 	/* OR of parallel vacuum flags.  See vacuum.h for flags. */
 	uint8		amparallelvacuumoptions;
 	/* type of data stored in index, or InvalidOid if variable */
diff --git a/src/include/utils/rel.h b/src/include/utils/rel.h
index b4faa1c123..1bde51ab25 100644
--- a/src/include/utils/rel.h
+++ b/src/include/utils/rel.h
@@ -159,6 +159,7 @@ typedef struct RelationData
 	Bitmapset  *rd_keyattr;		/* cols that can be ref'd by foreign keys */
 	Bitmapset  *rd_pkattr;		/* cols included in primary key */
 	Bitmapset  *rd_idattr;		/* included in replica identity index */
+	Bitmapset  *rd_hotblockingattr;	/* cols blocking HOT update */
 
 	PublicationActions *rd_pubactions;	/* publication actions */
 
diff --git a/src/include/utils/relcache.h b/src/include/utils/relcache.h
index d2c17575f6..bf575bedf8 100644
--- a/src/include/utils/relcache.h
+++ b/src/include/utils/relcache.h
@@ -58,7 +58,8 @@ typedef enum IndexAttrBitmapKind
 	INDEX_ATTR_BITMAP_ALL,
 	INDEX_ATTR_BITMAP_KEY,
 	INDEX_ATTR_BITMAP_PRIMARY_KEY,
-	INDEX_ATTR_BITMAP_IDENTITY_KEY
+	INDEX_ATTR_BITMAP_IDENTITY_KEY,
+	INDEX_ATTR_BITMAP_HOT_BLOCKING
 } IndexAttrBitmapKind;
 
 extern Bitmapset *RelationGetIndexAttrBitmap(Relation relation,
diff --git a/src/test/modules/dummy_index_am/dummy_index_am.c b/src/test/modules/dummy_index_am/dummy_index_am.c
index 5365b0639e..9d409faff5 100644
--- a/src/test/modules/dummy_index_am/dummy_index_am.c
+++ b/src/test/modules/dummy_index_am/dummy_index_am.c
@@ -298,6 +298,7 @@ dihandler(PG_FUNCTION_ARGS)
 	amroutine->amcanparallel = false;
 	amroutine->amcaninclude = false;
 	amroutine->amusemaintenanceworkmem = false;
+	amroutine->amhotblocking = true;
 	amroutine->amparallelvacuumoptions = VACUUM_OPTION_NO_PARALLEL;
 	amroutine->amkeytype = InvalidOid;
 
diff --git a/src/test/regress/expected/brin.out b/src/test/regress/expected/brin.out
index e53d6e4885..08dab22b1c 100644
--- a/src/test/regress/expected/brin.out
+++ b/src/test/regress/expected/brin.out
@@ -567,3 +567,52 @@ SELECT * FROM brintest_3 WHERE b < '0';
 
 DROP TABLE brintest_3;
 RESET enable_seqscan;
+-- test BRIN index doesn't block HOT update
+CREATE TABLE brin_hot (
+        id  integer PRIMARY KEY,
+        val integer NOT NULL
+) WITH (autovacuum_enabled = off, fillfactor = 70);
+INSERT INTO brin_hot SELECT *, 0 FROM generate_series(1, 235);
+CREATE INDEX val_brin ON brin_hot using brin(val);
+CREATE FUNCTION wait_for_hot_stats() RETURNS void AS $$
+DECLARE
+        start_time timestamptz := clock_timestamp();
+        updated bool;
+BEGIN
+        -- we don't want to wait forever; loop will exit after 30 seconds
+        FOR i IN 1 .. 300 LOOP
+                SELECT (pg_stat_get_tuples_hot_updated('brin_hot'::regclass::oid) > 0) INTO updated;
+                EXIT WHEN updated;
+
+                -- wait a little
+                PERFORM pg_sleep_for('100 milliseconds');
+                -- reset stats snapshot so we can test again
+                PERFORM pg_stat_clear_snapshot();
+        END LOOP;
+        -- report time waited in postmaster log (where it won't change test output)
+        RAISE log 'wait_for_hot_stats delayed % seconds',
+          EXTRACT(epoch FROM clock_timestamp() - start_time);
+END
+$$ LANGUAGE plpgsql;
+UPDATE brin_hot SET val = -3 WHERE id = 42;
+-- We can't just call wait_for_hot_stats() at this point, because we only
+-- transmit stats when the session goes idle, and we probably didn't
+-- transmit the last couple of counts yet thanks to the rate-limiting logic
+-- in pgstat_report_stat().  But instead of waiting for the rate limiter's
+-- timeout to elapse, let's just start a new session.  The old one will
+-- then send its stats before dying.
+\c -
+SELECT wait_for_hot_stats();
+ wait_for_hot_stats 
+--------------------
+ 
+(1 row)
+
+SELECT pg_stat_get_tuples_hot_updated('brin_hot'::regclass::oid);
+ pg_stat_get_tuples_hot_updated 
+--------------------------------
+                              1
+(1 row)
+
+DROP TABLE brin_hot;
+DROP FUNCTION wait_for_hot_stats();
diff --git a/src/test/regress/sql/brin.sql b/src/test/regress/sql/brin.sql
index 3bd866d947..740e0b77a5 100644
--- a/src/test/regress/sql/brin.sql
+++ b/src/test/regress/sql/brin.sql
@@ -509,3 +509,49 @@ SELECT * FROM brintest_3 WHERE b < '0';
 
 DROP TABLE brintest_3;
 RESET enable_seqscan;
+
+-- test BRIN index doesn't block HOT update
+CREATE TABLE brin_hot (
+        id  integer PRIMARY KEY,
+        val integer NOT NULL
+) WITH (autovacuum_enabled = off, fillfactor = 70);
+
+INSERT INTO brin_hot SELECT *, 0 FROM generate_series(1, 235);
+CREATE INDEX val_brin ON brin_hot using brin(val);
+
+CREATE FUNCTION wait_for_hot_stats() RETURNS void AS $$
+DECLARE
+        start_time timestamptz := clock_timestamp();
+        updated bool;
+BEGIN
+        -- we don't want to wait forever; loop will exit after 30 seconds
+        FOR i IN 1 .. 300 LOOP
+                SELECT (pg_stat_get_tuples_hot_updated('brin_hot'::regclass::oid) > 0) INTO updated;
+                EXIT WHEN updated;
+
+                -- wait a little
+                PERFORM pg_sleep_for('100 milliseconds');
+                -- reset stats snapshot so we can test again
+                PERFORM pg_stat_clear_snapshot();
+        END LOOP;
+        -- report time waited in postmaster log (where it won't change test output)
+        RAISE log 'wait_for_hot_stats delayed % seconds',
+          EXTRACT(epoch FROM clock_timestamp() - start_time);
+END
+$$ LANGUAGE plpgsql;
+
+UPDATE brin_hot SET val = -3 WHERE id = 42;
+
+-- We can't just call wait_for_hot_stats() at this point, because we only
+-- transmit stats when the session goes idle, and we probably didn't
+-- transmit the last couple of counts yet thanks to the rate-limiting logic
+-- in pgstat_report_stat().  But instead of waiting for the rate limiter's
+-- timeout to elapse, let's just start a new session.  The old one will
+-- then send its stats before dying.
+\c -
+
+SELECT wait_for_hot_stats();
+SELECT pg_stat_get_tuples_hot_updated('brin_hot'::regclass::oid);
+
+DROP TABLE brin_hot;
+DROP FUNCTION wait_for_hot_stats();
-- 
2.31.1

