Take RI fast-path snapshot after locking referenced relation

ri_FastPathCheck() acquired its scan snapshot before opening the
referenced relation.  If it then waited for the relation lock in READ
COMMITTED mode, a referenced row committed during the wait would not be
visible to the old snapshot.  The check could consequently report a
foreign key violation even though the referenced row existed.

The batched path acquires its scan snapshot after opening the relations,
so it does not have the lock-wait problem.  The per-row path remains
reachable during constraint validation and checks re-entered from a
batch flush.

Take the snapshot after opening the referenced relation and reloading
the constraint information.  This also agrees with the SPI path, which
acquires the referenced-relation lock before selecting the snapshot
used for the check.

Also make the scan snapshot active for the duration of the check, so
STABLE cast and equality functions use the same snapshot as the index
scan.  Otherwise, a STABLE cast can miss rows inserted by earlier AFTER
triggers and cause a valid foreign key to be rejected.  Add a
regression test comparing this case with the SPI path.

Finally, also advance the command counter after acquiring the
referenced-relation lock and before reloading the constraint
information, so local invalidations are processed under the lock, as in
the SPI path.

Add an injection-point test that pauses a per-row check before locking
the referenced relation, commits the missing referenced row in another
session, and verifies that validation succeeds.  Force per-row
validation by using a role without SELECT privilege on the referenced
table, so the test exercises this path even with batching enabled.

Add isolation-test coverage for the visibility of a referenced row
committed after the referencing transaction has executed an earlier
command.  A later command can see such a row in READ COMMITTED, but not
in REPEATABLE READ or SERIALIZABLE.

Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Reviewed-by: Amit Kapila <amit.kapila16@gmail.com>
Discussion: https://postgr.es/m/CA+HiwqEhm+_=bs=2wavAJz-UqC+1KebD31++mapJQQGweE8iQQ@mail.gmail.com
Backpatch-through: 19

src/backend/utils/adt/ri_triggers.c

@@ -2876,23 +2876,43 @@ ri_FastPathCheck(RI_ConstraintInfo *riinfo,
 	int			saved_sec_context;
 	Snapshot	snapshot;
 
-	/*
-	 * Advance the command counter so the snapshot sees the effects of prior
-	 * triggers in this statement.  Mirrors what the SPI path does in
-	 * ri_PerformCheck().
-	 */
-	CommandCounterIncrement();
-	snapshot = RegisterSnapshot(GetTransactionSnapshot());
-
 	INJECTION_POINT("ri-before-pk-lock", NULL);
 
 	pk_rel = table_open(riinfo->pk_relid, RowShareLock);
 
+	/*
+	 * Advance the command counter so the check sees the effects of prior
+	 * triggers in this statement, as SPI does when executing the query issued
+	 * by ri_PerformCheck().  Do this after locking the referenced relation
+	 * and before reloading the constraint information, so local invalidations
+	 * are processed under the lock.
+	 */
+	CommandCounterIncrement();
+
 	/* Re-read the constraint under that lock; see ri_FastPathGetEntry(). */
 	riinfo = ri_LoadConstraintInfo(riinfo->constraint_id);
 
 	idx_rel = index_open(riinfo->conindid, AccessShareLock);
 
+	/*
+	 * Only now take the snapshot the scan will use.  Acquiring it before
+	 * table_open() would let an unbounded amount of time pass while we wait
+	 * for the lock, during which another transaction can commit the very row
+	 * we are about to look for.  The scan would not see it and the check
+	 * would report a violation for a key that exists.
+	 *
+	 * The SPI path does not have this problem: for this check it passes
+	 * InvalidSnapshot, so SPI takes the snapshot after the
+	 * referenced-relation lock has been acquired.
+	 *
+	 * Make this snapshot active too, as SPI does.  STABLE cast and equality
+	 * functions use the active snapshot, so leaving the outer query's
+	 * snapshot active could hide changes made by earlier triggers even though
+	 * the index scan can see them.
+	 */
+	snapshot = RegisterSnapshot(GetTransactionSnapshot());
+	PushActiveSnapshot(snapshot);
+
 	slot = table_slot_create(pk_rel, NULL);
 
 	GetUserIdAndSecContext(&saved_userid, &saved_sec_context);
@@ -2929,6 +2949,7 @@ ri_FastPathCheck(RI_ConstraintInfo *riinfo,
 	index_endscan(scandesc);
 	ExecDropSingleTupleTableSlot(slot);
 	UnregisterSnapshot(snapshot);
+	PopActiveSnapshot();
 
 	if (!found)
 		ri_ReportViolation(riinfo, pk_rel, fk_rel,

src/test/isolation/expected/fk-snapshot-2.out

@@ -59,3 +59,47 @@ step s1c: COMMIT;
 step s2ins: <... completed>
 ERROR:  could not serialize access due to concurrent delete
 step s2c: COMMIT;
+
+starting permutation: s1rr s2rr s2sel s1ins s1c s2ins2 s2c
+step s1rr: BEGIN ISOLATION LEVEL REPEATABLE READ;
+step s2rr: BEGIN ISOLATION LEVEL REPEATABLE READ;
+step s2sel: SELECT count(*) FROM parent;
+count
+-----
+    1
+(1 row)
+
+step s1ins: INSERT INTO parent VALUES (2);
+step s1c: COMMIT;
+step s2ins2: INSERT INTO child VALUES (2, 2);
+ERROR:  insert or update on table "child" violates foreign key constraint "child_parent_id_fkey"
+step s2c: COMMIT;
+
+starting permutation: s1ser s2ser s2sel s1ins s1c s2ins2 s2c
+step s1ser: BEGIN ISOLATION LEVEL SERIALIZABLE;
+step s2ser: BEGIN ISOLATION LEVEL SERIALIZABLE;
+step s2sel: SELECT count(*) FROM parent;
+count
+-----
+    1
+(1 row)
+
+step s1ins: INSERT INTO parent VALUES (2);
+step s1c: COMMIT;
+step s2ins2: INSERT INTO child VALUES (2, 2);
+ERROR:  insert or update on table "child" violates foreign key constraint "child_parent_id_fkey"
+step s2c: COMMIT;
+
+starting permutation: s1rc s2rc s2sel s1ins s1c s2ins2 s2c
+step s1rc: BEGIN ISOLATION LEVEL READ COMMITTED;
+step s2rc: BEGIN ISOLATION LEVEL READ COMMITTED;
+step s2sel: SELECT count(*) FROM parent;
+count
+-----
+    1
+(1 row)
+
+step s1ins: INSERT INTO parent VALUES (2);
+step s1c: COMMIT;
+step s2ins2: INSERT INTO child VALUES (2, 2);
+step s2c: COMMIT;

src/test/isolation/specs/fk-snapshot-2.spec

@@ -21,6 +21,7 @@ step s1rr	{ BEGIN ISOLATION LEVEL REPEATABLE READ; }
 step s1ser	{ BEGIN ISOLATION LEVEL SERIALIZABLE; }
 step s1del	{ DELETE FROM parent WHERE parent_id = 1; }
 step s1c	{ COMMIT; }
+step s1ins	{ INSERT INTO parent VALUES (2); }
 
 session s2
 step s2rc	{ BEGIN ISOLATION LEVEL READ COMMITTED; }
@@ -28,6 +29,8 @@ step s2rr	{ BEGIN ISOLATION LEVEL REPEATABLE READ; }
 step s2ser	{ BEGIN ISOLATION LEVEL SERIALIZABLE; }
 step s2ins	{ INSERT INTO child VALUES (1, 1); }
 step s2c	{ COMMIT; }
+step s2sel	{ SELECT count(*) FROM parent; }
+step s2ins2	{ INSERT INTO child VALUES (2, 2); }
 
 # Violates referential integrity unless we use a crosscheck snapshot,
 # which is up-to-date compared with the transaction's snapshot.
@@ -48,3 +51,15 @@ permutation s1ser s2ser s2ins s1del s2c s1c
 # We raise a concurrent update error
 # which is good enough:
 permutation s1ser s2ser s1del s2ins s1c s2c
+
+# A parent row committed by another transaction after this one took its
+# snapshot.  RI_FKey_check passes detectNewRows = false, so the check runs
+# under the transaction snapshot rather than a current one, and the row is
+# correctly not visible: referencing it is a violation.  Only the parent-side
+# checks need to see rows committed since the snapshot.
+permutation s1rr s2rr s2sel s1ins s1c s2ins2 s2c
+permutation s1ser s2ser s2sel s1ins s1c s2ins2 s2c
+
+# The same order in READ COMMITTED, where the check's snapshot is a fresh one
+# and the parent row is visible.
+permutation s1rc s2rc s2sel s1ins s1c s2ins2 s2c

src/test/modules/injection_points/Makefile

@@ -20,6 +20,7 @@ ISOLATION = basic \
 	    repack_temporal_multirange \
 	    repack_toast \
 	    ri_fastpath_reindex \
+	    ri_fastpath_snapshot \
 	    syscache-update-pruned \
 	    wait_cleanup \
 	    heap_lock_update

src/test/modules/injection_points/expected/ri_fastpath_snapshot.out

@@ -0,0 +1,19 @@
+Parsed test spec with 2 sessions
+
+starting permutation: validate ins wake checked validated
+step validate: ALTER TABLE ri_snapshot_fk VALIDATE CONSTRAINT ri_snapshot_fkey; <waiting ...>
+step ins: INSERT INTO ri_snapshot_pk VALUES (500);
+step wake: 
+    DO $$BEGIN
+        PERFORM injection_points_detach('ri-before-pk-lock');
+        PERFORM injection_points_wakeup('ri-before-pk-lock');
+    END$$;
+
+step validate: <... completed>
+step checked: 
+step validated: SELECT convalidated FROM pg_constraint WHERE conname = 'ri_snapshot_fkey';
+convalidated
+------------
+t           
+(1 row)
+

src/test/modules/injection_points/meson.build

@@ -51,6 +51,7 @@ tests += {
       'repack_temporal_multirange',
       'repack_toast',
       'ri_fastpath_reindex',
+      'ri_fastpath_snapshot',
       'syscache-update-pruned',
       'wait_cleanup',
       'heap_lock_update',

src/test/modules/injection_points/specs/ri_fastpath_snapshot.spec

@@ -0,0 +1,50 @@
+# Verify that the per-row RI probe sees a referenced row committed while
+# the probe is parked at ri-before-pk-lock.
+#
+# Lack of SELECT permission on the referenced table makes RI_Initial_Check
+# fall back to per-row validation.
+
+setup
+{
+    CREATE EXTENSION injection_points;
+    CREATE ROLE regress_ri_snapshot;
+    CREATE TABLE ri_snapshot_pk (id int PRIMARY KEY);
+    CREATE TABLE ri_snapshot_fk (pid int);
+    INSERT INTO ri_snapshot_fk VALUES (500);
+    ALTER TABLE ri_snapshot_fk ADD CONSTRAINT ri_snapshot_fkey
+        FOREIGN KEY (pid) REFERENCES ri_snapshot_pk NOT VALID;
+    ALTER TABLE ri_snapshot_fk OWNER TO regress_ri_snapshot;
+}
+
+teardown
+{
+    DROP TABLE ri_snapshot_fk, ri_snapshot_pk;
+    DROP ROLE regress_ri_snapshot;
+    DROP EXTENSION injection_points;
+}
+
+session s1
+setup
+{
+    DO $$BEGIN
+        PERFORM injection_points_set_local();
+        PERFORM injection_points_attach('ri-before-pk-lock', 'wait');
+    END$$;
+    SET ROLE regress_ri_snapshot;
+    SET default_transaction_isolation = 'read committed';
+}
+step validate { ALTER TABLE ri_snapshot_fk VALIDATE CONSTRAINT ri_snapshot_fkey; }
+step checked { }
+
+session s2
+step ins { INSERT INTO ri_snapshot_pk VALUES (500); }
+step wake
+{
+    DO $$BEGIN
+        PERFORM injection_points_detach('ri-before-pk-lock');
+        PERFORM injection_points_wakeup('ri-before-pk-lock');
+    END$$;
+}
+step validated { SELECT convalidated FROM pg_constraint WHERE conname = 'ri_snapshot_fkey'; }
+
+permutation validate ins wake checked validated

src/test/regress/expected/foreign_key.out

@@ -3682,6 +3682,75 @@ NOTICE:  fp_auto_pk called
 NOTICE:  fp_auto_pk called
 DROP TABLE fp_fk_cci, fp_pk_cci;
 DROP FUNCTION fp_auto_pk;
+-- A STABLE cast used by an FK check must see changes made by earlier AFTER
+-- triggers, using the check's snapshot rather than the outer query's snapshot.
+-- Compare the per-row fast path with a partitioned-parent SPI check.
+BEGIN;
+CREATE SCHEMA ri_snapshot;
+SET LOCAL search_path = ri_snapshot, pg_catalog;
+CREATE TYPE lookup_key AS (v int);
+CREATE TABLE lookup_rows (v int);
+CREATE FUNCTION lookup_key_to_int(k lookup_key) RETURNS int
+LANGUAGE plpgsql STABLE AS $$
+DECLARE
+    result int;
+BEGIN
+    SELECT r.v INTO result FROM lookup_rows AS r WHERE r.v = k.v;
+    RETURN coalesce(result, -1);
+END;
+$$;
+CREATE CAST (lookup_key AS int)
+    WITH FUNCTION lookup_key_to_int(lookup_key) AS IMPLICIT;
+CREATE TABLE pk_fast (v int PRIMARY KEY);
+CREATE TABLE pk_spi (v int PRIMARY KEY) PARTITION BY RANGE (v);
+CREATE TABLE pk_spi_p PARTITION OF pk_spi
+    FOR VALUES FROM (MINVALUE) TO (MAXVALUE);
+INSERT INTO pk_fast VALUES (1);
+INSERT INTO pk_spi VALUES (1);
+CREATE TABLE fk_fast (k lookup_key REFERENCES pk_fast(v));
+CREATE TABLE fk_spi (k lookup_key REFERENCES pk_spi(v));
+CREATE FUNCTION add_lookup_row() RETURNS trigger LANGUAGE plpgsql AS $$
+BEGIN
+    INSERT INTO lookup_rows VALUES ((NEW.k).v);
+    RETURN NEW;
+END;
+$$;
+-- Sort before the RI trigger, so the lookup row is inserted first.
+CREATE TRIGGER "AAA_lookup" AFTER INSERT ON fk_fast
+    FOR EACH ROW EXECUTE FUNCTION add_lookup_row();
+CREATE TRIGGER "AAA_lookup" AFTER INSERT ON fk_spi
+    FOR EACH ROW EXECUTE FUNCTION add_lookup_row();
+-- Invoke the fast-path check from another FK's cast.  Even with batching
+-- enabled, a check nested inside the end-of-batch flush takes the per-row path.
+CREATE TYPE driver_key AS (v int);
+CREATE FUNCTION driver_key_to_int(k driver_key) RETURNS int
+LANGUAGE plpgsql VOLATILE AS $$
+BEGIN
+    INSERT INTO fk_fast VALUES (ROW(k.v)::lookup_key);
+    RETURN k.v;
+END;
+$$;
+CREATE CAST (driver_key AS int)
+    WITH FUNCTION driver_key_to_int(driver_key) AS IMPLICIT;
+CREATE TABLE driver (k driver_key REFERENCES pk_fast(v));
+-- Both checks must succeed.  Without an active snapshot for the per-row
+-- check, its STABLE cast misses the lookup row and returns -1 instead of 1.
+INSERT INTO driver VALUES (ROW(1)::driver_key);
+SELECT count(*) FROM fk_fast;
+ count 
+-------
+     1
+(1 row)
+
+DELETE FROM lookup_rows;
+INSERT INTO fk_spi VALUES (ROW(1)::lookup_key);
+SELECT count(*) FROM fk_spi;
+ count 
+-------
+     1
+(1 row)
+
+ROLLBACK;
 -- Multi-column FK: exercises batched per-row probing with composite keys
 CREATE TABLE fp_pk_multi (a int, b int, PRIMARY KEY (a, b));
 INSERT INTO fp_pk_multi SELECT i, i FROM generate_series(1, 100) i;

src/test/regress/sql/foreign_key.sql

@@ -2654,6 +2654,71 @@ INSERT INTO fp_fk_cci VALUES (1), (2), (3);
 DROP TABLE fp_fk_cci, fp_pk_cci;
 DROP FUNCTION fp_auto_pk;
 
+-- A STABLE cast used by an FK check must see changes made by earlier AFTER
+-- triggers, using the check's snapshot rather than the outer query's snapshot.
+-- Compare the per-row fast path with a partitioned-parent SPI check.
+BEGIN;
+CREATE SCHEMA ri_snapshot;
+SET LOCAL search_path = ri_snapshot, pg_catalog;
+
+CREATE TYPE lookup_key AS (v int);
+CREATE TABLE lookup_rows (v int);
+CREATE FUNCTION lookup_key_to_int(k lookup_key) RETURNS int
+LANGUAGE plpgsql STABLE AS $$
+DECLARE
+    result int;
+BEGIN
+    SELECT r.v INTO result FROM lookup_rows AS r WHERE r.v = k.v;
+    RETURN coalesce(result, -1);
+END;
+$$;
+CREATE CAST (lookup_key AS int)
+    WITH FUNCTION lookup_key_to_int(lookup_key) AS IMPLICIT;
+
+CREATE TABLE pk_fast (v int PRIMARY KEY);
+CREATE TABLE pk_spi (v int PRIMARY KEY) PARTITION BY RANGE (v);
+CREATE TABLE pk_spi_p PARTITION OF pk_spi
+    FOR VALUES FROM (MINVALUE) TO (MAXVALUE);
+INSERT INTO pk_fast VALUES (1);
+INSERT INTO pk_spi VALUES (1);
+CREATE TABLE fk_fast (k lookup_key REFERENCES pk_fast(v));
+CREATE TABLE fk_spi (k lookup_key REFERENCES pk_spi(v));
+
+CREATE FUNCTION add_lookup_row() RETURNS trigger LANGUAGE plpgsql AS $$
+BEGIN
+    INSERT INTO lookup_rows VALUES ((NEW.k).v);
+    RETURN NEW;
+END;
+$$;
+-- Sort before the RI trigger, so the lookup row is inserted first.
+CREATE TRIGGER "AAA_lookup" AFTER INSERT ON fk_fast
+    FOR EACH ROW EXECUTE FUNCTION add_lookup_row();
+CREATE TRIGGER "AAA_lookup" AFTER INSERT ON fk_spi
+    FOR EACH ROW EXECUTE FUNCTION add_lookup_row();
+
+-- Invoke the fast-path check from another FK's cast.  Even with batching
+-- enabled, a check nested inside the end-of-batch flush takes the per-row path.
+CREATE TYPE driver_key AS (v int);
+CREATE FUNCTION driver_key_to_int(k driver_key) RETURNS int
+LANGUAGE plpgsql VOLATILE AS $$
+BEGIN
+    INSERT INTO fk_fast VALUES (ROW(k.v)::lookup_key);
+    RETURN k.v;
+END;
+$$;
+CREATE CAST (driver_key AS int)
+    WITH FUNCTION driver_key_to_int(driver_key) AS IMPLICIT;
+CREATE TABLE driver (k driver_key REFERENCES pk_fast(v));
+
+-- Both checks must succeed.  Without an active snapshot for the per-row
+-- check, its STABLE cast misses the lookup row and returns -1 instead of 1.
+INSERT INTO driver VALUES (ROW(1)::driver_key);
+SELECT count(*) FROM fk_fast;
+DELETE FROM lookup_rows;
+INSERT INTO fk_spi VALUES (ROW(1)::lookup_key);
+SELECT count(*) FROM fk_spi;
+ROLLBACK;
+
 -- Multi-column FK: exercises batched per-row probing with composite keys
 CREATE TABLE fp_pk_multi (a int, b int, PRIMARY KEY (a, b));
 INSERT INTO fp_pk_multi SELECT i, i FROM generate_series(1, 100) i;