wait on lock take

fdbserver/DDTxnProcessor.actor.cpp

@@ -705,12 +705,12 @@ struct DDMockTxnProcessorImpl {
 		std::sort(params.destinationTeam.begin(), params.destinationTeam.end());
 		std::sort(params.healthyDestinations.begin(), params.healthyDestinations.end());
 
-		self->rawStartMovement(params, tssMapping);
+		wait(self->rawStartMovement(params, tssMapping));
 		ASSERT(tssMapping.empty());
 
 		wait(checkFetchingState(self, params.destinationTeam, params.keys));
 
-		self->rawFinishMovement(params, tssMapping);
+		wait(self->rawFinishMovement(params, tssMapping));
 		if (!params.dataMovementComplete.isSet())
 			params.dataMovementComplete.send(Void());
 		return Void();
@@ -888,10 +888,14 @@ Future<std::vector<ProcessData>> DDMockTxnProcessor::getWorkers() const {
 	return Future<std::vector<ProcessData>>();
 }
 
-void DDMockTxnProcessor::rawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping) {
+ACTOR Future<Void> rawStartMovement(std::shared_ptr<MockGlobalState> mgs,
+                                    MoveKeysParams params,
+                                    std::map<UID, StorageServerInterface> tssMapping) {
 	// There won’t be parallel rawStart or rawFinish in mock world due to the fact the following *mock* transaction code
 	// will always finish without coroutine switch.
 	ASSERT(params.startMoveKeysParallelismLock->activePermits() == 0);
+	wait(params.startMoveKeysParallelismLock->take(TaskPriority::DataDistributionLaunch));
+	state FlowLock::Releaser releaser(*params.startMoveKeysParallelismLock);
 
 	std::vector<ShardsAffectedByTeamFailure::Team> destTeams;
 	destTeams.emplace_back(params.destinationTeam, true);
@@ -920,13 +924,22 @@ void DDMockTxnProcessor::rawStartMovement(MoveKeysParams& params, std::map<UID,
 		server.setShardStatus(params.keys, MockShardStatus::INFLIGHT, mgs->restrictSize);
 		server.signalFetchKeys(params.keys, randomRangeSize);
 	}
+	return Void();
 }
 
-void DDMockTxnProcessor::rawFinishMovement(MoveKeysParams& params,
+Future<Void> DDMockTxnProcessor::rawStartMovement(MoveKeysParams& params,
-                                           const std::map<UID, StorageServerInterface>& tssMapping) {
+                                                  std::map<UID, StorageServerInterface>& tssMapping) {
+	return ::rawStartMovement(mgs, params, tssMapping);
+}
+
+ACTOR Future<Void> rawFinishMovement(std::shared_ptr<MockGlobalState> mgs,
+                                     MoveKeysParams params,
+                                     std::map<UID, StorageServerInterface> tssMapping) {
 	// There won’t be parallel rawStart or rawFinish in mock world due to the fact the following *mock* transaction code
 	// will always finish without coroutine switch.
 	ASSERT(params.finishMoveKeysParallelismLock->activePermits() == 0);
+	wait(params.finishMoveKeysParallelismLock->take(TaskPriority::DataDistributionLaunch));
+	state FlowLock::Releaser releaser(*params.finishMoveKeysParallelismLock);
 
 	// get source and dest teams
 	auto [destTeams, srcTeams] = mgs->shardMapping->getTeamsForFirstShard(params.keys);
@@ -953,4 +966,10 @@ void DDMockTxnProcessor::rawFinishMovement(MoveKeysParams& params,
 	}
 	mgs->shardMapping->finishMove(params.keys);
 	mgs->shardMapping->defineShard(params.keys); // coalesce for merge
+	return Void();
+}
+
+Future<Void> DDMockTxnProcessor::rawFinishMovement(MoveKeysParams& params,
+                                                   const std::map<UID, StorageServerInterface>& tssMapping) {
+	return ::rawFinishMovement(mgs, params, tssMapping);
 }

fdbserver/include/fdbserver/DDTxnProcessor.h

@@ -292,9 +292,9 @@ public:
 	Future<std::vector<ProcessData>> getWorkers() const override;
 
 protected:
-	void rawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping);
+	Future<Void> rawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping);
 
-	void rawFinishMovement(MoveKeysParams& params, const std::map<UID, StorageServerInterface>& tssMapping);
+	Future<Void> rawFinishMovement(MoveKeysParams& params, const std::map<UID, StorageServerInterface>& tssMapping);
 };
 
 #endif // FOUNDATIONDB_DDTXNPROCESSOR_H

fdbserver/workloads/IDDTxnProcessorApiCorrectness.actor.cpp

@@ -80,12 +80,12 @@ void verifyInitDataEqual(Reference<InitialDataDistribution> real, Reference<Init
 class DDMockTxnProcessorTester : public DDMockTxnProcessor {
 public:
 	explicit DDMockTxnProcessorTester(std::shared_ptr<MockGlobalState> mgs = nullptr) : DDMockTxnProcessor(mgs) {}
-	void testRawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping) {
+	Future<Void> testRawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping) {
-		rawStartMovement(params, tssMapping);
+		return rawStartMovement(params, tssMapping);
 	}
 
-	void testRawFinishMovement(MoveKeysParams& params, const std::map<UID, StorageServerInterface>& tssMapping) {
+	Future<Void> testRawFinishMovement(MoveKeysParams& params, const std::map<UID, StorageServerInterface>& tssMapping) {
-		rawFinishMovement(params, tssMapping);
+		return rawFinishMovement(params, tssMapping);
 	}
 };
 
@@ -94,12 +94,12 @@ public:
 	explicit DDTxnProcessorTester(Database cx) : DDTxnProcessor(cx) {}
 
 	Future<Void> testRawStartMovement(MoveKeysParams& params, std::map<UID, StorageServerInterface>& tssMapping) {
-		return this->rawStartMovement(params, tssMapping);
+		return rawStartMovement(params, tssMapping);
 	}
 
 	Future<Void> testRawFinishMovement(MoveKeysParams& params,
 	                                   const std::map<UID, StorageServerInterface>& tssMapping) {
-		return this->rawFinishMovement(params, tssMapping);
+		return rawFinishMovement(params, tssMapping);
 	}
 };
 
@@ -279,7 +279,7 @@ struct IDDTxnProcessorApiWorkload : TestWorkload {
 			wait(store(params.lock, self->real->takeMoveKeysLock(UID())));
 			try {
 				// test start
-				self->mock->testRawStartMovement(params, emptyTssMapping);
+				wait(self->mock->testRawStartMovement(params, emptyTssMapping));
 				wait(self->real->testRawStartMovement(params, emptyTssMapping));
 
 				// test finish or started but cancelled movement
@@ -288,7 +288,7 @@ struct IDDTxnProcessorApiWorkload : TestWorkload {
 					break;
 				}
 
-				self->mock->testRawFinishMovement(params, emptyTssMapping);
+				wait(self->mock->testRawFinishMovement(params, emptyTssMapping));
 				wait(self->real->testRawFinishMovement(params, emptyTssMapping));
 				break;
 			} catch (Error& e) {

tests/CMakeLists.txt

@@ -171,6 +171,7 @@ if(WITH_PYTHON)
   add_fdb_test(TEST_FILES fast/GetEstimatedRangeSize.toml)
   add_fdb_test(TEST_FILES fast/GetMappedRange.toml)
   add_fdb_test(TEST_FILES fast/IDDTxnProcessorRawStartMovement.toml)
+  add_fdb_test(TEST_FILES fast/IDDTxnProcessorMoveKeys.toml IGNORE)
   add_fdb_test(TEST_FILES fast/PrivateEndpoints.toml)
   add_fdb_test(TEST_FILES fast/ProtocolVersion.toml)
   add_fdb_test(TEST_FILES fast/RandomSelector.toml)

tests/fast/IDDTxnProcessorMoveKeys.toml

@@ -0,0 +1,13 @@
+[configuration]
+generateFearless = false # prevent generating remote dc because in MGS there's no region setting yet
+disableTss = true # There's no TSS in MGS this prevent the DD operate TSS mapping
+
+[[knobs]]
+max_added_sources_multiplier = 0 # set to 0 because it's impossible to make sure SS and mock SS will finish fetch keys at the same time.
+
+[[test]]
+testTitle = 'IDDTxnProcessorMoveKeys'
+
+    [[test.workload]]
+    testName = 'IDDTxnProcessorApiCorrectness'
+    testDuration = 50.0

tests/fast/IDDTxnProcessorRawStartMovement.toml

@@ -6,7 +6,7 @@ disableTss = true # There's no TSS in MGS this prevent the DD operate TSS mappin
 max_added_sources_multiplier = 0 # set to 0 because it's impossible to make sure SS and mock SS will finish fetch keys at the same time.
 
 [[test]]
-testTitle = 'IDDTxnProcessorApiCorrectness'
+testTitle = 'IDDTxnProcessorRawStartMovement'
 
     [[test.workload]]
     testName = 'IDDTxnProcessorApiCorrectness'