Client: Implement screenshot captures on demand

For now they're blocking, since the time taken isn't that long. It does skip a few frames, but offloading the encoding step would be more complex than just doing it on the main thread and eating the dropped frames.

It can also timeout the surface texture and crash the renderer, but that's a separate problem and will be tackled soon (TM).

Core/NativeClient/Renderer.lua

@@ -1,4 +1,5 @@
 local bit = require("bit")
+local console = require("console")
 local etrace = require("etrace")
 local ffi = require("ffi")
 local interop = require("interop")
@@ -66,6 +67,7 @@ local Renderer = {
 		},
 	},
 	DEBUG_DISCARDED_BACKGROUND_PIXELS = false, -- This is really slow (disk I/O); don't enable unless necessary
+	SCREENSHOT_OUTPUT_DIRECTORY = "Screenshots",
 	numWidgetTransformsUsedThisFrame = 0,
 	errorStrings = {
 		INVALID_VERTEX_BUFFER = "Cannot upload geometry with invalid vertex buffer",
@@ -224,7 +226,11 @@ function Renderer:RenderNextFrame(deltaTime)
 		for materialIndex, meshes in pairs(meshesByMaterial) do
 			local material = self.supportedMaterials[materialIndex]
 			-- Should skip this if there aren't any meshes (wasteful to switch for no reason)?
+			if self.isCapturingScreenshot then
+				RenderPassEncoder:SetPipeline(renderPass, material.offlineRenderingPipeline.wgpuPipeline)
+			else
 				RenderPassEncoder:SetPipeline(renderPass, material.surfaceRenderingPipeline.wgpuPipeline)
+			end
 			for _, mesh in ipairs(meshes) do
 				for index, animation in ipairs(mesh.keyframeAnimations) do
 					animation:UpdateWithDeltaTime(deltaTime / 10E5)
@@ -241,8 +247,11 @@ function Renderer:RenderNextFrame(deltaTime)
 	do
 		local uiRenderPass = self:BeginUserInterfaceRenderPass(commandEncoder, nextTextureView)
 		RenderPassEncoder:SetBindGroup(uiRenderPass, 0, self.cameraViewportUniform.bindGroup, 0, nil)
-
+		if self.isCapturingScreenshot then
+			RenderPassEncoder:SetPipeline(uiRenderPass, UserInterfaceMaterial.offlineRenderingPipeline.wgpuPipeline)
+		else
 			RenderPassEncoder:SetPipeline(uiRenderPass, UserInterfaceMaterial.surfaceRenderingPipeline.wgpuPipeline)
+		end
 		self.numWidgetTransformsUsedThisFrame = 0
 		rml.bindings.rml_context_update(self.rmlContext)
 		-- NO MORE CHANGES here before rendering the updated state!
@@ -255,15 +264,42 @@ function Renderer:RenderNextFrame(deltaTime)
 	local uiRenderPassTime = uv.hrtime() - uiRenderPassStartTime
 
 	local commandSubmitStartTime = uv.hrtime()
-	self:SubmitCommandBuffer(commandEncoder)
+	CommandEncoder:Submit(commandEncoder, self.wgpuDevice)
 	local commandSubmissionTime = uv.hrtime() - commandSubmitStartTime
 
+	if self.isCapturingScreenshot then
+		local rgbaImageBytes, width, height = self.screenshotTexture:DownloadPixelBuffer(self.wgpuDevice)
+		-- This assumes the buffer read is blocking (which is suboptimal); streamline later, via events
+		self:SaveCapturedScreenshot(rgbaImageBytes, width, height)
+		self.isCapturingScreenshot = false
+	end
+
 	self.backingSurface:PresentNextFrame()
 
 	local totalFrameTime = worldRenderPassTime + uiRenderPassTime + commandSubmissionTime
 	return totalFrameTime, worldRenderPassTime, uiRenderPassTime, commandSubmissionTime
 end
 
+function Renderer:SaveCapturedScreenshot(rgbaImageBytes, width, height)
+	console.startTimer("[Renderer] SaveCapturedScreenshot")
+
+	local screenshotFileName = format("RagLite_Screenshot_%s.jpg", os.date("%Y%m%d%H%M%S"))
+	C_FileSystem.MakeDirectoryTree(Renderer.SCREENSHOT_OUTPUT_DIRECTORY)
+	local screenshotFilePath = path.join(Renderer.SCREENSHOT_OUTPUT_DIRECTORY, screenshotFileName)
+
+	local imageFileContents = C_ImageProcessing.EncodeJPG(rgbaImageBytes, width, height)
+
+	C_FileSystem.WriteFile(screenshotFilePath, imageFileContents)
+	printf(
+		"Screenshot taken: %s (raw size: %s, encoded size: %s)",
+		screenshotFileName,
+		string.filesize(#rgbaImageBytes),
+		string.filesize(#imageFileContents)
+	)
+
+	console.stopTimer("[Renderer] SaveCapturedScreenshot")
+end
+
 local rmlEventNames = {
 	[ffi.C.ERROR_EVENT] = "UNKNOWN_RENDER_COMMAND",
 	[ffi.C.GEOMETRY_RENDER_EVENT] = "GEOMETRY_RENDER_EVENT",
@@ -518,16 +554,6 @@ function Renderer:DrawWidget(renderPass, compiledWidgetGeometry, offsetU, offset
 	)
 end
 
-function Renderer:SubmitCommandBuffer(commandEncoder)
-	local commandBufferDescriptor = new("WGPUCommandBufferDescriptor")
-	local commandBuffer = CommandEncoder:Finish(commandEncoder, commandBufferDescriptor)
-
-	-- The WebGPU API expects an array here, but currently this renderer only supports a single buffer (to keep things simple)
-	local queue = Device:GetQueue(self.wgpuDevice)
-	local commandBuffers = new("WGPUCommandBuffer[1]", commandBuffer)
-	Queue:Submit(queue, 1, commandBuffers)
-end
-
 function Renderer:UploadMeshGeometry(mesh)
 	local positions = mesh.vertexPositions
 	local colors = mesh.vertexColors

Core/NativeClient/WebGPU/Buffer.lua

@@ -5,6 +5,7 @@ local webgpu = require("webgpu")
 local Buffer = {
 	VERTEX_BUFFER_FLAGS = bit.bor(ffi.C.WGPUBufferUsage_CopyDst, ffi.C.WGPUBufferUsage_Vertex),
 	INDEX_BUFFER_FLAGS = bit.bor(ffi.C.WGPUBufferUsage_CopyDst, ffi.C.WGPUBufferUsage_Index),
+	READBACK_BUFFER_FLAGS = bit.bor(ffi.C.WGPUBufferUsage_MapRead, ffi.C.WGPUBufferUsage_CopyDst),
 }
 
 local ALIGNMENT_IN_BYTES = 4

Core/NativeClient/WebGPU/CommandEncoder.lua

@@ -1,3 +1,6 @@
+local Device = require("Core.NativeClient.WebGPU.Device")
+local Queue = require("Core.NativeClient.WebGPU.Queue")
+
 local webgpu = require("webgpu")
 
 local CommandEncoder = {}
@@ -10,6 +13,15 @@ function CommandEncoder:Finish(wgpuCommandEncoder, wgpuCommandBufferDescriptor)
 	return webgpu.bindings.wgpu_command_encoder_finish(wgpuCommandEncoder, wgpuCommandBufferDescriptor)
 end
 
+function CommandEncoder:Submit(commandEncoder, wgpuDevice)
+	local commandBufferDescriptor = new("WGPUCommandBufferDescriptor")
+	local commandBuffer = self:Finish(commandEncoder, commandBufferDescriptor)
+
+	local queue = Device:GetQueue(wgpuDevice)
+	local commandBuffers = new("WGPUCommandBuffer[1]", commandBuffer)
+	Queue:Submit(queue, 1, commandBuffers)
+end
+
 CommandEncoder.__call = CommandEncoder.Construct
 CommandEncoder.__index = CommandEncoder
 setmetatable(CommandEncoder, CommandEncoder)

Core/NativeClient/WebGPU/RenderTargets/ScreenshotCaptureTexture.lua

@@ -1,8 +1,17 @@
 local bit = require("bit")
+local console = require("console")
 local ffi = require("ffi")
 local webgpu = require("webgpu")
 
-local ScreenshotCaptureTexture = {}
+local Buffer = require("Core.NativeClient.WebGPU.Buffer")
+local CommandEncoder = require("Core.NativeClient.WebGPU.CommandEncoder")
+local Device = require("Core.NativeClient.WebGPU.Device")
+
+local format = string.format
+
+local ScreenshotCaptureTexture = {
+	OUTPUT_TEXTURE_FORMAT = ffi.C.WGPUTextureFormat_RGBA8UnormSrgb,
+}
 
 function ScreenshotCaptureTexture:Construct(wgpuDevice, width, height)
 	local textureDescriptor = new("WGPUTextureDescriptor")
@@ -32,9 +41,100 @@ function ScreenshotCaptureTexture:Construct(wgpuDevice, width, height)
 		clearValue = new("WGPUColor", { 0, 0, 0, 0 }),
 	})
 
-	setmetatable(instance, self)
+	local pixelBufferDescriptor = new("WGPUBufferDescriptor")
+	pixelBufferDescriptor.mappedAtCreation = false
+	pixelBufferDescriptor.usage = Buffer.READBACK_BUFFER_FLAGS
+	pixelBufferDescriptor.size = 4 * width * height -- RGBA
+	local pixelBuffer = webgpu.bindings.wgpu_device_create_buffer(wgpuDevice, pixelBufferDescriptor)
+
+	local instance = {
+		width = width,
+		height = height,
+		wgpuDevice = wgpuDevice,
+		wgpuTexture = texture,
+		wgpuTextureView = textureView,
+		wgpuTextureDescriptor = textureDescriptor,
+		wgpuTextureViewDescriptor = textureViewDescriptor,
+		colorAttachment = renderPassColorAttachment,
+		readbackBuffer = pixelBuffer,
+		readbackBufferDescriptor = pixelBufferDescriptor,
+	}
+
+	local inheritanceLookupMetatable = {
+		__index = self,
+	}
+	setmetatable(instance, inheritanceLookupMetatable)
+
 	return instance
 end
+function ScreenshotCaptureTexture:DownloadPixelBuffer(wgpuDevice)
+	console.startTimer("[ScreenshotCaptureTexture] DownloadPixelBuffer")
+
+	local rgbaPixelBufferSize = tonumber(self.readbackBufferDescriptor.size)
+	self.rgbaPixelBuffer = buffer.new(rgbaPixelBufferSize)
+
+	local commandEncoderDescriptor = new("WGPUCommandEncoderDescriptor")
+	local commandEncoder = Device:CreateCommandEncoder(wgpuDevice, commandEncoderDescriptor)
+
+	local source = new("WGPUImageCopyTexture", {
+		mipLevel = 0,
+		origin = { 0, 0, 0 },
+		aspect = ffi.C.WGPUTextureAspect_All,
+		texture = self.wgpuTexture,
+	})
+
+	local destination = new("WGPUImageCopyBuffer", {
+		buffer = self.readbackBuffer,
+		layout = {
+			bytesPerRow = 4 * self.width,
+			offset = 0,
+			rowsPerImage = self.height,
+		},
+	})
+
+	webgpu.bindings.wgpu_command_encoder_copy_texture_to_buffer(
+		commandEncoder,
+		source,
+		destination,
+		new("WGPUExtent3D", { self.width, self.height, 1 })
+	)
+
+	CommandEncoder:Submit(commandEncoder, self.wgpuDevice)
+
+	local function onAsyncMapRequestCompleted(status)
+		if status ~= ffi.C.WGPUBufferMapAsyncStatus_Success then
+			error(format("Failed to map screenshot readback buffer (status: %s)", status))
+			return
+		end
+
+		local pixelData = webgpu.bindings.wgpu_buffer_get_const_mapped_range(
+			self.readbackBuffer,
+			0,
+			self.readbackBufferDescriptor.size
+		)
+
+		self.rgbaPixelBuffer:putcdata(pixelData, rgbaPixelBufferSize)
+
+		-- Can safely unmap here since the image data was copied (not ideal, can probably avoid this copy?)
+		webgpu.bindings.wgpu_buffer_unmap(self.readbackBuffer)
+
+		console.stopTimer("[ScreenshotCaptureTexture] DownloadPixelBuffer")
+	end
+
+	webgpu.bindings.wgpu_buffer_map_async(
+		self.readbackBuffer,
+		ffi.C.WGPUMapMode_Read,
+		0,
+		self.readbackBufferDescriptor.size,
+		onAsyncMapRequestCompleted,
+		nil
+	)
+
+	-- Should probably do this asynchronously and not block the render loop (optimize later)
+	webgpu.bindings.wgpu_device_poll(self.wgpuDevice, true, nil)
+
+	return self.rgbaPixelBuffer, self.width, self.height
+end
 
 class("ScreenshotCaptureTexture", ScreenshotCaptureTexture)
 

Tests/NativeClient/Renderer.spec.lua

@@ -668,6 +668,41 @@ describe("Renderer", function()
 			end, expectedErrorMessage)
 		end)
 	end)
+
+	describe("SaveCapturedScreenshot", function()
+		before(function()
+			Renderer.SCREENSHOT_OUTPUT_DIRECTORY = "TemporaryScreenshotsDir"
+		end)
+
+		after(function()
+			local screenshots = C_FileSystem.ReadDirectoryTree(Renderer.SCREENSHOT_OUTPUT_DIRECTORY)
+			for screenshotFile, isFile in pairs(screenshots) do
+				C_FileSystem.Delete(screenshotFile)
+			end
+			C_FileSystem.Delete(Renderer.SCREENSHOT_OUTPUT_DIRECTORY)
+			Renderer.SCREENSHOT_OUTPUT_DIRECTORY = "Screenshots"
+		end)
+
+		it("should create the configured screenshots directory if it doesn't yet exist", function()
+			assertFalse(C_FileSystem.Exists(Renderer.SCREENSHOT_OUTPUT_DIRECTORY))
+			Renderer:SaveCapturedScreenshot("\255\254\253\0", 1, 1)
+			assertTrue(C_FileSystem.Exists(Renderer.SCREENSHOT_OUTPUT_DIRECTORY))
+		end)
+
+		it("should save the provided image in the configured screenshots directory", function()
+			Renderer:SaveCapturedScreenshot("\000\000\000\000	", 1, 1)
+			local screenshots = C_FileSystem.ReadDirectoryTree(Renderer.SCREENSHOT_OUTPUT_DIRECTORY)
+			assertEquals(table.count(screenshots), 1)
+			for screenshotFile, isFile in pairs(screenshots) do
+				local fileContents = C_FileSystem.ReadFile(screenshotFile)
+				local rgbaImageBytes, width, height = C_ImageProcessing.DecodeFileContents(fileContents)
+				assertEquals(width, 1)
+				assertEquals(height, 1)
+				local expectedImageBytes = "\000\000\000\255" -- JPG doesn't support transparency
+				assertEquals(rgbaImageBytes, expectedImageBytes)
+			end
+		end)
+	end)
 end)
 
 VirtualGPU:Disable()