pytorch

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)
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
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.
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.
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.
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)
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

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

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.
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