Invalidate AppendOnlyHash cache entry for AT_SetDistributedBy cases

ALTER TABLE commands that are tagged as AT_SetDistributedBy require a
gather motion and does its own variation of creating a temporary table
for CTAS (basically bypassing the usual ATRewriteTable which actually
does do AppendOnlyHash cache entry invalidation). Without the
AppendOnlyHash cache entry invalidation, the entry could have
invisible leaks in its AOSegfileStatus array that will be stuck in
state AOSEG_STATE_AWAITING_DROP. These leaks will persist until the
user evicts the cache entry by not using the table to allow another AO
table to cache itself in that slot or by restarting the database. We
fix this issue by invalidating the cache entry at the end of
AT_SetDistributedBy ALTER TABLE cases.

src/backend/commands/tablecmds.c

@@ -13013,6 +13013,7 @@ ATExecSetDistributedBy(Relation rel, Node *node, AlterTableCmd *cmd)
 	RangeVar   *tmprv;
 	Oid			tmprelid;
 	Oid			tarrelid = RelationGetRelid(rel);
+	char		tarrelstorage = rel->rd_rel->relstorage;
 	List	   *oid_map = NIL;
 	bool        rand_pol = false;
 	bool        rep_pol = false;
@@ -13672,6 +13673,25 @@ ATExecSetDistributedBy(Relation rel, Node *node, AlterTableCmd *cmd)
 		performDeletion(&object, DROP_RESTRICT);
 	}
 
+	if (relstorage_is_ao(tarrelstorage) && IS_QUERY_DISPATCHER())
+	{
+		/*
+		 * Drop the shared memory hash table entry for this table if it
+		 * exists. We must do so since before the rewrite we probably have few
+		 * non-zero segfile entries for this table while after the rewrite
+		 * only segno zero will be full and the others will be empty. By
+		 * dropping the hash entry we force refreshing the entry from the
+		 * catalog the next time a write into this AO table comes along.
+		 *
+		 * Note that ALTER already took an exclusive lock on the old relation
+		 * so we are guaranteed to not drop the hash entry from under any
+		 * concurrent operation.
+		 */
+		LWLockAcquire(AOSegFileLock, LW_EXCLUSIVE);
+		AORelRemoveHashEntry(tarrelid);
+		LWLockRelease(AOSegFileLock);
+	}
+
 l_distro_fini:
 
 	/* MPP-6929: metadata tracking */

src/test/isolation2/expected/reorganize_after_ao_vacuum_skip_drop.out

+-- Ensure segfiles in AOSEG_STATE_AWAITING_DROP are not leaked in
+-- AppendOnlyHash after doing an AT_SetDistributedBy operation which
+-- rewrites the relation differently than other ALTER operations.
+
+CREATE TABLE reorganize_after_ao_vacuum_skip_drop (a INT, b INT) WITH (appendonly=true);
+CREATE
+INSERT INTO reorganize_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 10) AS i;
+INSERT 10
+
+DELETE FROM reorganize_after_ao_vacuum_skip_drop;
+DELETE 10
+
+-- We should see all aosegs in state 1
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+segno|state
+-----+-----
+1    |1    
+(1 row)
+
+-- VACUUM while another session holds lock
+1: BEGIN;
+BEGIN
+1: SELECT COUNT(*) FROM reorganize_after_ao_vacuum_skip_drop;
+count
+-----
+0    
+(1 row)
+2: VACUUM reorganize_after_ao_vacuum_skip_drop;
+VACUUM
+1: END;
+END
+
+-- We should see an aoseg in state 2 (AOSEG_STATE_AWAITING_DROP)
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+segno|state
+-----+-----
+1    |2    
+2    |1    
+(2 rows)
+
+-- The AO relation should be rewritten and AppendOnlyHash entry invalidated
+1: ALTER TABLE reorganize_after_ao_vacuum_skip_drop SET WITH (reorganize=true);
+ALTER
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+segno|state
+-----+-----
+(0 rows)
+
+-- Check if insert goes into segno 1 instead of segno 2. If it did not
+-- go into segno 1, there was a leak in the AppendOnlyHash entry.
+1: INSERT INTO reorganize_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 100) AS i;
+INSERT 100
+0U: SELECT segno, tupcount > 0, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+segno|?column?|state
+-----+--------+-----
+1    |t       |1    
+(1 row)

src/test/isolation2/isolation2_schedule

@@ -61,6 +61,7 @@ test: uao/vacuum_self_serializable3_row
 test: uao/vacuum_while_insert_row
 test: uao/vacuum_while_vacuum_row
 test: uao/vacuum_cleanup_row
+test: reorganize_after_ao_vacuum_skip_drop
 
 # Tests on Append-Optimized tables (column-oriented).
 test: uao/alter_while_vacuum_column uao/alter_while_vacuum2_column

src/test/isolation2/sql/reorganize_after_ao_vacuum_skip_drop.sql

+-- Ensure segfiles in AOSEG_STATE_AWAITING_DROP are not leaked in
+-- AppendOnlyHash after doing an AT_SetDistributedBy operation which
+-- rewrites the relation differently than other ALTER operations.
+
+CREATE TABLE reorganize_after_ao_vacuum_skip_drop (a INT, b INT) WITH (appendonly=true);
+INSERT INTO reorganize_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 10) AS i;
+
+DELETE FROM reorganize_after_ao_vacuum_skip_drop;
+
+-- We should see all aosegs in state 1
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+
+-- VACUUM while another session holds lock
+1: BEGIN;
+1: SELECT COUNT(*) FROM reorganize_after_ao_vacuum_skip_drop;
+2: VACUUM reorganize_after_ao_vacuum_skip_drop;
+1: END;
+
+-- We should see an aoseg in state 2 (AOSEG_STATE_AWAITING_DROP)
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+
+-- The AO relation should be rewritten and AppendOnlyHash entry invalidated
+1: ALTER TABLE reorganize_after_ao_vacuum_skip_drop SET WITH (reorganize=true);
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');
+
+-- Check if insert goes into segno 1 instead of segno 2. If it did not
+-- go into segno 1, there was a leak in the AppendOnlyHash entry.
+1: INSERT INTO reorganize_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 100) AS i;
+0U: SELECT segno, tupcount > 0, state FROM gp_toolkit.__gp_aoseg_name('reorganize_after_ao_vacuum_skip_drop');