From c67253a0e3e89a89bd95c0304c4237c5b7634b7e Mon Sep 17 00:00:00 2001
From: nezuo <hinckleytmicah@gmail.com>
Date: Mon, 23 Oct 2023 18:01:57 -0700
Subject: [PATCH] Call beforeCloseCallback in a promise

---
 src/Document.lua      | 56 +++++++++++++++++++++++--------------------
 src/Document.test.lua | 22 +++++++++++------
 2 files changed, 45 insertions(+), 33 deletions(-)

diff --git a/src/Document.lua b/src/Document.lua
index c56a567..0b6cd3a 100644
--- a/src/Document.lua
+++ b/src/Document.lua
@@ -1,5 +1,6 @@
 local Compression = require(script.Parent.Compression)
 local freezeDeep = require(script.Parent.freezeDeep)
+local Promise = require(script.Parent.Parent.Promise)
 
 --[=[
 	@class Document
@@ -81,7 +82,7 @@ end
 	:::
 
 	:::warning
-	Throws an error or yields if the beforeClose callback does.
+	If the beforeClose callback errors, the returned promise will reject and the data will not be saved.
 	:::
 
 	@return Promise<()>
@@ -89,42 +90,45 @@ end
 function Document:close()
 	assert(not self.closed and not self.callingBeforeClose, "Cannot close a closed document")
 
-	local function close()
-		self.closed = true
-
-		self.collection.openDocuments[self.key] = nil
-
-		self.collection.autoSave:removeDocument(self)
-	end
+	local promise = Promise.resolve()
 
 	if self.beforeCloseCallback ~= nil then
 		self.callingBeforeClose = true
 
-		local ok, err = pcall(self.beforeCloseCallback)
-
-		if not ok then
-			close()
-			error(`beforeClose callback threw error: {tostring(err)}`)
-		end
+		promise = Promise.new(function(resolve, reject)
+			local ok, err = pcall(self.beforeCloseCallback)
 
-		self.callingBeforeClose = false
+			if not ok then
+				reject(`beforeClose callback threw error: {tostring(err)}`)
+			else
+				resolve()
+			end
+		end)
 	end
 
-	close()
+	return promise
+		:finally(function()
+			self.closed = true
 
-	return self.collection.data:save(self.collection.dataStore, self.key, function(value)
-		if value.lockId ~= self.lockId then
-			return "fail", "The session lock was stolen"
-		end
+			self.collection.openDocuments[self.key] = nil
 
-		local scheme, compressed = Compression.compress(self.data)
+			self.collection.autoSave:removeDocument(self)
+		end)
+		:andThen(function()
+			return self.collection.data:save(self.collection.dataStore, self.key, function(value)
+				if value.lockId ~= self.lockId then
+					return "fail", "The session lock was stolen"
+				end
 
-		value.compressionScheme = scheme
-		value.data = compressed
-		value.lockId = nil
+				local scheme, compressed = Compression.compress(self.data)
 
-		return "succeed", value
-	end)
+				value.compressionScheme = scheme
+				value.data = compressed
+				value.lockId = nil
+
+				return "succeed", value
+			end)
+		end)
 end
 
 --[=[
diff --git a/src/Document.test.lua b/src/Document.test.lua
index a6ecfbf..017ea1f 100644
--- a/src/Document.test.lua
+++ b/src/Document.test.lua
@@ -44,11 +44,14 @@ return function(x)
 		local pendingSave = document:save()
 		local pendingClose = document:close() -- This should override the pending save.
 
+		context.dataStoreService.budget.budgets[Enum.DataStoreRequestType.UpdateAsync] = 1
+
 		context.dataStoreService.yield:stopYield()
 
-		assert(pendingSave == pendingClose, "save and close didn't merge")
+		local values = Promise.all({ ongoingSave, pendingSave }):expect()
 
-		local values = Promise.all({ ongoingSave, pendingSave, pendingClose }):expect()
+		-- If the save and close don't merge, this will throw because there isn't enough budget to close.
+		pendingClose:now("save and close didn't merge"):expect()
 
 		-- save and close should never resolve with a value.
 		-- It's checked in this test to make sure it works with save merging.
@@ -221,7 +224,7 @@ return function(x)
 			end)
 
 			shouldThrow(function()
-				document:close()
+				document:close():expect()
 			end, "beforeClose callback threw error")
 		end)
 
@@ -233,11 +236,11 @@ return function(x)
 			end)
 
 			shouldThrow(function()
-				document:close()
+				document:close():expect()
 			end, "beforeClose callback threw error")
 		end)
 
-		x.test("closes document even if document error", function(context)
+		x.test("closes document even if beforeClose errors", function(context)
 			local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS)
 
 			local promise = collection:load("document")
@@ -248,14 +251,19 @@ return function(x)
 			end)
 
 			shouldThrow(function()
-				document:close()
+				document:close():expect()
 			end)
 
-			assert(collection:load("document") ~= promise, "collection:load should return a new promise")
+			local secondPromise = collection:load("document")
+
+			assert(secondPromise ~= promise, "collection:load should return a new promise")
 
 			shouldThrow(function()
 				document:write({ foo = "baz" })
 			end, "Cannot write to a closed document")
+
+			-- Ignore the could not acquire lock error.
+			secondPromise:catch(function() end)
 		end)
 
 		x.test("saves new data", function(context)