123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200 |
- #pragma once
- #include <c10/util/Exception.h>
- #include <c10/util/Optional.h>
- #include <cerrno>
- #include <cstdio>
- #include <cstdlib>
- #include <cstring>
- #include <string>
- #include <utility>
- #include <vector>
- #if !defined(_WIN32)
- #include <unistd.h>
- #else // defined(_WIN32)
- #include <Windows.h>
- #include <fileapi.h>
- #endif // defined(_WIN32)
- namespace c10 {
- namespace detail {
- // Creates the filename pattern passed to and completed by `mkstemp`.
- // Returns std::vector<char> because `mkstemp` needs a (non-const) `char*` and
- // `std::string` only provides `const char*` before C++17.
- #if !defined(_WIN32)
- inline std::vector<char> make_filename(std::string name_prefix) {
- // The filename argument to `mkstemp` needs "XXXXXX" at the end according to
- // http://pubs.opengroup.org/onlinepubs/009695399/functions/mkstemp.html
- static const std::string kRandomPattern = "XXXXXX";
- // We see if any of these environment variables is set and use their value, or
- // else default the temporary directory to `/tmp`.
- static const char* env_variables[] = {"TMPDIR", "TMP", "TEMP", "TEMPDIR"};
- std::string tmp_directory = "/tmp";
- for (const char* variable : env_variables) {
- if (const char* path = getenv(variable)) {
- tmp_directory = path;
- break;
- }
- }
- std::vector<char> filename;
- filename.reserve(
- tmp_directory.size() + name_prefix.size() + kRandomPattern.size() + 2);
- filename.insert(filename.end(), tmp_directory.begin(), tmp_directory.end());
- filename.push_back('/');
- filename.insert(filename.end(), name_prefix.begin(), name_prefix.end());
- filename.insert(filename.end(), kRandomPattern.begin(), kRandomPattern.end());
- filename.push_back('\0');
- return filename;
- }
- #endif // !defined(_WIN32)
- } // namespace detail
- struct TempFile {
- #if !defined(_WIN32)
- TempFile() : fd(-1) {}
- TempFile(std::string name, int fd) : fd(fd), name(std::move(name)) {}
- TempFile(const TempFile&) = delete;
- TempFile(TempFile&& other) noexcept
- : fd(other.fd), name(std::move(other.name)) {
- other.fd = -1;
- other.name.clear();
- }
- TempFile& operator=(const TempFile&) = delete;
- TempFile& operator=(TempFile&& other) noexcept {
- fd = other.fd;
- name = std::move(other.name);
- other.fd = -1;
- other.name.clear();
- return *this;
- }
- ~TempFile() {
- if (fd >= 0) {
- unlink(name.c_str());
- close(fd);
- }
- }
- int fd;
- #endif // !defined(_WIN32)
- std::string name;
- };
- struct TempDir {
- TempDir() = default;
- explicit TempDir(std::string name) : name(std::move(name)) {}
- TempDir(const TempDir&) = delete;
- TempDir(TempDir&& other) noexcept : name(std::move(other.name)) {
- other.name.clear();
- }
- TempDir& operator=(const TempDir&) = delete;
- TempDir& operator=(TempDir&& other) noexcept {
- name = std::move(other.name);
- other.name.clear();
- return *this;
- }
- ~TempDir() {
- if (!name.empty()) {
- #if !defined(_WIN32)
- rmdir(name.c_str());
- #else // defined(_WIN32)
- RemoveDirectoryA(name.c_str());
- #endif // defined(_WIN32)
- }
- }
- std::string name;
- };
- /// Attempts to return a temporary file or returns `nullopt` if an error
- /// occurred.
- ///
- /// The file returned follows the pattern
- /// `<tmp-dir>/<name-prefix><random-pattern>`, where `<tmp-dir>` is the value of
- /// the `"TMPDIR"`, `"TMP"`, `"TEMP"` or
- /// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`;
- /// `<name-prefix>` is the value supplied to this function, and
- /// `<random-pattern>` is a random sequence of numbers.
- /// On Windows, `name_prefix` is ignored and `tmpnam` is used.
- inline c10::optional<TempFile> try_make_tempfile(
- std::string name_prefix = "torch-file-") {
- #if defined(_WIN32)
- return TempFile{std::tmpnam(nullptr)};
- #else
- std::vector<char> filename = detail::make_filename(std::move(name_prefix));
- const int fd = mkstemp(filename.data());
- if (fd == -1) {
- return c10::nullopt;
- }
- // Don't make the string from string(filename.begin(), filename.end(), or
- // there will be a trailing '\0' at the end.
- return TempFile(filename.data(), fd);
- #endif // defined(_WIN32)
- }
- /// Like `try_make_tempfile`, but throws an exception if a temporary file could
- /// not be returned.
- inline TempFile make_tempfile(std::string name_prefix = "torch-file-") {
- if (auto tempfile = try_make_tempfile(std::move(name_prefix))) {
- return std::move(*tempfile);
- }
- TORCH_CHECK(false, "Error generating temporary file: ", std::strerror(errno));
- }
- /// Attempts to return a temporary directory or returns `nullopt` if an error
- /// occurred.
- ///
- /// The directory returned follows the pattern
- /// `<tmp-dir>/<name-prefix><random-pattern>/`, where `<tmp-dir>` is the value
- /// of the `"TMPDIR"`, `"TMP"`, `"TEMP"` or
- /// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`;
- /// `<name-prefix>` is the value supplied to this function, and
- /// `<random-pattern>` is a random sequence of numbers.
- /// On Windows, `name_prefix` is ignored and `tmpnam` is used.
- inline c10::optional<TempDir> try_make_tempdir(
- std::string name_prefix = "torch-dir-") {
- #if defined(_WIN32)
- while (true) {
- const char* dirname = std::tmpnam(nullptr);
- if (!dirname) {
- return c10::nullopt;
- }
- if (CreateDirectoryA(dirname, NULL)) {
- return TempDir(dirname);
- }
- if (GetLastError() != ERROR_ALREADY_EXISTS) {
- return c10::nullopt;
- }
- }
- return c10::nullopt;
- #else
- std::vector<char> filename = detail::make_filename(std::move(name_prefix));
- const char* dirname = mkdtemp(filename.data());
- if (!dirname) {
- return c10::nullopt;
- }
- return TempDir(dirname);
- #endif // defined(_WIN32)
- }
- /// Like `try_make_tempdir`, but throws an exception if a temporary directory
- /// could not be returned.
- inline TempDir make_tempdir(std::string name_prefix = "torch-dir-") {
- if (auto tempdir = try_make_tempdir(std::move(name_prefix))) {
- return std::move(*tempdir);
- }
- TORCH_CHECK(
- false, "Error generating temporary directory: ", std::strerror(errno));
- }
- } // namespace c10
|