Drop-in Nevermore Session Locked Store (Not battle tested)

MemoryStoreUtils2.lua

--[=[
	@class MemoryStoreUtils2

	Utils for working with MemoryStore.
]=]

local require = require(script.Parent.loader).load(script)

local DEBUG_MAP = false

local Promise = require("Promise")

local MemoryStoreUtils2 = {}

function MemoryStoreUtils2.promiseStealFlagAtomic(map: MemoryStoreSortedMap, key: string, expirationSeconds: number)
	assert(typeof(map) == "Instance" and map:IsA("MemoryStoreSortedMap"), "Bad map")
	assert(typeof(key) == "string", "Bad key")
	assert(typeof(expirationSeconds) == "number", "Bad expirationSeconds")

	return Promise.spawn(function(resolve, reject)
		if DEBUG_MAP then
			print(("[MemoryStoreUtils.promiseStealFlagAtomic] - Stealing flag %q"):format(key))
		end

		local didWeStealTheFlag: boolean? = nil

		local ok, err = pcall(function()
			map:UpdateAsync(key, function(oldValue: boolean?)
				if oldValue == nil then
					-- Not locked! Write immediately to claim it as our own!
					didWeStealTheFlag = true
					return true
				else
					-- Locked! No-op.
					didWeStealTheFlag = false
					return nil
				end
			end, expirationSeconds)
		end)

		if not ok then
			if DEBUG_MAP then
				warn(("Failed to write map due to %q"):format(err or "nil"))
			end
			return reject(err or "[MemoryStoreUtils.promiseStealFlagAtomic] - Failed atomically steal flag on map")
		end
		return resolve(didWeStealTheFlag)
	end)
end

function MemoryStoreUtils2.promiseWriteFlag(map: MemoryStoreSortedMap, key: string, expirationSeconds: number)
	assert(typeof(map) == "Instance" and map:IsA("MemoryStoreSortedMap"), "Bad map")
	assert(typeof(key) == "string", "Bad key")
	assert(typeof(expirationSeconds) == "number", "Bad expirationSeconds")

	return Promise.spawn(function(resolve, reject)
		if DEBUG_MAP then
			print(("[MemoryStoreUtils.promiseWriteFlag] - Writing %q"):format(key))
		end

		local ok, err = pcall(function()
			map:SetAsync(key, true, expirationSeconds)
		end)

		if not ok then
			if DEBUG_MAP then
				warn(("Failed to write map due to %q"):format(err or "nil"))
			end
			return reject(err or "[MemoryStoreUtils.promiseWriteFlag] - Failed to write flag on map")
		end
		return resolve()
	end)
end

function MemoryStoreUtils2.promiseFreeFlag(map: MemoryStoreSortedMap, key: string)
	assert(typeof(map) == "Instance" and map:IsA("MemoryStoreSortedMap"), "Bad map")
	assert(typeof(key) == "string", "Bad key")

	return Promise.spawn(function(resolve, reject)
		if DEBUG_MAP then
			print(("[MemoryStoreUtils.promiseFreeFlag] - Freeing %q"):format(key))
		end

		local ok, err = pcall(function()
			map:RemoveAsync(key)
		end)

		if not ok then
			if DEBUG_MAP then
				warn(("Failed to free on map due to %q"):format(err or "nil"))
			end
			return reject(err or "[MemoryStoreUtils.promiseFreeFlag] - Failed to free on map")
		end
		return resolve()
	end)
end

return MemoryStoreUtils2

SessionedDataStore.lua

--[=[
	Wraps the datastore object to provide async cached loading and saving. See [SessionedDataStoreStage] for more API.

	Has the following features
	* Automatic saving every 5 minutes
	* Jitter (doesn't save all at the same time)
	* De-duplication (only updates data it needs)
	* Battle tested across multiple top games.

	```lua
	local playerMoneyValue = Instance.new("IntValue")
	playerMoneyValue.Value = 0

	local dataStore = SessionedDataStore.new(DataStoreService:GetDataStore("test"), "test-store")
	dataStore:Load("money", 0):Then(function(money)
		playerMoneyValue.Value = money
		dataStore:StoreOnValueChange("money", playerMoneyValue)
	end):Catch(function()
		-- TODO: Notify player
	end)
	```

	To use a datastore for a player, it's recommended you use the [PlayerDataStoreService]. This looks
	something like this. See [ServiceBag] for more information on service initialization.

	```lua
	local serviceBag = ServiceBag.new()
	local playerDataStoreService = serviceBag:GetService(require("PlayerDataStoreService"))

	serviceBag:Init()
	serviceBag:Start()

	local topMaid = Maid.new()

	local function handlePlayer(player)
		local maid = Maid.new()

		local playerMoneyValue = Instance.new("IntValue")
		playerMoneyValue.Name = "Money"
		playerMoneyValue.Value = 0
		playerMoneyValue.Parent = player

		maid:GivePromise(playerDataStoreService:PromiseDataStore(Players)):Then(function(dataStore)
			maid:GivePromise(dataStore:Load("money", 0))
				:Then(function(money)
					playerMoneyValue.Value = money
					maid:GiveTask(dataStore:StoreOnValueChange("money", playerMoneyValue))
				end)
		end)

		topMaid[player] = maid
	end
	Players.PlayerAdded:Connect(handlePlayer)
	Players.PlayerRemoving:Connect(function(player)
		topMaid[player] = nil
	end)
	for _, player in pairs(Players:GetPlayers()) do
		task.spawn(handlePlayer, player)
	end
	```

	@server
	@class SessionedDataStore
]=]

-- https://github.com/Qualadore/Tutorial-preventing-duplication-and-data-loss-from-trading.

local require = require(script.Parent.loader).load(script)

local MemoryStoreService = game:GetService("MemoryStoreService")

local DataStoreDeleteToken = require("DataStoreDeleteToken")
local DataStorePromises = require("DataStorePromises")
local DataStoreStage = require("DataStoreStage")
local Maid = require("Maid")
local MemoryStoreUtils2 = require("MemoryStoreUtils2")
local Promise = require("Promise")
local Signal = require("Signal")

local DEBUG_WRITING = false

local AUTO_SAVE_TIME = 60 * 5
local CHECK_DIVISION = 15
local JITTER = 20 -- Randomly assign jitter so if a ton of players join at once we don't hit the datastore at once

local SESSION_LOCK_TIME = 60 -- How long does the lock last?
local SESSION_LOCK_CHECK_SHORTWAVE_TIME = 1.5 -- How often should we check whether the lock has expired? Use this time the first few attempts.
local SESSION_LOCK_CHECK_BACKOFF_TIME = 6 -- How often do we check if the lock has cleared after a few attempts?
local SESSION_LOCK_ATTEMPTS_BEFORE_BACKOFF = 2
local SESSION_LOCK_ATTEMPTS_BEFORE_STEAL = 5
local SESSION_LOCK_WRITE_TIME = 30 -- Once stolen, how often should we refresh the lock?

local SessionedDataStore = setmetatable({}, DataStoreStage)
SessionedDataStore.ClassName = "SessionedDataStore"
SessionedDataStore.__index = SessionedDataStore

--[=[
	Constructs a new SessionedDataStore. See [SessionedDataStoreStage] for more API.
	@param robloxDataStore SessionedDataStore
	@param key string
	@return SessionedDataStore
]=]
function SessionedDataStore.new(robloxDataStore: DataStore, key)
	local self = setmetatable(DataStoreStage.new(), SessionedDataStore)

	assert(typeof(key) == "string" and #key > 0, "Bad key")
	assert(typeof(robloxDataStore) == "Instance" and robloxDataStore:IsA("DataStore"), "Bad robloxDataStore")

	self._key = key
	self._robloxDataStore = robloxDataStore

	-- TODO: Try to recover DataStore scope for our key. For now, this is sufficient for deduplication.
	local memoryStoreMapName = ("%sSessionLock"):format(robloxDataStore.Name)
	self._lockMemoryStoreMap = MemoryStoreService:GetSortedMap(memoryStoreMapName)

	--[=[
		Prop that fires when saving. Promise will resolve once saving is complete.
		@prop Saving Signal<Promise>
		@within SessionedDataStore
	]=]
	self.Saving = Signal.new() -- :Fire(promise)
	self._maid:GiveTask(self.Saving)

	task.spawn(function()
		while self.Destroy do
			for _ = 1, CHECK_DIVISION do
				task.wait(AUTO_SAVE_TIME / CHECK_DIVISION)
				if not self.Destroy then
					break
				end
			end

			if not self.Destroy then
				break
			end

			-- Apply additional jitter on auto-save
			task.wait(math.random(1, JITTER))

			if not self.Destroy then
				break
			end

			self:Save()
		end
	end)

	-- TODO: Free flag on destruction. May be a bad idea?
	self._maid:GiveTask(function()
		if self:DidAcquireExclusiveLock() then
			-- We don't use BindToClose here.
			-- When destroying this instance, we'll probably also be saving data.
			-- In which case, the save will hold the server open.
			-- DataStores are slow, MemoryStores are fast. We can probably write the key in that time.
			MemoryStoreUtils2.promiseFreeFlag(self._lockMemoryStoreMap, self._key)
		end
	end)

	return self
end

--[=[
	Returns the full path for the datastore
	@return string
]=]
function SessionedDataStore:GetFullPath()
	return ("RobloxDataStore@%s"):format(self._key)
end

--[=[
	Returns whether the datastore failed.
	@return boolean
]=]
function SessionedDataStore:DidLoadFail()
	if not self._loadPromise then
		return false
	end

	if self._loadPromise:IsRejected() then
		return true
	end

	return false
end

function SessionedDataStore:DidAcquireExclusiveLock()
	if not self._exclusivityPromise then
		return false
	end

	if self._exclusivityPromise:IsRejected() then
		return false
	end

	return true
end

--[=[
	Returns whether the datastore has loaded successfully.

	@return Promise<boolean>
]=]
function SessionedDataStore:PromiseLoadSuccessful()
	return self._maid:GivePromise(self:_promiseLoad()):Then(function()
		return true
	end, function()
		return false
	end)
end

--[=[
	Saves all stored data.
	@return Promise
]=]
function SessionedDataStore:Save()
	if not self:DidAcquireExclusiveLock() then
		warn("[SessionedDataStore] - Not saving, didn't acquire exclusive lock")
		return Promise.rejected("Didn't acquire sesion lock, not saving")
	end

	if self:DidLoadFail() then
		warn("[SessionedDataStore] - Not saving, failed to load")
		return Promise.rejected("Load not successful, not saving")
	end

	if DEBUG_WRITING then
		print("[SessionedDataStore.Save] - Starting save routine")
	end

	-- Avoid constructing promises for every callback down the datastore
	-- upon save.
	return (self:_promiseInvokeSavingCallbacks() or Promise.resolved()):Then(function()
		if not self:HasWritableData() then
			-- Nothing to save, don't update anything
			if DEBUG_WRITING then
				print("[SessionedDataStore.Save] - Not saving, nothing staged")
			end
			return nil
		else
			return self:_saveData(self:GetNewWriter())
		end
	end)
end

--[=[
	Loads data. This returns the originally loaded data.
	This makes no guarantee about the availability of saving this data. It merely says we have it available.
	@param keyName string
	@param defaultValue any?
	@return any?
]=]
function SessionedDataStore:Load(keyName, defaultValue)
	return self:_promiseLoad():Then(function(data)
		return self:_afterLoadGetAndApplyStagedData(keyName, data, defaultValue)
	end)
end

function SessionedDataStore:_saveData(writer)
	local maid = Maid.new()

	local promise = Promise.new()
	promise:Resolve(maid:GivePromise(DataStorePromises.updateAsync(self._robloxDataStore, self._key, function(data)
		if promise:IsRejected() then
			-- Cancel if we have another request
			return nil
		end

		data = writer:WriteMerge(data or {})
		assert(data ~= DataStoreDeleteToken, "Cannot delete from UpdateAsync")

		if DEBUG_WRITING then
			print("[SessionedDataStore] - Writing", game:GetService("HttpService"):JSONEncode(data))
		end

		return data
	end)):Catch(function(err)
		-- Might be caused by Maid rejecting state
		warn("[SessionedDataStore] - Failed to UpdateAsync data", err)
		return Promise.rejected(err)
	end))

	self._maid._saveMaid = maid

	if self.Saving.Destroy then
		self.Saving:Fire(promise)
	end

	return promise
end

function SessionedDataStore:_promiseLoad()
	if self._loadPromise then
		return self._loadPromise
	end

	self._loadPromise = self:_promiseExclusivity()
		:Then(function()
			return self._maid:GivePromise(DataStorePromises.getAsync(self._robloxDataStore, self._key))
		end)
		:Then(function(data)
			if data == nil then
				return {}
			elseif type(data) == "table" then
				return data
			else
				return Promise.rejected("Failed to load data. Wrong type '" .. type(data) .. "'")
			end
		end, function(err)
			-- Log:
			warn("[SessionedDataStore] - Failed to GetAsync data", err)
			return Promise.rejected(err)
		end)

	return self._loadPromise
end

function SessionedDataStore:_promiseExclusivity()
	if self._exclusivityPromise then
		return self._exclusivityPromise
	end

	local promise = Promise.new()

	local maid = Maid.new()
	promise:Finally(function()
		maid:Destroy()
	end)

	-- Attempt to get the lock.
	maid:GiveTask(task.spawn(function()
		local attempts = 0

		-- TODO: Reject after 'n' failures to claim the lock?
		-- Try to steal the flag.
		while true do
			maid._flagPromise =
				MemoryStoreUtils2.promiseStealFlagAtomic(self._lockMemoryStoreMap, self._key, SESSION_LOCK_TIME)
			local isOk, res = maid._flagPromise:Yield()

			attempts += 1

			-- TODO: Handle promise rejection from maid cleanup.
			if not isOk then
				warn(("[SessionedDataStore] Failed to query memory store. %q"):format(res))
			end

			-- If promise is ok, 'res' is true when we stole the flag on the MemoryStore.
			-- Otherwise 'res' is probably a string errorcode or something.
			if isOk and res then
				-- All good! If we got the lock immediately, there was no session lock.
				-- Tell consumers this, so they can decide if their data is stale and they want to reload.
				return promise:Resolve(attempts == 1)
			end

			-- We couldn't get exclusivity this time. Try again after a delay.
			if attempts < SESSION_LOCK_ATTEMPTS_BEFORE_BACKOFF then
				-- This is typical. Sometimes we reject straight away, no big deal.
				task.wait(SESSION_LOCK_CHECK_SHORTWAVE_TIME)
			elseif attempts < SESSION_LOCK_ATTEMPTS_BEFORE_STEAL then
				-- Now we're a bit more concerned.
				warn(("[SessionedDataStore] Lock on %q hasn't expired, after %i attempts!"):format(self._key, attempts))
				task.wait(SESSION_LOCK_CHECK_BACKOFF_TIME)
			else
				-- Guh. Whatever. Steal the lock, that other server either crashed or is taking too long.
				warn(("[SessionedDataStore] Stealing lock for %q - this is taking too long."):format(self._key))
				return promise:Resolve(false)
			end
		end
	end))

	self._exclusivityPromise = self._maid:GivePromise(promise)

	-- Once we've got the exclusivity, aim to keep it.
	self._exclusivityPromise:Tap(function(_didStoreStartUnlocked: boolean)
		self._maid:GiveTask(task.defer(function()
			while true do
				-- We just got it. Wait upfront before refreshing.
				task.wait(SESSION_LOCK_WRITE_TIME)

				local isOk = self._maid
					:GivePromise(
						MemoryStoreUtils2.promiseWriteFlag(self._lockMemoryStoreMap, self._key, SESSION_LOCK_TIME)
					)
					:Yield()

				if not isOk then
					warn(("[SessionedDataStore] Failed to write flag %q"):format(self._key))
				end
			end
		end))
	end)

	return self._exclusivityPromise
end

return SessionedDataStore