pass /MockGlobalState/initialAsEmptyDatabaseMGS

fdbrpc/include/fdbrpc/RangeMap.h

@@ -216,6 +216,19 @@ public:
 
 	Future<Void> clearAsync() { return map.clearAsync(); }
 
+	Metric getMetric(const_iterator x) const { return map.getMetric(x); }
+	Metric getMetric(iterator x) const { return getMetric(const_iterator{ x }); }
+
+	Metric sumTo(const_iterator to) const { return map.sumTo(to); }
+	Metric sumTo(iterator to) const { return sumTo(const_iterator{ to }); }
+
+	Metric sumRange(const_iterator begin, const_iterator end) const { return map.sumRange(begin, end); }
+	Metric sumRange(iterator begin, iterator end) const { return map.sumRange(begin, end); }
+	template <class KeyCompatible>
+	Metric sumRange(const KeyCompatible& begin, const KeyCompatible& end) const {
+		return map.sumRange(begin, end);
+	}
+
 protected:
 	Map<Key, Val, pair_type, Metric> map;
 	const MetricFunc mf;

fdbserver/DDTxnProcessor.actor.cpp

@@ -504,6 +504,7 @@ DDMockTxnProcessor::getServerListAndProcessClasses() {
 	return res;
 }
 
+// reconstruct team combination from shardMapping
 std::set<std::vector<UID>> DDMockTxnProcessor::getAllTeamsInRegion(bool primary) const {
 	auto teams = mgs->shardMapping->getAllTeams();
 	std::set<std::vector<UID>> res;
@@ -516,7 +517,8 @@ std::set<std::vector<UID>> DDMockTxnProcessor::getAllTeamsInRegion(bool primary)
 }
 
 inline void transformTeamsToServerIds(std::vector<ShardsAffectedByTeamFailure::Team>& teams,
-                                std::vector<UID>& primaryIds, std::vector<UID>& remoteIds) {
+                                      std::vector<UID>& primaryIds,
+                                      std::vector<UID>& remoteIds) {
 	std::set<UID> primary, remote;
 	for (auto& team : teams) {
 		team.primary ? primary.insert(team.servers.begin(), team.servers.end())
@@ -526,6 +528,7 @@ inline void transformTeamsToServerIds(std::vector<ShardsAffectedByTeamFailure::T
 	remoteIds = std::vector<UID>(remote.begin(), remote.end());
 }
 
+// reconstruct DDShardInfos from shardMapping
 std::vector<DDShardInfo> DDMockTxnProcessor::getDDShardInfos() const {
 	std::vector<DDShardInfo> res;
 	res.reserve(mgs->shardMapping->getNumberOfShards());
@@ -566,6 +569,14 @@ Future<Reference<InitialDataDistribution>> DDMockTxnProcessor::getInitialDataDis
 	return res;
 }
 
+inline void removeFailedServerFromTeams(std::vector<ShardsAffectedByTeamFailure::Team>& teams, const UID& serverID) {
+	for (auto& team : teams) {
+		team.removeServer(serverID);
+	}
+	teams.erase(std::remove_if(teams.begin(), teams.end(), [](const auto& team) { return team.servers.empty(); }),
+	            teams.end());
+}
+
 // Remove the server from shardMapping and set serverKeysFalse to the server's serverKeys list.
 // Changes to keyServer and serverKey must happen symmetrically in this function.
 // If serverID is the last source server for a shard, the shard will be erased, and then be assigned
@@ -575,6 +586,22 @@ Future<Void> DDMockTxnProcessor::removeKeysFromFailedServer(const UID& serverID,
                                                             const MoveKeysLock& lock,
                                                             const DDEnabledState* ddEnabledState) const {
 	auto& mss = mgs->allServers.at(serverID);
+	auto allRange = mss.getAllRanges();
+	for (auto it = allRange.begin(); it != allRange.end(); ++it) {
+		KeyRangeRef curRange = it->range();
+		auto [curTeams, prevTeams] = mgs->shardMapping->getTeamsFor(curRange);
+		bool inFlight = !prevTeams.empty();
+		// remove failed server from source and destination teams
+		removeFailedServerFromTeams(curTeams, serverID);
+		removeFailedServerFromTeams(prevTeams, serverID);
+
+		if (inFlight) {
+
+		} else {
+			if (curTeams.empty()) {
+			}
+		}
+	}
 
 	return Void();
 }
@@ -583,7 +610,7 @@ Future<Void> DDMockTxnProcessor::removeStorageServer(const UID& serverID,
                                                      const Optional<UID>& tssPairID,
                                                      const MoveKeysLock& lock,
                                                      const DDEnabledState* ddEnabledState) const {
-
+	ASSERT(mgs->shardMapping->getNumberOfShards(serverID) == 0);
 	mgs->allServers.erase(serverID);
 	return Void();
 }

fdbserver/MockGlobalState.cpp

@@ -20,15 +20,71 @@
 
 #include "fdbserver/MockGlobalState.h"
 
+bool MockStorageServer::allShardStatusEqual(KeyRangeRef range, MockShardStatus status) {
+	auto ranges = serverKeys.intersectingRanges(range);
+	ASSERT(!ranges.empty()); // at least the range is allKeys
+
+	for (auto it = ranges.begin(); it != ranges.end(); ++it) {
+		if (it->cvalue().status != status)
+			return false;
+	}
+	return true;
+}
+
 void MockGlobalState::initialAsEmptyDatabaseMGS(const DatabaseConfiguration& conf, uint64_t defaultDiskSpace) {
 	ASSERT(conf.storageTeamSize > 0);
 	configuration = conf;
 	std::vector<UID> serverIds;
-	for(int i = 1; i <= conf.storageTeamSize; ++ i) {
+	for (int i = 1; i <= conf.storageTeamSize; ++i) {
 		UID id = indexToUID(i);
 		serverIds.push_back(id);
 		allServers[id] = MockStorageServer(id, defaultDiskSpace);
-		allServers[id].serverKeys.insert(allKeys, true);
+		allServers[id].serverKeys.insert(allKeys, { MockShardStatus::COMPLETED, 0 });
 	}
-	shardMapping->assignRangeToTeams(allKeys, {Team(serverIds, true)});
+	shardMapping->assignRangeToTeams(allKeys, { Team(serverIds, true) });
+}
+
+bool MockGlobalState::serverIsSourceForShard(const UID& serverId, KeyRangeRef shard, bool inFlightShard) {
+	if (allServers.find(serverId) == allServers.end())
+		return false;
+
+	// check serverKeys
+	auto& mss = allServers.at(serverId);
+	if (!mss.allShardStatusEqual(shard, MockShardStatus::COMPLETED)) {
+		return false;
+	}
+
+	// check keyServers
+	auto teams = shardMapping->getTeamsFor(shard);
+	if (inFlightShard) {
+		return std::any_of(teams.second.begin(), teams.second.end(), [&serverId](const Team& team) {
+			return team.hasServer(serverId);
+		});
+	}
+	return std::any_of(
+	    teams.first.begin(), teams.first.end(), [&serverId](const Team& team) { return team.hasServer(serverId); });
+}
+
+bool MockGlobalState::serverIsDestForShard(const UID& serverId, KeyRangeRef shard) {
+	return false;
+}
+
+TEST_CASE("/MockGlobalState/initialAsEmptyDatabaseMGS/SimpleThree") {
+	BasicTestConfig testConfig;
+	testConfig.simpleConfig = true;
+	testConfig.minimumReplication = 3;
+	testConfig.logAntiQuorum = 0;
+	DatabaseConfiguration dbConfig = generateNormalDatabaseConfiguration(testConfig);
+	TraceEvent("UnitTestDbConfig").detail("Config", dbConfig.toString());
+
+	auto mgs = std::make_shared<MockGlobalState>();
+	mgs->initialAsEmptyDatabaseMGS(dbConfig);
+	for (int i = 1; i <= dbConfig.storageTeamSize; ++i) {
+		auto id = MockGlobalState::indexToUID(i);
+		std::cout << "Check server " << i << "\n";
+		ASSERT(mgs->serverIsSourceForShard(id, allKeys));
+		ASSERT(mgs->allServers.at(id).serverKeys.sumRange(allKeys.begin, allKeys.end) == 0);
+	}
+
+	return Void();
 }

fdbserver/SimulatedCluster.actor.cpp

@@ -82,7 +82,7 @@ bool destructed = false;
 
 // Configuration details specified in workload test files that change the simulation
 // environment details
-class TestConfig: public BasicTestConfig {
+class TestConfig : public BasicTestConfig {
 	class ConfigBuilder {
 		using value_type = toml::basic_value<toml::discard_comments>;
 		using base_variant = std::variant<int, float, double, bool, std::string, std::vector<int>, ConfigDBType>;
@@ -452,6 +452,9 @@ public:
 			}
 		}
 	}
+
+	TestConfig() = default;
+	TestConfig(const BasicTestConfig& config): BasicTestConfig(config){}
 };
 
 template <class T>
@@ -1852,7 +1855,9 @@ void SimulationConfig::setTss(const TestConfig& testConfig) {
 }
 
 void setConfigDB(TestConfig const& testConfig) {
-	g_simulator.configDBType = testConfig.getConfigDBType();
+	if(g_pSimulator) {
+		g_simulator.configDBType = testConfig.getConfigDBType();
+	}
 }
 
 // Generates and sets an appropriate configuration for the database according to
@@ -2605,9 +2610,8 @@ ACTOR void setupAndRun(std::string dataFolder,
 	ASSERT(false);
 }
 
-DatabaseConfiguration generateNormalDatabaseConfiguration(const BasicTestConfig& testConfig, uint64_t defaultDiskSpace) {
-	TestConfig config;
-	config.BasicTestConfig::operator=(testConfig);
+DatabaseConfiguration generateNormalDatabaseConfiguration(const BasicTestConfig& testConfig) {
+	TestConfig config(testConfig);
 	SimulationConfig simConf(config);
 	return simConf.db;
 }

fdbserver/include/fdbserver/MockGlobalState.h

@@ -28,14 +28,29 @@
 #include "SimulatedCluster.h"
 #include "ShardsAffectedByTeamFailure.h"
 
+
+template <class Metric>
+struct ShardSizeMetric {
+	template <typename pair_type>
+	Metric operator()(pair_type const& p) const {
+		return Metric(p.value.shardSize);
+	}
+};
+
+enum class MockShardStatus { EMPTY = 0, COMPLETED, INFLIGHT };
+
 class MockStorageServer {
 public:
+	struct ShardInfo {
+		MockShardStatus status;
+		uint64_t shardSize;
+	};
 	// control plane statistics associated with a real storage server
 	uint64_t usedDiskSpace = 0, availableDiskSpace;
 
 	// In-memory counterpart of the `serverKeys` in system keyspace
-	// the value bool is equal to "[[serverKeysTrue]]" |" [[serverKeysFalse]]" and metrics uint64_t is the shard size
-	KeyRangeMap<bool, uint64_t> serverKeys;
+	// the value ShardStatus is [InFlight, Completed, Empty] and metrics uint64_t is the shard size
+	KeyRangeMap<ShardInfo, uint64_t, ShardSizeMetric<uint64_t>> serverKeys;
 
 	// sampled metrics
 	StorageServerMetrics metrics;
@@ -50,6 +65,10 @@ public:
 	  : usedDiskSpace(usedDiskSpace), availableDiskSpace(availableDiskSpace), id(id) {
 		ssi.uniqueID = id;
 	}
+
+	decltype(serverKeys)::Ranges getAllRanges() { return serverKeys.ranges(); }
+
+	bool allShardStatusEqual(KeyRangeRef range, MockShardStatus status);
 };
 
 class MockGlobalState {
@@ -70,6 +89,11 @@ public:
 	static UID indexToUID(uint64_t a) { return UID(a, a); }
 	void initialAsEmptyDatabaseMGS(const DatabaseConfiguration& conf,
 	                               uint64_t defaultDiskSpace = 1000LL * 1024 * 1024 * 1024);
+
+	// check methods
+	bool serverIsSourceForShard(const UID& serverId, KeyRangeRef shard, bool inFlightShard = false);
+	bool serverIsDestForShard(const UID& serverId, KeyRangeRef shard);
+	bool allShardRemovedFromServer(const UID& serverId);
 };
 
 #endif // FOUNDATIONDB_MOCKGLOBALSTATE_H

fdbserver/include/fdbserver/ShardsAffectedByTeamFailure.h

@@ -49,6 +49,14 @@ public:
 		bool operator==(const Team& r) const { return servers == r.servers && primary == r.primary; }
 		bool operator!=(const Team& r) const { return !(*this == r); }
 
+		bool hasServer(const UID& id) const {
+			return std::find(servers.begin(), servers.end(), id) != servers.end();
+		}
+
+		void removeServer(const UID& id) {
+			servers.erase(std::remove(servers.begin(), servers.end(), id), servers.end());
+		}
+
 		std::string toString() const { return describe(servers); };
 	};
 
@@ -117,6 +125,8 @@ public:
 	auto getAllRanges() const -> decltype(shard_teams)::ConstRanges;
 	// get total shards count
 	size_t getNumberOfShards() const;
+	// void removeFailedServerForRange(KeyRangeRef keys, const UID& serverID, const std::pair<std::vector<Team>,
+	// std::vector<Team>>& involvedTeams);
 };
 
 #endif // FOUNDATIONDB_SHARDSAFFECTEDBYTEAMFAILURE_H