1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168
|
#include <c10/util/Exception.h>
#include <c10/util/env.h>
#include <c10/util/error.h>
#include <c10/util/tempfile.h>
#include <fmt/format.h>
#if !defined(_WIN32)
#include <unistd.h>
#include <cerrno>
#else // defined(_WIN32)
#include <Windows.h>
#include <fcntl.h>
#include <fileapi.h>
#include <io.h>
#endif // defined(_WIN32)
// Creates the filename pattern passed to and completed by `mkstemp`.
#if !defined(_WIN32)
static std::string make_filename(std::string_view name_prefix) {
// The filename argument to `mkstemp` needs "XXXXXX" at the end according to
// http://pubs.opengroup.org/onlinepubs/009695399/functions/mkstemp.html
constexpr const char* kRandomPattern = "XXXXXX";
// We see if any of these environment variables is set and use their value, or
// else default the temporary directory to `/tmp`.
std::string tmp_directory = "/tmp";
for (const char* variable : {"TMPDIR", "TMP", "TEMP", "TEMPDIR"}) {
auto path_opt = c10::utils::get_env(variable);
if (path_opt.has_value()) {
tmp_directory = path_opt.value();
break;
}
}
return fmt::format("{}/{}{}", tmp_directory, name_prefix, kRandomPattern);
}
#else
static std::string make_filename() {
char name[L_tmpnam_s]{};
auto res = tmpnam_s(name, L_tmpnam_s);
if (res != 0) {
TORCH_WARN("Error generating temporary file");
return "";
}
return name;
}
#endif // !defined(_WIN32)
namespace c10 {
/// Attempts to return a temporary file or returns `nullopt` if an error
/// occurred.
std::optional<TempFile> try_make_tempfile(std::string_view name_prefix) {
#if defined(_WIN32)
auto filename = make_filename();
#else
auto filename = make_filename(name_prefix);
#endif
if (filename.empty()) {
return std::nullopt;
}
#if defined(_WIN32)
return TempFile(std::move(filename));
#else
const int fd = mkstemp(filename.data());
if (fd == -1) {
return std::nullopt;
}
return TempFile(std::move(filename), fd);
#endif // defined(_WIN32)
}
/// Like `try_make_tempfile`, but throws an exception if a temporary file could
/// not be returned.
TempFile make_tempfile(std::string_view name_prefix) {
if (auto tempfile = try_make_tempfile(name_prefix)) {
return std::move(*tempfile);
}
TORCH_CHECK(
false, "Error generating temporary file: ", c10::utils::str_error(errno));
}
/// Attempts to return a temporary directory or returns `nullopt` if an error
/// occurred.
std::optional<TempDir> try_make_tempdir(std::string_view name_prefix) {
#if defined(_WIN32)
for (int i = 0; i < 10; i++) {
auto dirname = make_filename();
if (dirname.empty()) {
return std::nullopt;
}
if (CreateDirectoryA(dirname.c_str(), nullptr)) {
return TempDir(dirname);
}
if (GetLastError() == ERROR_SUCCESS) {
return std::nullopt;
}
}
return std::nullopt;
#else
auto filename = make_filename(name_prefix);
const char* dirname = mkdtemp(filename.data());
if (!dirname) {
return std::nullopt;
}
return TempDir(dirname);
#endif // defined(_WIN32)
}
#if defined(_WIN32)
bool TempFile::open() {
if (fd != -1) {
return false;
}
auto err = _sopen_s(
&fd,
name.c_str(),
_O_CREAT | _O_TEMPORARY | _O_EXCL | _O_BINARY | _O_RDWR,
_SH_DENYNO,
_S_IREAD | _S_IWRITE);
if (err != 0) {
fd = -1;
return false;
}
return true;
}
#endif
TempFile::~TempFile() {
if (!name.empty()) {
#if !defined(_WIN32)
if (fd >= 0) {
unlink(name.c_str());
close(fd);
}
#else
if (fd >= 0) {
_close(fd);
}
#endif
}
}
TempDir::~TempDir() {
if (!name.empty()) {
#if !defined(_WIN32)
rmdir(name.c_str());
#else // defined(_WIN32)
RemoveDirectoryA(name.c_str());
#endif // defined(_WIN32)
}
}
/// Like `try_make_tempdir`, but throws an exception if a temporary directory
/// could not be returned.
TempDir make_tempdir(std::string_view name_prefix) {
if (auto tempdir = try_make_tempdir(name_prefix)) {
return std::move(*tempdir);
}
#if !defined(_WIN32)
TORCH_CHECK(
false,
"Error generating temporary directory: ",
c10::utils::str_error(errno));
#else // defined(_WIN32)
TORCH_CHECK(false, "Error generating temporary directory");
#endif // defined(_WIN32)
}
} // namespace c10
|