From 39eff8c5696206839addf8bcd649ec92f42aaa78 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew.noyes@snowflake.com>
Date: Tue, 3 Aug 2021 09:43:42 -0700
Subject: [PATCH] Add trace_partial_file_suffix network option

---
 .../unit/trace_partial_file_suffix_test.cpp   |  5 ++--
 fdbclient/NativeAPI.actor.cpp                 |  7 +++++-
 fdbclient/NativeAPI.actor.h                   |  1 +
 fdbclient/vexillographer/fdb.options          |  3 +++
 flow/FileTraceLogWriter.cpp                   | 23 +++++++++++++++----
 flow/FileTraceLogWriter.h                     |  3 +++
 flow/Trace.cpp                                | 12 +++++++---
 flow/Trace.h                                  |  3 ++-
 8 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/bindings/c/test/unit/trace_partial_file_suffix_test.cpp b/bindings/c/test/unit/trace_partial_file_suffix_test.cpp
index d3be3c79e7..6462d84f5b 100644
--- a/bindings/c/test/unit/trace_partial_file_suffix_test.cpp
+++ b/bindings/c/test/unit/trace_partial_file_suffix_test.cpp
@@ -53,12 +53,11 @@ int main(int argc, char** argv) {
 	fdb_check(fdb_select_api_version(710));
 
 	std::string file_identifier = "trace_partial_file_suffix_test" + std::to_string(std::random_device{}());
-	// std::string trace_partial_file_suffix = ".tmp";
-	std::string trace_partial_file_suffix = "";
+	std::string trace_partial_file_suffix = ".tmp";
 
 	set_net_opt(FDBNetworkOption::FDB_NET_OPTION_TRACE_ENABLE, "");
 	set_net_opt(FDBNetworkOption::FDB_NET_OPTION_TRACE_FILE_IDENTIFIER, file_identifier);
-	// set_net_opt(FDBNetworkOption::FDB_NET_OPTION_TRACE_PARTIAL_FILE_SUFFIX, trace_partial_file_suffix);
+	set_net_opt(FDBNetworkOption::FDB_NET_OPTION_TRACE_PARTIAL_FILE_SUFFIX, trace_partial_file_suffix);
 
 	fdb_check(fdb_setup_network());
 	std::thread network_thread{ &fdb_run_network };
diff --git a/fdbclient/NativeAPI.actor.cpp b/fdbclient/NativeAPI.actor.cpp
index 0d05f33c58..0a6a474249 100644
--- a/fdbclient/NativeAPI.actor.cpp
+++ b/fdbclient/NativeAPI.actor.cpp
@@ -1698,7 +1698,8 @@ Database Database::createDatabase(Reference<ClusterConnectionFile> connFile,
 			              networkOptions.traceDirectory.get(),
 			              "trace",
 			              networkOptions.traceLogGroup,
-			              networkOptions.traceFileIdentifier);
+			              networkOptions.traceFileIdentifier,
+			              networkOptions.tracePartialFileSuffix);
 
 			TraceEvent("ClientStart")
 			    .detail("SourceVersion", getSourceVersion())
@@ -1856,6 +1857,10 @@ void setNetworkOption(FDBNetworkOptions::Option option, Optional<StringRef> valu
 			throw invalid_option_value();
 		}
 		break;
+	case FDBNetworkOptions::TRACE_PARTIAL_FILE_SUFFIX:
+		validateOptionValuePresent(value);
+		networkOptions.tracePartialFileSuffix = value.get().toString();
+		break;
 	case FDBNetworkOptions::KNOB: {
 		validateOptionValuePresent(value);
 
diff --git a/fdbclient/NativeAPI.actor.h b/fdbclient/NativeAPI.actor.h
index e671bac8c1..3636b30c46 100644
--- a/fdbclient/NativeAPI.actor.h
+++ b/fdbclient/NativeAPI.actor.h
@@ -68,6 +68,7 @@ struct NetworkOptions {
 	std::string traceFormat;
 	std::string traceClockSource;
 	std::string traceFileIdentifier;
+	std::string tracePartialFileSuffix;
 	Optional<bool> logClientInfo;
 	Reference<ReferencedObject<Standalone<VectorRef<ClientVersionRef>>>> supportedVersions;
 	bool runLoopProfilingEnabled;
diff --git a/fdbclient/vexillographer/fdb.options b/fdbclient/vexillographer/fdb.options
index e0f908a69c..27fa43aa8f 100644
--- a/fdbclient/vexillographer/fdb.options
+++ b/fdbclient/vexillographer/fdb.options
@@ -57,6 +57,9 @@ description is not currently required but encouraged.
     <Option name="trace_file_identifier" code="36"
             paramType="String" paramDescription="The identifier that will be part of all trace file names"
             description="Once provided, this string will be used to replace the port/PID in the log file names." />
+    <Option name="trace_partial_file_suffix" code="39"
+            paramType="String" paramDesciption="Append this suffix to partially written log files. When a log file is complete, rename it to remove the suffix."
+            description="" />
     <Option name="knob" code="40"
             paramType="String" paramDescription="knob_name=knob_value"
             description="Set internal tuning or debugging knobs"/>
diff --git a/flow/FileTraceLogWriter.cpp b/flow/FileTraceLogWriter.cpp
index e57f7502f6..18dfc52a07 100644
--- a/flow/FileTraceLogWriter.cpp
+++ b/flow/FileTraceLogWriter.cpp
@@ -87,11 +87,13 @@ FileTraceLogWriter::FileTraceLogWriter(std::string const& directory,
                                        std::string const& processName,
                                        std::string const& basename,
                                        std::string const& extension,
+                                       std::string const& tracePartialFileSuffix,
                                        uint64_t maxLogsSize,
                                        std::function<void()> const& onError,
                                        Reference<ITraceLogIssuesReporter> const& issues)
-  : directory(directory), processName(processName), basename(basename), extension(extension), maxLogsSize(maxLogsSize),
-    traceFileFD(-1), index(0), issues(issues), onError(onError) {}
+  : directory(directory), processName(processName), basename(basename), extension(extension),
+    tracePartialFileSuffix(tracePartialFileSuffix), maxLogsSize(maxLogsSize), traceFileFD(-1), index(0), issues(issues),
+    onError(onError) {}
 
 void FileTraceLogWriter::addref() {
 	ReferenceCounted<FileTraceLogWriter>::addref();
@@ -158,7 +160,8 @@ void FileTraceLogWriter::open() {
 	// log10(index) < 10
 	UNSTOPPABLE_ASSERT(indexWidth < 10);
 
-	auto finalname = format("%s.%d.%d.%s", basename.c_str(), indexWidth, index, extension.c_str());
+	finalname =
+	    format("%s.%d.%d.%s%s", basename.c_str(), indexWidth, index, extension.c_str(), tracePartialFileSuffix.c_str());
 	while ((traceFileFD = __open(finalname.c_str(), TRACEFILE_FLAGS, TRACEFILE_MODE)) == -1) {
 		lastError(errno);
 		if (errno == EEXIST) {
@@ -166,7 +169,12 @@ void FileTraceLogWriter::open() {
 			indexWidth = unsigned(::floor(log10f(float(index))));
 
 			UNSTOPPABLE_ASSERT(indexWidth < 10);
-			finalname = format("%s.%d.%d.%s", basename.c_str(), indexWidth, index, extension.c_str());
+			finalname = format("%s.%d.%d.%s%s",
+			                   basename.c_str(),
+			                   indexWidth,
+			                   index,
+			                   extension.c_str(),
+			                   tracePartialFileSuffix.c_str());
 		} else {
 			fprintf(stderr,
 			        "ERROR: could not create trace log file `%s' (%d: %s)\n",
@@ -178,7 +186,7 @@ void FileTraceLogWriter::open() {
 
 			int errorNum = errno;
 			onMainThreadVoid(
-			    [finalname, errorNum] {
+			    [finalname = finalname, errorNum] {
 				    TraceEvent(SevWarnAlways, "TraceFileOpenError")
 				        .detail("Filename", finalname)
 				        .detail("ErrorCode", errorNum)
@@ -201,6 +209,11 @@ void FileTraceLogWriter::close() {
 		while (__close(traceFileFD))
 			threadSleep(0.1);
 	}
+	traceFileFD = -1;
+	if (!tracePartialFileSuffix.empty()) {
+		renameFile(finalname, finalname.substr(0, finalname.size() - tracePartialFileSuffix.size()));
+	}
+	finalname = "";
 }
 
 void FileTraceLogWriter::roll() {
diff --git a/flow/FileTraceLogWriter.h b/flow/FileTraceLogWriter.h
index 2d28a076b4..4f22998f7b 100644
--- a/flow/FileTraceLogWriter.h
+++ b/flow/FileTraceLogWriter.h
@@ -51,6 +51,8 @@ private:
 	std::string processName;
 	std::string basename;
 	std::string extension;
+	std::string finalname;
+	std::string tracePartialFileSuffix;
 
 	uint64_t maxLogsSize;
 	int traceFileFD;
@@ -66,6 +68,7 @@ public:
 	                   std::string const& processName,
 	                   std::string const& basename,
 	                   std::string const& extension,
+	                   std::string const& tracePartialFileSuffix,
 	                   uint64_t maxLogsSize,
 	                   std::function<void()> const& onError,
 	                   Reference<ITraceLogIssuesReporter> const& issues);
diff --git a/flow/Trace.cpp b/flow/Trace.cpp
index 5b10dee691..bd4e4a36f7 100644
--- a/flow/Trace.cpp
+++ b/flow/Trace.cpp
@@ -115,6 +115,7 @@ private:
 	std::string directory;
 	std::string processName;
 	Optional<NetworkAddress> localAddress;
+	std::string tracePartialFileSuffix;
 
 	Reference<IThreadPool> writer;
 	uint64_t rollsize;
@@ -288,13 +289,15 @@ public:
 	          std::string const& timestamp,
 	          uint64_t rs,
 	          uint64_t maxLogsSize,
-	          Optional<NetworkAddress> na) {
+	          Optional<NetworkAddress> na,
+	          std::string const& tracePartialFileSuffix) {
 		ASSERT(!writer && !opened);
 
 		this->directory = directory;
 		this->processName = processName;
 		this->logGroup = logGroup;
 		this->localAddress = na;
+		this->tracePartialFileSuffix = tracePartialFileSuffix;
 
 		basename = format("%s/%s.%s.%s",
 		                  directory.c_str(),
@@ -306,6 +309,7 @@ public:
 		    processName,
 		    basename,
 		    formatter->getExtension(),
+		    tracePartialFileSuffix,
 		    maxLogsSize,
 		    [this]() { barriers->triggerAll(); },
 		    issues));
@@ -715,7 +719,8 @@ void openTraceFile(const NetworkAddress& na,
                    std::string directory,
                    std::string baseOfBase,
                    std::string logGroup,
-                   std::string identifier) {
+                   std::string identifier,
+                   std::string tracePartialFileSuffix) {
 	if (g_traceLog.isOpen())
 		return;
 
@@ -739,7 +744,8 @@ void openTraceFile(const NetworkAddress& na,
 	                format("%lld", time(nullptr)),
 	                rollsize,
 	                maxLogsSize,
-	                !g_network->isSimulated() ? na : Optional<NetworkAddress>());
+	                !g_network->isSimulated() ? na : Optional<NetworkAddress>(),
+	                tracePartialFileSuffix);
 
 	uncancellable(recurring(&flushTraceFile, FLOW_KNOBS->TRACE_FLUSH_INTERVAL, TaskPriority::FlushTrace));
 	g_traceBatch.dump();
diff --git a/flow/Trace.h b/flow/Trace.h
index ac9da282bc..467422b9cc 100644
--- a/flow/Trace.h
+++ b/flow/Trace.h
@@ -564,7 +564,8 @@ void openTraceFile(const NetworkAddress& na,
                    std::string directory = ".",
                    std::string baseOfBase = "trace",
                    std::string logGroup = "default",
-                   std::string identifier = "");
+                   std::string identifier = "",
+                   std::string tracePartialFileSuffix = "");
 void initTraceEventMetrics();
 void closeTraceFile();
 bool traceFileIsOpen();