Merge pull request #35 from nezuo/remove-close-twice-error

Don't error when calling close again

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Lapis Changelog
 
 ## Unreleased Changes
+* `Document:close` no longer errors when called again and instead returns the original promise. ([#35])
+  * This is so it won't error when called from `PlayerRemoving` if `game:BindToClose` happens to run first.
+
+[#35]: https://github.com/nezuo/lapis/pull/35
 
 ## 0.2.8 - December 27, 2023
 * Removed internal compression code since compression is no longer planned ([#31])

src/Document.lua

@@ -1,14 +1,18 @@
 local freezeDeep = require(script.Parent.freezeDeep)
 local Promise = require(script.Parent.Parent.Promise)
 
-local function runCallback(name, callback)
+local function runCallback(document, name, callback)
 	if callback == nil then
 		return Promise.resolve()
 	end
 
+	document.callingCallback = name
+
 	return Promise.new(function(resolve, reject)
 		local ok, message = pcall(callback)
 
+		document.callingCallback = nil
+
 		if not ok then
 			reject(`{name} callback threw error: {message}`)
 		else
@@ -32,7 +36,6 @@ function Document.new(collection, key, validate, lockId, data, userIds)
 		data = data,
 		userIds = userIds,
 		closed = false,
-		callingCloseCallbacks = false,
 	}, Document)
 end
 
@@ -71,6 +74,8 @@ end
 	@param userId number
 ]=]
 function Document:addUserId(userId)
+	assert(not self.closed, "Cannot add user id to a closed document")
+
 	if table.find(self.userIds, userId) == nil then
 		table.insert(self.userIds, userId)
 	end
@@ -84,6 +89,8 @@ end
 	@param userId number
 ]=]
 function Document:removeUserId(userId)
+	assert(not self.closed, "Cannot remove user id from a closed document")
+
 	local index = table.find(self.userIds, userId)
 
 	if index ~= nil then
@@ -105,9 +112,10 @@ end
 	@return Promise<()>
 ]=]
 function Document:save()
-	assert(not self.closed and not self.callingCloseCallbacks, "Cannot save a closed document")
+	assert(not self.closed, "Cannot save a closed document")
+	assert(self.callingCallback == nil, `Cannot save in {self.callingCallback} callback`)
 
-	return runCallback("beforeSave", self.beforeSaveCallback):andThen(function()
+	return runCallback(self, "beforeSave", self.beforeSaveCallback):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"
@@ -124,9 +132,7 @@ end
 	Saves the document and removes the session lock. The document is unusable after calling. If a save is currently in
 	progress it will close the document instead.
 
-	:::warning
-	Throws an error if the document was closed.
-	:::
+	If called again, it will return the promise from the original call.
 
 	:::warning
 	If the beforeSave or beforeClose callbacks error, the returned promise will reject and the data will not be saved.
@@ -135,31 +141,33 @@ end
 	@return Promise<()>
 ]=]
 function Document:close()
-	assert(not self.closed and not self.callingCloseCallbacks, "Cannot close a closed document")
+	assert(self.callingCallback == nil, `Cannot close in {self.callingCallback} callback`)
 
-	self.callingCloseCallbacks = true
+	if self.closePromise == nil then
+		self.closePromise = runCallback(self, "beforeSave", self.beforeSaveCallback)
+			:andThenCall(runCallback, self, "beforeClose", self.beforeCloseCallback)
+			:finally(function()
+				self.closed = true
 
-	return runCallback("beforeSave", self.beforeSaveCallback)
-		:andThenCall(runCallback, "beforeClose", self.beforeCloseCallback)
-		:finally(function()
-			self.closed = true
+				self.collection.openDocuments[self.key] = nil
 
-			self.collection.openDocuments[self.key] = nil
-
-			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
+				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.data = self.data
-				value.lockId = nil
+					value.data = self.data
+					value.lockId = nil
 
-				return "succeed", value, self.userIds
+					return "succeed", value, self.userIds
+				end)
 			end)
-		end)
+	end
+
+	return self.closePromise
 end
 
 --[=[

src/Document.test.lua

@@ -101,7 +101,7 @@ return function(x)
 		end, "foo must be a string")
 	end)
 
-	x.test("throws when writing/saving/closing a closed document", function(context)
+	x.test("methods throw when called on a closed document", function(context)
 		local document = context.lapis.createCollection("5", DEFAULT_OPTIONS):load("doc"):expect()
 
 		local promise = document:close()
@@ -115,12 +115,24 @@ return function(x)
 		end, "Cannot save a closed document")
 
 		shouldThrow(function()
-			document:close()
-		end, "Cannot close a closed document")
+			document:addUserId(1234)
+		end, "Cannot add user id to a closed document")
+
+		shouldThrow(function()
+			document:removeUserId(1234)
+		end, "Cannot remove user id from a closed document")
 
 		promise:expect()
 	end)
 
+	x.test("close returns first promise when called again", function(context)
+		local document = context.lapis.createCollection("col", DEFAULT_OPTIONS):load("doc"):expect()
+
+		local promise = document:close()
+
+		assertEqual(promise, document:close())
+	end)
+
 	x.test("loads with default data", function(context)
 		local document = context.lapis.createCollection("o", DEFAULT_OPTIONS):load("a"):expect()
 

src/init.test.lua

@@ -268,8 +268,8 @@ return function(x)
 		-- Verify each document has been closed.
 		for _, document in { one, two, three } do
 			shouldThrow(function()
-				document:close():expect()
-			end, "Cannot close a closed document")
+				document:save():expect()
+			end, "Cannot save a closed document")
 		end
 
 		context.dataStoreService.yield:stopYield()