Make getTeamByServers O(1) in time (#12938) (#13012) (#13244)

* Make getTeamByServers O(1) in time (#12938)

* Make getTeamByServers O(1) in time

* address feedback

* address feedback

fdbclient/ServerKnobs.cpp

@@ -1021,6 +1021,7 @@ void ServerKnobs::initialize(Randomize randomize, ClientKnobs* clientKnobs, IsSi
 	init( DD_PREFER_LOW_READ_UTIL_TEAM,                          true );
 	init( DD_TRACE_MOVE_BYTES_AVERAGE_INTERVAL,                   120);
 	init( MOVING_WINDOW_SAMPLE_SIZE,                         10000000); // 10MB
+	init( DD_TEAMS_BY_SERVER_IDS_CONSISTENCY_CHECK_PROB_SIM,       0.0 ); if( isSimulated ) DD_TEAMS_BY_SERVER_IDS_CONSISTENCY_CHECK_PROB_SIM = (deterministicRandom()->random01() * 0.4) + 0.1;
 
 	//Storage Server
 	init( STORAGE_LOGGING_DELAY,                                 5.0 );

fdbclient/include/fdbclient/ServerKnobs.h

@@ -1074,6 +1074,9 @@ public:
 	// Rolling window duration over which the average bytes moved by DD is calculated for the 'MovingData' trace event.
 	double DD_TRACE_MOVE_BYTES_AVERAGE_INTERVAL;
 	int64_t MOVING_WINDOW_SAMPLE_SIZE;
+	// Probability of running the consistency check between teams and teamsByServerIDs
+	// in getTeamByServers (simulation only).
+	double DD_TEAMS_BY_SERVER_IDS_CONSISTENCY_CHECK_PROB_SIM;
 
 	// Storage Server
 	double STORAGE_LOGGING_DELAY;

fdbserver/DDTeamCollection.actor.cpp

@@ -235,13 +235,11 @@ public:
 
 	// Find the team with the exact storage servers as req.src.
 	static void getTeamByServers(DDTeamCollection* self, GetTeamRequest req) {
-		const std::string servers = TCTeamInfo::serversToString(req.src);
+		getTeamByServersConsistencyCheckInSim(self);
 		Optional<Reference<IDataDistributionTeam>> res;
-		for (const auto& team : self->teams) {
-			if (team->getServerIDsStr() == servers) {
-				res = team;
-				break;
-			}
+		auto it = self->teamsByServerIDs.find(TCTeamInfo::serversToString(req.src));
+		if (it != self->teamsByServerIDs.end()) {
+			res = it->second;
 		}
 		req.reply.send(std::make_pair(res, false));
 	}
@@ -381,6 +379,34 @@ public:
 		}
 	}
 
+	// Probabilistic consistency check between teams and teamsByServerIDs
+	// Run only in simulation with a probability of DD_TEAMS_BY_SERVER_IDS_CONSISTENCY_CHECK_PROB_SIM
+	// We may need to tune this knob if simulation runs too slowly (in real-time) and results in
+	// ExternalTimeout in Joshua
+	static void getTeamByServersConsistencyCheckInSim(DDTeamCollection* self) {
+		// This check can be expensive in prod so only run it in simulation
+		if (!g_network->isSimulated()) {
+			return;
+		}
+
+		if (deterministicRandom()->random01() < SERVER_KNOBS->DD_TEAMS_BY_SERVER_IDS_CONSISTENCY_CHECK_PROB_SIM) {
+			std::unordered_map<std::string, Reference<TCTeamInfo>> expected;
+			for (const auto& team : self->teams) {
+				expected[team->getServerIDsStr()] = team;
+			}
+			ASSERT(expected.size() == self->teamsByServerIDs.size());
+			for (const auto& [key, value] : expected) {
+				auto it = self->teamsByServerIDs.find(key);
+				ASSERT(it != self->teamsByServerIDs.end());
+				ASSERT(it->second == value);
+			}
+			TraceEvent("TeamByServerIDsConsistencyCheckPassed")
+			    .suppressFor(5.0)
+			    .detail("TeamsSize", self->teams.size())
+			    .detail("MapSize", self->teamsByServerIDs.size());
+		}
+	}
+
 	// Return a threshold of team queue size which guarantees at least DD_LONG_STORAGE_QUEUE_TEAM_MAJORITY_PERCENTILE
 	// portion of teams that have longer storage queues
 	// A team storage queue size is defined as the longest storage queue size among all SSes of the team
@@ -4993,6 +5019,7 @@ void DDTeamCollection::addTeam(const std::vector<Reference<TCServerInfo>>& newTe
 
 	// For a good team, we add it to teams and create machine team for it when necessary
 	teams.push_back(teamInfo);
+	teamsByServerIDs[teamInfo->getServerIDsStr()] = teamInfo;
 	for (auto& server : newTeamServers) {
 		server->addTeam(teamInfo);
 	}
@@ -5918,6 +5945,10 @@ void DDTeamCollection::addServer(StorageServerInterface newServer,
 
 bool DDTeamCollection::removeTeam(Reference<TCTeamInfo> team) {
 	TraceEvent("RemovedServerTeam", distributorId).detail("Team", team->getDesc());
+	auto it = teamsByServerIDs.find(team->getServerIDsStr());
+	if (it != teamsByServerIDs.end()) {
+		teamsByServerIDs.erase(it);
+	}
 	bool found = false;
 	for (int t = 0; t < teams.size(); t++) {
 		if (teams[t] == team) {

fdbserver/include/fdbserver/DDTeamCollection.h

@@ -684,7 +684,12 @@ public:
 	std::map<Standalone<StringRef>, Reference<TCMachineInfo>> machine_info;
 	std::vector<Reference<TCMachineTeamInfo>> machineTeams; // all machine teams
 
+	// IMPORTANT: teams and teamsByServerIDs MUST be consistent, so any time we
+	// mutate teams, we must also mutate teamsByServerIDs
 	std::vector<Reference<TCTeamInfo>> teams;
+	// O(1) hash map from server ID string to team information
+	// Currently used by getTeamByServers
+	std::unordered_map<std::string, Reference<TCTeamInfo>> teamsByServerIDs;
 
 	std::vector<DDTeamCollection*> teamCollections;
 	AsyncTrigger printDetailedTeamsInfo;