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