Use SerializationService as a fallback for when patch application fails (#1030)

plugin/src/ServeSession.lua

@@ -1,11 +1,15 @@
 local StudioService = game:GetService("StudioService")
 local RunService = game:GetService("RunService")
+local ChangeHistoryService = game:GetService("ChangeHistoryService")
+local SerializationService = game:GetService("SerializationService")
+local Selection = game:GetService("Selection")
 
 local Packages = script.Parent.Parent.Packages
 local Log = require(Packages.Log)
 local Fmt = require(Packages.Fmt)
 local t = require(Packages.t)
 local Promise = require(Packages.Promise)
+local Timer = require(script.Parent.Timer)
 
 local ChangeBatcher = require(script.Parent.ChangeBatcher)
 local encodePatchUpdate = require(script.Parent.ChangeBatcher.encodePatchUpdate)
@@ -95,6 +99,8 @@ function ServeSession.new(options)
 		__changeBatcher = changeBatcher,
 		__statusChangedCallback = nil,
 		__connections = connections,
+		__precommitCallbacks = {},
+		__postcommitCallbacks = {},
 	}
 
 	setmetatable(self, ServeSession)
@@ -125,12 +131,46 @@ function ServeSession:setConfirmCallback(callback)
 	self.__userConfirmCallback = callback
 end
 
+--[=[
+	Hooks a function to run before patch application.
+	The provided function is called with the incoming patch and an InstanceMap
+	as parameters.
+]=]
 function ServeSession:hookPrecommit(callback)
-	return self.__reconciler:hookPrecommit(callback)
+	table.insert(self.__precommitCallbacks, callback)
+	Log.trace("Added precommit callback: {}", callback)
+
+	return function()
+		-- Remove the callback from the list
+		for i, cb in self.__precommitCallbacks do
+			if cb == callback then
+				table.remove(self.__precommitCallbacks, i)
+				Log.trace("Removed precommit callback: {}", callback)
+				break
+			end
+		end
+	end
 end
 
+--[=[
+	Hooks a function to run after patch application.
+	The provided function is called with the applied patch, the current
+	InstanceMap, and a PatchSet containing any unapplied changes.
+]=]
 function ServeSession:hookPostcommit(callback)
-	return self.__reconciler:hookPostcommit(callback)
+	table.insert(self.__postcommitCallbacks, callback)
+	Log.trace("Added postcommit callback: {}", callback)
+
+	return function()
+		-- Remove the callback from the list
+		for i, cb in self.__postcommitCallbacks do
+			if cb == callback then
+				table.remove(self.__postcommitCallbacks, i)
+				Log.trace("Removed postcommit callback: {}", callback)
+				break
+			end
+		end
+	end
 end
 
 function ServeSession:start()
@@ -206,6 +246,169 @@ function ServeSession:__onActiveScriptChanged(activeScript)
 	self.__apiContext:open(scriptId)
 end
 
+function ServeSession:__replaceInstances(idList)
+	if #idList == 0 then
+		return true, PatchSet.newEmpty()
+	end
+	-- It would be annoying if selection went away, so we try to preserve it.
+	local selection = Selection:Get()
+	local selectionMap = {}
+	for i, instance in selection do
+		selectionMap[instance] = i
+	end
+
+	-- TODO: Should we do this in multiple requests so we can more granularly mark failures?
+	local modelSuccess, replacements = self.__apiContext
+		:serialize(idList)
+		:andThen(function(response)
+			Log.debug("Deserializing results from serialize endpoint")
+			local objects = SerializationService:DeserializeInstancesAsync(response.modelContents)
+			if not objects[1] then
+				return Promise.reject("Serialize endpoint did not deserialize into any Instances")
+			end
+			if #objects[1]:GetChildren() ~= #idList then
+				return Promise.reject("Serialize endpoint did not return the correct number of Instances")
+			end
+
+			local instanceMap = {}
+			for _, item in objects[1]:GetChildren() do
+				instanceMap[item.Name] = item.Value
+			end
+			return instanceMap
+		end)
+		:await()
+
+	local refSuccess, refPatch = self.__apiContext
+		:refPatch(idList)
+		:andThen(function(response)
+			return response.patch
+		end)
+		:await()
+
+	if not (modelSuccess and refSuccess) then
+		return false
+	end
+
+	for id, replacement in replacements do
+		local oldInstance = self.__instanceMap.fromIds[id]
+		self.__instanceMap:insert(id, replacement)
+		Log.trace("Swapping Instance {} out via api/models/ endpoint", id)
+		local oldParent = oldInstance.Parent
+		for _, child in oldInstance:GetChildren() do
+			child.Parent = replacement
+		end
+
+		replacement.Parent = oldParent
+		-- ChangeHistoryService doesn't like it if an Instance has been
+		-- Destroyed. So, we have to accept the potential memory hit and
+		-- just set the parent to `nil`.
+		oldInstance.Parent = nil
+
+		if selectionMap[oldInstance] then
+			-- This is a bit funky, but it saves the order of Selection
+			-- which might matter for some use cases.
+			selection[selectionMap[oldInstance]] = replacement
+		end
+	end
+
+	local patchApplySuccess, unappliedPatch = pcall(self.__reconciler.applyPatch, self.__reconciler, refPatch)
+	if patchApplySuccess then
+		Selection:Set(selection)
+		return true, unappliedPatch
+	else
+		error(unappliedPatch)
+	end
+end
+
+function ServeSession:__applyPatch(patch)
+	local patchTimestamp = DateTime.now():FormatLocalTime("LTS", "en-us")
+	local historyRecording = ChangeHistoryService:TryBeginRecording("Rojo: Patch " .. patchTimestamp)
+	if not historyRecording then
+		-- There can only be one recording at a time
+		Log.debug("Failed to begin history recording for " .. patchTimestamp .. ". Another recording is in progress.")
+	end
+
+	Timer.start("precommitCallbacks")
+	-- Precommit callbacks must be serial in order to obey the contract that
+	-- they execute before commit
+	for _, callback in self.__precommitCallbacks do
+		local success, err = pcall(callback, patch, self.__instanceMap)
+		if not success then
+			Log.warn("Precommit hook errored: {}", err)
+		end
+	end
+	Timer.stop()
+
+	local patchApplySuccess, unappliedPatch = pcall(self.__reconciler.applyPatch, self.__reconciler, patch)
+	if not patchApplySuccess then
+		if historyRecording then
+			ChangeHistoryService:FinishRecording(historyRecording, Enum.FinishRecordingOperation.Commit)
+		end
+		-- This might make a weird stack trace but the only way applyPatch can
+		-- fail is if a bug occurs so it's probably fine.
+		error(unappliedPatch)
+	end
+
+	if PatchSet.isEmpty(unappliedPatch) then
+		if historyRecording then
+			ChangeHistoryService:FinishRecording(historyRecording, Enum.FinishRecordingOperation.Commit)
+		end
+		return
+	end
+
+	local addedIdList = PatchSet.addedIdList(unappliedPatch)
+	local updatedIdList = PatchSet.updatedIdList(unappliedPatch)
+
+	local actualUnappliedPatches = PatchSet.newEmpty()
+	if Settings:get("enableSyncFallback") then
+		Log.debug("ServeSession:__replaceInstances(unappliedPatch.added)")
+		Timer.start("ServeSession:__replaceInstances(unappliedPatch.added)")
+		local addSuccess, unappliedAddedRefs = self:__replaceInstances(addedIdList)
+		Timer.stop()
+
+		Log.debug("ServeSession:__replaceInstances(unappliedPatch.updated)")
+		Timer.start("ServeSession:__replaceInstances(unappliedPatch.updated)")
+		local updateSuccess, unappliedUpdateRefs = self:__replaceInstances(updatedIdList)
+		Timer.stop()
+
+		if addSuccess then
+			table.clear(unappliedPatch.added)
+			PatchSet.assign(actualUnappliedPatches, unappliedAddedRefs)
+		end
+		if updateSuccess then
+			table.clear(unappliedPatch.updated)
+			PatchSet.assign(actualUnappliedPatches, unappliedUpdateRefs)
+		end
+	else
+		Log.debug("Skipping ServeSession:__replaceInstances because of setting")
+	end
+	PatchSet.assign(actualUnappliedPatches, unappliedPatch)
+
+	if not PatchSet.isEmpty(actualUnappliedPatches) then
+		Log.debug(
+			"Could not apply all changes requested by the Rojo server:\n{}",
+			PatchSet.humanSummary(self.__instanceMap, unappliedPatch)
+		)
+	end
+
+	Timer.start("postcommitCallbacks")
+	-- Postcommit callbacks can be called with spawn since regardless of firing order, they are
+	-- guaranteed to be called after the commit
+	for _, callback in self.__postcommitCallbacks do
+		task.spawn(function()
+			local success, err = pcall(callback, patch, self.__instanceMap, actualUnappliedPatches)
+			if not success then
+				Log.warn("Postcommit hook errored: {}", err)
+			end
+		end)
+	end
+	Timer.stop()
+
+	if historyRecording then
+		ChangeHistoryService:FinishRecording(historyRecording, Enum.FinishRecordingOperation.Commit)
+	end
+end
+
 function ServeSession:__initialSync(serverInfo)
 	return self.__apiContext:read({ serverInfo.rootInstanceId }):andThen(function(readResponseBody)
 		-- Tell the API Context that we're up-to-date with the version of
@@ -280,15 +483,7 @@ function ServeSession:__initialSync(serverInfo)
 
 			return self.__apiContext:write(inversePatch)
 		elseif userDecision == "Accept" then
-			local unappliedPatch = self.__reconciler:applyPatch(catchUpPatch)
-
-			if not PatchSet.isEmpty(unappliedPatch) then
-				Log.debug(
-					"Could not apply all changes requested by the Rojo server:\n{}",
-					PatchSet.humanSummary(self.__instanceMap, unappliedPatch)
-				)
-			end
-
+			self:__applyPatch(catchUpPatch)
 			return Promise.resolve()
 		else
 			return Promise.reject("Invalid user decision: " .. userDecision)
@@ -311,14 +506,7 @@ function ServeSession:__mainSyncLoop()
 					Log.trace("Serve session {} retrieved {} messages", tostring(self), #messages)
 
 					for _, message in messages do
-						local unappliedPatch = self.__reconciler:applyPatch(message)
-
-						if not PatchSet.isEmpty(unappliedPatch) then
-							Log.debug(
-								"Could not apply all changes requested by the Rojo server:\n{}",
-								PatchSet.humanSummary(self.__instanceMap, unappliedPatch)
-							)
-						end
+						self:__applyPatch(message)
 					end
 				end)
 				:await()