Synchronize static members of KZGlobalService

src/kz/global/kz_global.cpp

@@ -105,11 +105,20 @@ void KZGlobalService::Cleanup()
 
 void KZGlobalService::UpdateGlobalCache()
 {
-	if (KZGlobalService::currentMap.has_value())
+	u16 currentMapID = 0;
+
 	{
-		return;
+		std::unique_lock currentMapLock(KZGlobalService::currentMapMutex);
+
+		if (!KZGlobalService::currentMap.has_value())
+		{
+			return;
+		}
+
+		currentMapID = KZGlobalService::currentMap->id;
 	}
-	KZ::API::events::WantWorldRecordsForCache data(KZGlobalService::currentMap->id);
+
+	KZ::API::events::WantWorldRecordsForCache data(currentMapID);
 
 	auto callback = [](const KZ::API::events::WorldRecordsForCache &records)
 	{
@@ -163,7 +172,10 @@ void KZGlobalService::OnActivateServer()
 
 	auto callback = [currentMapName](const KZ::API::events::MapInfo &mapInfo)
 	{
-		KZGlobalService::currentMap = mapInfo.map;
+		{
+			std::unique_lock handshakeLock(KZGlobalService::handshakeLock);
+			KZGlobalService::currentMap = mapInfo.map;
+		}
 
 		if (!mapInfo.map.has_value())
 		{
@@ -182,6 +194,30 @@ void KZGlobalService::OnActivateServer()
 	KZGlobalService::SendMessage("map-change", data, callback);
 }
 
+void KZGlobalService::OnServerGamePostSimulate()
+{
+	std::unordered_map<u64, StoredCallback> readyCallbacks;
+
+	{
+		std::unique_lock callbacksLock(KZGlobalService::callbacksMutex);
+
+		for (auto it = KZGlobalService::callbacks.cbegin(); it != KZGlobalService::callbacks.cend();)
+		{
+			auto curr = it++;
+
+			if (curr->second.payload.has_value())
+			{
+				readyCallbacks.insert(KZGlobalService::callbacks.extract(curr));
+			}
+		}
+	}
+
+	for (auto &[_, callback] : readyCallbacks)
+	{
+		callback.callback(*callback.payload);
+	}
+}
+
 void KZGlobalService::OnPlayerAuthorized()
 {
 	KZ::API::events::PlayerJoin data;
@@ -201,10 +237,21 @@ void KZGlobalService::OnPlayerAuthorized()
 
 	KZGlobalService::SendMessage("player-join", data, callback);
 
-	if (KZGlobalService::currentMap.has_value())
+	u16 currentMapID = 0;
+
+	{
+		std::unique_lock currentMapLock(KZGlobalService::currentMapMutex);
+
+		if (KZGlobalService::currentMap.has_value())
+		{
+			currentMapID = KZGlobalService::currentMap->id;
+		}
+	}
+
+	if (currentMapID != 0)
 	{
 		KZ::API::events::WantPlayerRecords data;
-		data.mapId = KZGlobalService::currentMap->id;
+		data.mapId = currentMapID;
 		data.playerId = this->player->GetSteamId64();
 
 		auto callback = [player = this->player](const KZ::API::events::PlayerRecords &pbs)
@@ -270,10 +317,14 @@ bool KZGlobalService::SubmitRecord(u16 filterID, f64 time, u32 teleports, std::s
 	{
 		return false;
 	}
-	if (!KZGlobalService::currentMap.has_value())
+
 	{
-		META_CONPRINTF("[KZ::Global] Cannot submit record on non-global map.\n");
-		return false;
+		std::unique_lock currentMapLock(KZGlobalService::currentMapMutex);
+		if (!KZGlobalService::currentMap.has_value())
+		{
+			META_CONPRINTF("[KZ::Global] Cannot submit record on non-global map.\n");
+			return false;
+		}
 	}
 
 	KZ::API::events::NewRecord data;
@@ -411,12 +462,14 @@ void KZGlobalService::OnWebSocketMessage(const ix::WebSocketMessagePtr &message)
 				return;
 			}
 
-			auto callback = KZGlobalService::callbacks.find(messageId);
-
-			if (callback != KZGlobalService::callbacks.end())
 			{
-				callback->second(payload);
-				KZGlobalService::callbacks.erase(messageId);
+				std::unique_lock callbacksLock(KZGlobalService::callbacksMutex);
+				auto callback = KZGlobalService::callbacks.find(messageId);
+
+				if (callback != KZGlobalService::callbacks.end())
+				{
+					callback->second.payload = std::make_optional(payload);
+				}
 			}
 
 			break;
@@ -496,7 +549,8 @@ void KZGlobalService::CompleteWebSocketHandshake(const ix::WebSocketMessagePtr &
 
 	if (ack.map.has_value())
 	{
-		ack.map.swap(KZGlobalService::currentMap);
+		std::unique_lock currentMapLock(KZGlobalService::currentMapMutex);
+		KZGlobalService::currentMap.swap(ack.map);
 	}
 
 	META_CONPRINTF("[KZ::Global] Completed handshake!\n");
@@ -546,7 +600,7 @@ void KZGlobalService::SendMessage(const char *event, const T &data, CallbackFunc
 	payload.Set("event", event);
 	payload.Set("data", data);
 
-	KZGlobalService::callbacks[messageId] = [messageId, callback](Json responseJson)
+	auto callbackToStore = [messageId, callback](Json responseJson)
 	{
 		if (!responseJson.IsValid())
 		{
@@ -567,5 +621,10 @@ void KZGlobalService::SendMessage(const char *event, const T &data, CallbackFunc
 		callback(responseData);
 	};
 
+	{
+		std::unique_lock callbacksLock(KZGlobalService::callbacksMutex);
+		KZGlobalService::callbacks[messageId] = {callbackToStore, std::nullopt};
+	}
+
 	KZGlobalService::apiSocket->send(payload.ToString());
 }

src/kz/global/kz_global.h

@@ -63,6 +63,7 @@ class KZGlobalService : public KZBaseService
 	static void RestoreConVars();
 
 	static void OnActivateServer();
+	static void OnServerGamePostSimulate();
 
 	void OnPlayerAuthorized();
 	void OnClientDisconnect();
@@ -150,15 +151,37 @@ class KZGlobalService : public KZBaseService
 	 */
 	static inline u64 nextMessageId = 1;
 
+	/**
+	 * Protects `currentMap`.
+	 */
+	static inline std::mutex currentMapMutex {};
+
 	/**
 	 * The API's view of the current map.
 	 */
 	static inline std::optional<KZ::API::Map> currentMap {};
 
+	/**
+	 * Protects `callbacks`.
+	 */
+	static inline std::mutex callbacksMutex {};
+
+	struct StoredCallback
+	{
+		Callback<const Json &> callback;
+
+		/**
+		 * The payload received by the WebSocket.
+		 *
+		 * If this is `std::nullopt` then we haven't received a response yet and cannot execute the callback.
+		 */
+		std::optional<Json> payload;
+	};
+
 	/**
 	 * Callbacks to execute when the API sends messages with specific IDs.
 	 */
-	static inline std::unordered_map<u64, Callback<const Json &>> callbacks {};
+	static inline std::unordered_map<u64, StoredCallback> callbacks {};
 
 	/**
 	 * Called bx IXWebSocket whenever we receive a message.

src/utils/hooks.cpp

@@ -686,6 +686,7 @@ static_function CServerSideClientBase *Hook_ConnectClientPost(const char *pszNam
 static_function void Hook_ServerGamePostSimulate(const EventServerGamePostSimulate_t *)
 {
 	ProcessTimers();
+	KZGlobalService::OnServerGamePostSimulate();
 }
 
 static_function void Hook_BuildGameSessionManifest(const EventBuildGameSessionManifest_t *msg)