[YSQL] #1199: Enable FOR SHARE and FOR KEY SHARE row locking in YSQL …

…SELECT statements

Summary:
Support row-level explicit locking in `SELECT` statements for YSQL tables.
The new supported cases are:
- `SELECT .. FOR SHARE`
- `SELECT .. FOR KEY SHARE`

Additionally, since the `FOR KEY SHARE` case is also used internally for referential integrity
checking (i.e. foreign keys), this diff also enables DMLs on tables with foreign key references
regardless of isolation level (previously only worked in `SERIALIZABLE` isolation).

Implementation is done by passing the row mark as flag in the DocDB read request and taking
appropriate lock for those requests that have that corresponding row mark.

Limitations (that will be addressed in follow-up commits):
- The `FOR UPDATE` and `FOR NO KEY UPDATE` are not yet supported
- The locking for `FOR KEY SHARE` is stronger than strictly required (strong read lock instead of
   weak read lock). This is still correct but less efficient (will conflict more often than required).
- We don't support having different row mark settings within one DocDB (read) batch. This is fine
  for now because we do not batch such requests together, but later we might.

The current implementation makes data visible between two transactions when it shouldn't. This is because we are currently not detecting read conflicts. #2523 tracks this issue.

Test Plan:
Java test testForeignKeyConflicts has been converted to two tests: testForeignKeyConflictsWithSerializableIsolation and testForeignKeyConflictsWithSnapshotIsolation/

```postgres=# create table t(k int primary key, s text);
CREATE TABLE
postgres=# select * from t where k = 1 for key share;
 k | s
---+---
(0 rows)

postgres=# insert into t(k, s) values (1, 'a');
INSERT 0 1
postgres=# select * from t where k = 1 for key share;
 k | s
---+---
 1 | a
(1 row)

postgres=# select * from t for key share;
 k | s
---+---
 1 | a
(1 row)
```

```postgres=# create table t(k int primary key, s text);
CREATE TABLE
postgres=#  insert into t(k, s) values (1, 'a');
INSERT 0 1
postgres=# begin;
BEGIN
postgres=# insert into t(k,s) values (2, 'b');
INSERT 0 1
postgres=# select * from t for key share;
 k | s
---+---
 1 | a
 2 | b
(2 rows)

postgres=# select * from t where k = 1 for key share;
 k | s
---+---
 1 | a
(1 row)

postgres=# abort;
ROLLBACK
postgres=# select * from t where k = 1 for key share;
 k | s
---+---
 1 | a
(1 row)

postgres=# select * from t for key share;
 k | s
---+---
 1 | a
(1 row)
```

```
postgres=# create table t(id int primary key, s text);
CREATE TABLE
postgres=# create table s(id int primary key , t_id int references t(id));
CREATE TABLE
postgres=# insert into t(id, s) values (1, 'a');
INSERT 0 1
postgres=# insert into t(id, s) values (2, 'b');
INSERT 0 1
postgres=# insert into s(id, t_id) values (100, 1);
INSERT 0 1
postgres=# select s.id, t.s from s, t where s.t_id = t.id;
 id  | s
-----+---
 100 | a
(1 row)

postgres=# begin;
BEGIN
postgres=# delete from t where id = 2;
DELETE 1
postgres=# select * from t;
 id | s
----+---
  1 | a
(1 row)

postgres=# commit;
COMMIT
postgres=# select * from t;
 id | s
----+---
  1 | a
(1 row)
```

Reviewers: mihnea, mikhail

Reviewed By: mikhail

Subscribers: bogdan, yql

Differential Revision: https://phabricator.dev.yugabyte.com/D7007

java/yb-pgsql/src/test/java/org/yb/pgsql/TestPgForeignKey.java

@@ -61,14 +61,23 @@ private void checkRows(Statement statement,
   }
 
   @Test
-  public void testForeignKeyConflicts() throws Exception {
+  public void testForeignKeyConflictsWithSerializableIsolation() throws Exception {
+    testForeignKeyConflicts(Connection.TRANSACTION_SERIALIZABLE);
+  }
+
+  @Test
+  public void testForeignKeyConflictsWithSnapshotIsolation() throws Exception {
+    testForeignKeyConflicts(Connection.TRANSACTION_REPEATABLE_READ);
+  }
+
+  private void testForeignKeyConflicts(int pgIsolationLevel) throws Exception {
 
     Set<Row> expectedPkRows = new HashSet<>();
     Set<Row> expectedFkRows = new HashSet<>();
 
     // Set up the tables.
     try (Statement statement = connection.createStatement()) {
-      connection.setTransactionIsolation(Connection.TRANSACTION_SERIALIZABLE);
+      connection.setTransactionIsolation(pgIsolationLevel);
 
       // Setup pk table.
       statement.execute("create table pk(id int primary key)");
@@ -86,8 +95,19 @@ public void testForeignKeyConflicts() throws Exception {
       }
       checkRows(statement, "fk", expectedFkRows);
 
-      try (Connection connection1 = createConnectionSerializableNoAutoCommit();
-           Connection connection2 = createConnectionSerializableNoAutoCommit()) {
+      IsolationLevel isolationLevel = IsolationLevel.REPEATABLE_READ;
+      if (pgIsolationLevel == Connection.TRANSACTION_SERIALIZABLE) {
+        isolationLevel = IsolationLevel.SERIALIZABLE;
+      }
+
+      try (Connection connection1 = newConnectionBuilder()
+              .setIsolationLevel(isolationLevel)
+              .setAutoCommit(AutoCommit.DISABLED)
+              .connect();
+           Connection connection2 = newConnectionBuilder()
+                   .setIsolationLevel(isolationLevel)
+                   .setAutoCommit(AutoCommit.DISABLED)
+                   .connect()) {
 
         // Test update/delete conflicts.
         for (int id = 1; id < 20; id += 2) {

src/postgres/src/backend/access/index/indexam.c

@@ -589,7 +589,10 @@ index_getnext_tid(IndexScanDesc scan, ScanDirection direction)
 	 * scan->xs_recheck and possibly scan->xs_itup/scan->xs_hitup, though we
 	 * pay no attention to those fields here.
 	 */
-	found = scan->indexRelation->rd_amroutine->amgettuple(scan, direction);
+	if (IsYBBackedRelation(scan->indexRelation))
+		found = ybcingettuple(scan, direction);
+	else
+		found = scan->indexRelation->rd_amroutine->amgettuple(scan, direction);
 
 	/* Reset kill flag immediately for safety */
 	scan->kill_prior_tuple = false;

src/postgres/src/backend/executor/execMain.c

@@ -959,15 +959,17 @@ InitPlan(QueryDesc *queryDesc, int eflags)
 		{
 			case ROW_MARK_EXCLUSIVE:
 			case ROW_MARK_NOKEYEXCLUSIVE:
-			case ROW_MARK_SHARE:
-			case ROW_MARK_KEYSHARE:
 				relation = heap_open(relid, RowShareLock);
-				if (IsYBRelation(relation)) 
+				if (IsYBRelation(relation))
 				{
 					YBRaiseNotSupported("SELECT locking option only supported for temporary tables",
 										1199);
 				}
 				break;
+			case ROW_MARK_SHARE:
+			case ROW_MARK_KEYSHARE:
+				relation = heap_open(relid, RowShareLock);
+				break;
 			case ROW_MARK_REFERENCE:
 				relation = heap_open(relid, AccessShareLock);
 				break;

src/postgres/src/backend/executor/execUtils.c

@@ -176,6 +176,7 @@ CreateExecutorState(void)
 	estate->yb_exec_params.limit_count = -1;
 	estate->yb_exec_params.limit_offset = 0;
 	estate->yb_exec_params.limit_use_default = true;
+	estate->yb_exec_params.rowmark = -1;
 
 	return estate;
 }

src/postgres/src/backend/executor/nodeIndexscan.c

@@ -133,6 +133,14 @@ IndexNext(IndexScanState *node)
 	 */
 	if (IsYugaByteEnabled()) {
 		scandesc->yb_exec_params = &estate->yb_exec_params;
+		// Add row marks.
+		scandesc->yb_exec_params->rowmark = -1;
+		ListCell   *l;
+		foreach(l, estate->es_rowMarks) {
+			ExecRowMark *erm = (ExecRowMark *) lfirst(l);
+			scandesc->yb_exec_params->rowmark = erm->markType;
+			break;
+		}
 	}
 
 	/*

src/postgres/src/backend/executor/nodeLockRows.c

@@ -60,6 +60,28 @@ ExecLockRows(PlanState *pstate)
 lnext:
 	slot = ExecProcNode(outerPlan);
 
+	int n_yb_relations = 0;
+	int n_relations = 0;
+	foreach(lc, node->lr_arowMarks)
+	{
+		ExecAuxRowMark *aerm = (ExecAuxRowMark *) lfirst(lc);
+		ExecRowMark *erm = aerm->rowmark;
+		if (IsYBBackedRelation(erm->relation))
+		{
+			n_yb_relations++;
+		}
+		n_relations++;
+	}
+
+	if (n_yb_relations == n_relations)
+		return slot;
+
+	if (n_yb_relations > 0)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("Mixing Yugabyte relations and not Yugabyte relations with "
+						"row locks is not supported")));
+
 	if (TupIsNull(slot))
 		return NULL;
 

src/postgres/src/backend/executor/ybc_fdw.c

@@ -270,6 +270,15 @@ ybcBeginForeignScan(ForeignScanState *node, int eflags)
 	ResourceOwnerRememberYugaByteStmt(CurrentResourceOwner, ybc_state->handle);
 	ybc_state->stmt_owner = CurrentResourceOwner;
 	ybc_state->exec_params = &estate->yb_exec_params;
+
+	ybc_state->exec_params->rowmark = -1;
+	ListCell   *l;
+	foreach(l, estate->es_rowMarks) {
+		ExecRowMark *erm = (ExecRowMark *) lfirst(l);
+		ybc_state->exec_params->rowmark = erm->markType;
+		break;
+	}
+
 	ybc_state->is_exec_done = false;
 
 	/* Set the current syscatalog version (will check that we are up to date) */

src/postgres/src/backend/optimizer/prep/preptlist.c

@@ -143,7 +143,13 @@ preprocess_targetlist(PlannerInfo *root)
 
 		if (rc->allMarkTypes & ~(1 << ROW_MARK_COPY))
 		{
-			if (target_relation && IsYBBackedRelation(target_relation)) 
+			bool is_yb_relation = false;
+			if (!target_relation)
+				is_yb_relation = IsYBRelationById(getrelid(rc->rti, range_table));
+			else
+				is_yb_relation = IsYBBackedRelation(target_relation);
+
+			if (is_yb_relation)
 			{
 				/* Need to fetch YB TID */
 				var = makeVar(rc->rti,
@@ -153,8 +159,8 @@ preprocess_targetlist(PlannerInfo *root)
 								InvalidOid,
 								0);
 				snprintf(resname, sizeof(resname), "ybctid%u", rc->rowmarkId);
-			} 
-			else 
+			}
+			else
 			{
 				/* Need to fetch TID */
 				var = makeVar(rc->rti,

src/postgres/src/backend/utils/adt/ri_triggers.c

@@ -244,25 +244,6 @@ static void ri_ReportViolation(const RI_ConstraintInfo *riinfo,
 				   int queryno) pg_attribute_noreturn();
 
 
-/*
- * In YB mode we currently only support foreign key DMLs
- * in serializable mode (for YB tables).
- * TODO To be removed when we support all cases.
- */
-static bool IsYBForeignKeyConstraint(Relation pk_relation)
-{
-	if (IsYBRelation(pk_relation))
-	{
-		if (!IsolationIsSerializable())
-		{
-			YBRaiseNotSupported("Operation only supported in SERIALIZABLE "
-								"isolation level", 1199);
-		}
-		return true;
-	}
-	return false;
-}
-
 /* ----------
  * RI_FKey_check -
  *
@@ -451,14 +432,7 @@ RI_FKey_check(TriggerData *trigdata)
 			queryoids[i] = fk_type;
 		}
 
-		/*
-		 * TODO In YB mode we currently only allow foreign key DMLs
-		 * in YB serializable mode -- so no need for key share here
-		 */
-		if (!IsYBForeignKeyConstraint(pk_rel))
-		{
-			appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
-		}
+		appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
 
 		/* Prepare and save the plan */
 		qplan = ri_PlanCheck(querybuf.data, riinfo->nkeys, queryoids,
@@ -594,15 +568,8 @@ ri_Check_Pk_Match(Relation pk_rel, Relation fk_rel,
 			queryoids[i] = pk_type;
 		}
 
-		/*
-		 * TODO In YB mode we currently only allow foreign key DMLs
-		 * in YB serializable mode -- so no need for key share here
-		 */
-		if (!IsYBForeignKeyConstraint(pk_rel))
-		{
-			appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
-		}
-
+		appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
+
 		/* Prepare and save the plan */
 		qplan = ri_PlanCheck(querybuf.data, riinfo->nkeys, queryoids,
 							 &qkey, fk_rel, pk_rel, true);
@@ -867,10 +834,7 @@ ri_restrict(TriggerData *trigdata, bool is_no_action)
 				* TODO In YB mode we currently only allow foreign key DMLs
 				* in YB serializable mode -- so no need for key share here
 				*/
-				if (!IsYBForeignKeyConstraint(pk_rel))
-				{
-					appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
-				}
+				appendStringInfoString(&querybuf, " FOR KEY SHARE OF x");
 
 				/* Prepare and save the plan */
 				qplan = ri_PlanCheck(querybuf.data, riinfo->nkeys, queryoids,

src/postgres/src/test/regress/expected/yb_foreign_key.out

@@ -1,8 +1,9 @@
 --
 -- FOREIGN KEY (YB-added tests)
 --
--- TODO - YB only supports foreign keys in serializable isolation currently.
-SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE;
+-- TODO: Run this test with REPEATABLE READ isolation level:
+--       https://github.com/yugabyte/yugabyte-db/issues/2604
+--
 -- MATCH FULL
 --
 -- First test, check and cascade

src/postgres/src/test/regress/expected/yb_pg_foreign_key.out

@@ -1,8 +1,8 @@
 --
 -- FOREIGN KEY
 --
--- TODO - YB only supports foreign keys in serializable isolation currently.
-SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE;
+-- TODO: Run this test with REPEATABLE READ isolation level:                                                                                                                                                            --       https://github.com/yugabyte/yugabyte-db/issues/2604
+--
 -- MATCH FULL
 --
 -- First test, check and cascade

src/postgres/src/test/regress/sql/yb_foreign_key.sql

@@ -1,10 +1,9 @@
 --
 -- FOREIGN KEY (YB-added tests)
 --
-
--- TODO - YB only supports foreign keys in serializable isolation currently.
-SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE;
-
+-- TODO: Run this test with REPEATABLE READ isolation level:
+--       https://github.com/yugabyte/yugabyte-db/issues/2604
+--
 -- MATCH FULL
 --
 -- First test, check and cascade
@@ -114,4 +113,4 @@ SELECT * FROM ITABLE ORDER BY ptest1, ptest2;
 SELECT * FROM FKTABLE ORDER BY ftest1, ftest2, ftest3;
 
 DROP TABLE ITABLE CASCADE;
-DROP TABLE FKTABLE;
+DROP TABLE FKTABLE;

src/postgres/src/test/regress/sql/yb_pg_foreign_key.sql

@@ -1,10 +1,8 @@
 --
 -- FOREIGN KEY
 --
-
--- TODO - YB only supports foreign keys in serializable isolation currently.
-SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL SERIALIZABLE;
-
+-- TODO: Run this test with REPEATABLE READ isolation level:                                                                                                                                                            --       https://github.com/yugabyte/yugabyte-db/issues/2604
+--
 -- MATCH FULL
 --
 -- First test, check and cascade
@@ -547,8 +545,8 @@ ptest3) REFERENCES pktable);
 DROP TABLE PKTABLE;
 
 -- TODO: YugaByte does not yet support table inheritance.
--- Leaving the first failing statement uncommented so that this test 
--- will fail when the feature is implemented (the full test should 
+-- Leaving the first failing statement uncommented so that this test
+-- will fail when the feature is implemented (the full test should
 -- be uncommented then).
 
 --
@@ -657,8 +655,8 @@ drop table pktable_base;
 --
 
 -- TODO YB does not yet support deferrable foreign key constraints.
--- Leaving the first failing statement uncommented so that this test 
--- will fail when the feature is implemented (the full test should 
+-- Leaving the first failing statement uncommented so that this test
+-- will fail when the feature is implemented (the full test should
 -- be uncommented then).
 
 -- deferrable, explicitly deferred
@@ -855,8 +853,8 @@ DROP TABLE pktable, fktable;
 -- cause the on-INSERT RI trigger not to be fired.
 
 -- TODO YugaByte does not support deferrable constraints yet.
--- Leaving the first failing statement uncommented so that this test 
--- will fail when the feature is implemented (the full test should 
+-- Leaving the first failing statement uncommented so that this test
+-- will fail when the feature is implemented (the full test should
 -- be uncommented then).
 
 CREATE TEMP TABLE pktable (
@@ -1128,8 +1126,8 @@ drop table pktable2, fktable2;
 --
 
 -- TODO YugaByte does not support partitioned tables yet.
--- Leaving the first failing statement uncommented so that this test 
--- will fail when the feature is implemented (the full test should 
+-- Leaving the first failing statement uncommented so that this test
+-- will fail when the feature is implemented (the full test should
 -- be uncommented then).
 
 -- partitioned table in the referenced side are not allowed

src/yb/common/common.proto

@@ -366,6 +366,28 @@ enum IsolationLevel {
   SERIALIZABLE_ISOLATION = 2;
 }
 
+// This enum matches enum RowMarkType defined in src/include/nodes/plannodes.h.
+enum RowMarkType
+{
+  // Obtain exclusive tuple lock. Not supported yet.
+  ROW_MARK_EXCLUSIVE = 0;
+
+  // Obtain no-key exclusive tuple lock. Not supported yet.
+  ROW_MARK_NOKEYEXCLUSIVE = 1;
+
+  // Obtain shared tuple lock.
+  ROW_MARK_SHARE = 2;
+
+  // Obtain keyshare tuple lock.
+  ROW_MARK_KEYSHARE = 3;
+
+  // Not supported. Used for postgres compatibility.
+  ROW_MARK_REFERENCE = 4;
+
+  // Not supported. Used for postgres compatibility.
+  ROW_MARK_COPY = 5;
+}
+
 enum TransactionStatus {
   CREATED = 1;
   PENDING = 2;