From 83e33b4dbe580a0879987925d62b845e97397c0f Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Wed, 1 Feb 2023 09:53:00 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 506345436 --- mediapipe/tasks/cc/core/BUILD | 13 ++-- .../tasks/cc/core/external_file_handler.cc | 77 +++++++++++++++---- 2 files changed, 72 insertions(+), 18 deletions(-) diff --git a/mediapipe/tasks/cc/core/BUILD b/mediapipe/tasks/cc/core/BUILD index d440271df..e5bc18306 100644 --- a/mediapipe/tasks/cc/core/BUILD +++ b/mediapipe/tasks/cc/core/BUILD @@ -40,16 +40,19 @@ cc_library( srcs = ["external_file_handler.cc"], hdrs = ["external_file_handler.h"], deps = [ - "//mediapipe/framework/port:integral_types", - "//mediapipe/framework/port:status", - "//mediapipe/tasks/cc:common", - "//mediapipe/tasks/cc/core/proto:external_file_cc_proto", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", - ], + "//mediapipe/framework/port:integral_types", + "//mediapipe/framework/port:status", + "//mediapipe/tasks/cc:common", + "//mediapipe/tasks/cc/core/proto:external_file_cc_proto", + ] + select({ + "//mediapipe:windows": ["@bazel_tools//tools/cpp/runfiles"], + "//conditions:default": [], + }), ) cc_library( diff --git a/mediapipe/tasks/cc/core/external_file_handler.cc b/mediapipe/tasks/cc/core/external_file_handler.cc index ff30bea72..a95b8e744 100644 --- a/mediapipe/tasks/cc/core/external_file_handler.cc +++ b/mediapipe/tasks/cc/core/external_file_handler.cc @@ -37,12 +37,17 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/match.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "mediapipe/framework/port/status_macros.h" #include "mediapipe/tasks/cc/common.h" #include "mediapipe/tasks/cc/core/proto/external_file.pb.h" +#ifdef _WIN32 +#include "tools/cpp/runfiles/runfiles.h" +#endif // _WIN32 + namespace mediapipe { namespace tasks { namespace core { @@ -50,13 +55,21 @@ namespace { using ::absl::StatusCode; +#ifndef O_BINARY +#ifdef _O_BINARY +#define O_BINARY _O_BINARY +#else +#define O_BINARY 0 // If this isn't defined, the platform doesn't need it. +#endif // _O_BINARY +#endif // O_BINARY + // Gets the offset aligned to page size for mapping given files into memory by // file descriptor correctly, as according to mmap(2), the offset used in mmap // must be a multiple of sysconf(_SC_PAGE_SIZE). int64 GetPageSizeAlignedOffset(int64 offset) { #ifdef _WIN32 // mmap is not used on Windows - return -1; + return 0; #else int64 aligned_offset = offset; int64 page_size = sysconf(_SC_PAGE_SIZE); @@ -64,7 +77,7 @@ int64 GetPageSizeAlignedOffset(int64 offset) { aligned_offset = offset / page_size * page_size; } return aligned_offset; -#endif +#endif // _WIN32 } } // namespace @@ -83,6 +96,24 @@ ExternalFileHandler::CreateFromExternalFile( return handler; } +absl::StatusOr PathToResourceAsFile(std::string path) { +#ifndef _WIN32 + return path; +#else + if (absl::StartsWith(path, "./")) { + path = "mediapipe" + path.substr(1); + } + + std::string error; + std::unique_ptr<::bazel::tools::cpp::runfiles::Runfiles> runfiles( + ::bazel::tools::cpp::runfiles::Runfiles::Create("", &error)); + if (!runfiles) { + return absl::InternalError("Unable to initialize runfiles: " + error); + } + return runfiles->Rlocation(path); +#endif // _WIN32 +} + absl::Status ExternalFileHandler::MapExternalFile() { if (!external_file_.file_content().empty()) { return absl::OkStatus(); @@ -101,12 +132,6 @@ absl::Status ExternalFileHandler::MapExternalFile() { return absl::OkStatus(); } -// TODO: Add Windows support -#ifdef _WIN32 - return CreateStatusWithPayload(StatusCode::kFailedPrecondition, - "File loading is not yet supported on Windows", - MediaPipeTasksStatus::kFileReadError); -#else if (external_file_.file_name().empty() && !external_file_.has_file_descriptor_meta()) { return CreateStatusWithPayload( @@ -118,7 +143,9 @@ absl::Status ExternalFileHandler::MapExternalFile() { // Obtain file descriptor, offset and size. int fd = -1; if (!external_file_.file_name().empty()) { - owned_fd_ = open(external_file_.file_name().c_str(), O_RDONLY); + ASSIGN_OR_RETURN(std::string file_name, + PathToResourceAsFile(external_file_.file_name())); + owned_fd_ = open(file_name.c_str(), O_RDONLY | O_BINARY); if (owned_fd_ < 0) { const std::string error_message = absl::StrFormat( "Unable to open file at %s", external_file_.file_name()); @@ -149,6 +176,12 @@ absl::Status ExternalFileHandler::MapExternalFile() { } fd = owned_fd_; } else { +#ifdef _WIN32 + return CreateStatusWithPayload( + StatusCode::kFailedPrecondition, + "File descriptors are not supported on Windows.", + MediaPipeTasksStatus::kFileReadError); +#else fd = external_file_.file_descriptor_meta().fd(); if (fd < 0) { return CreateStatusWithPayload( @@ -158,6 +191,7 @@ absl::Status ExternalFileHandler::MapExternalFile() { } buffer_offset_ = external_file_.file_descriptor_meta().offset(); buffer_size_ = external_file_.file_descriptor_meta().length(); +#endif // _WIN32 } // Get actual file size. Always use 0 as offset to lseek(2) to get the actual // file size, as SEEK_END returns the size of the file *plus* offset. @@ -189,22 +223,37 @@ absl::Status ExternalFileHandler::MapExternalFile() { buffer_size_ + buffer_offset_, file_size), MediaPipeTasksStatus::kInvalidArgumentError); } + // If buffer_offset_ is not multiple of sysconf(_SC_PAGE_SIZE), align with // extra leading bytes and adjust buffer_size_ to account for the extra // leading bytes. buffer_aligned_offset_ = GetPageSizeAlignedOffset(buffer_offset_); buffer_aligned_size_ = buffer_size_ + buffer_offset_ - buffer_aligned_offset_; + +#ifdef _WIN32 + buffer_ = malloc(file_size); + // Return the file pointer back to the beginning of the file + lseek(fd, 0L, SEEK_SET); + buffer_size_ = read(fd, buffer_, file_size); + if (buffer_size_ <= 0) { + free(buffer_); + buffer_ = nullptr; + } +#else // Map into memory. buffer_ = mmap(/*addr=*/nullptr, buffer_aligned_size_, PROT_READ, MAP_SHARED, fd, buffer_aligned_offset_); if (buffer_ == MAP_FAILED) { + buffer_ = nullptr; + } +#endif // _WIN32 + if (!buffer_) { return CreateStatusWithPayload( StatusCode::kUnknown, absl::StrFormat("Unable to map file to memory buffer, errno=%d", errno), MediaPipeTasksStatus::kFileMmapError); } return absl::OkStatus(); -#endif } absl::string_view ExternalFileHandler::GetFileContent() { @@ -223,11 +272,13 @@ absl::string_view ExternalFileHandler::GetFileContent() { } ExternalFileHandler::~ExternalFileHandler() { -#ifndef _WIN32 - if (buffer_ != MAP_FAILED) { + if (buffer_) { +#ifdef _WIN32 + free(buffer_); +#else munmap(buffer_, buffer_aligned_size_); +#endif // _WIN32 } -#endif if (owned_fd_ >= 0) { close(owned_fd_); }