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
21namespace c10 {
22namespace 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)
27inline 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
59struct 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
92struct 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.
130inline 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.
148inline 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.
165inline 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.
193inline 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