diff --git a/src/backend/distributed/master/master_node_protocol.c b/src/backend/distributed/master/master_node_protocol.c index 035db256d..6976317e9 100644 --- a/src/backend/distributed/master/master_node_protocol.c +++ b/src/backend/distributed/master/master_node_protocol.c @@ -69,6 +69,7 @@ int ShardReplicationFactor = 1; /* desired replication factor for shards */ int ShardMaxSize = 1048576; /* maximum size in KB one shard can grow to */ int ShardPlacementPolicy = SHARD_PLACEMENT_ROUND_ROBIN; int NextShardId = 0; +int NextPlacementId = 0; static List * GetTableReplicaIdentityCommand(Oid relationId); static Datum WorkerNodeGetDatum(WorkerNode *workerNode, TupleDesc tupleDescriptor); @@ -294,12 +295,19 @@ GetNextShardId() { text *sequenceName = NULL; Oid sequenceId = InvalidOid; - Datum sequenceIdDatum = NULL; + Datum sequenceIdDatum = 0; Oid savedUserId = InvalidOid; int savedSecurityContext = 0; Datum shardIdDatum = 0; uint64 shardId = 0; + /* + * In regression tests, we would like to generate shard IDs consistently + * even if the tests run in parallel. Instead of the sequence, we can use + * the next_shard_id GUC to specify which shard ID the current session should + * generate next. The GUC is automatically increased by 1 every time a new + * shard ID is generated. + */ if (NextShardId > 0) { shardId = NextShardId; @@ -364,14 +372,33 @@ master_get_new_placementid(PG_FUNCTION_ARGS) uint64 GetNextPlacementId(void) { - text *sequenceName = cstring_to_text(PLACEMENTID_SEQUENCE_NAME); - Oid sequenceId = ResolveRelationId(sequenceName); - Datum sequenceIdDatum = ObjectIdGetDatum(sequenceId); + text *sequenceName = NULL; + Oid sequenceId = InvalidOid; + Datum sequenceIdDatum = 0; Oid savedUserId = InvalidOid; int savedSecurityContext = 0; Datum placementIdDatum = 0; uint64 placementId = 0; + /* + * In regression tests, we would like to generate placement IDs consistently + * even if the tests run in parallel. Instead of the sequence, we can use + * the next_placement_id GUC to specify which shard ID the current session + * should generate next. The GUC is automatically increased by 1 every time + * a new placement ID is generated. + */ + if (NextPlacementId > 0) + { + placementId = NextPlacementId; + NextPlacementId += 1; + + return placementId; + } + + sequenceName = cstring_to_text(PLACEMENTID_SEQUENCE_NAME); + sequenceId = ResolveRelationId(sequenceName); + sequenceIdDatum = ObjectIdGetDatum(sequenceId); + GetUserIdAndSecContext(&savedUserId, &savedSecurityContext); SetUserIdAndSecContext(CitusExtensionOwner(), SECURITY_LOCAL_USERID_CHANGE); diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index eef9897bb..5cc4bc79d 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -828,6 +828,21 @@ RegisterCitusConfigVariables(void) GUC_NO_SHOW_ALL, NULL, NULL, NULL); + DefineCustomIntVariable( + "citus.next_placement_id", + gettext_noop("Set the next placement ID to use in placement creation."), + gettext_noop("Placement IDs are normally generated using a sequence. If " + "next_placement_id is set to a non-zero value, placement IDs will " + "instead be generated by incrementing from the value of " + "this GUC and this will be reflected in the GUC. This is " + "mainly useful to ensure consistent placement IDs when running " + "tests in parallel."), + &NextPlacementId, + 0, 0, INT_MAX, + PGC_USERSET, + GUC_NO_SHOW_ALL, + NULL, NULL, NULL); + DefineCustomIntVariable( "citus.max_task_string_size", gettext_noop("Sets the maximum size (in bytes) of a worker task call string."), diff --git a/src/include/distributed/master_protocol.h b/src/include/distributed/master_protocol.h index 9875196dc..c37a23b42 100644 --- a/src/include/distributed/master_protocol.h +++ b/src/include/distributed/master_protocol.h @@ -93,6 +93,7 @@ extern int ShardReplicationFactor; extern int ShardMaxSize; extern int ShardPlacementPolicy; extern int NextShardId; +extern int NextPlacementId; extern bool IsCoordinator(void);