Fix visimap consults for unique checks during UPDATEs

This fixes #17183.

When consulting the visimap during an UPDATE for the purposes of
uniqueness checks, we used to refer to the visimap from the delete half
of the update.

This is the wrong structure to look at as this structure is not meant to
be consulted while deletes are in flight (we haven't reached
end-of-delete where visibility info from the visimapDelete structure
flows into the catalog).

Instead, we should be consulting the visimapDelete structure attached to
the deleteDesc. This structure can handle visimap queries for tuples
that have visimap changes that haven't yet been persisted to the catalog
table.

The effect of not using this structure meant running into issues such
as: "attempted to update invisible tuple" when we would attempt to
persist a dirty visimap entry in AppendOnlyVisimap_IsVisible() with a
call to AppendOnlyVisimap_Store(). The dirty entry is one which was
introduced by the delete half of the update. Our existing test did not
trip this issue because the update did not need a swap-out of the
current entry. With enough data, however, the issue reproduces, as
evidenced in #17183.

Co-authored-by: Ashwin Agrawal <aashwin@vmware.com>
Reviewed-by: Haolin Wang <whaolin@vmware.com>

src/backend/access/aocs/aocsam_handler.c

@@ -326,8 +326,8 @@ aoco_dml_finish(Relation relation, CmdType operation)
 		state->uniqueCheckDesc->blockDirectory = NULL;
 
 		/*
-		 * If this fetch is a part of an update, then we have been reusing the
-		 * visimap used by the delete half of the update, which would have
+		 * If this fetch is a part of an UPDATE, then we have been reusing the
+		 * visimapDelete used by the delete half of the UPDATE, which would have
 		 * already been cleaned up above. Clean up otherwise.
 		 */
 		if (!had_delete_desc)
@@ -336,6 +336,7 @@ aoco_dml_finish(Relation relation, CmdType operation)
 			pfree(state->uniqueCheckDesc->visimap);
 		}
 		state->uniqueCheckDesc->visimap = NULL;
+		state->uniqueCheckDesc->visiMapDelete = NULL;
 
 		pfree(state->uniqueCheckDesc);
 		state->uniqueCheckDesc = NULL;
@@ -499,18 +500,29 @@ get_or_create_unique_check_desc(Relation relation, Snapshot snapshot)
 													  relation,
 													  relation->rd_att->natts, /* numColGroups */
 													  snapshot);
+
 		/*
-		 * If this is part of an update, we need to reuse the visimap used by
-		 * the delete half of the update. This is to avoid spurious conflicts
-		 * when the key's previous and new value are identical. Using the
-		 * visimap from the delete half ensures that the visimap can recognize
-		 * any tuples deleted by us prior to this insert, within this command.
+		 * If this is part of an UPDATE, we need to reuse the visimapDelete
+		 * support structure from the delete half of the update. This is to
+		 * avoid spurious conflicts when the key's previous and new value are
+		 * identical. Using it ensures that we can recognize any tuples deleted
+		 * by us prior to this insert, within this command.
+		 *
+		 * Note: It is important that we reuse the visimapDelete structure and
+		 * not the visimap structure. This is because, when a uniqueness check
+		 * is performed as part of an UPDATE, visimap changes aren't persisted
+		 * yet (they are persisted at dml_finish() time, see
+		 * AppendOnlyVisimapDelete_Finish()). So, if we use the visimap
+		 * structure, we would not necessarily see all the changes.
 		 */
 		if (state->deleteDesc)
-			uniqueCheckDesc->visimap = &state->deleteDesc->visibilityMap;
+		{
+			uniqueCheckDesc->visiMapDelete = &state->deleteDesc->visiMapDelete;
+			uniqueCheckDesc->visimap = NULL;
+		}
 		else
 		{
-			/* Initialize the visimap */
+			/* COPY/INSERT: Initialize the visimap */
 			uniqueCheckDesc->visimap = palloc0(sizeof(AppendOnlyVisimap));
 			AppendOnlyVisimap_Init_forUniqueCheck(uniqueCheckDesc->visimap,
 												  relation,
@@ -994,8 +1006,9 @@ aoco_index_unique_check(Relation rel,
 	if (TransactionIdIsValid(snapshot->xmin) || TransactionIdIsValid(snapshot->xmax))
 		return true;
 
-	/* Now, consult the visimap */
-	visible = AppendOnlyVisimap_UniqueCheck(uniqueCheckDesc->visimap,
+	/* Now, perform a visibility check against the visimap infrastructure */
+	visible = AppendOnlyVisimap_UniqueCheck(uniqueCheckDesc->visiMapDelete,
+											uniqueCheckDesc->visimap,
 											aoTupleId,
 											snapshot);
 

src/backend/access/appendonly/appendonly_visimap.c

@@ -828,6 +828,36 @@ AppendOnlyVisimapDelete_WriteBackStashedEntries(AppendOnlyVisimapDelete *visiMap
 	}
 }
 
+/*
+ * Checks if the given tuple id is visible according to the visimapDelete
+ * support structure.
+ * A positive result is a necessary but not sufficient condition for
+ * a tuple to be visible to the user.
+ *
+ * Loads the entry for the tuple id before checking the bit.
+ */
+bool
+AppendOnlyVisimapDelete_IsVisible(AppendOnlyVisimapDelete *visiMapDelete,
+								  AOTupleId *aoTupleId)
+{
+	AppendOnlyVisimap *visiMap;
+
+	Assert(visiMapDelete);
+	Assert(aoTupleId);
+
+	elogif(Debug_appendonly_print_visimap, LOG,
+		   "Append-only visi map delete: IsVisible check "
+		   "(tupleId) = %s",
+		   AOTupleIdToString(aoTupleId));
+
+	visiMap = visiMapDelete->visiMap;
+	Assert(visiMap);
+
+	AppendOnlyVisimapDelete_LoadTuple(visiMapDelete, aoTupleId);
+
+	return AppendOnlyVisimapEntry_IsVisible(&visiMap->visimapEntry, aoTupleId);
+}
+
 
 /*
  * Finishes the delete operation.
@@ -928,8 +958,8 @@ AppendOnlyVisimap_Finish_forUniquenessChecks(
 {
 	AppendOnlyVisimapStore *visimapStore = &visiMap->visimapStore;
 	/*
-	 * The snapshot was either reset to NULL in between calls or already cleaned
-	 * up (if this was part of an update command)
+	 * The snapshot was never set or reset to NULL in between calls to
+	 * AppendOnlyVisimap_UniqueCheck().
 	 */
 	Assert(visimapStore->snapshot == InvalidSnapshot);
 

src/backend/access/appendonly/appendonlyam_handler.c

@@ -297,7 +297,7 @@ appendonly_dml_finish(Relation relation, CmdType operation)
 
 		/*
 		 * If this fetch is a part of an update, then we have been reusing the
-		 * visimap used by the delete half of the update, which would have
+		 * visimapDelete used by the delete half of the update, which would have
 		 * already been cleaned up above. Clean up otherwise.
 		 */
 		if (!had_delete_desc)
@@ -306,6 +306,7 @@ appendonly_dml_finish(Relation relation, CmdType operation)
 			pfree(state->uniqueCheckDesc->visimap);
 		}
 		state->uniqueCheckDesc->visimap = NULL;
+		state->uniqueCheckDesc->visiMapDelete = NULL;
 
 		pfree(state->uniqueCheckDesc);
 		state->uniqueCheckDesc = NULL;
@@ -465,17 +466,27 @@ get_or_create_unique_check_desc(Relation relation, Snapshot snapshot)
 													  snapshot);
 
 		/*
-		 * If this is part of an update, we need to reuse the visimap used by
-		 * the delete half of the update. This is to avoid spurious conflicts
-		 * when the key's previous and new value are identical. Using the
-		 * visimap from the delete half ensures that the visimap can recognize
-		 * any tuples deleted by us prior to this insert, within this command.
+		 * If this is part of an UPDATE, we need to reuse the visimapDelete
+		 * support structure from the delete half of the update. This is to
+		 * avoid spurious conflicts when the key's previous and new value are
+		 * identical. Using it ensures that we can recognize any tuples deleted
+		 * by us prior to this insert, within this command.
+		 *
+		 * Note: It is important that we reuse the visimapDelete structure and
+		 * not the visimap structure. This is because, when a uniqueness check
+		 * is performed as part of an UPDATE, visimap changes aren't persisted
+		 * yet (they are persisted at dml_finish() time, see
+		 * AppendOnlyVisimapDelete_Finish()). So, if we use the visimap
+		 * structure, we would not necessarily see all the changes.
 		 */
 		if (state->deleteDesc)
-			uniqueCheckDesc->visimap = &state->deleteDesc->visibilityMap;
+		{
+			uniqueCheckDesc->visiMapDelete = &state->deleteDesc->visiMapDelete;
+			uniqueCheckDesc->visimap = NULL;
+		}
 		else
 		{
-			/* Initialize the visimap */
+			/* COPY/INSERT: Initialize the visimap */
 			uniqueCheckDesc->visimap = palloc0(sizeof(AppendOnlyVisimap));
 			AppendOnlyVisimap_Init_forUniqueCheck(uniqueCheckDesc->visimap,
 												  relation,
@@ -761,8 +772,9 @@ appendonly_index_unique_check(Relation rel,
 	if (TransactionIdIsValid(snapshot->xmin) || TransactionIdIsValid(snapshot->xmax))
 		return true;
 
-	/* Now, consult the visimap */
-	visible = AppendOnlyVisimap_UniqueCheck(uniqueCheckDesc->visimap,
+	/* Now, perform a visibility check against the visimap infrastructure */
+	visible = AppendOnlyVisimap_UniqueCheck(uniqueCheckDesc->visiMapDelete,
+											uniqueCheckDesc->visimap,
 											aoTupleId,
 											snapshot);
 

src/include/access/appendonly_visimap.h

@@ -160,6 +160,10 @@ void AppendOnlyVisimapDelete_Finish(
 
 void
 AppendOnlyVisimapDelete_LoadTuple(AppendOnlyVisimapDelete *visiMapDelete,
+								   AOTupleId *aoTupleId);
+
+bool
+AppendOnlyVisimapDelete_IsVisible(AppendOnlyVisimapDelete *visiMapDelete,
 								  AOTupleId *aoTupleId);
 
 /*
@@ -168,27 +172,59 @@ AppendOnlyVisimapDelete_LoadTuple(AppendOnlyVisimapDelete *visiMapDelete,
  * During a uniqueness check, look up the visimap to see if a tuple was deleted
  * by a *committed* transaction.
  *
- * Note: We need to use the passed in per-tuple snapshot to perform the block
- * directory lookup. See AppendOnlyVisimap_Init_forUniqueCheck() for details on
- * why we can't set up the metadata snapshot at init time.
- * If this is part of an update, we are reusing the visimap from the delete half
- * of the update, so better restore its snapshot once we are done.
+ * If this uniqueness check is part of an UPDATE, we consult the visiMapDelete
+ * structure. Otherwise, we consult the visiMap structure. Only one of these
+ * arguments should be supplied.
  */
-static inline bool AppendOnlyVisimap_UniqueCheck(
-											AppendOnlyVisimap *visiMap,
-											AOTupleId *aoTupleId,
-											Snapshot appendOnlyMetaDataSnapshot)
+static inline bool AppendOnlyVisimap_UniqueCheck(AppendOnlyVisimapDelete *visiMapDelete,
+												 AppendOnlyVisimap *visiMap,
+												 AOTupleId *aoTupleId,
+												 Snapshot appendOnlyMetaDataSnapshot)
 {
-	Snapshot save_snapshot;
-	bool visible;
+	Snapshot          save_snapshot;
+	bool              visible;
 
 	Assert(appendOnlyMetaDataSnapshot->snapshot_type == SNAPSHOT_DIRTY ||
-			appendOnlyMetaDataSnapshot->snapshot_type == SNAPSHOT_SELF);
+			   appendOnlyMetaDataSnapshot->snapshot_type == SNAPSHOT_SELF);
+
+	if (visiMapDelete)
+	{
+		/* part of an UPDATE */
+		Assert(!visiMap);
+		Assert(visiMapDelete->visiMap);
+
+		/* Save the snapshot used for the delete half of the UPDATE */
+		save_snapshot = visiMapDelete->visiMap->visimapStore.snapshot;
+
+		/*
+		 * Replace with per-tuple snapshot meant for uniqueness checks. See
+		 * AppendOnlyVisimap_Init_forUniqueCheck() for details on why we can't
+		 * set up the metadata snapshot at init time.
+		 */
+		visiMapDelete->visiMap->visimapStore.snapshot = appendOnlyMetaDataSnapshot;
+
+		visible = AppendOnlyVisimapDelete_IsVisible(visiMapDelete, aoTupleId);
+
+		/* Restore the snapshot used for the delete half of the UPDATE */
+		visiMapDelete->visiMap->visimapStore.snapshot = save_snapshot;
+	}
+	else
+	{
+		/* part of a COPY/INSERT */
+		Assert(visiMap);
+
+		/*
+		 * Set up per-tuple snapshot meant for uniqueness checks. See
+		 * AppendOnlyVisimap_Init_forUniqueCheck() for details on why we can't
+		 * set up the metadata snapshot at init time.
+		 */
+		visiMap->visimapStore.snapshot = appendOnlyMetaDataSnapshot;
+
+		visible = AppendOnlyVisimap_IsVisible(visiMap, aoTupleId);
+
+		visiMap->visimapStore.snapshot = NULL; /* be a good citizen */
+	}
 
-	save_snapshot = visiMap->visimapStore.snapshot;
-	visiMap->visimapStore.snapshot = appendOnlyMetaDataSnapshot;
-	visible = AppendOnlyVisimap_IsVisible(visiMap, aoTupleId);
-	visiMap->visimapStore.snapshot = save_snapshot;
 	return visible;
 }
 

src/include/cdb/cdbaocsam.h

@@ -288,7 +288,10 @@ typedef struct AOCSDeleteDescData *AOCSDeleteDesc;
 typedef struct AOCSUniqueCheckDescData
 {
 	AppendOnlyBlockDirectory *blockDirectory;
+	/* visimap to check for deleted tuples as part of INSERT/COPY */
 	AppendOnlyVisimap 		 *visimap;
+	/* visimap support structure to check for deleted tuples as part of UPDATE */
+	AppendOnlyVisimapDelete  *visiMapDelete;
 } AOCSUniqueCheckDescData;
 
 typedef struct AOCSUniqueCheckDescData *AOCSUniqueCheckDesc;

src/include/cdb/cdbappendonlyam.h

@@ -402,7 +402,10 @@ typedef struct AppendOnlyDeleteDescData *AppendOnlyDeleteDesc;
 typedef struct AppendOnlyUniqueCheckDescData
 {
 	AppendOnlyBlockDirectory *blockDirectory;
+	/* visimap to check for deleted tuples as part of INSERT/COPY */
 	AppendOnlyVisimap 		 *visimap;
+	/* visimap support structure to check for deleted tuples as part of UPDATE */
+	AppendOnlyVisimapDelete  *visiMapDelete;
 } AppendOnlyUniqueCheckDescData;
 
 typedef struct AppendOnlyUniqueCheckDescData *AppendOnlyUniqueCheckDesc;

src/test/regress/input/uao_dml/uao_dml_unique_index_update.source

@@ -97,6 +97,18 @@ SELECT * FROM uao_unique_index_update;
 
 DROP TABLE uao_unique_index_update;
 
+-- Case 8_1: Update where pre-update key = post-update key
+-- but updates span multiple visimap rows---------------------------------------
+CREATE TABLE uao_unique_index_update_1 (a INT unique, b INT) DISTRIBUTED REPLICATED;
+INSERT INTO uao_unique_index_update_1 SELECT a, 1 FROM generate_series (1, 32768) a;
+UPDATE uao_unique_index_update_1 SET b = 2;
+-- If we were to select the contents of the visimap relation on any of the
+-- segments we would see 2 entries at this point, one with first_row_no = 0 and
+-- the other with first_row_no = 32768.
+UPDATE uao_unique_index_update_1 SET b = 3;
+SELECT DISTINCT b FROM uao_unique_index_update_1;
+DROP TABLE uao_unique_index_update_1;
+
 -- Case 9: Updating tx inserts a key that already exists------------------------
 CREATE TABLE uao_unique_index_update (a INT unique);
 INSERT INTO uao_unique_index_update VALUES (1), (2);

src/test/regress/output/uao_dml/uao_dml_unique_index_update.source

@@ -134,6 +134,22 @@ SELECT * FROM uao_unique_index_update;
 (1 row)
 
 DROP TABLE uao_unique_index_update;
+-- Case 8_1: Update where pre-update key = post-update key
+-- but updates span multiple visimap rows---------------------------------------
+CREATE TABLE uao_unique_index_update_1 (a INT unique, b INT) DISTRIBUTED REPLICATED;
+INSERT INTO uao_unique_index_update_1 SELECT a, 1 FROM generate_series (1, 32768) a;
+UPDATE uao_unique_index_update_1 SET b = 2;
+-- If we were to select the contents of the visimap relation on any of the
+-- segments we would see 2 entries at this point, one with first_row_no = 0 and
+-- the other with first_row_no = 32768.
+UPDATE uao_unique_index_update_1 SET b = 3;
+SELECT DISTINCT b FROM uao_unique_index_update_1;
+ b 
+---
+ 3
+(1 row)
+
+DROP TABLE uao_unique_index_update_1;
 -- Case 9: Updating tx inserts a key that already exists------------------------
 CREATE TABLE uao_unique_index_update (a INT unique);
 INSERT INTO uao_unique_index_update VALUES (1), (2);