chore(PE-8188): add patch balances work

.gitignore

@@ -17,3 +17,4 @@ coverage
 
 mainnet/**/outputs/*
 mainnet/outputs/*
+._*

src/balances.lua

@@ -96,4 +96,39 @@ function balances.walletHasSufficientBalance(wallet, quantity)
 	return Balances[wallet] ~= nil and Balances[wallet] >= quantity
 end
 
+---@param oldBalances table<string, number> A table of addresses and their balances
+---@param newBalances table<string, number> A table of addresses and their balances
+---@return table<string, boolean> affectedBalancesAddresses table of addresses that have had balance changes
+function balances.patchBalances(oldBalances, newBalances)
+	assert(type(oldBalances) == "table", "Old balances must be a table")
+	assert(type(newBalances) == "table", "New balances must be a table")
+	local affectedBalancesAddresses = {}
+	for address, _ in pairs(oldBalances) do
+		if Balances[address] ~= oldBalances[address] then
+			affectedBalancesAddresses[address] = true
+		end
+	end
+	for address, _ in pairs(newBalances) do
+		if oldBalances[address] ~= newBalances[address] then
+			affectedBalancesAddresses[address] = true
+		end
+	end
+
+	--- For simplicity we always include the protocol balance in the patch message
+	--- this also prevents us from sending an empty patch message and deleting the entire hyperbeam balances table
+	local patchMessage = { device = "[email protected]", balances = { [ao.id] = Balances[ao.id] or 0 } }
+	for address, _ in pairs(affectedBalancesAddresses) do
+		patchMessage.balances[address] = Balances[address] or 0
+	end
+
+	-- only send the patch message if there are affected balances, otherwise we'll end up deleting the entire hyperbeam balances table
+	if next(patchMessage.balances) == nil then
+		return {}
+	else
+		ao.send(patchMessage)
+	end
+
+	return affectedBalancesAddresses
+end
+
 return balances

src/main.lua

@@ -90,6 +90,8 @@ local ActionMap = {
 	ApprovePrimaryNameRequest = "Approve-Primary-Name-Request",
 	PrimaryNames = "Primary-Names",
 	PrimaryName = "Primary-Name",
+	-- Hyperbeam Patch Balances
+	PatchHyperbeamBalances = "Patch-Hyperbeam-Balances",
 }
 
 --- @param msg ParsedMessage
@@ -471,6 +473,8 @@ local function addEventingHandler(handlerName, pattern, handleFn, critical, prin
 	critical = critical or false
 	printEvent = printEvent == nil and true or printEvent
 	Handlers.add(handlerName, pattern, function(msg)
+		-- Store the old balances to compare after the handler has run for patching state
+		local oldBalances = utils.deepCopy(Balances)
 		-- add an ARIOEvent to the message if it doesn't exist
 		msg.ioEvent = msg.ioEvent or ARIOEvent(msg)
 		-- global handler for all eventing errors, so we can log them and send a notice to the sender for non critical errors and discard the memory on critical errors
@@ -494,6 +498,9 @@ local function addEventingHandler(handlerName, pattern, handleFn, critical, prin
 			error(errorWithEvent, 0) -- 0 ensures not to include this line number in the error message
 		end
 
+		-- Send patch message to HB
+		balances.patchBalances(oldBalances, Balances)
+
 		msg.ioEvent:addField("Handler-Memory-KiB-Used", collectgarbage("count"), false)
 		collectgarbage("collect")
 		msg.ioEvent:addField("Final-Memory-KiB-Used", collectgarbage("count"), false)
@@ -2137,17 +2144,12 @@ end)
 -- Pagination handlers
 
 addEventingHandler("paginatedRecords", function(msg)
-        return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
+	return msg.Action == "Paginated-Records" or msg.Action == ActionMap.Records
 end, function(msg)
-        local page = utils.parsePaginationTags(msg)
-        local result = arns.getPaginatedRecords(
-                page.cursor,
-                page.limit,
-                page.sortBy or "startTimestamp",
-                page.sortOrder,
-                page.filters
-        )
-        Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
+	local page = utils.parsePaginationTags(msg)
+	local result =
+		arns.getPaginatedRecords(page.cursor, page.limit, page.sortBy or "startTimestamp", page.sortOrder, page.filters)
+	Send(msg, { Target = msg.From, Action = "Records-Notice", Data = json.encode(result) })
 end)
 
 addEventingHandler("paginatedGateways", function(msg)
@@ -2691,4 +2693,18 @@ addEventingHandler("allPaginatedGatewayVaults", utils.hasMatchingTag("Action", "
 	Send(msg, { Target = msg.From, Action = "All-Gateway-Vaults-Notice", Data = json.encode(result) })
 end)
 
+addEventingHandler(
+	ActionMap.PatchHyperbeamBalances,
+	utils.hasMatchingTag("Action", ActionMap.PatchHyperbeamBalances),
+	function(msg)
+		assert(msg.From == Owner, "Only the owner can trigger " .. ActionMap.PatchHyperbeamBalances)
+		local patchMessage = { device = "[email protected]", balances = utils.deepCopy(Balances) }
+		ao.send(patchMessage)
+		return Send(msg, {
+			Target = msg.From,
+			Action = ActionMap.PatchHyperbeamBalances .. "-Notice",
+		})
+	end
+)
+
 return main