Invalidate AppendOnlyHash cache entry at end of TRUNCATE

TRUNCATE will rewrite the relation by creating a temporary table and
swapping it with the real relation. For AO, this includes the
auxiliary tables which is concerning for the AO relation's pg_aoseg
table which holds information that a AO segment file is available for
write or waiting to be compacted/dropped. Since we do not currently
invalidate the AppendOnlyHash cache entry, 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 TRUNCATE
on AO relations.

src/backend/commands/tablecmds.c

@@ -1770,6 +1770,21 @@ ExecuteTruncate(TruncateStmt *stmt)
 	{
 		Relation	rel = (Relation) lfirst(cell);
 
+		if (RelationIsAppendOptimized(rel) && 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.
+			 */
+			LWLockAcquire(AOSegFileLock, LW_EXCLUSIVE);
+			AORelRemoveHashEntry(RelationGetRelid(rel));
+			LWLockRelease(AOSegFileLock);
+		}
+
 		heap_close(rel, NoLock);
 	}
 }

src/test/isolation2/expected/truncate_after_ao_vacuum_skip_drop.out

+-- Ensure segfiles in AOSEG_STATE_AWAITING_DROP are not leaked in
+-- AppendOnlyHash after doing a TRUNCATE.
+
+CREATE TABLE truncate_after_ao_vacuum_skip_drop (a INT, b INT) WITH (appendonly=true);
+CREATE
+INSERT INTO truncate_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 10) AS i;
+INSERT 10
+
+DELETE FROM truncate_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('truncate_after_ao_vacuum_skip_drop');
+segno|state
+-----+-----
+1    |1    
+(1 row)
+
+-- VACUUM while another session holds lock
+1: BEGIN;
+BEGIN
+1: SELECT COUNT(*) FROM truncate_after_ao_vacuum_skip_drop;
+count
+-----
+0    
+(1 row)
+2: VACUUM truncate_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('truncate_after_ao_vacuum_skip_drop');
+segno|state
+-----+-----
+1    |2    
+2    |1    
+(2 rows)
+
+-- The AO relation should be rewritten and AppendOnlyHash entry invalidated
+1: TRUNCATE truncate_after_ao_vacuum_skip_drop;
+TRUNCATE
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('truncate_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 truncate_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('truncate_after_ao_vacuum_skip_drop');
+segno|?column?|state
+-----+--------+-----
+1    |t       |1    
+(1 row)

src/test/isolation2/isolation2_schedule

@@ -61,7 +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
+test: reorganize_after_ao_vacuum_skip_drop truncate_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/truncate_after_ao_vacuum_skip_drop.sql

+-- Ensure segfiles in AOSEG_STATE_AWAITING_DROP are not leaked in
+-- AppendOnlyHash after doing a TRUNCATE.
+
+CREATE TABLE truncate_after_ao_vacuum_skip_drop (a INT, b INT) WITH (appendonly=true);
+INSERT INTO truncate_after_ao_vacuum_skip_drop SELECT i as a, i as b FROM generate_series(1, 10) AS i;
+
+DELETE FROM truncate_after_ao_vacuum_skip_drop;
+
+-- We should see all aosegs in state 1
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('truncate_after_ao_vacuum_skip_drop');
+
+-- VACUUM while another session holds lock
+1: BEGIN;
+1: SELECT COUNT(*) FROM truncate_after_ao_vacuum_skip_drop;
+2: VACUUM truncate_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('truncate_after_ao_vacuum_skip_drop');
+
+-- The AO relation should be rewritten and AppendOnlyHash entry invalidated
+1: TRUNCATE truncate_after_ao_vacuum_skip_drop;
+0U: SELECT segno, state FROM gp_toolkit.__gp_aoseg_name('truncate_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 truncate_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('truncate_after_ao_vacuum_skip_drop');