Replace the output file safely in tracy-capture.

The old pre-flight truncated the output at startup with -f, so an
interrupted capture lost the trace, and the probe fopen followed a
symlink. The existing file is now moved to FILE~ for the duration of
the capture and restored if the trace cannot be written.

capture/CMakeLists.txt

@@ -16,6 +16,7 @@ include(${CMAKE_CURRENT_LIST_DIR}/../cmake/config.cmake)
 include(${CMAKE_CURRENT_LIST_DIR}/../cmake/vendor.cmake)
 include(${CMAKE_CURRENT_LIST_DIR}/../cmake/server.cmake)
 include(${CMAKE_CURRENT_LIST_DIR}/../cmake/GitRef.cmake)
+include(${CMAKE_CURRENT_LIST_DIR}/../cmake/util.cmake)
 
 set(PROGRAM_FILES
     src/capture.cpp
@@ -24,7 +25,7 @@ set(PROGRAM_FILES
 
 add_executable(${PROJECT_NAME} ${PROGRAM_FILES} ${COMMON_FILES} ${SERVER_FILES})
 add_git_ref(${PROJECT_NAME})
-target_link_libraries(${PROJECT_NAME} PRIVATE TracyServer TracyGetOpt)
+target_link_libraries(${PROJECT_NAME} PRIVATE TracyServer TracyUtil TracyGetOpt)
 set_property(DIRECTORY ${CMAKE_CURRENT_LIST_DIR} PROPERTY VS_STARTUP_PROJECT ${PROJECT_NAME})
 
 install(TARGETS ${PROJECT_NAME} DESTINATION ${CMAKE_INSTALL_BINDIR})

capture/src/capture.cpp

@@ -11,7 +11,6 @@
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <sys/stat.h>
 
 #include "../../server/TracyFileWrite.hpp"
 #include "../../server/TracyPrint.hpp"
@@ -19,6 +18,7 @@
 #include "../../server/TracyWorker.hpp"
 #include "../../public/common/TracyVersion.hpp"
 #include "GitRef.hpp"
+#include "CaptureFileBackup.hpp"
 
 #include "CaptureOutput.hpp"
 
@@ -96,21 +96,18 @@ int main( int argc, char** argv )
 
     if( !address || !output ) Usage();
 
-    struct stat st;
-    if( stat( output, &st ) == 0 && !overwrite )
+    const char* prepError = nullptr;
+    const auto prep = tracy::PrepareOutputFile( output, overwrite, &prepError );
+    if( prep == tracy::OutputPrep::Exists )
     {
         printf( "Output file %s already exists! Use -f to force overwrite.\n", output );
         return 4;
     }
-
-    FILE* test = fopen( output, "wb" );
-    if( !test )
+    if( prep == tracy::OutputPrep::Unusable )
     {
-        printf( "Cannot open output file %s for writing!\n", output );
+        printf( "Cannot use output file: %s!\n", prepError );
         return 5;
     }
-    fclose( test );
-    unlink( output );
 
     printf( "Connecting to %s:%i...", address, port );
     fflush( stdout );
@@ -169,10 +166,12 @@ int main( int argc, char** argv )
         f->Finish();
         const auto stats = f->GetCompressionStatistics();
         printf( "Trace size %s (%.2f%% ratio)\n", tracy::MemSizeToString( stats.second ), 100.f * stats.second / stats.first );
+        tracy::DiscardOutputBackup();
     }
     else
     {
         AnsiPrintf( ANSI_RED ANSI_BOLD, " failed!\n");
+        tracy::RestoreOutputBackup();
     }
 
     return 0;

cmake/util.cmake

@@ -0,0 +1,10 @@
+set(TRACY_UTIL_DIR ${CMAKE_CURRENT_LIST_DIR}/../util)
+
+set(TRACY_UTIL_SOURCES
+    CaptureFileBackup.cpp
+)
+
+list(TRANSFORM TRACY_UTIL_SOURCES PREPEND "${TRACY_UTIL_DIR}/")
+
+add_library(TracyUtil STATIC EXCLUDE_FROM_ALL ${TRACY_UTIL_SOURCES})
+target_include_directories(TracyUtil PUBLIC ${TRACY_UTIL_DIR})

util/CaptureFileBackup.cpp

@@ -0,0 +1,92 @@
+#include <cstdio>
+#include <cstdlib>
+#include <filesystem>
+#include <system_error>
+
+#include "CaptureFileBackup.hpp"
+
+namespace tracy
+{
+
+static std::filesystem::path s_backupPath, s_backupOf;
+
+static char s_errorBuf[512];
+
+OutputPrep PrepareOutputFile( const char* path, bool overwrite, const char** error )
+{
+    std::error_code ec;
+    const std::filesystem::path out( path );
+
+    const auto st = std::filesystem::symlink_status( out, ec );
+    const bool absent = st.type() == std::filesystem::file_type::not_found ||
+                        ec == std::errc::no_such_file_or_directory;
+    if( ec && !absent )
+    {
+        snprintf( s_errorBuf, sizeof( s_errorBuf ), "cannot access output path %s: %s", path, ec.message().c_str() );
+        *error = s_errorBuf;
+        return OutputPrep::Unusable;
+    }
+    ec.clear();
+    const bool exists = !absent;
+    if( exists && !std::filesystem::is_regular_file( st ) )
+    {
+        snprintf( s_errorBuf, sizeof( s_errorBuf ), "output path %s exists and is not a regular file", path );
+        *error = s_errorBuf;
+        return OutputPrep::Unusable;
+    }
+    if( exists && !overwrite )
+    {
+        *error = nullptr;
+        return OutputPrep::Exists;
+    }
+
+    FILE* probe = fopen( path, "ab" );
+    if( !probe )
+    {
+        snprintf( s_errorBuf, sizeof( s_errorBuf ), "cannot open output file %s for writing", path );
+        *error = s_errorBuf;
+        return OutputPrep::Unusable;
+    }
+    fclose( probe );
+
+    if( exists )
+    {
+        auto backup = out;
+        backup += "~";
+        std::filesystem::rename( out, backup, ec );
+        if( ec )
+        {
+            snprintf( s_errorBuf, sizeof( s_errorBuf ), "cannot move existing %s out of the way: %s", path, ec.message().c_str() );
+            *error = s_errorBuf;
+            return OutputPrep::Unusable;
+        }
+        s_backupOf = out;
+        s_backupPath = backup;
+        atexit( RestoreOutputBackup );
+    }
+    else
+    {
+        std::filesystem::remove( out, ec );
+    }
+    *error = nullptr;
+    return OutputPrep::Ok;
+}
+
+void RestoreOutputBackup()
+{
+    if( s_backupOf.empty() ) return;
+    std::error_code ec;
+    std::filesystem::rename( s_backupPath, s_backupOf, ec );
+    if( ec ) fprintf( stderr, "Could not restore %s from %s: %s\n", s_backupOf.string().c_str(), s_backupPath.string().c_str(), ec.message().c_str() );
+    s_backupOf.clear();
+}
+
+void DiscardOutputBackup()
+{
+    if( s_backupOf.empty() ) return;
+    std::error_code ec;
+    std::filesystem::remove( s_backupPath, ec );
+    s_backupOf.clear();
+}
+
+}

util/CaptureFileBackup.hpp

@@ -0,0 +1,23 @@
+#ifndef __CAPTUREFILEBACKUP_HPP__
+#define __CAPTUREFILEBACKUP_HPP__
+
+#include <stdint.h>
+
+namespace tracy
+{
+
+enum class OutputPrep : uint8_t
+{
+    Ok,
+    Exists,
+    Unusable
+};
+
+OutputPrep PrepareOutputFile( const char* path, bool overwrite, const char** error );
+
+void RestoreOutputBackup();
+void DiscardOutputBackup();
+
+}
+
+#endif