diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/llvm/include/llvm/Support/raw_ostream.h @@ -410,6 +410,12 @@ uint64_t pos = 0; + /// If set, the stream tries to lock underlying file before writing to it. + bool LockForUpdate = false; + + /// Time in milliseconds during which the stream tries to lock file. + unsigned LockTimeout = 1000; + /// See raw_ostream::write_impl. void write_impl(const char *Ptr, size_t Size) override; @@ -476,6 +482,15 @@ void enable_colors(bool enable) override { ColorEnabled = enable; } + /// Requires locking file before write to it. + void lockForUpdate(bool X) { LockForUpdate = X; } + bool lockForUpdate() const { return LockForUpdate; } + + /// Sets time in milliseconds during which the stream tries to lock + /// underlying file. + void lockTimeout(unsigned X) { LockTimeout = X; } + unsigned lockTimeout() const { return LockTimeout; } + std::error_code error() const { return EC; } /// Return the value of the flag in this raw_fd_ostream indicating whether an diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp --- a/llvm/lib/Support/raw_ostream.cpp +++ b/llvm/lib/Support/raw_ostream.cpp @@ -704,6 +704,12 @@ MaxWriteSize = 1024 * 1024 * 1024; #endif + if (LockForUpdate) + if (std::error_code EC = llvm::sys::fs::lockFile(FD, LockTimeout)) { + error_detected(EC); + return; + } + do { size_t ChunkSize = std::min(Size, MaxWriteSize); ssize_t ret = ::write(FD, Ptr, ChunkSize); @@ -735,6 +741,10 @@ Ptr += ret; Size -= ret; } while (Size > 0); + + if (LockForUpdate) + if (std::error_code EC = llvm::sys::fs::unlockFile(FD)) + error_detected(EC); } void raw_fd_ostream::close() { diff --git a/llvm/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp --- a/llvm/unittests/Support/raw_ostream_test.cpp +++ b/llvm/unittests/Support/raw_ostream_test.cpp @@ -7,7 +7,9 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/SmallString.h" +#include "llvm/Support/Errc.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/FileUtilities.h" #include "llvm/Support/Format.h" #include "llvm/Support/raw_ostream.h" #include "gtest/gtest.h" @@ -356,4 +358,41 @@ { raw_fd_ostream("-", EC, sys::fs::OpenFlags::OF_None); } { raw_fd_ostream("-", EC, sys::fs::OpenFlags::OF_None); } } + +#ifdef _WIN32 +// Windows refuses lock request if file region is already locked by the same +// process. POSIX system in this case updates the existing lock. +TEST(raw_fd_ostreamTest, LockForUpdate) { + int FD; + std::error_code EC; + SmallString<64> TempPath; + EC = sys::fs::createTemporaryFile("test", "temp", FD, TempPath); + ASSERT_FALSE(EC); + FileRemover Cleanup(TempPath); + + raw_fd_ostream Stream(TempPath, EC); + ASSERT_FALSE(EC); + Stream.lockForUpdate(true); + Stream.write("ABCDE", 5); + Stream.flush(); + ASSERT_FALSE(Stream.error()); + Stream.write("FGHIJ", 5); + Stream.flush(); + ASSERT_FALSE(Stream.error()); + + EC = sys::fs::lockFile(FD); + ASSERT_FALSE(EC); + Stream.write("12345", 5); + Stream.flush(); + ASSERT_EQ(errc::no_lock_available, Stream.error()); + + Stream.clear_error(); + EC = sys::fs::unlockFile(FD); + ASSERT_FALSE(EC); + Stream.write("abcde", 5); + Stream.flush(); + ASSERT_FALSE(Stream.error()); +} +#endif + }