Merge pull request #13 from nezuo/fix-save-merging

Fix save merging

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ## Unreleased Changes
 * Remove write cooldown throttling since write cooldowns [were removed](https://devforum.roblox.com/t/removal-of-6s-cool-down-for-data-stores/2436230)
+* Fix save merging algorithm ([#13])
+
+[#13]: https://github.com/nezuo/lapis/pull/13
 
 ## 0.2.1 - June 10, 2023
 * Move TestEZ and DataStoreServiceMock to dev dependencies

src/AutoSave.lua

@@ -1,7 +1,5 @@
 local RunService = game:GetService("RunService")
 
-local Promise = require(script.Parent.Parent.Promise)
-
 local UPDATE_INTERVAL = 5 * 60
 
 local AutoSave = {}
@@ -41,16 +39,7 @@ function AutoSave:start()
 			self.documents[#self.documents]:close()
 		end
 
-		local promises = {}
-
-		-- This will wait for documents that closed before BindToClose was called.
-		for _, pendingSaves in self.data:getPendingSaves() do
-			for _, pendingSave in pendingSaves do
-				table.insert(promises, pendingSave.promise)
-			end
-		end
-
-		Promise.allSettled(promises):await()
+		self.data:waitForOngoingSaves():await()
 	end)
 end
 

src/Data/init.lua

@@ -8,20 +8,41 @@ function Data.new(config)
 	return setmetatable({
 		config = config,
 		throttle = Throttle.new(config),
-		pendingSaves = {},
+		ongoingSaves = {},
 	}, Data)
 end
 
-function Data:getPendingSave(dataStore, key)
-	if self.pendingSaves[dataStore] == nil or self.pendingSaves[dataStore][key] == nil then
+function Data:waitForOngoingSave(dataStore, key)
+	if self.ongoingSaves[dataStore] == nil or self.ongoingSaves[dataStore][key] == nil then
 		return Promise.resolve()
 	end
 
-	return self.pendingSaves[dataStore][key].promise
+	local ongoingSave = self.ongoingSaves[dataStore][key]
+
+	return Promise.allSettled({
+		ongoingSave.promise,
+		if ongoingSave.pendingSave ~= nil then ongoingSave.pendingSave.promise else nil,
+	})
+end
+
+function Data:waitForOngoingSaves()
+	local promises = {}
+
+	for _, ongoingSave in self.ongoingSaves do
+		table.insert(
+			promises,
+			Promise.allSettled({
+				ongoingSave.promise,
+				if ongoingSave.pendingSave ~= nil then ongoingSave.pendingSave.promise else nil,
+			})
+		)
+	end
+
+	return Promise.allSettled(promises)
 end
 
 function Data:load(dataStore, key, transform)
-	return self:getPendingSave(dataStore, key):andThen(function()
+	return self:waitForOngoingSave(dataStore, key):andThen(function()
 		local attempts = self.config:get("loadAttempts")
 		local retryDelay = self.config:get("loadRetryDelay")
 
@@ -30,43 +51,45 @@ function Data:load(dataStore, key, transform)
 end
 
 function Data:save(dataStore, key, transform)
-	if self.pendingSaves[dataStore] == nil then
-		self.pendingSaves[dataStore] = {}
+	if self.ongoingSaves[dataStore] == nil then
+		self.ongoingSaves[dataStore] = {}
 	end
 
-	local pendingSave = self.pendingSaves[dataStore][key]
-
-	if pendingSave ~= nil then
-		pendingSave.transform = transform
-
-		return pendingSave.promise
-	else
-		self.pendingSaves[dataStore][key] = { transform = transform }
+	local ongoingSave = self.ongoingSaves[dataStore][key]
 
+	if ongoingSave == nil then
 		local attempts = self.config:get("saveAttempts")
+		local promise = self.throttle:updateAsync(dataStore, key, transform, attempts):finally(function()
+			self.ongoingSaves[dataStore][key] = nil
 
-		local promise = self.throttle
-			:updateAsync(dataStore, key, function(...)
-				return self.pendingSaves[dataStore][key].transform(...)
-			end, attempts)
-			:finally(function()
-				self.pendingSaves[dataStore][key] = nil
-
-				if next(self.pendingSaves[dataStore]) == nil then
-					self.pendingSaves[dataStore] = nil
-				end
-			end)
+			if next(self.ongoingSaves[dataStore]) == nil then
+				self.ongoingSaves[dataStore] = nil
+			end
+		end)
 
 		if promise:getStatus() == Promise.Status.Started then
-			self.pendingSaves[dataStore][key].promise = promise
+			self.ongoingSaves[dataStore][key] = { promise = promise }
 		end
 
 		return promise
-	end
-end
+	elseif ongoingSave.pendingSave == nil then
+		local pendingSave = { transform = transform }
+
+		local function save()
+			return self:save(dataStore, key, pendingSave.transform)
+		end
+
+		-- promise:finally(save) can't be used because if the ongoingSave promise rejects, so will the promise returned from finally.
+		pendingSave.promise = ongoingSave.promise:andThen(save, save)
 
-function Data:getPendingSaves()
-	return self.pendingSaves
+		ongoingSave.pendingSave = pendingSave
+
+		return pendingSave.promise
+	else
+		ongoingSave.pendingSave.transform = transform
+
+		return ongoingSave.pendingSave.promise
+	end
 end
 
 return Data

src/Document.spec.lua

@@ -10,25 +10,46 @@ local DEFAULT_OPTIONS = {
 }
 
 return function()
-	itFIXME("combines save and close requests", function(context)
-		local document = context.lapis.createCollection("fff", DEFAULT_OPTIONS):load("doc"):expect()
+	it("it should not merge close into save when save is running", function(context)
+		local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("doc"):expect()
 
-		document:write({
-			foo = "updated value",
-		})
+		-- It's not safe to merge saves when UpdateAsync is running.
+		-- This will yield the UpdateAsync call until stopYield is called.
+		context.dataStoreService.yield:startYield()
 
 		local save = document:save()
+		document:write({ foo = "new" })
 		local close = document:close()
 
+		context.dataStoreService.yield:stopYield()
+
 		Promise.all({ save, close }):expect()
 
-		expect(save).to.equal(close)
+		local saved = context.read("collection", "doc")
 
-		local saved = context.read("fff", "doc")
+		-- If data.foo == "bar", that means the close was merged with the save when it wasn't safe to.
+		expect(saved.data.foo).to.equal("new")
+	end)
+
+	it("it should merge pending saves", function(context)
+		local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("doc"):expect()
+
+		context.dataStoreService.yield:startYield()
+
+		local ongoingSave = document:save()
+
+		local pendingSave = document:save()
+		local pendingClose = document:close() -- This should override the pending save.
+
+		context.dataStoreService.yield:stopYield()
+
+		expect(pendingSave).to.equal(pendingClose)
+
+		Promise.all({ ongoingSave, pendingSave, pendingClose }):expect()
+
+		local saved = context.read("collection", "doc")
 
-		expect(saved).to.be.a("table")
 		expect(saved.lockId).never.to.be.ok()
-		expect(saved.data.foo).to.equal("updated value")
 	end)
 
 	it("saves data", function(context)