resgroup: move UnassignResGroup into AtEOXact_ResGroup

* Do UnassignResGroup within prepareTransaction too, prepareTransaction()
will put QE out of any transactions temporarily until the second commit
command arrives, so any failures in this gap will cause leaks of resource
group including slots etc.
* Clean code, move UnassignResGroup() into AtEOXact_ResGroup() so resource
group related codes will not spread across prepare/commit/abort functions.
* Do not call callback functions in PrepareTransaction because the transaction
is not trully commited.

src/backend/access/transam/xact.c

@@ -2740,7 +2740,7 @@ CommitTransaction(void)
 	/* All relations that are in the vacuum process are being commited now. */
 	ResetVacuumRels();
 
-	/* Process resource group related callbacks */
+	/* Process resource group commit processing */
 	AtEOXact_ResGroup(true);
 
 	/* Check we've released all buffer pins */
@@ -2837,10 +2837,6 @@ CommitTransaction(void)
 	RESUME_INTERRUPTS();
 
 	freeGangsForPortal(NULL);
-
-	/* Release resource group slot at the end of a transaction */
-	if (ShouldUnassignResGroup())
-		UnassignResGroup();
 }
 
 
@@ -2891,6 +2887,9 @@ PrepareTransaction(void)
 			break;
 	}
 
+	/* detach from current resouce group */
+	AtPrepare_ResGroup();
+
 	/* Now we can shut down the deferred-trigger manager */
 	AfterTriggerEndXact(true);
 
@@ -3046,9 +3045,6 @@ PrepareTransaction(void)
 						 RESOURCE_RELEASE_BEFORE_LOCKS,
 						 true, true);
 
-	/* Process resource group related callbacks */
-	AtEOXact_ResGroup(true);
-
 	/* Check we've released all buffer pins */
 	AtEOXact_Buffers(true);
 
@@ -3214,6 +3210,9 @@ AbortTransaction(void)
 	if (Gp_role == GP_ROLE_DISPATCH && IsResQueueEnabled())
 		AtAbort_ResScheduler();
 		
+	/* Perform any Resource Group abort processing. */
+	AtEOXact_ResGroup(false);
+
 	/* Perform any AO table abort processing */
 	AtAbort_AppendOnly();
 
@@ -3273,7 +3272,6 @@ AbortTransaction(void)
 		ResourceOwnerRelease(TopTransactionResourceOwner,
 							 RESOURCE_RELEASE_BEFORE_LOCKS,
 							 false, true);
-		AtEOXact_ResGroup(false);
 		AtEOXact_Buffers(false);
 		AtEOXact_RelationCache(false);
 		AtEOXact_Inval(false);
@@ -3393,10 +3391,6 @@ CleanupTransaction(void)
 	s->state = TRANS_DEFAULT;
 
 	finishDistributedTransactionContext("CleanupTransaction", true);
-
-	/* Release resource group slot at the end of a transaction */
-	if (ShouldUnassignResGroup())
-		UnassignResGroup();
 }
 
 /*

src/backend/commands/resgroupcmds.c

@@ -109,7 +109,7 @@ static void registerResourceGroupCallback(ResourceGroupCallback callback, void *
  * Note the callback functions would be removed as being processed.
  */
 void
-AtEOXact_ResGroup(bool isCommit)
+HandleResGroupDDLCallbacks(bool isCommit)
 {
 	if (ResourceGroup_callback.callback == NULL)
 		return;

src/backend/utils/resgroup/resgroup.c

@@ -403,6 +403,24 @@ AllocResGroupEntry(Oid groupId, const ResGroupCaps *caps)
 	LWLockRelease(ResGroupLock);
 }
 
+void
+AtEOXact_ResGroup(bool isCommit)
+{
+	HandleResGroupDDLCallbacks(isCommit);
+
+	/* Release resource group slot at the end of a transaction */
+	if (ShouldUnassignResGroup())
+		UnassignResGroup();
+}
+
+void
+AtPrepare_ResGroup(void)
+{
+	/* Release resource group slot */
+	if (ShouldUnassignResGroup())
+		UnassignResGroup();
+}
+
 /*
  * Load the resource groups in shared memory. Note this
  * can only be done after enough setup has been done. This uses
@@ -2521,6 +2539,24 @@ ResGroupGetMemInfo(int *memLimit, int *slotQuota, int *sharedQuota)
 	*sharedQuota = groupGetMemSharedExpected(caps);
 }
 
+int64
+ResGroupGetSelfCapability(const char *prop)
+{
+	if (!strcmp(prop, "memory_limit"))
+		return self->caps.memLimit;
+	if (!strcmp(prop, "concurrency"))
+		return self->caps.concurrency;
+	if (!strcmp(prop, "cpu_rate_limit"))
+		return self->caps.cpuRateLimit;
+	if (!strcmp(prop, "memory_shared_quota"))
+		return self->caps.memSharedQuota;
+	if (!strcmp(prop, "memory_spill_ratio"))
+		return self->caps.memSpillRatio;
+
+	elog(ERROR, "unknow property %s", prop);
+	return -1;
+}
+
 /*
  * Validate the consistency of the resgroup information in self.
  *

src/include/commands/resgroupcmds.h

@@ -30,5 +30,8 @@ extern void GetResGroupCapabilities(Relation rel,
 									Oid groupId,
 									ResGroupCaps *resgroupCaps);
 extern void AtEOXact_ResGroup(bool isCommit);
+extern void GetResGroupCapabilities(Relation rel, Oid groupId,
+									ResGroupCaps *resgroupCaps);
+extern void HandleResGroupDDLCallbacks(bool isCommit);
 
 #endif   /* RESGROUPCMDS_H */

src/include/utils/resgroup.h

@@ -95,6 +95,8 @@ typedef enum
 /*
  * Functions in resgroup.c
  */
+extern void AtEOXact_ResGroup(bool isCommit);
+extern void AtPrepare_ResGroup(void);
 
 /* Shared memory and semaphores */
 extern Size ResGroupShmemSize(void);
@@ -140,6 +142,7 @@ extern int32 ResGroupGetMaxChunksPerQuery(void);
 
 /* test helper function */
 extern void ResGroupGetMemInfo(int *memLimit, int *slotQuota, int *sharedQuota);
+extern int64 ResGroupGetSelfCapability(const char *prop);
 
 extern int64 ResourceGroupGetQueryMemoryLimit(void);
 

src/test/isolation2/input/resgroup/resgroup_alter_memory.source

@@ -495,3 +495,60 @@ DROP ROLE role1_memory_test;
 DROP ROLE role2_memory_test;
 DROP RESOURCE GROUP rg1_memory_test;
 DROP RESOURCE GROUP rg2_memory_test;
+
+--
+-- Test PrepareTransaction report an error
+--
+CREATE OR REPLACE FUNCTION getSelfRGCapability(cstring) RETURNS int AS
+'@abs_builddir@/../regress/regress@DLSUFFIX@', 'getSelfRGCapability'
+LANGUAGE C READS SQL DATA;
+CREATE RESOURCE GROUP rg_test_group WITH (cpu_rate_limit=5, memory_limit=5);
+CREATE ROLE rg_test_role RESOURCE GROUP rg_test_group;
+
+SET debug_dtm_action = "fail_begin_command";
+SET debug_dtm_action_target = "protocol";
+SET debug_dtm_action_protocol = "prepare";
+SET debug_dtm_action_segment = 0;
+
+-- ALTER should fail and the memory_limit in both catalog and share memory are
+-- still 5%
+ALTER RESOURCE GROUP rg_test_group set memory_limit 1;
+
+RESET debug_dtm_action;
+RESET debug_dtm_action_target;
+RESET debug_dtm_action_protocol;
+RESET debug_dtm_action_segment;
+
+-- should still be 5% on both QD and QE
+select memory_limit from gp_toolkit.gp_resgroup_config where groupname = 'rg_test_group';
+SET ROLE rg_test_role;
+select getSelfRGCapability('memory_limit');
+select getSelfRGCapability('memory_limit') from gp_dist_random('gp_id');
+
+SET ROLE none;
+
+--
+-- Test error happen on commit_prepare, DDL success after retry
+--
+SET debug_dtm_action = "fail_begin_command";
+SET debug_dtm_action_target = "protocol";
+SET debug_dtm_action_protocol = "commit_prepared";
+SET debug_dtm_action_segment = 0;
+
+-- ALTER should success
+ALTER RESOURCE GROUP rg_test_group set memory_limit 4;
+
+RESET debug_dtm_action;
+RESET debug_dtm_action_target;
+RESET debug_dtm_action_protocol;
+RESET debug_dtm_action_segment;
+
+-- should still be 4% on both QD and QE
+select memory_limit from gp_toolkit.gp_resgroup_config where groupname = 'rg_test_group';
+SET ROLE rg_test_role;
+select getSelfRGCapability('memory_limit');
+select getSelfRGCapability('memory_limit') from gp_dist_random('gp_id');
+
+SET ROLE none;
+DROP ROLE rg_test_role;
+DROP RESOURCE GROUP rg_test_group;

src/test/isolation2/output/resgroup/resgroup_alter_memory.source

@@ -862,3 +862,113 @@ DROP RESOURCE GROUP rg1_memory_test;
 DROP
 DROP RESOURCE GROUP rg2_memory_test;
 DROP
+
+--
+-- Test PrepareTransaction report an error
+--
+CREATE OR REPLACE FUNCTION getSelfRGCapability(cstring) RETURNS int AS '@abs_builddir@/../regress/regress@DLSUFFIX@', 'getSelfRGCapability' LANGUAGE C READS SQL DATA;
+CREATE
+CREATE RESOURCE GROUP rg_test_group WITH (cpu_rate_limit=5, memory_limit=5);
+CREATE
+CREATE ROLE rg_test_role RESOURCE GROUP rg_test_group;
+CREATE
+
+SET debug_dtm_action = "fail_begin_command";
+SET
+SET debug_dtm_action_target = "protocol";
+SET
+SET debug_dtm_action_protocol = "prepare";
+SET
+SET debug_dtm_action_segment = 0;
+SET
+
+-- ALTER should fail and the memory_limit in both catalog and share memory are
+-- still 5%
+ALTER RESOURCE GROUP rg_test_group set memory_limit 1;
+ERROR:  The distributed transaction 'Prepare' broadcast failed to one or more segments for gid = 1509599543-0000000983. (cdbtm.c:698)
+
+RESET debug_dtm_action;
+RESET
+RESET debug_dtm_action_target;
+RESET
+RESET debug_dtm_action_protocol;
+RESET
+RESET debug_dtm_action_segment;
+RESET
+
+-- should still be 5% on both QD and QE
+select memory_limit from gp_toolkit.gp_resgroup_config where groupname = 'rg_test_group';
+memory_limit
+------------
+5           
+(1 row)
+SET ROLE rg_test_role;
+SET
+select getSelfRGCapability('memory_limit');
+getselfrgcapability
+-------------------
+5                 
+(1 row)
+select getSelfRGCapability('memory_limit') from gp_dist_random('gp_id');
+getselfrgcapability
+-------------------
+5                 
+5                 
+5                 
+(3 rows)
+
+SET ROLE none;
+SET
+
+--
+-- Test error happen on commit_prepare, DDL success after retry
+--
+SET debug_dtm_action = "fail_begin_command";
+SET
+SET debug_dtm_action_target = "protocol";
+SET
+SET debug_dtm_action_protocol = "commit_prepared";
+SET
+SET debug_dtm_action_segment = 0;
+SET
+
+-- ALTER should success
+ALTER RESOURCE GROUP rg_test_group set memory_limit 4;
+ALTER
+
+RESET debug_dtm_action;
+RESET
+RESET debug_dtm_action_target;
+RESET
+RESET debug_dtm_action_protocol;
+RESET
+RESET debug_dtm_action_segment;
+RESET
+
+-- should still be 4% on both QD and QE
+select memory_limit from gp_toolkit.gp_resgroup_config where groupname = 'rg_test_group';
+memory_limit
+------------
+4           
+(1 row)
+SET ROLE rg_test_role;
+SET
+select getSelfRGCapability('memory_limit');
+getselfrgcapability
+-------------------
+4                 
+(1 row)
+select getSelfRGCapability('memory_limit') from gp_dist_random('gp_id');
+getselfrgcapability
+-------------------
+4                 
+4                 
+4                 
+(3 rows)
+
+SET ROLE none;
+SET
+DROP ROLE rg_test_role;
+DROP
+DROP RESOURCE GROUP rg_test_group;
+DROP

src/test/regress/regress_gp.c

@@ -68,6 +68,7 @@ extern Datum userdata_project(PG_FUNCTION_ARGS);
 extern Datum checkResourceQueueMemoryLimits(PG_FUNCTION_ARGS);
 extern Datum repeatPalloc(PG_FUNCTION_ARGS);
 extern Datum resGroupPalloc(PG_FUNCTION_ARGS);
+extern Datum getSelfRGCapability(PG_FUNCTION_ARGS);
 
 /* Gang management test support */
 extern Datum gangRaiseInfo(PG_FUNCTION_ARGS);
@@ -604,6 +605,21 @@ repeatPalloc(PG_FUNCTION_ARGS)
 	PG_RETURN_INT32(0);
 }
 
+PG_FUNCTION_INFO_V1(getSelfRGCapability);
+Datum
+getSelfRGCapability(PG_FUNCTION_ARGS)
+{
+	char *cap = PG_GETARG_CSTRING(0);
+	int64 ret = -1;
+
+	if (!IsResGroupEnabled())
+		PG_RETURN_INT64(-1);
+
+	ret = ResGroupGetSelfCapability(cap);
+
+	PG_RETURN_INT64(ret);
+}
+
 PG_FUNCTION_INFO_V1(resGroupPalloc);
 Datum
 resGroupPalloc(PG_FUNCTION_ARGS)