From 0bd500eb2936b8e9457fa8d9a92f11b91686222f Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Fri, 6 Dec 2024 08:52:56 -0800 Subject: [PATCH] Add memory mapping and locking to file helpers PiperOrigin-RevId: 703514385 --- mediapipe/framework/BUILD | 6 + mediapipe/framework/deps/BUILD | 27 ++++ mediapipe/framework/deps/file_helpers.cc | 186 +++++++++++++++++++++- mediapipe/framework/deps/file_helpers.h | 8 + mediapipe/framework/deps/mlock_helpers.cc | 53 ++++++ mediapipe/framework/deps/mlock_helpers.h | 11 ++ mediapipe/framework/deps/mmapped_file.h | 46 ++++++ mediapipe/framework/port/BUILD | 1 + 8 files changed, 337 insertions(+), 1 deletion(-) create mode 100644 mediapipe/framework/deps/mlock_helpers.cc create mode 100644 mediapipe/framework/deps/mlock_helpers.h create mode 100644 mediapipe/framework/deps/mmapped_file.h diff --git a/mediapipe/framework/BUILD b/mediapipe/framework/BUILD index f5fc73815c..7ecac979ba 100644 --- a/mediapipe/framework/BUILD +++ b/mediapipe/framework/BUILD @@ -681,9 +681,15 @@ cc_library( hdrs = ["resources.h"], visibility = ["//visibility:public"], deps = [ + "//mediapipe/framework/deps:mlock_helpers", + "//mediapipe/framework/deps:mmapped_file", + "//mediapipe/framework/port:file_helpers", + "//mediapipe/framework/port:logging", + "//mediapipe/framework/port:status", "//mediapipe/framework/tool:status_util", "//mediapipe/util:resource_util", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log:absl_log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/mediapipe/framework/deps/BUILD b/mediapipe/framework/deps/BUILD index 1b87352367..da9835037d 100644 --- a/mediapipe/framework/deps/BUILD +++ b/mediapipe/framework/deps/BUILD @@ -117,6 +117,15 @@ cc_library( visibility = ["//visibility:public"], ) +cc_library( + name = "mmapped_file", + hdrs = ["mmapped_file.h"], + deps = [ + "//mediapipe/framework/port:logging", + "@com_google_absl//absl/status", + ], +) + cc_library( name = "file_helpers", srcs = ["file_helpers.cc"], @@ -124,6 +133,24 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":file_path", + ":mmapped_file", + ":platform_strings", + "//mediapipe/framework/formats:unique_fd", + "//mediapipe/framework/port:status", + "@com_google_absl//absl/base:config", + "@com_google_absl//absl/cleanup", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + ], +) + +cc_library( + name = "mlock_helpers", + srcs = ["mlock_helpers.cc"], + hdrs = ["mlock_helpers.h"], + visibility = ["//visibility:public"], + deps = [ ":platform_strings", "//mediapipe/framework/port:status", "@com_google_absl//absl/status", diff --git a/mediapipe/framework/deps/file_helpers.cc b/mediapipe/framework/deps/file_helpers.cc index b889867203..d444b452a4 100644 --- a/mediapipe/framework/deps/file_helpers.cc +++ b/mediapipe/framework/deps/file_helpers.cc @@ -17,27 +17,36 @@ #ifdef _WIN32 #include #include +#include #else #include +#include +#include #endif // _WIN32 #include #include #include #include +#include #include +#include +#include "absl/base/config.h" +#include "absl/cleanup/cleanup.h" // IWYU pragma: keep #include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "mediapipe/framework/deps/file_path.h" +#include "mediapipe/framework/deps/mmapped_file.h" #include "mediapipe/framework/deps/platform_strings.h" // IWYU pragma: keep +#include "mediapipe/framework/formats/unique_fd.h" #include "mediapipe/framework/port/status_macros.h" namespace mediapipe { namespace file { namespace { - // Helper class that returns all entries (files, directories) in a directory, // except "." and "..". Example usage: // @@ -197,6 +206,181 @@ absl::Status AppendStringToFile(absl::string_view path, return absl::OkStatus(); } +#ifdef _WIN32 +class WindowsMMap : public MemoryMappedFile { + public: + WindowsMMap(std::string path, const void* base_address, size_t length, + HANDLE file_handle, HANDLE mapping_handle) + : MemoryMappedFile(std::move(path), base_address, length), + file_handle_(file_handle), + mapping_handle_(mapping_handle) {} + + virtual absl::Status Close() override; + + private: + const HANDLE file_handle_; + const HANDLE mapping_handle_; +}; + +absl::StatusOr> MMapFile( + absl::string_view path) { + std::string name_string = std::string(path); + const HANDLE file_handle = CreateFile( + /*lpFileName=*/Utf8ToNative(name_string).c_str(), + /*dwDesiredAccess=*/GENERIC_READ, + /*dwShareMode=*/FILE_SHARE_READ, + /*lpSecurityAttributes=*/NULL, + /*dwCreationDisposition=*/OPEN_EXISTING, + /*dwFlagsAndAttributes=*/FILE_ATTRIBUTE_NORMAL, + /*hTemplateFile=*/NULL); + if (file_handle == INVALID_HANDLE_VALUE) { + return absl::UnavailableError( + absl::StrCat("Failed to open the file '", path, + "' for reading: ", FormatLastError())); + } + absl::Cleanup file_closer = [file_handle] { CloseHandle(file_handle); }; + + // We're calling `CreateFileMappingA` regardless of `UNICODE` because we don't + // pass the `lpName` string parameter. + const HANDLE mapping_handle = CreateFileMappingA( + /*hFile=*/file_handle, + /*lpFileMappingAttributes=*/NULL, + /*flProtect=*/PAGE_READONLY, + /*dwMaximumSizeHigh=*/0, // If `dwMaximumSize{Low,High} are zero, + /*dwMaximumSizeLow=*/0, // the maximum mapping size is the file size. + /*lpName=*/NULL); + if (mapping_handle == INVALID_HANDLE_VALUE) { + return absl::UnavailableError( + absl::StrCat("Failed to create a memory mapping for the file '", path, + "': ", FormatLastError())); + } + absl::Cleanup mapping_closer = [mapping_handle] { + CloseHandle(mapping_handle); + }; + + const LPVOID base_address = MapViewOfFile( + /*hFileMappingObject=*/mapping_handle, + /*dwDesiredAccess=*/FILE_MAP_READ, + /*dwFileOffsetHigh=*/0, + /*dwFileOffsetLow=*/0, + /*dwNumberOfBytesToMap=*/0 // Extends to the file end. + ); + if (base_address == NULL) { + return absl::UnavailableError(absl::StrCat( + "Failed to memory-map the file '", path, "': ", FormatLastError())); + } + + LARGE_INTEGER large_length; + const BOOL success = GetFileSizeEx(file_handle, &large_length); + if (!success) { + return absl::UnavailableError( + absl::StrCat("Failed to determine the size of the file '", path, + "': ", FormatLastError())); + } + const size_t length = static_cast(large_length.QuadPart); + + std::move(file_closer).Cancel(); + std::move(mapping_closer).Cancel(); + + return std::make_unique(std::move(name_string), base_address, + length, file_handle, mapping_handle); +} + +absl::Status WindowsMMap::Close() { + BOOL success = UnmapViewOfFile(BaseAddress()); + if (!success) { + return absl::UnavailableError(absl::StrCat( + "Failed to unmap the file '", Path(), "': ", FormatLastError())); + } + success = CloseHandle(mapping_handle_); + if (!success) { + return absl::UnavailableError( + absl::StrCat("Failed to close the memory mapping for file '", Path(), + "': " << FormatLastError())); + } + success = CloseHandle(file_handle_); + if (!success) { + return absl::UnavailableError(absl::StrCat( + "Failed to close the file '", Path(), "': ", FormatLastError())); + } + return absl::OkStatus(); +} +#elif ABSL_HAVE_MMAP +class PosixMMap : public MemoryMappedFile { + public: + PosixMMap(std::string path, const void* base_address, size_t length, + UniqueFd&& fd) + : MemoryMappedFile(path, base_address, length), + unique_fd_(std::move(fd)) {} + + absl::Status Close() override; + + private: + UniqueFd unique_fd_; +}; + +absl::StatusOr> MMapFile( + absl::string_view path) { + std::string name_string = std::string(path); + const int fd = open(name_string.c_str(), O_RDONLY); + if (fd < 0) { + return absl::UnavailableError(absl::StrCat( + "Couldn't open file '", path, "' for reading: ", FormatLastError())); + } + UniqueFd unique_fd(fd); + + struct stat file_stat; + const int status = fstat(unique_fd.Get(), &file_stat); + if (status < 0) { + return absl::UnavailableError( + absl::StrCat("Couldn't stat file '", path, "': ", FormatLastError())); + } + size_t length = file_stat.st_size; + + const void* base_address = + mmap(nullptr, length, PROT_READ, /*flags=*/0, unique_fd.Get(), + /*offset=*/0); + if (base_address == nullptr) { + return absl::UnavailableError(absl::StrCat( + "Couldn't map file '", path, "' into memory: ", FormatLastError())); + } + + return std::make_unique(std::move(name_string), base_address, + length, std::move(unique_fd)); +} + +absl::Status PosixMMap::Close() { + int status = munmap(const_cast(BaseAddress()), Length()); + if (status < 0) { + return absl::UnavailableError(absl::StrCat( + "Couldn't unmap file '", Path(), "' from memory: ", FormatLastError())); + } + + status = close(unique_fd_.Release()); + if (status < 0) { + return absl::UnavailableError(absl::StrCat("Couldn't close file '", Path(), + "': ", FormatLastError())); + } + return absl::OkStatus(); +} +#else // _WIN32 / ABSL_HAVE_MMAP +absl::StatusOr> MMapFile( + absl::string_view path) { + return absl::UnavailableError(absl::StrCat( + "No supported memory-mapping mechanism is provided for file '", path, + "'")); +} + +absl::Status LockMemory(const void* base_address, size_t length) { + return absl::UnavailableError("Locking memory unsupported"); +} + +absl::Status UnlockMemory(const void* base_address, size_t length) { + return absl::UnavailableError( + "Shouldn't attempt unlocking memory where locking is not supported"); +} +#endif // _WIN32 / ABSL_HAVE_MMAP + absl::Status MatchInTopSubdirectories(const std::string& parent_directory, const std::string& file_name, std::vector* results) { diff --git a/mediapipe/framework/deps/file_helpers.h b/mediapipe/framework/deps/file_helpers.h index 146ea5daf1..6724fa5427 100644 --- a/mediapipe/framework/deps/file_helpers.h +++ b/mediapipe/framework/deps/file_helpers.h @@ -16,7 +16,9 @@ #define MEDIAPIPE_FRAMEWORK_DEPS_FILE_HELPERS_H_ #include "absl/status/status.h" +#include "absl/status/statusor.h" #include "absl/strings/match.h" +#include "mediapipe/framework/deps/mmapped_file.h" namespace mediapipe { namespace file { @@ -29,6 +31,12 @@ absl::Status SetContents(absl::string_view file_name, absl::Status AppendStringToFile(absl::string_view file_name, absl::string_view contents); +absl::StatusOr> MMapFile( + absl::string_view path); + +absl::Status LockMemory(const void* base_address, size_t length); +absl::Status UnlockMemory(const void* base_address, size_t length); + absl::Status MatchInTopSubdirectories(const std::string& parent_directory, const std::string& file_name, std::vector* results); diff --git a/mediapipe/framework/deps/mlock_helpers.cc b/mediapipe/framework/deps/mlock_helpers.cc new file mode 100644 index 0000000000..63a994f6f4 --- /dev/null +++ b/mediapipe/framework/deps/mlock_helpers.cc @@ -0,0 +1,53 @@ +#include "mediapipe/framework/deps/mlock_helpers.h" + +#include + +#ifdef _WIN32 +#include +#else +#include +#endif + +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "mediapipe/framework/deps/platform_strings.h" + +namespace mediapipe { +#ifdef _WIN32 +absl::Status LockMemory(const void* base_address, size_t length) { + BOOL status = VirtualLock(const_cast(base_address), length); + if (!status) { + return absl::UnavailableError( + absl::StrCat("Failed to lock pages in memory: ", FormatLastError())); + } + return absl::OkStatus(); +} + +absl::Status UnlockMemory(const void* base_address, size_t length) { + BOOL status = VirtualUnlock(const_cast(base_address), length); + if (!status) { + return absl::UnavailableError( + absl::StrCat("Failed to unlock memory pages: ", FormatLastError())); + } + return absl::OkStatus(); +} +#else // _WIN32 +absl::Status LockMemory(const void* base_address, size_t length) { + int status = mlock(base_address, length); + if (status < 0) { + return absl::UnavailableError( + absl::StrCat("Failed to lock pages in memory: ", FormatLastError())); + } + return absl::OkStatus(); +} + +absl::Status UnlockMemory(const void* base_address, size_t length) { + int status = munlock(base_address, length); + if (status < 0) { + return absl::UnavailableError( + absl::StrCat("Failed to unlock memory pages: ", FormatLastError())); + } + return absl::OkStatus(); +} +#endif // _WIN32 +} // namespace mediapipe diff --git a/mediapipe/framework/deps/mlock_helpers.h b/mediapipe/framework/deps/mlock_helpers.h new file mode 100644 index 0000000000..1f4a1ab585 --- /dev/null +++ b/mediapipe/framework/deps/mlock_helpers.h @@ -0,0 +1,11 @@ +#ifndef MEDIAPIPE_FRAMEWORK_DEPS_MLOCK_HELPERS_H_ +#define MEDIAPIPE_FRAMEWORK_DEPS_MLOCK_HELPERS_H_ +#include "absl/status/status.h" + +namespace mediapipe { +// Uses `mlock`/`VirtualLock` to pin memory pages. +absl::Status LockMemory(const void* base_address, size_t length); +// Unlocks a previously locked memory region. +absl::Status UnlockMemory(const void* base_address, size_t length); +} // namespace mediapipe +#endif // MEDIAPIPE_FRAMEWORK_DEPS_MLOCK_HELPERS_H_ diff --git a/mediapipe/framework/deps/mmapped_file.h b/mediapipe/framework/deps/mmapped_file.h new file mode 100644 index 0000000000..3409762f9f --- /dev/null +++ b/mediapipe/framework/deps/mmapped_file.h @@ -0,0 +1,46 @@ +// Copyright 2024 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include "absl/status/status.h" +#include "mediapipe/framework/port/logging.h" + +#ifndef MEDIAPIPE_FRAMEWORK_DEPS_MMAPPED_FILE_H_ +#define MEDIAPIPE_FRAMEWORK_DEPS_MMAPPED_FILE_H_ +namespace mediapipe { +namespace file { +class MemoryMappedFile { + public: + MemoryMappedFile(std::string path, const void* base_address, size_t length) + : path_(std::move(path)), base_address_(base_address), length_(length) {} + + virtual absl::Status Close() = 0; + + virtual ~MemoryMappedFile() = default; + + const std::string& Path() const { return path_; } + const void* BaseAddress() const { return base_address_; } + size_t Length() const { return length_; } + + private: + std::string path_; + const void* base_address_; + size_t length_; +}; +} // namespace file +} // namespace mediapipe +#endif // MEDIAPIPE_FRAMEWORK_DEPS_MMAPPED_FILE_H_ diff --git a/mediapipe/framework/port/BUILD b/mediapipe/framework/port/BUILD index 31c8223044..985a4655f0 100644 --- a/mediapipe/framework/port/BUILD +++ b/mediapipe/framework/port/BUILD @@ -113,6 +113,7 @@ cc_library( ":status", "//mediapipe/framework/deps:file_helpers", "//mediapipe/framework/deps:file_path", + "//mediapipe/framework/deps:mmapped_file", ], )