CoopAllTheThings/hook/src/vk_capture.hpp

// Shared Vulkan swap-chain read-back used by BOTH early-presence paths: the inline hook
// (hook/src/vk_hook.cpp) and the implicit layer (vk_layer/coop_vk_layer.cpp). They differ only in
// how they get into the dispatch chain; the capture itself -- copy the presented image into a shared
// keyed-mutex D3D11 texture for the host to sample -- is identical, so it lives here once.
//
// Performance contract (the whole reason this is a separate component): the read-back must NOT
// stall the game's present thread. Doing the GPU copy + a synchronous CPU read of the mapped
// staging buffer + the swizzle + the D3D upload on the present thread drops a 144 FPS game to a
// few FPS -- especially when the staging buffer is plain HOST_COHERENT (on a discrete GPU that's
// write-combined/uncached, where a scattered CPU read runs at PCIe latency: hundreds of ms for
// one 1080p frame). This component avoids both halves:
//   * the present thread only records + submits the copy (sub-millisecond) and returns immediately;
//   * a dedicated reaper thread waits the copy's fence, reads the staging buffer, swizzles and
//     uploads -- off the critical path; and
//   * the staging buffer prefers HOST_CACHED memory so the CPU read is cached, not uncached PCIe.
// A ring of in-flight slots lets capture keep pace with the game; if the reaper falls behind the
// present thread simply skips a frame (the game never waits). Proven by tests/vk_capture_perf_test.
#pragma once

#include <atomic>
#include <condition_variable>
#include <cstdint>
#include <deque>
#include <functional>
#include <mutex>
#include <thread>
#include <vector>

#include <windows.h>

#include <d3d11.h>
#include <dxgi1_2.h>

#define VK_NO_PROTOTYPES
#include <vulkan/vulkan.h>

#include "shared_video_texture.hpp"

namespace coop::hook {

class VkCapture {
  public:
	// Device entry points the read-back needs (resolved by the caller via the real
	// vkGetDeviceProcAddr; GetPhysicalDeviceMemoryProperties is instance-level).
	struct Fns {
		PFN_vkGetDeviceQueue GetDeviceQueue;
		PFN_vkCreateCommandPool CreateCommandPool;
		PFN_vkDestroyCommandPool DestroyCommandPool;
		PFN_vkAllocateCommandBuffers AllocateCommandBuffers;
		PFN_vkBeginCommandBuffer BeginCommandBuffer;
		PFN_vkEndCommandBuffer EndCommandBuffer;
		PFN_vkResetCommandBuffer ResetCommandBuffer;
		PFN_vkCmdPipelineBarrier CmdPipelineBarrier;
		PFN_vkCmdCopyImageToBuffer CmdCopyImageToBuffer;
		PFN_vkQueueSubmit QueueSubmit;
		PFN_vkCreateFence CreateFence;
		PFN_vkDestroyFence DestroyFence;
		PFN_vkWaitForFences WaitForFences;
		PFN_vkResetFences ResetFences;
		PFN_vkGetFenceStatus GetFenceStatus;
		PFN_vkCreateSemaphore CreateSemaphore;
		PFN_vkDestroySemaphore DestroySemaphore;
		PFN_vkCreateBuffer CreateBuffer;
		PFN_vkDestroyBuffer DestroyBuffer;
		PFN_vkGetBufferMemoryRequirements GetBufferMemoryRequirements;
		PFN_vkAllocateMemory AllocateMemory;
		PFN_vkFreeMemory FreeMemory;
		PFN_vkBindBufferMemory BindBufferMemory;
		PFN_vkMapMemory MapMemory;
		PFN_vkUnmapMemory UnmapMemory;
		PFN_vkInvalidateMappedMemoryRanges InvalidateMappedMemoryRanges;
		PFN_vkDeviceWaitIdle DeviceWaitIdle;
		PFN_vkGetPhysicalDeviceMemoryProperties GetPhysicalDeviceMemoryProperties;
	};

	VkCapture() = default;
	~VkCapture();
	VkCapture(const VkCapture&) = delete;
	VkCapture& operator=(const VkCapture&) = delete;

	// Bind to the game's device + queue family and start the reaper thread. `pid` names the shared
	// texture (video_share_name). `on_frame(w,h)` runs on the reaper thread after each frame is
	// published (the caller does its own IPC / stat bookkeeping there). Idempotent-ish: call once.
	void init(VkPhysicalDevice phys, VkDevice device, std::uint32_t queue_family, const Fns& fns, unsigned long pid,
	          std::function<void(std::uint32_t, std::uint32_t)> on_frame);

	bool active() const { return m_device != VK_NULL_HANDLE; }

	// Present-thread entry. `image` is the swap-chain image being presented (in PRESENT_SRC layout).
	// If a capture was queued, returns true and sets `out_sem` to a semaphore the real present must
	// wait on (our copy consumes `wait`/`wait_count` and signals `out_sem`, so present still orders
	// after rendering). Returns false to present `image` unchanged (unsupported format, no free slot,
	// or not initialised). Never waits on the GPU and never touches the staging buffer's bytes.
	bool present(VkImage image, VkFormat fmt, std::uint32_t w, std::uint32_t h, const VkSemaphore* wait,
	             std::uint32_t wait_count, VkSemaphore& out_sem);

	// Stop the reaper thread, drain the device (DeviceWaitIdle), and free every Vulkan/D3D resource.
	// Must NOT be called from DllMain (it joins a thread). Safe if never initialised, and safe to call
	// more than once.
	void shutdown();

	std::uint64_t frames_published() const { return m_published.load(std::memory_order_relaxed); }

	// Test seam: copy the most recently published RGBA frame out (tightly packed w*4). False if none.
	bool last_frame(std::vector<unsigned char>& out, std::uint32_t& w, std::uint32_t& h);

  private:
	static constexpr int kSlots = 4; // in-flight copies; also the present-semaphore reuse slack

	struct Slot {
		VkCommandBuffer cmd = VK_NULL_HANDLE;
		VkFence fence = VK_NULL_HANDLE;
		VkSemaphore present_sem = VK_NULL_HANDLE;
		VkBuffer staging = VK_NULL_HANDLE;
		VkDeviceMemory mem = VK_NULL_HANDLE;
		void* mapped = nullptr;
		VkDeviceSize size = 0;
		bool coherent = true; // if false, Invalidate before the CPU reads the mapping
		std::uint32_t w = 0;
		std::uint32_t h = 0;
		VkFormat fmt = VK_FORMAT_UNDEFINED;
		std::atomic<bool> busy{false}; // true from present() submit until the reaper finishes it
	};

	bool ensure_slot_pool();
	bool ensure_staging(Slot& s, std::uint32_t w, std::uint32_t h);
	void free_slots();
	bool find_readback_memory(std::uint32_t type_bits, std::uint32_t& out_index, bool& out_coherent);

	bool ensure_d3d();
	void release_d3d();

	void reaper_main();
	void reap_slot(Slot& s);

	// --- Vulkan side (present thread creates/records; reaper reads the mapping) ---
	VkPhysicalDevice m_phys = VK_NULL_HANDLE;
	VkDevice m_device = VK_NULL_HANDLE;
	std::uint32_t m_qfam = 0;
	VkQueue m_queue = VK_NULL_HANDLE;
	Fns m_fns{};
	VkCommandPool m_pool = VK_NULL_HANDLE;
	Slot m_slots[kSlots];
	int m_next = 0;

	// --- D3D11 shared texture (reaper thread only) ---
	ID3D11Device* m_d3d = nullptr;
	ID3D11DeviceContext* m_d3d_ctx = nullptr;
	SharedVideoTexture m_shared;
	std::vector<unsigned char> m_rgba; // reaper scratch (swizzled frame)
	unsigned long m_pid = 0;
	std::function<void(std::uint32_t, std::uint32_t)> m_on_frame;

	// --- reaper thread + queue ---
	std::thread m_reaper;
	std::mutex m_q_mutex;
	std::condition_variable m_q_cv;
	std::deque<int> m_pending; // slot indices awaiting read-back
	bool m_stop = false;

	std::atomic<std::uint64_t> m_published{0};

	// last published frame (for the test seam)
	std::mutex m_last_mutex;
	std::vector<unsigned char> m_last;
	std::uint32_t m_last_w = 0;
	std::uint32_t m_last_h = 0;
};

} // namespace coop::hook