diff --git a/src/Document.lua b/src/Document.lua index 7f4e902..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 @@ -15,6 +16,7 @@ function Document.new(collection, key, validate, lockId, data) lockId = lockId, data = data, closed = false, + callingBeforeClose = false, }, Document) end @@ -55,7 +57,7 @@ end @return Promise<()> ]=] function Document:save() - assert(not self.closed, "Cannot save a closed document") + assert(not self.closed and not self.callingBeforeClose, "Cannot save a closed document") return self.collection.data:save(self.collection.dataStore, self.key, function(value) if value.lockId ~= self.lockId then @@ -79,30 +81,70 @@ end Throws an error if the document was closed. ::: + :::warning + If the beforeClose callback errors, the returned promise will reject and the data will not be saved. + ::: + @return Promise<()> ]=] function Document:close() - assert(not self.closed, "Cannot close a closed document") + assert(not self.closed and not self.callingBeforeClose, "Cannot close a closed document") - self.closed = true + local promise = Promise.resolve() - self.collection.openDocuments[self.key] = nil + if self.beforeCloseCallback ~= nil then + self.callingBeforeClose = true - self.collection.autoSave:removeDocument(self) + promise = Promise.new(function(resolve, reject) + local ok, err = pcall(self.beforeCloseCallback) - 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 + if not ok then + reject(`beforeClose callback threw error: {tostring(err)}`) + else + resolve() + end + end) + end - local scheme, compressed = Compression.compress(self.data) + return promise + :finally(function() + self.closed = true - value.compressionScheme = scheme - value.data = compressed - value.lockId = nil + self.collection.openDocuments[self.key] = nil - return "succeed", value - 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 + + local scheme, compressed = Compression.compress(self.data) + + value.compressionScheme = scheme + value.data = compressed + value.lockId = nil + + return "succeed", value + end) + end) +end + +--[=[ + Sets a callback that is run inside `document:close` before it saves. The document can be read and written to in the + callback. + + :::warning + Throws an error if it was called previously. + ::: + + @param callback () -> () +]=] +function Document:beforeClose(callback) + assert(self.beforeCloseCallback == nil, "Document:beforeClose can only be called once") + + self.beforeCloseCallback = callback end return Document diff --git a/src/Document.test.lua b/src/Document.test.lua index 6cde4a7..017ea1f 100644 --- a/src/Document.test.lua +++ b/src/Document.test.lua @@ -10,6 +10,7 @@ local DEFAULT_OPTIONS = { } return function(x) + local assertEqual = x.assertEqual local shouldThrow = x.shouldThrow x.test("it should not merge close into save when save is running", function(context) @@ -43,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. @@ -200,4 +204,80 @@ return function(x) promise:expect() end) + + x.nested("Document:beforeClose", function() + x.test("throws when setting twice", function(context) + local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect() + + document:beforeClose(function() end) + + shouldThrow(function() + document:beforeClose(function() end) + end, "Document:beforeClose can only be called once") + end) + + x.test("throws when calling close in callback", function(context) + local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect() + + document:beforeClose(function() + document:close() + end) + + shouldThrow(function() + document:close():expect() + end, "beforeClose callback threw error") + end) + + x.test("throws when calling save in callback", function(context) + local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect() + + document:beforeClose(function() + document:save() + end) + + shouldThrow(function() + document:close():expect() + end, "beforeClose callback threw error") + end) + + x.test("closes document even if beforeClose errors", function(context) + local collection = context.lapis.createCollection("collection", DEFAULT_OPTIONS) + + local promise = collection:load("document") + local document = promise:expect() + + document:beforeClose(function() + error("error") + end) + + shouldThrow(function() + document:close():expect() + end) + + 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) + local document = context.lapis.createCollection("collection", DEFAULT_OPTIONS):load("document"):expect() + + document:beforeClose(function() + document:read() -- This checks that read doesn't error in the callback. + + document:write({ foo = "new" }) + end) + + document:close():expect() + + assertEqual(context.read("collection", "document").data.foo, "new") + end) + end) end diff --git a/src/init.lua b/src/init.lua index be73f57..845b1d6 100644 --- a/src/init.lua +++ b/src/init.lua @@ -33,6 +33,7 @@ export type Document = { write: (self: Document, T) -> (), save: (self: Document) -> PromiseTypes.TypedPromise<()>, close: (self: Document) -> PromiseTypes.TypedPromise<()>, + beforeClose: (self: Document, callback: () -> ()) -> (), } --[=[ diff --git a/wally.lock b/wally.lock index f109279..92ee559 100644 --- a/wally.lock +++ b/wally.lock @@ -14,10 +14,10 @@ dependencies = [] [[package]] name = "nezuo/lapis" -version = "0.2.4" -dependencies = [["Promise", "evaera/promise@4.0.0"], ["DataStoreServiceMock", "nezuo/data-store-service-mock@0.3.0"], ["Midori", "nezuo/midori@0.1.2"]] +version = "0.2.5" +dependencies = [["Promise", "evaera/promise@4.0.0"], ["DataStoreServiceMock", "nezuo/data-store-service-mock@0.3.0"], ["Midori", "nezuo/midori@0.1.3"]] [[package]] name = "nezuo/midori" -version = "0.1.2" +version = "0.1.3" dependencies = [] diff --git a/wally.toml b/wally.toml index 1b80dfa..23a108d 100644 --- a/wally.toml +++ b/wally.toml @@ -10,5 +10,5 @@ realm = "server" Promise = "evaera/promise@4.0.0" [dev-dependencies] -Midori = "nezuo/midori@0.1.2" +Midori = "nezuo/midori@0.1.3" DataStoreServiceMock = "nezuo/data-store-service-mock@0.3.0"