Improvements to interal_os utilities.

Subprocess now supports Pipe redirects.

PiperOrigin-RevId: 738088103
Change-Id: Iadaf5a30f7ad16fc2fbe66be0f6c16d618a68c06

Author: laramiel

tensorstore/internal/BUILD

@@ -418,6 +418,7 @@ tensorstore_cc_library(
     hdrs = ["flat_cord_builder.h"],
     deps = [
         "//tensorstore/internal/os:memory_region",
+        "//tensorstore/util:span",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings:cord",

tensorstore/internal/flat_cord_builder.h

@@ -26,6 +26,7 @@
 #include "absl/log/absl_check.h"
 #include "absl/strings/cord.h"
 #include "tensorstore/internal/os/memory_region.h"
+#include "tensorstore/util/span.h"
 
 namespace tensorstore {
 namespace internal {
@@ -75,6 +76,10 @@ class FlatCordBuilder {
     inuse_ = size;
   }
 
+  tensorstore::span<char> available_span() {
+    return tensorstore::span(region_.data() + inuse_, available());
+  }
+
   /// Append data to the builder.
   void Append(std::string_view sv) {
     if (ABSL_PREDICT_FALSE(sv.empty())) return;
@@ -83,7 +88,12 @@ class FlatCordBuilder {
     inuse_ += sv.size();
   }
 
-  absl::Cord Build() && { return std::move(region_).as_cord(); }
+  absl::Cord Build() && {
+    if (inuse_ == region_.size()) {
+      return std::move(region_).as_cord();
+    }
+    return std::move(region_).as_cord().Subcord(0, inuse_);
+  }
 
  private:
   internal_os::MemoryRegion region_;

tensorstore/internal/os/BUILD

@@ -77,6 +77,7 @@ tensorstore_cc_library(
         "//tensorstore/internal/tracing",
         "//tensorstore/util:quote_string",
         "//tensorstore/util:result",
+        "//tensorstore/util:span",
         "//tensorstore/util:status",
         "@com_google_absl//absl/base:config",
         "@com_google_absl//absl/base:core_headers",
@@ -128,6 +129,7 @@ tensorstore_cc_test(
     deps = [
         ":file_util",
         "//tensorstore/internal/testing:scoped_directory",
+        "//tensorstore/util:span",
         "//tensorstore/util:status_testutil",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
@@ -170,6 +172,7 @@ tensorstore_cc_test(
         ":file_util",
         "//tensorstore/internal/testing:scoped_directory",
         "//tensorstore/util:result",
+        "//tensorstore/util:span",
         "//tensorstore/util:status",
         "//tensorstore/util:status_testutil",
         "@com_google_absl//absl/log:absl_check",
@@ -191,12 +194,13 @@ tensorstore_cc_library(
     tags = ["msvc"],
     deps = [
         ":error_code",
+        ":file_util",
         ":include_windows",
         ":wstring",
         "//tensorstore/util:result",
         "//tensorstore/util:status",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/status",
@@ -212,21 +216,20 @@ tensorstore_cc_test(
         "msvc",
     ],
     deps = [
+        ":file_util",
         ":subprocess",
         "//tensorstore/internal:env",
         "//tensorstore/internal:path",
         "//tensorstore/internal/testing:scoped_directory",
         "//tensorstore/util:result",
-        "//tensorstore/util:status",
+        "//tensorstore/util:span",
         "//tensorstore/util:status_testutil",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest",
-        "@com_google_riegeli//riegeli/bytes:fd_reader",
-        "@com_google_riegeli//riegeli/bytes:read_all",
     ],
 )
 
@@ -286,7 +289,6 @@ tensorstore_cc_library(
     hdrs = ["memory_region.h"],
     deps = [
         "//tensorstore/util:result",
-        "@com_google_absl//absl/debugging:leak_check",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/strings:cord",

tensorstore/internal/os/file_lister_test.cc

@@ -28,6 +28,7 @@
 #include "tensorstore/internal/os/file_util.h"
 #include "tensorstore/internal/testing/scoped_directory.h"
 #include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
 #include "tensorstore/util/status.h"
 #include "tensorstore/util/status_testutil.h"
 
@@ -39,10 +40,9 @@ using ::tensorstore::internal_os::FsyncDirectory;
 using ::tensorstore::internal_os::FsyncFile;
 using ::tensorstore::internal_os::MakeDirectory;
 using ::tensorstore::internal_os::OpenDirectoryDescriptor;
-using ::tensorstore::internal_os::OpenExistingFileForReading;
 using ::tensorstore::internal_os::OpenFileWrapper;
 using ::tensorstore::internal_os::OpenFlags;
-using ::tensorstore::internal_os::ReadFromFile;
+using ::tensorstore::internal_os::PReadFromFile;
 using ::tensorstore::internal_os::RecursiveFileList;
 using ::tensorstore::internal_os::WriteToFile;
 using ::tensorstore::internal_testing::ScopedTemporaryDirectory;
@@ -191,8 +191,8 @@ TEST(RecursiveFileListEntryTest, DeleteWithOpenFile) {
 
   {
     // Open file; it should be deleted.
-    auto f =
-        OpenExistingFileForReading(absl::StrCat(tmpdir.path(), "/read.txt"));
+    auto f = OpenFileWrapper(absl::StrCat(tmpdir.path(), "/read.txt"),
+                             OpenFlags::DefaultRead);
     EXPECT_THAT(f, IsOk());
 
     std::vector<std::string> files;
@@ -213,7 +213,8 @@ TEST(RecursiveFileListEntryTest, DeleteWithOpenFile) {
                 IsOk());
 
     char buf[16];
-    EXPECT_THAT(ReadFromFile(f->get(), buf, 3, 0), IsOkAndHolds(3));
+    EXPECT_THAT(PReadFromFile(f->get(), tensorstore::span(buf, 3), 0),
+                IsOkAndHolds(3));
   }
 
   std::vector<std::string> files;

tensorstore/internal/os/file_lock.cc

@@ -53,7 +53,7 @@ absl::Status FileLock::Delete() && {
   assert(fd_ != FileDescriptorTraits::Invalid());
 
   auto fd = std::exchange(fd_, FileDescriptorTraits::Invalid());
-  auto status = internal_os::DeleteOpenFile(fd, lock_path_);
+  auto status = DeleteOpenFile(fd, lock_path_);
   Unlock(fd);
   FileDescriptorTraits::Close(fd);
   return MaybeAnnotateStatus(std::move(status), "Failed to clean lock file");
@@ -70,13 +70,13 @@ Result<FileLock> AcquireFileLock(std::string lock_path) {
   using private_t = FileLock::private_t;
   TENSORSTORE_ASSIGN_OR_RETURN(
       UniqueFileDescriptor fd,
-      internal_os::OpenFileWrapper(lock_path, OpenFlags::DefaultWrite));
+      OpenFileWrapper(lock_path, OpenFlags::DefaultWrite));
   FileInfo a, b;
   FileInfo* info = &a;
 
   // Is this a network filesystem?
-  TENSORSTORE_RETURN_IF_ERROR(internal_os::GetFileInfo(fd.get(), info));
-  if (!internal_os::IsRegularFile(*info)) {
+  TENSORSTORE_RETURN_IF_ERROR(GetFileInfo(fd.get(), info));
+  if (!IsRegularFile(*info)) {
     return absl::FailedPreconditionError(
         absl::StrCat("Not a regular file: ", lock_path));
   }
@@ -85,20 +85,19 @@ Result<FileLock> AcquireFileLock(std::string lock_path) {
   while (true) {
     // Acquire lock.
     TENSORSTORE_ASSIGN_OR_RETURN(
-        auto unlock_fn, internal_os::AcquireFdLock(fd.get()),
+        auto unlock_fn, AcquireFdLock(fd.get()),
         MaybeAnnotateStatus(_, absl::StrCat("Failed to acquire lock on file: ",
                                             QuoteString(lock_path))));
 
     // Reopening the file should give the same value since the lock is held.
     TENSORSTORE_ASSIGN_OR_RETURN(
         UniqueFileDescriptor other_fd,
-        internal_os::OpenFileWrapper(lock_path, OpenFlags::DefaultWrite));
+        OpenFileWrapper(lock_path,
+                        OpenFlags::DefaultWrite | OpenFlags::CloseOnExec));
 
     FileInfo* other_info = info == &a ? &b : &a;
-    TENSORSTORE_RETURN_IF_ERROR(
-        internal_os::GetFileInfo(other_fd.get(), other_info));
-    if (internal_os::GetDeviceId(a) == internal_os::GetDeviceId(b) &&
-        internal_os::GetFileId(a) == internal_os::GetFileId(b)) {
+    TENSORSTORE_RETURN_IF_ERROR(GetFileInfo(other_fd.get(), other_info));
+    if (GetDeviceId(a) == GetDeviceId(b) && GetFileId(a) == GetFileId(b)) {
       // Lock was acquired successfully.
       return FileLock(private_t(), std::move(lock_path), fd.release(),
                       std::move(unlock_fn));
@@ -122,10 +121,11 @@ Result<FileLock> AcquireExclusiveFile(std::string lock_path,
 
   // Determine whether the lock file is stale.
   auto detect_stale_lock = [&]() mutable {
-    auto read_fd = OpenFileWrapper(lock_path, OpenFlags::OpenReadOnly);
+    auto read_fd = OpenFileWrapper(
+        lock_path, OpenFlags::OpenReadOnly | OpenFlags::CloseOnExec);
     if (read_fd.ok()) {
       TENSORSTORE_RETURN_IF_ERROR(GetFileInfo(read_fd->get(), &info));
-      if (!internal_os::IsRegularFile(info)) {
+      if (!IsRegularFile(info)) {
         // A lock file must be a regular file, not a symlink or directory.
         return absl::FailedPreconditionError(
             absl::StrCat("Not a regular file: ", lock_path));
@@ -148,11 +148,11 @@ Result<FileLock> AcquireExclusiveFile(std::string lock_path,
     n++;
     if (m > 1000) m = 1000;
 
-    auto fd = internal_os::OpenFileWrapper(
-        lock_path,
-        OpenFlags::Create | OpenFlags::Exclusive | OpenFlags::OpenReadWrite);
+    auto fd = OpenFileWrapper(
+        lock_path, OpenFlags::Create | OpenFlags::Exclusive |
+                       OpenFlags::OpenReadWrite | OpenFlags::CloseOnExec);
     if (fd.ok()) {
-      TENSORSTORE_RETURN_IF_ERROR(internal_os::GetFileInfo(fd->get(), &info));
+      TENSORSTORE_RETURN_IF_ERROR(GetFileInfo(fd->get(), &info));
       return FileLock(private_t{}, std::move(lock_path), fd->release(),
                       std::nullopt);
     }

tensorstore/internal/os/file_util.cc

@@ -21,17 +21,18 @@
 #include <string>
 
 #include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
 #include "tensorstore/util/status.h"
 
 namespace tensorstore {
 namespace internal_os {
 
 Result<std::string> ReadAllToString(const std::string& path) {
-  TENSORSTORE_ASSIGN_OR_RETURN(
-      auto fd, OpenFileWrapper(path, internal_os::OpenFlags::OpenReadOnly));
+  TENSORSTORE_ASSIGN_OR_RETURN(auto fd,
+                               OpenFileWrapper(path, OpenFlags::OpenReadOnly));
 
-  internal_os::FileInfo info;
-  TENSORSTORE_RETURN_IF_ERROR(internal_os::GetFileInfo(fd.get(), &info));
+  FileInfo info;
+  TENSORSTORE_RETURN_IF_ERROR(GetFileInfo(fd.get(), &info));
 
   // Handle the case where the file is empty.
   std::string result(internal_os::GetSize(info), 0);
@@ -42,7 +43,8 @@ Result<std::string> ReadAllToString(const std::string& path) {
   size_t offset = 0;
   TENSORSTORE_ASSIGN_OR_RETURN(
       auto read,
-      internal_os::ReadFromFile(fd.get(), &result[0], result.size(), 0));
+      PReadFromFile(fd.get(), tensorstore::span(&result[0], result.size()),
+                    offset));
   offset += read;
 
   while (true) {
@@ -53,7 +55,8 @@ Result<std::string> ReadAllToString(const std::string& path) {
       char buffer[4096];
       TENSORSTORE_ASSIGN_OR_RETURN(
           read,
-          internal_os::ReadFromFile(fd.get(), buffer, sizeof(buffer), offset));
+          PReadFromFile(fd.get(), tensorstore::span(buffer, sizeof(buffer)),
+                        offset));
       if (read > 0) {
         // Amortized resize; double the size of the buffer.
         if (read > result.size()) {
@@ -63,8 +66,10 @@ Result<std::string> ReadAllToString(const std::string& path) {
       }
     } else {
       TENSORSTORE_ASSIGN_OR_RETURN(
-          read, internal_os::ReadFromFile(fd.get(), &result[offset],
-                                          result.size() - offset, offset));
+          read, PReadFromFile(
+                    fd.get(),
+                    tensorstore::span(&result[offset], result.size() - offset),
+                    offset));
     }
     if (read == 0) {
       result.resize(offset);

tensorstore/internal/os/file_util.h

@@ -24,14 +24,14 @@
 
 #include <string>
 #include <string_view>
-#include <utility>
 
 #include "absl/status/status.h"
 #include "absl/strings/cord.h"
 #include "absl/time/time.h"
 #include "tensorstore/internal/os/memory_region.h"
 #include "tensorstore/internal/os/unique_handle.h"
 #include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
 
 // Include system headers last to reduce impact of macros.
 #ifndef _WIN32
@@ -119,8 +119,10 @@ enum class OpenFlags : int {
   Create = O_CREAT,
   Append = O_APPEND,
   Exclusive = O_EXCL,
+  CloseOnExec = O_CLOEXEC,
 
-  DefaultWrite = O_CREAT | O_WRONLY,
+  DefaultRead = O_RDONLY | O_CLOEXEC,
+  DefaultWrite = O_CREAT | O_WRONLY | O_CLOEXEC,
 };
 
 inline OpenFlags operator|(OpenFlags a, OpenFlags b) {
@@ -136,20 +138,22 @@ inline OpenFlags operator&(OpenFlags a, OpenFlags b) {
 Result<UniqueFileDescriptor> OpenFileWrapper(const std::string& path,
                                              OpenFlags flags);
 
-inline Result<UniqueFileDescriptor> OpenExistingFileForReading(
-    const std::string& path) {
-  return OpenFileWrapper(path, OpenFlags::OpenReadOnly);
-}
+/// Reads from an open file.
+///
+/// \param fd Open file descriptor.
+/// \param buffer[out] Pointer to memory where data will be stored.
+/// \returns Number of bytes read or a failure absl::Status code.
+Result<ptrdiff_t> ReadFromFile(FileDescriptor fd,
+                               tensorstore::span<char> buffer);
 
 /// Reads from an open file.
 ///
 /// \param fd Open file descriptor.
-/// \param buf[out] Pointer to memory where data will be stored.
-/// \param count Maximum number of bytes to read.
+/// \param buffer[out] Pointer to memory where data will be stored.
 /// \param offset Byte offset within file at which to start reading.
 /// \returns Number of bytes read or a failure absl::Status code.
-Result<ptrdiff_t> ReadFromFile(FileDescriptor fd, void* buf, size_t count,
-                               int64_t offset);
+Result<ptrdiff_t> PReadFromFile(FileDescriptor fd,
+                                tensorstore::span<char> buffer, int64_t offset);
 
 /// Reads the entire file into a string.
 ///

tensorstore/internal/os/file_util_posix.cc

@@ -57,6 +57,7 @@
 #include "tensorstore/internal/tracing/logged_trace_span.h"
 #include "tensorstore/util/quote_string.h"
 #include "tensorstore/util/result.h"
+#include "tensorstore/util/span.h"
 #include "tensorstore/util/status.h"
 
 #if ABSL_HAVE_MMAP
@@ -257,15 +258,36 @@ Result<UniqueFileDescriptor> OpenFileWrapper(const std::string& path,
   }
 }
 
-Result<ptrdiff_t> ReadFromFile(FileDescriptor fd, void* buf, size_t count,
-                               int64_t offset) {
+Result<ptrdiff_t> ReadFromFile(FileDescriptor fd,
+                               tensorstore::span<char> buffer) {
   LoggedTraceSpan tspan(__func__, detail_logging.Level(1),
-                        {{"fd", fd}, {"count", count}, {"offset", offset}});
+                        {{"fd", fd}, {"count", buffer.size()}});
   ssize_t n;
   do {
     PotentiallyBlockingRegion region;
-    n = ::pread(fd, buf, count, static_cast<off_t>(offset));
-  } while ((n < 0) && (errno == EINTR || errno == EAGAIN));
+    n = ::read(fd, buffer.data(), buffer.size());
+  } while ((n < 0) && (errno == EINTR));
+  if (n >= 0) {
+    return n;
+  } else if (errno == EAGAIN || errno == EWOULDBLOCK) {
+    return 0;
+  }
+  auto status = StatusFromOsError(errno, "Failed to read from file");
+  return std::move(tspan).EndWithStatus(std::move(status));
+}
+
+Result<ptrdiff_t> PReadFromFile(FileDescriptor fd,
+                                tensorstore::span<char> buffer,
+                                int64_t offset) {
+  LoggedTraceSpan tspan(
+      __func__, detail_logging.Level(1),
+      {{"fd", fd}, {"count", buffer.size()}, {"offset", offset}});
+  ssize_t n;
+  do {
+    PotentiallyBlockingRegion region;
+    n = ::pread(fd, buffer.data(), buffer.size(), static_cast<off_t>(offset));
+  } while ((n < 0) &&
+           (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK));
   if (n >= 0) {
     return n;
   }