tempfile.h 5.9 KB

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