fix: DataStores now use one message per a place instead of many, which cuts down on subscription costs

Author: Quenty

pnpm-lock.yaml

@@ -81,6 +81,7 @@ local PromiseMaidUtils = require("PromiseMaidUtils")
 local PromiseRetryUtils = require("PromiseRetryUtils")
 local PromiseUtils = require("PromiseUtils")
 local Rx = require("Rx")
+local ServiceBag = require("ServiceBag")
 local Signal = require("Signal")
 local Symbol = require("Symbol")
 local ValueObject = require("ValueObject")
@@ -179,16 +180,41 @@ function DataStore.SetSessionLockingEnabled(self: DataStore, sessionLockingEnabl
 			self._sessionLockingEnabledHelper = DataStoreLockHelper.new(self)
 			self._maid._sessionLockingEnabledHelper = self._sessionLockingEnabledHelper
 		end
+	else
+		self._sessionLockingEnabledHelper = nil
+		self._maid._sessionLockingEnabledHelper = nil
+	end
+end
+
+--[=[
+	Sets session messaging enabled.
 
-		if not self._sessionMessagingEnabledHelper then
-			self._sessionMessagingEnabledHelper = DataStoreMessageHelper.new(self)
+	Currently this only works in conjunction with session locking, and allows for a session lock
+	to gracefully close when requested by another session.
+
+	@param isEnabled boolean
+	@param serviceBag ServiceBag -- Required when enabling
+]=]
+function DataStore.SetSessionMessagingEnabled(
+	self: DataStore,
+	isEnabled: boolean,
+	serviceBag: ServiceBag.ServiceBag
+): ()
+	if isEnabled then
+		assert(serviceBag, "Must provide serviceBag when enabling session messaging")
+
+		if
+			self._sessionMessagingEnabledHelper == nil
+			or self._sessionMessagingEnabledHelper:GetServiceBag() ~= serviceBag
+		then
+			self._maid._sessionMessagingEnabledHelper = nil
+			self._sessionMessagingEnabledHelper = nil
+
+			self._sessionMessagingEnabledHelper = DataStoreMessageHelper.new(serviceBag, self)
 			self._maid._sessionMessagingEnabledHelper = self._sessionMessagingEnabledHelper
 		end
 	else
-		self._sessionLockingEnabledHelper = nil
 		self._sessionMessagingEnabledHelper = nil
-
-		self._maid._sessionLockingEnabledHelper = nil
 		self._maid._sessionMessagingEnabledHelper = nil
 	end
 end
@@ -618,7 +644,11 @@ function DataStore._promiseGetAsyncNoCache(self: DataStore): Promise.Promise<()>
 							if self._sessionMessagingEnabledHelper and tryMessagingServiceSessionClose then
 								-- Gracefully kick to avoid losing memory
 								self._sessionMessagingEnabledHelper
-									:PromiseCloseSessionGraceful(lockResult.blockingSession.SessionId)
+									:PromiseCloseSessionGraceful(
+										lockResult.blockingSession.PlaceId,
+										lockResult.blockingSession.JobId,
+										lockResult.blockingSession.SessionId
+									)
 									:Then(function()
 										-- Give enough time for Roblox to replicate changes
 										-- We probably could bump back to the loop but this has slightly better error messages

src/datastore/src/Server/DataStore.lua

@@ -1,5 +1,4 @@
 --!strict
-
 --[=[
     @class DataStoreMessageHelper
 ]=]
@@ -8,8 +7,10 @@ local require = require(script.Parent.loader).load(script)
 
 local BaseObject = require("BaseObject")
 local MessagingServiceUtils = require("MessagingServiceUtils")
+local PlaceMessagingService = require("PlaceMessagingService")
 local Promise = require("Promise")
 local PromiseMaidUtils = require("PromiseMaidUtils")
+local ServiceBag = require("ServiceBag")
 
 local DEBUG_LOG = false
 
@@ -21,6 +22,8 @@ export type DataStoreMessageHelper =
 	typeof(setmetatable(
 		{} :: {
 			_dataStore: any,
+			_serviceBag: ServiceBag.ServiceBag,
+			_placeMessagingService: PlaceMessagingService.PlaceMessagingService,
 			_sessionClosedNotifications: { [string]: Promise.Promise<()> },
 		},
 		{} :: typeof({ __index = DataStoreMessageHelper })
@@ -38,30 +41,39 @@ export type CloseSessionComplete = {
 
 export type DataStoreMessage = CloseSessionRequest | CloseSessionComplete
 
-function DataStoreMessageHelper.new(dataStore: any): DataStoreMessageHelper
+function DataStoreMessageHelper.new(serviceBag: ServiceBag.ServiceBag, dataStore: any): DataStoreMessageHelper
 	local self: DataStoreMessageHelper = setmetatable(BaseObject.new() :: any, DataStoreMessageHelper)
 
+	self._serviceBag = assert(serviceBag, "No serviceBag")
+	self._placeMessagingService = self._serviceBag:GetService(PlaceMessagingService :: any)
+
 	self._dataStore = assert(dataStore, "No dataStore")
 	self._sessionClosedNotifications = {}
 
-	-- Explicitly don't give to maid so wecan disconnect the subscription before destroying
-	MessagingServiceUtils.promiseSubscribe(self:_getActiveStoreTopic(self._dataStore:GetSessionId()), function(data)
-		self:_handleSubscription(data)
-	end):Then(function(subscription)
-		if not self.Destroy then
-			subscription:Disconnect()
-		end
-
-		self._maid:GiveTask(function()
-			subscription:Disconnect()
-		end)
-	end)
+	self._maid:GiveTask(
+		self._placeMessagingService
+			:ObserveMessages(self:_getActiveStoreTopic(self._dataStore:GetSessionId()))
+			:Subscribe(function(data, metadata)
+				self:_handleSubscription(data, metadata)
+			end)
+	)
 
 	return self
 end
 
+--[=[
+	Returns the service bag used by this helper. (Used for comparison)
+
+	@return ServiceBag
+]=]
+function DataStoreMessageHelper.GetServiceBag(self: DataStoreMessageHelper): ServiceBag.ServiceBag
+	return self._serviceBag
+end
+
 function DataStoreMessageHelper.PromiseCloseSessionGraceful(
 	self: DataStoreMessageHelper,
+	placeId: number,
+	jobId: string,
 	sessionId: string
 ): Promise.Promise<()>
 	local promise: any = self._sessionClosedNotifications[sessionId]
@@ -84,7 +96,7 @@ function DataStoreMessageHelper.PromiseCloseSessionGraceful(
 		end))
 	end)
 
-	return self:PromiseMessage(sessionId, {
+	return self:PromiseMessage(placeId, jobId, sessionId, {
 		type = "close-session",
 		requesterSessionId = self._dataStore:GetSessionId(),
 	}):Then(function()
@@ -94,6 +106,8 @@ end
 
 function DataStoreMessageHelper.PromiseMessage(
 	self: DataStoreMessageHelper,
+	placeId: number,
+	jobId: string,
 	sessionId: string,
 	message: DataStoreMessage
 ): Promise.Promise<()>
@@ -103,14 +117,15 @@ function DataStoreMessageHelper.PromiseMessage(
 		print("[DataStoreMessageHelper] - Sending message:", MessagingServiceUtils.toHumanReadable(message))
 	end
 
-	return self._maid:GivePromise(MessagingServiceUtils.promisePublish(self:_getActiveStoreTopic(sessionId), message))
+	local topic = self:_getActiveStoreTopic(sessionId)
+	return self._maid:GivePromise(self._placeMessagingService:SendMessage(placeId, jobId, topic, message))
 end
 
 function DataStoreMessageHelper._handleSubscription(
 	self: DataStoreMessageHelper,
-	subscriptionData: MessagingServiceUtils.SubscriptionData
-)
-	local data = subscriptionData.Data
+	data: any,
+	metadata: PlaceMessagingService.PlacePacketMetadata
+): ()
 	if type(data) ~= "table" or type(data.type) ~= "string" then
 		warn(`[DataStoreMessageHelper] - Received malformed message: {MessagingServiceUtils.toHumanReadable(data)}`)
 		return
@@ -129,13 +144,10 @@ function DataStoreMessageHelper._handleSubscription(
 
 			closeSessionPromise:Then(function()
 				-- We could have GCed by now, but try to send off a notification to the requester
-				MessagingServiceUtils.promisePublish(
-					topic,
-					{
-						type = "close-session-complete",
-						senderId = senderId,
-					} :: CloseSessionComplete
-				)
+				self._placeMessagingService:SendMessageToAddress(metadata.from, topic, {
+					type = "close-session-complete",
+					senderId = senderId,
+				})
 			end)
 		end
 	elseif data.type == "close-session-complete" then

src/datastore/src/Server/DataStoreMessageHelper.lua

@@ -60,6 +60,7 @@ local Maid = require("Maid")
 local PendingPromiseTracker = require("PendingPromiseTracker")
 local Promise = require("Promise")
 local PromiseUtils = require("PromiseUtils")
+local ServiceBag = require("ServiceBag")
 
 local PlayerDataStoreManager = setmetatable({}, BaseObject)
 PlayerDataStoreManager.ClassName = "PlayerDataStoreManager"
@@ -73,6 +74,7 @@ export type PlayerDataStoreManager =
 	typeof(setmetatable(
 		{} :: {
 			_robloxDataStore: any,
+			_serviceBag: ServiceBag.ServiceBag,
 			_keyGenerator: KeyGenerator,
 			_datastores: { [PlayerUserId]: DataStore.DataStore },
 			_removing: { [PlayerUserId]: boolean },
@@ -93,6 +95,7 @@ export type PlayerDataStoreManager =
 	@return PlayerDataStoreManager
 ]=]
 function PlayerDataStoreManager.new(
+	serviceBag: ServiceBag.ServiceBag,
 	robloxDataStore: DataStore,
 	keyGenerator: KeyGenerator,
 	skipBindingToClose: boolean?
@@ -101,8 +104,9 @@ function PlayerDataStoreManager.new(
 
 	assert(type(skipBindingToClose) == "boolean" or skipBindingToClose == nil, "Bad skipBindingToClose")
 
-	self._robloxDataStore = robloxDataStore or error("No robloxDataStore")
-	self._keyGenerator = keyGenerator or error("No keyGenerator")
+	self._robloxDataStore = assert(robloxDataStore, "No robloxDataStore")
+	self._keyGenerator = assert(keyGenerator, "No keyGenerator")
+	self._serviceBag = assert(serviceBag, "No serviceBag")
 
 	self._maid._savingConns = Maid.new()
 
@@ -220,9 +224,10 @@ function PlayerDataStoreManager._createDataStore(
 
 	local maid = Maid.new()
 
-	-- TODO: Destroy DataStore after cleanup
+	-- DataStore is cleaned up very carefully in _removePlayerDataStore
 	local datastore = DataStore.new(self._robloxDataStore, self:_getKey(userId))
 	datastore:SetSessionLockingEnabled(true)
+	datastore:SetSessionMessagingEnabled(true, self._serviceBag)
 	datastore:SetUserIdList({ userId })
 
 	maid:GivePromise(datastore:PromiseSessionLockingFailed()):Then(function()
@@ -260,12 +265,16 @@ function PlayerDataStoreManager._createDataStore(
 	return datastore
 end
 
-function PlayerDataStoreManager._removePlayerDataStore(self: PlayerDataStoreManager, userId: PlayerUserId)
+function PlayerDataStoreManager._removePlayerDataStore(self: PlayerDataStoreManager, userId: PlayerUserId): ()
 	local datastore = self._datastores[userId]
 	if not datastore then
 		return
 	end
 
+	if self._removing[userId] then
+		return
+	end
+
 	self._removing[userId] = true
 
 	local removingPromises: { Promise.Promise<any?> } = {}

src/datastore/src/Server/PlayerDataStoreManager.lua

@@ -41,9 +41,10 @@ function PlayerDataStoreService.Init(self: PlayerDataStoreService, serviceBag: S
 
 	-- External
 	self._bindToCloseService = self._serviceBag:GetService(require("BindToCloseService"))
+	self._serviceBag:GetService(require("PlaceMessagingService"))
 
+	-- State
 	self._promiseStarted = self._maid:Add(Promise.new())
-
 	self._dataStoreName = "PlayerData"
 	self._dataStoreScope = "SaveData"
 end
@@ -140,7 +141,7 @@ function PlayerDataStoreService.PromiseManager(
 			return DataStorePromises.promiseDataStore(self._dataStoreName, self._dataStoreScope)
 		end)
 		:Then(function(dataStore)
-			local manager = self._maid:Add(PlayerDataStoreManager.new(dataStore, function(player)
+			local manager = self._maid:Add(PlayerDataStoreManager.new(self._serviceBag, dataStore, function(player)
 				if type(player) == "number" then
 					return tostring(player)
 				else

src/datastore/src/Server/PlayerDataStoreService.lua

@@ -30,7 +30,11 @@
   ],
   "dependencies": {
     "@quenty/loader": "workspace:*",
-    "@quenty/promise": "workspace:*"
+    "@quenty/maid": "workspace:*",
+    "@quenty/promise": "workspace:*",
+    "@quenty/rx": "workspace:*",
+    "@quenty/servicebag": "workspace:*",
+    "@quenty/statestack": "workspace:*"
   },
   "devDependencies": {
     "@quenty/loader": "workspace:*"