Optimize lock for select-for-update and similar queries

Add GUC 'enable_lock_optimization' to enable lock optimization for
select-for-update and similar queries. If the lock for query can be
optimized, we use tuple lock instead of relation lock to improve
concurrency performance.

Author: HuSen8891

src/backend/parser/analyze.c

@@ -70,6 +70,14 @@
 #include "parser/parse_func.h"
 #include "utils/lsyscache.h"
 
+/*
+ * GUC parameter
+ *
+ * Enable locking optimization for extended query protocol to avoid
+ * ExclusiveLock in case of select-for-update and similar queries.
+ */
+bool enableLockOptimization = false;
+
 /* Working state for transformSetOperationTree_internal */
 typedef struct
 {
@@ -132,7 +140,6 @@ static bool test_raw_expression_coverage(Node *node, void *context);
 static int get_distkey_by_name(char *key, IntoClause *into, Query *qry, bool *found);
 static void setQryDistributionPolicy(ParseState *pstate, IntoClause *into, Query *qry);
 
-static bool checkCanOptSelectLockingClause(SelectStmt *stmt);
 static bool queryNodeSearch(Node *node, void *context);
 static void sanity_check_on_conflict_update_set_distkey(GpPolicy  *policy, List *onconflict_set);
 static void sanity_check_on_conflict_update(Oid relid, List *on_conflict_set, Node *on_conflict_where);
@@ -4080,7 +4087,7 @@ setQryDistributionPolicy(ParseState *pstate, IntoClause *into, Query *qry)
  * can behave like Postgres. We have to know it before
  * we acquire any locks on the tables.
  */
-static bool
+bool
 checkCanOptSelectLockingClause(SelectStmt *stmt)
 {
 	QueryNodeSearchContext ctx = {false};

src/backend/tcop/postgres.c

@@ -2171,6 +2171,22 @@ exec_parse_message(const char *query_string,	/* string to execute */
 			 */
 			if (IS_SINGLENODE())
 				((SelectStmt *)raw_parse_tree->stmt)->disableLockingOptimization = false;
+			else if (enableLockOptimization)
+			{
+				SelectStmt *stmt = (SelectStmt *)raw_parse_tree->stmt;
+
+				/*
+				 * If GUC 'enableLockOptimization' is on, we try to optimize the lock for
+				 * select-for-update and similar queries.
+				 *
+				 * If the lock cannot be optimized, use the default way.
+				 */
+				stmt->disableLockingOptimization = false;
+				if (!checkCanOptSelectLockingClause(stmt))
+				{
+					((SelectStmt *)raw_parse_tree->stmt)->disableLockingOptimization = true;
+				}
+			}
 			else
 				((SelectStmt *)raw_parse_tree->stmt)->disableLockingOptimization = true;
 		}
@@ -2472,6 +2488,26 @@ exec_bind_message(StringInfo input_message)
 
 	portal->is_extended_query = true;
 
+	/*
+	 * If GUC 'enableLockOptimization' is on, we try to optimize the lock for
+	 * select-for-update and similar queries.
+	 *
+	 * If the lock for query can be optimized, we use tuple lock instead of
+	 * relation lock to improve Concurrency Performance. LockRows is executed
+	 * on segment, but reader gangs cannot execute LockRows, so we need to
+	 * dispatch query plan to writer gangs, that's why we reset is_extended_query
+	 * to false here.
+	 */
+	if (enableLockOptimization && IsA(psrc->raw_parse_tree->stmt, SelectStmt))
+	{
+		SelectStmt *stmt = (SelectStmt *)psrc->raw_parse_tree->stmt;
+
+		if (checkCanOptSelectLockingClause(stmt))
+		{
+			portal->is_extended_query = false;
+		}
+	}
+
 	/*
 	 * Prepare to copy stuff into the portal's memory context.  We do all this
 	 * copying first, because it could possibly fail (out-of-memory) and we
@@ -2973,6 +3009,45 @@ exec_execute_message(const char *portal_name, int64 max_rows)
 	if (max_rows <= 0)
 		max_rows = FETCH_ALL;
 
+	/*
+	 * If the lock for select-for-update and similar queries is optimized, we should
+	 * fetch all rows here.
+	 *
+	 * Since we optimize the lock for query, the query plan is dispatched to writer
+	 * gangs to execute. If we do not fetch all rows here, the writer gangs are occupied
+	 * and can not execute other query plans.
+	 */
+	if (enableLockOptimization)
+	{
+		CachedPlanSource *psrc;
+
+		if (portal->prepStmtName)
+		{
+			PreparedStatement *pstmt;
+
+			pstmt = FetchPreparedStatement(portal->prepStmtName, true);
+			psrc = pstmt->plansource;
+		}
+		else
+		{
+			/* special-case the unnamed statement */
+			psrc = unnamed_stmt_psrc;
+			if (!psrc)
+				ereport(ERROR,
+						(errcode(ERRCODE_UNDEFINED_PSTATEMENT),
+						errmsg("unnamed prepared statement does not exist")));
+		}
+
+		if (IsA(psrc->raw_parse_tree->stmt, SelectStmt) &&
+			checkCanOptSelectLockingClause((SelectStmt *)psrc->raw_parse_tree->stmt) &&
+			max_rows != FETCH_ALL)
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_INTERNAL_ERROR),
+					errmsg("Should fetch all rows for query if lock optimization is enabled")));
+		}
+	}
+
 	completed = PortalRun(portal,
 						  max_rows,
 						  true, /* always top level */

src/backend/utils/misc/guc.c

@@ -67,6 +67,7 @@
 #include "optimizer/optimizer.h"
 #include "optimizer/paths.h"
 #include "optimizer/planmain.h"
+#include "parser/analyze.h"
 #include "parser/parse_expr.h"
 #include "parser/parse_type.h"
 #include "parser/parser.h"
@@ -2209,6 +2210,13 @@ static struct config_bool ConfigureNamesBool[] =
 		NULL, NULL, NULL
 	},
 
+	{
+		{"enable_lock_optimization", PGC_SIGHUP, CUSTOM_OPTIONS},
+		&enableLockOptimization,
+		false,
+		NULL, NULL, NULL
+	},
+
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, false, NULL, NULL, NULL

src/include/parser/analyze.h

@@ -17,6 +17,9 @@
 #include "parser/parse_node.h"
 #include "utils/queryjumble.h"
 
+/* GUC parameter */
+extern PGDLLIMPORT bool enableLockOptimization;
+
 /* Hook for plugins to get control at end of parse analysis */
 typedef void (*post_parse_analyze_hook_type) (ParseState *pstate,
 											  Query *query,
@@ -44,7 +47,7 @@ extern void CheckSelectLocking(Query *qry, LockClauseStrength strength);
 extern void applyLockingClause(Query *qry, Index rtindex,
 							   LockClauseStrength strength,
 							   LockWaitPolicy waitPolicy, bool pushedDown);
-
+extern bool checkCanOptSelectLockingClause(SelectStmt *stmt);
 extern List *BuildOnConflictExcludedTargetlist(Relation targetrel,
 											   Index exclRelIndex);
 

src/include/utils/unsync_guc_name.h

@@ -117,6 +117,7 @@
 		"enable_incremental_sort",
 		"enable_indexonlyscan",
 		"enable_indexscan",
+		"enable_lock_optimization",
 		"enable_material",
 		"enable_memoize",
 		"enable_mergejoin",

src/test/regress/expected/rangefuncs_cdb.out

@@ -19,6 +19,7 @@ name not in ('enable_parallel',
  enable_incremental_sort        | off
  enable_indexonlyscan           | on
  enable_indexscan               | on
+ enable_lock_optimization       | off
  enable_material                | on
  enable_memoize                 | on
  enable_mergejoin               | off
@@ -32,7 +33,7 @@ name not in ('enable_parallel',
  enable_seqscan                 | on
  enable_sort                    | on
  enable_tidscan                 | on
-(22 rows)
+(23 rows)
 
 -- start_ignore
 create schema rangefuncs_cdb;

src/test/regress/expected/sysviews.out

@@ -117,6 +117,7 @@ name not in ('enable_parallel',
  enable_incremental_sort        | off
  enable_indexonlyscan           | on
  enable_indexscan               | on
+ enable_lock_optimization       | off
  enable_material                | on
  enable_memoize                 | on
  enable_mergejoin               | off
@@ -130,7 +131,7 @@ name not in ('enable_parallel',
  enable_seqscan                 | on
  enable_sort                    | on
  enable_tidscan                 | on
-(22 rows)
+(23 rows)
 
 -- Test that the pg_timezone_names and pg_timezone_abbrevs views are
 -- more-or-less working.  We can't test their contents in any great detail

src/test/singlenode_regress/expected/sysviews.out

@@ -116,6 +116,7 @@ name not in ('enable_parallel',
  enable_incremental_sort        | off
  enable_indexonlyscan           | on
  enable_indexscan               | on
+ enable_lock_optimization       | off
  enable_material                | on
  enable_memoize                 | on
  enable_mergejoin               | off
@@ -129,7 +130,7 @@ name not in ('enable_parallel',
  enable_seqscan                 | on
  enable_sort                    | on
  enable_tidscan                 | on
-(22 rows)
+(23 rows)
 
 -- Test that the pg_timezone_names and pg_timezone_abbrevs views are
 -- more-or-less working.  We can't test their contents in any great detail