1 | #pragma once |
2 | |
3 | #include <c10/util/Exception.h> |
4 | #include <c10/util/Optional.h> |
5 | |
6 | #include <cerrno> |
7 | #include <cstdio> |
8 | #include <cstdlib> |
9 | #include <cstring> |
10 | #include <string> |
11 | #include <utility> |
12 | #include <vector> |
13 | |
14 | #if !defined(_WIN32) |
15 | #include <unistd.h> |
16 | #else // defined(_WIN32) |
17 | #include <Windows.h> |
18 | #include <fileapi.h> |
19 | #endif // defined(_WIN32) |
20 | |
21 | namespace c10 { |
22 | namespace detail { |
23 | // Creates the filename pattern passed to and completed by `mkstemp`. |
24 | // Returns std::vector<char> because `mkstemp` needs a (non-const) `char*` and |
25 | // `std::string` only provides `const char*` before C++17. |
26 | #if !defined(_WIN32) |
27 | inline std::vector<char> make_filename(std::string name_prefix) { |
28 | // The filename argument to `mkstemp` needs "XXXXXX" at the end according to |
29 | // http://pubs.opengroup.org/onlinepubs/009695399/functions/mkstemp.html |
30 | static const std::string kRandomPattern = "XXXXXX" ; |
31 | |
32 | // We see if any of these environment variables is set and use their value, or |
33 | // else default the temporary directory to `/tmp`. |
34 | static const char* env_variables[] = {"TMPDIR" , "TMP" , "TEMP" , "TEMPDIR" }; |
35 | |
36 | std::string tmp_directory = "/tmp" ; |
37 | for (const char* variable : env_variables) { |
38 | if (const char* path = getenv(variable)) { |
39 | tmp_directory = path; |
40 | break; |
41 | } |
42 | } |
43 | |
44 | std::vector<char> filename; |
45 | filename.reserve( |
46 | tmp_directory.size() + name_prefix.size() + kRandomPattern.size() + 2); |
47 | |
48 | filename.insert(filename.end(), tmp_directory.begin(), tmp_directory.end()); |
49 | filename.push_back('/'); |
50 | filename.insert(filename.end(), name_prefix.begin(), name_prefix.end()); |
51 | filename.insert(filename.end(), kRandomPattern.begin(), kRandomPattern.end()); |
52 | filename.push_back('\0'); |
53 | |
54 | return filename; |
55 | } |
56 | #endif // !defined(_WIN32) |
57 | } // namespace detail |
58 | |
59 | struct TempFile { |
60 | #if !defined(_WIN32) |
61 | TempFile() : fd(-1) {} |
62 | TempFile(std::string name, int fd) : fd(fd), name(std::move(name)) {} |
63 | TempFile(const TempFile&) = delete; |
64 | TempFile(TempFile&& other) noexcept |
65 | : fd(other.fd), name(std::move(other.name)) { |
66 | other.fd = -1; |
67 | other.name.clear(); |
68 | } |
69 | |
70 | TempFile& operator=(const TempFile&) = delete; |
71 | TempFile& operator=(TempFile&& other) noexcept { |
72 | fd = other.fd; |
73 | name = std::move(other.name); |
74 | other.fd = -1; |
75 | other.name.clear(); |
76 | return *this; |
77 | } |
78 | |
79 | ~TempFile() { |
80 | if (fd >= 0) { |
81 | unlink(name.c_str()); |
82 | close(fd); |
83 | } |
84 | } |
85 | |
86 | int fd; |
87 | #endif // !defined(_WIN32) |
88 | |
89 | std::string name; |
90 | }; |
91 | |
92 | struct TempDir { |
93 | TempDir() = default; |
94 | explicit TempDir(std::string name) : name(std::move(name)) {} |
95 | TempDir(const TempDir&) = delete; |
96 | TempDir(TempDir&& other) noexcept : name(std::move(other.name)) { |
97 | other.name.clear(); |
98 | } |
99 | |
100 | TempDir& operator=(const TempDir&) = delete; |
101 | TempDir& operator=(TempDir&& other) noexcept { |
102 | name = std::move(other.name); |
103 | other.name.clear(); |
104 | return *this; |
105 | } |
106 | |
107 | ~TempDir() { |
108 | if (!name.empty()) { |
109 | #if !defined(_WIN32) |
110 | rmdir(name.c_str()); |
111 | #else // defined(_WIN32) |
112 | RemoveDirectoryA(name.c_str()); |
113 | #endif // defined(_WIN32) |
114 | } |
115 | } |
116 | |
117 | std::string name; |
118 | }; |
119 | |
120 | /// Attempts to return a temporary file or returns `nullopt` if an error |
121 | /// occurred. |
122 | /// |
123 | /// The file returned follows the pattern |
124 | /// `<tmp-dir>/<name-prefix><random-pattern>`, where `<tmp-dir>` is the value of |
125 | /// the `"TMPDIR"`, `"TMP"`, `"TEMP"` or |
126 | /// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`; |
127 | /// `<name-prefix>` is the value supplied to this function, and |
128 | /// `<random-pattern>` is a random sequence of numbers. |
129 | /// On Windows, `name_prefix` is ignored and `tmpnam` is used. |
130 | inline c10::optional<TempFile> try_make_tempfile( |
131 | std::string name_prefix = "torch-file-" ) { |
132 | #if defined(_WIN32) |
133 | return TempFile{std::tmpnam(nullptr)}; |
134 | #else |
135 | std::vector<char> filename = detail::make_filename(std::move(name_prefix)); |
136 | const int fd = mkstemp(filename.data()); |
137 | if (fd == -1) { |
138 | return c10::nullopt; |
139 | } |
140 | // Don't make the string from string(filename.begin(), filename.end(), or |
141 | // there will be a trailing '\0' at the end. |
142 | return TempFile(filename.data(), fd); |
143 | #endif // defined(_WIN32) |
144 | } |
145 | |
146 | /// Like `try_make_tempfile`, but throws an exception if a temporary file could |
147 | /// not be returned. |
148 | inline TempFile make_tempfile(std::string name_prefix = "torch-file-" ) { |
149 | if (auto tempfile = try_make_tempfile(std::move(name_prefix))) { |
150 | return std::move(*tempfile); |
151 | } |
152 | TORCH_CHECK(false, "Error generating temporary file: " , std::strerror(errno)); |
153 | } |
154 | |
155 | /// Attempts to return a temporary directory or returns `nullopt` if an error |
156 | /// occurred. |
157 | /// |
158 | /// The directory returned follows the pattern |
159 | /// `<tmp-dir>/<name-prefix><random-pattern>/`, where `<tmp-dir>` is the value |
160 | /// of the `"TMPDIR"`, `"TMP"`, `"TEMP"` or |
161 | /// `"TEMPDIR"` environment variable if any is set, or otherwise `/tmp`; |
162 | /// `<name-prefix>` is the value supplied to this function, and |
163 | /// `<random-pattern>` is a random sequence of numbers. |
164 | /// On Windows, `name_prefix` is ignored and `tmpnam` is used. |
165 | inline c10::optional<TempDir> try_make_tempdir( |
166 | std::string name_prefix = "torch-dir-" ) { |
167 | #if defined(_WIN32) |
168 | while (true) { |
169 | const char* dirname = std::tmpnam(nullptr); |
170 | if (!dirname) { |
171 | return c10::nullopt; |
172 | } |
173 | if (CreateDirectoryA(dirname, NULL)) { |
174 | return TempDir(dirname); |
175 | } |
176 | if (GetLastError() != ERROR_ALREADY_EXISTS) { |
177 | return c10::nullopt; |
178 | } |
179 | } |
180 | return c10::nullopt; |
181 | #else |
182 | std::vector<char> filename = detail::make_filename(std::move(name_prefix)); |
183 | const char* dirname = mkdtemp(filename.data()); |
184 | if (!dirname) { |
185 | return c10::nullopt; |
186 | } |
187 | return TempDir(dirname); |
188 | #endif // defined(_WIN32) |
189 | } |
190 | |
191 | /// Like `try_make_tempdir`, but throws an exception if a temporary directory |
192 | /// could not be returned. |
193 | inline TempDir make_tempdir(std::string name_prefix = "torch-dir-" ) { |
194 | if (auto tempdir = try_make_tempdir(std::move(name_prefix))) { |
195 | return std::move(*tempdir); |
196 | } |
197 | TORCH_CHECK( |
198 | false, "Error generating temporary directory: " , std::strerror(errno)); |
199 | } |
200 | } // namespace c10 |
201 | |