Merge branch 'main' into lm/restrict-content-type

1 files changed, 101 insertions, 0 deletions

diff --git a/src/zen/cmds/exec_cmd.h b/src/zen/cmds/exec_cmd.h
new file mode 100644
index 000000000..6311354c0
--- /dev/null
+++ b/src/zen/cmds/exec_cmd.h
@@ -0,0 +1,101 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
+
+#include "../zen.h"
+
+#include <zencompute/recordingreader.h>
+#include <zencore/compactbinarypackage.h>
+#include <zencore/guid.h>
+#include <zencore/iohash.h>
+
+#include <filesystem>
+#include <functional>
+#include <unordered_map>
+
+namespace zen {
+class CbPackage;
+class CbObject;
+struct IoHash;
+class ChunkResolver;
+}  // namespace zen
+
+#if ZEN_WITH_COMPUTE_SERVICES
+
+namespace zen::compute {
+class ComputeServiceSession;
+}
+
+namespace zen {
+
+/**
+ * Zen CLI command for executing functions from a recording
+ *
+ * Mostly for testing and debugging purposes
+ */
+
+class ExecCommand : public ZenCmdBase
+{
+public:
+	ExecCommand();
+	~ExecCommand();
+
+	static constexpr char Name[]		= "exec";
+	static constexpr char Description[] = "Execute functions from a recording";
+
+	virtual void			  Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) override;
+	virtual cxxopts::Options& Options() override { return m_Options; }
+
+private:
+	cxxopts::Options	  m_Options{Name, Description};
+	std::string			  m_HostName;
+	std::string			  m_OrchestratorUrl;
+	std::filesystem::path m_BeaconPath;
+	std::filesystem::path m_RecordingPath;
+	std::filesystem::path m_RecordingLogPath;
+	int					  m_Offset = 0;
+	int					  m_Stride = 1;
+	int					  m_Limit  = 0;
+	bool				  m_Quiet  = false;
+	std::string			  m_Mode{"http"};
+	std::filesystem::path m_OutputPath;
+	bool				  m_Binary = false;
+
+	struct FunctionDefinition
+	{
+		std::string FunctionName;
+		zen::Guid	FunctionVersion;
+		zen::Guid	BuildSystemVersion;
+		zen::IoHash WorkerId;
+	};
+
+	bool m_FunctionListEmittedOnce = false;
+	void EmitFunctionListOnce(const std::vector<FunctionDefinition>& FunctionList);
+	void EmitFunctionList(const std::vector<FunctionDefinition>& FunctionList);
+
+	std::unordered_map<zen::IoHash, zen::CbPackage> m_WorkerMap;
+	std::vector<FunctionDefinition>					m_FunctionList;
+	bool											m_VerboseLogging = false;
+	bool											m_QuietLogging	 = false;
+	bool											m_DumpActions	 = false;
+
+	zen::ChunkResolver*				   m_ChunkResolver	 = nullptr;
+	zen::compute::RecordingReaderBase* m_RecordingReader = nullptr;
+
+	void RegisterWorkerFunctionsFromDescription(const zen::CbObject& WorkerDesc, const zen::IoHash& WorkerId);
+
+	int ExecUsingSession(zen::compute::ComputeServiceSession& ComputeSession);
+
+	// Execution modes
+
+	int DumpWorkItems();
+	int HttpExecute();
+	int InProcessExecute();
+	int LocalMessagingExecute();
+	int BeaconExecute();
+	int BuildActionsLog();
+};
+
+}  // namespace zen
+
+#endif	// ZEN_WITH_COMPUTE_SERVICES