Lock RI fast-path rows as of the scan snapshot's command ID

ri_LockPKTuple() passed table_tuple_lock() the current command ID, read
at lock time.  ExecLockRows(), which locks the row for the SPI path's
SELECT ... FOR KEY SHARE, uses the estate's es_output_cid, fixed when
the query starts and equal to the command ID of the query's snapshot.
The two differ if user code run between taking the snapshot and locking
the row, such as the equality function called by the index scan,
advances the command counter.  If that code also updated the referenced
row, the lock reports TM_Invisible, and the check fails with "attempted
to lock invisible tuple", where the SPI path gets TM_SelfModified and
reports a foreign key violation.

To fix, use the snapshot's command ID in ri_LockPKTuple() to mirror
what ExecLockRows() does.

This is mostly hardening.  With FK values that need a cast now checked
through SPI, the only user code that can run there is the opclass's
equality function, and it would have to write to the referenced table.
That's unlikely, but the fast path should still behave like SPI.

Discussion: https://postgr.es/m/CA+HiwqG79XK1oObdZ2AwT660CeJ6s3Mn4LrFPCme-k4L2rF_ag@mail.gmail.com
Backpatch-through: 19

src/backend/utils/adt/ri_triggers.c

@@ -2955,8 +2955,16 @@ ri_LockPKTuple(Relation pk_rel, TupleTableSlot *slot, Snapshot snap,
 	if (!IsolationUsesXactSnapshot())
 		lockflags |= TUPLE_LOCK_FLAG_FIND_LAST_VERSION;
 
+	/*
+	 * Lock as of the command ID the scan's snapshot was taken with, as
+	 * ExecLockRows() uses the es_output_cid fixed when its query started.
+	 * User code run during the scan, such as an equality function, may have
+	 * advanced the current command ID since; with that, a row it updated
+	 * would look updated by an earlier command (TM_Invisible) rather than by
+	 * this one (TM_SelfModified).
+	 */
 	result = table_tuple_lock(pk_rel, &slot->tts_tid, snap,
-							  slot, GetCurrentCommandId(false),
+							  slot, snap->curcid,
 							  LockTupleKeyShare, LockWaitBlock,
 							  lockflags, &tmfd);
 

src/test/regress/expected/foreign_key.out

@@ -4150,3 +4150,33 @@ SELECT count(*) FROM fp_fk_ro_tmp;
 (1 row)
 
 DROP TABLE fp_fk_ro, fp_pk_ro, fp_fk_ro_tmp, fp_pk_ro_tmp;
+-- The row lock must use the command ID the check's snapshot was taken
+-- with.  An equality function that updates the referenced row advances the
+-- command counter before the lock; the row must then count as updated by
+-- the check's own command, so that the check reports a violation, as the
+-- SPI path does, rather than failing to lock an "invisible" tuple.  The
+-- equality function runs while the scan holds the index page locked; its
+-- UPDATE is a HOT update, since v isn't indexed, so it doesn't touch the
+-- index.  This is done in a transaction that is rolled back, so that other
+-- tests never see the operator class.
+BEGIN;
+CREATE TABLE fp_pk_cid (id int, v int DEFAULT 0);
+CREATE FUNCTION fp_cid_eq(int, int) RETURNS bool LANGUAGE plpgsql AS $$
+BEGIN
+    UPDATE fp_pk_cid SET v = v + 1 WHERE id OPERATOR(pg_catalog.=) $1;
+    RETURN $1 OPERATOR(pg_catalog.=) $2;
+END$$;
+CREATE OPERATOR =~= (LEFTARG = int, RIGHTARG = int, FUNCTION = fp_cid_eq);
+CREATE OPERATOR CLASS fp_cid_ops FOR TYPE int USING btree AS
+    OPERATOR 1 pg_catalog.<, OPERATOR 2 pg_catalog.<=, OPERATOR 3 =~=,
+    OPERATOR 4 pg_catalog.>=, OPERATOR 5 pg_catalog.>,
+    FUNCTION 1 btint4cmp(int, int);
+CREATE UNIQUE INDEX fp_pk_cid_id ON fp_pk_cid (id fp_cid_ops);
+INSERT INTO fp_pk_cid VALUES (1);
+CREATE TABLE fp_fk_cid (a int REFERENCES fp_pk_cid (id));
+SAVEPOINT fp_cid;
+INSERT INTO fp_fk_cid VALUES (1);	-- fails, as a violation
+ERROR:  insert or update on table "fp_fk_cid" violates foreign key constraint "fp_fk_cid_a_fkey"
+DETAIL:  Key (a)=(1) is not present in table "fp_pk_cid".
+ROLLBACK TO SAVEPOINT fp_cid;
+ROLLBACK;

src/test/regress/sql/foreign_key.sql

@@ -3086,3 +3086,32 @@ COMMIT;	-- succeeds
 SELECT count(*) FROM fp_fk_ro;
 SELECT count(*) FROM fp_fk_ro_tmp;
 DROP TABLE fp_fk_ro, fp_pk_ro, fp_fk_ro_tmp, fp_pk_ro_tmp;
+
+-- The row lock must use the command ID the check's snapshot was taken
+-- with.  An equality function that updates the referenced row advances the
+-- command counter before the lock; the row must then count as updated by
+-- the check's own command, so that the check reports a violation, as the
+-- SPI path does, rather than failing to lock an "invisible" tuple.  The
+-- equality function runs while the scan holds the index page locked; its
+-- UPDATE is a HOT update, since v isn't indexed, so it doesn't touch the
+-- index.  This is done in a transaction that is rolled back, so that other
+-- tests never see the operator class.
+BEGIN;
+CREATE TABLE fp_pk_cid (id int, v int DEFAULT 0);
+CREATE FUNCTION fp_cid_eq(int, int) RETURNS bool LANGUAGE plpgsql AS $$
+BEGIN
+    UPDATE fp_pk_cid SET v = v + 1 WHERE id OPERATOR(pg_catalog.=) $1;
+    RETURN $1 OPERATOR(pg_catalog.=) $2;
+END$$;
+CREATE OPERATOR =~= (LEFTARG = int, RIGHTARG = int, FUNCTION = fp_cid_eq);
+CREATE OPERATOR CLASS fp_cid_ops FOR TYPE int USING btree AS
+    OPERATOR 1 pg_catalog.<, OPERATOR 2 pg_catalog.<=, OPERATOR 3 =~=,
+    OPERATOR 4 pg_catalog.>=, OPERATOR 5 pg_catalog.>,
+    FUNCTION 1 btint4cmp(int, int);
+CREATE UNIQUE INDEX fp_pk_cid_id ON fp_pk_cid (id fp_cid_ops);
+INSERT INTO fp_pk_cid VALUES (1);
+CREATE TABLE fp_fk_cid (a int REFERENCES fp_pk_cid (id));
+SAVEPOINT fp_cid;
+INSERT INTO fp_fk_cid VALUES (1);	-- fails, as a violation
+ROLLBACK TO SAVEPOINT fp_cid;
+ROLLBACK;