zen copy-state command to copy a zenserver data directory without the bulk data (#510)

* zen copy-state command to copy a zenserver data directory without the bulk data

3 files changed, 206 insertions, 0 deletions

diff --git a/src/zen/cmds/admin_cmd.cpp b/src/zen/cmds/admin_cmd.cpp
index 8b8b416ba..b041aa46e 100644
--- a/src/zen/cmds/admin_cmd.cpp
+++ b/src/zen/cmds/admin_cmd.cpp
@@ -1,10 +1,12 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
 #include "admin_cmd.h"
+#include <zencore/filesystem.h>
 #include <zencore/logging.h>
 #include <zenhttp/formatters.h>
 #include <zenhttp/httpclient.h>
 #include <zenhttp/httpcommon.h>
+#include <zenutil/basicfile.h>
 
 ZEN_THIRD_PARTY_INCLUDES_START
 #include <cpr/cpr.h>
@@ -455,4 +457,186 @@ FlushCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
 	return 1;
 }
 
+//////////////////////////////////////////////////////////////////////////
+
+CopyStateCommand::CopyStateCommand()
+{
+	m_Options.add_options()("h,help", "Print help");
+	m_Options.add_option("", "", "data-path", "Zen server source data path", cxxopts::value(m_DataPath), "<data-path>");
+	m_Options.add_option("", "", "target-path", "Target data path", cxxopts::value(m_TargetPath), "<target-path>");
+	m_Options.add_option("",
+						 "",
+						 "skip-logs",
+						 "Only copy state index files, not log files (recommended to issue a zen flush before using this option)",
+						 cxxopts::value(m_SkipLogs)->default_value("false"),
+						 "<skip-logs>");
+	m_Options.parse_positional({"data-path", "target-path"});
+}
+
+CopyStateCommand::~CopyStateCommand() = default;
+
+static void
+Copy(const std::filesystem::path& Source, const std::filesystem::path& Target)
+{
+	CreateDirectories(Target.parent_path());
+	BasicFile SourceFile;
+	SourceFile.Open(Source, BasicFile::Mode::kRead);
+	BasicFile TargetFile;
+	TargetFile.Open(Target, BasicFile::Mode::kTruncate);
+	uint64_t			 Size	= SourceFile.FileSize();
+	uint64_t			 Offset = 0;
+	std::vector<uint8_t> Buffer(Min(size_t(Size), size_t(65536u)));
+	while (Offset < Size)
+	{
+		uint64_t CopyCount = Min<uint64_t>(Size - Offset, size_t(Buffer.size()));
+		SourceFile.Read(Buffer.data(), CopyCount, Offset);
+		TargetFile.Write(Buffer.data(), CopyCount, Offset);
+		Offset += CopyCount;
+	}
+	TargetFile.Flush();
+}
+
+static bool
+TryCopy(const std::filesystem::path& Source, const std::filesystem::path& Target)
+{
+	if (!std::filesystem::is_regular_file(Source))
+	{
+		return false;
+	}
+	Copy(Source, Target);
+	return true;
+}
+
+int
+CopyStateCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
+{
+	ZEN_UNUSED(GlobalOptions);
+
+	if (!ParseOptions(argc, argv))
+	{
+		return 0;
+	}
+
+	if (m_DataPath.empty())
+	{
+		throw OptionParseException("data path must be given");
+	}
+
+	if (!std::filesystem::is_directory(m_DataPath))
+	{
+		throw OptionParseException("data path must exist");
+	}
+
+	if (m_TargetPath.empty())
+	{
+		throw OptionParseException("target path must be given");
+	}
+
+	std::filesystem::path RootManifestPath		 = m_DataPath / "root_manifest";
+	std::filesystem::path TargetRootManifestPath = m_TargetPath / "root_manifest";
+
+	if (!TryCopy(RootManifestPath, TargetRootManifestPath))
+	{
+		throw OptionParseException("data path is invalid, missing root_manifest");
+	}
+
+	std::filesystem::path CachePath		  = m_DataPath / "cache";
+	std::filesystem::path TargetCachePath = m_TargetPath / "cache";
+
+	// Copy cache state
+	DirectoryContent CacheDirectoryContent;
+	GetDirectoryContent(CachePath, DirectoryContent::IncludeDirsFlag, CacheDirectoryContent);
+	for (const std::filesystem::path& NamespacePath : CacheDirectoryContent.Directories)
+	{
+		std::filesystem::path NamespaceName		  = NamespacePath.filename();
+		std::filesystem::path TargetNamespacePath = TargetCachePath / NamespaceName;
+
+		DirectoryContent CacheNamespaceContent;
+		GetDirectoryContent(NamespacePath, DirectoryContent::IncludeDirsFlag, CacheNamespaceContent);
+
+		for (const std::filesystem::path& BucketPath : CacheNamespaceContent.Directories)
+		{
+			std::filesystem::path BucketName	   = BucketPath.filename();
+			std::filesystem::path TargetBucketPath = TargetNamespacePath / BucketName;
+
+			std::filesystem::path ManifestPath		 = BucketPath / "zen_manifest";
+			std::filesystem::path TargetManifestPath = TargetBucketPath / "zen_manifest";
+			if (TryCopy(ManifestPath, TargetManifestPath))
+			{
+				if (!m_SkipLogs)
+				{
+					std::filesystem::path LogName		= fmt::format("{}.{}", BucketName.string(), "slog");
+					std::filesystem::path LogPath		= BucketPath / LogName;
+					std::filesystem::path TargetLogPath = TargetBucketPath / LogName;
+					Copy(LogPath, TargetLogPath);
+				}
+
+				std::filesystem::path IndexName		  = fmt::format("{}.{}", BucketName.string(), "uidx");
+				std::filesystem::path IndexPath		  = BucketPath / IndexName;
+				std::filesystem::path TargetIndexPath = TargetBucketPath / IndexName;
+				TryCopy(IndexPath, TargetIndexPath);
+			}
+		}
+	}
+
+	std::filesystem::path CasPath		= m_DataPath / "cas";
+	std::filesystem::path TargetCasPath = m_TargetPath / "cas";
+
+	{
+		std::filesystem::path UCasRootPath		 = CasPath / ".ucas_root";
+		std::filesystem::path TargetUCasRootPath = TargetCasPath / ".ucas_root";
+		Copy(UCasRootPath, TargetUCasRootPath);
+	}
+
+	if (!m_SkipLogs)
+	{
+		std::filesystem::path LogPath		= CasPath / "cas.ulog";
+		std::filesystem::path TargetLogPath = TargetCasPath / "cas.ulog";
+		Copy(LogPath, TargetLogPath);
+	}
+
+	{
+		std::filesystem::path IndexPath		  = CasPath / "cas.uidx";
+		std::filesystem::path TargetIndexPath = TargetCasPath / "cas.uidx";
+		TryCopy(IndexPath, TargetIndexPath);
+	}
+
+	{
+		std::filesystem::path SobsRootPath		 = CasPath / "sobs";
+		std::filesystem::path TargetSobsRootPath = TargetCasPath / "sobs";
+
+		{
+			std::filesystem::path SobsIndexPath		  = SobsRootPath / "sobs.uidx";
+			std::filesystem::path TargetSobsIndexPath = TargetSobsRootPath / "sobs.uidx";
+			TryCopy(SobsIndexPath, TargetSobsIndexPath);
+
+			if (!m_SkipLogs)
+			{
+				std::filesystem::path SobsLogPath		= SobsRootPath / "sobs.ulog";
+				std::filesystem::path TargetSobsLogPath = TargetSobsRootPath / "sobs.ulog";
+				Copy(SobsLogPath, TargetSobsLogPath);
+			}
+		}
+	}
+
+	{
+		std::filesystem::path TobsRootPath		 = CasPath / "tobs";
+		std::filesystem::path TargetTobsRootPath = TargetCasPath / "tobs";
+		{
+			std::filesystem::path TobsIndexPath		  = TobsRootPath / "tobs.uidx";
+			std::filesystem::path TargetTobsIndexPath = TargetTobsRootPath / "tobs.uidx";
+			TryCopy(TobsIndexPath, TargetTobsIndexPath);
+
+			if (!m_SkipLogs)
+			{
+				std::filesystem::path TobsLogPath		= TobsRootPath / "tobs.ulog";
+				std::filesystem::path TargetTobsLogPath = TargetTobsRootPath / "tobs.ulog";
+				Copy(TobsLogPath, TargetTobsLogPath);
+			}
+		}
+	}
+
+	return 0;
+}
+
 }  // namespace zen
diff --git a/src/zen/cmds/admin_cmd.h b/src/zen/cmds/admin_cmd.h
index 4b57eccdb..356f58363 100644
--- a/src/zen/cmds/admin_cmd.h
+++ b/src/zen/cmds/admin_cmd.h
@@ -4,6 +4,8 @@
 
 #include "../zen.h"
 
+#include <filesystem>
+
 namespace zen {
 
 /** Scrub storage
@@ -113,4 +115,22 @@ private:
 	std::string		 m_HostName;
 };
 
+/** Copy state
+ */
+class CopyStateCommand : public ZenCmdBase
+{
+public:
+	CopyStateCommand();
+	~CopyStateCommand();
+
+	virtual int				  Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) override;
+	virtual cxxopts::Options& Options() override { return m_Options; }
+
+private:
+	cxxopts::Options	  m_Options{"copy-state", "Copy zen server disk state"};
+	std::filesystem::path m_DataPath;
+	std::filesystem::path m_TargetPath;
+	bool				  m_SkipLogs;
+};
+
 }  // namespace zen
diff --git a/src/zen/zen.cpp b/src/zen/zen.cpp
index cfed591a8..241dbb599 100644
--- a/src/zen/zen.cpp
+++ b/src/zen/zen.cpp
@@ -200,6 +200,7 @@ main(int argc, char** argv)
 	CacheInfoCommand		 CacheInfoCmd;
 	CacheStatsCommand		 CacheStatsCmd;
 	CopyCommand				 CopyCmd;
+	CopyStateCommand		 CopyStateCmd;
 	CreateOplogCommand		 CreateOplogCmd;
 	DeleteOplogCommand		 DeleteOplogCmd;
 	CreateProjectCommand	 CreateProjectCmd;
@@ -252,6 +253,7 @@ main(int argc, char** argv)
 		{"cache-info",			&CacheInfoCmd,			"Info on cache, namespace or bucket"},
 		{"cache-stats",			&CacheStatsCmd,			"Stats on cache"},
 		{"copy",				&CopyCmd,				"Copy file(s)"},
+		{"copy-state",			&CopyStateCmd,			"Copy zen server disk state"},
 		{"dedup",				&DedupCmd,				"Dedup files"},
 		{"down",				&DownCmd,				"Bring zen server down"},
 		{"drop",				&DropCmd,				"Drop cache namespace or bucket"},