env.cc revision 667a023a417f36ff4eabaee77daf374dceaa5ddd 
1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7    http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include <deque>
17#include <vector>
18#if defined(__APPLE__)
19#include <mach-o/dyld.h>
20#endif
21#if defined(PLATFORM_WINDOWS)
22#include <windows.h>
23#define PATH_MAX MAX_PATH
24#else
25#include <unistd.h>
26#endif
27
28#include "tensorflow/core/lib/core/errors.h"
29#include "tensorflow/core/lib/gtl/map_util.h"
30#include "tensorflow/core/lib/gtl/stl_util.h"
31#include "tensorflow/core/lib/io/path.h"
32#include "tensorflow/core/platform/env.h"
33#include "tensorflow/core/platform/protobuf.h"
34
35namespace tensorflow {
36
37class FileSystemRegistryImpl : public FileSystemRegistry {
38 public:
39  Status Register(const string& scheme, Factory factory) override;
40  FileSystem* Lookup(const string& scheme) override;
41  Status GetRegisteredFileSystemSchemes(std::vector<string>* schemes) override;
42
43 private:
44  mutable mutex mu_;
45  mutable std::unordered_map<string, std::unique_ptr<FileSystem>> registry_
46      GUARDED_BY(mu_);
47};
48
49Status FileSystemRegistryImpl::Register(const string& scheme,
50                                        FileSystemRegistry::Factory factory) {
51  mutex_lock lock(mu_);
52  if (!registry_.emplace(string(scheme), std::unique_ptr<FileSystem>(factory()))
53           .second) {
54    return errors::AlreadyExists("File factory for ", scheme,
55                                 " already registered");
56  }
57  return Status::OK();
58}
59
60FileSystem* FileSystemRegistryImpl::Lookup(const string& scheme) {
61  mutex_lock lock(mu_);
62  const auto found = registry_.find(scheme);
63  if (found == registry_.end()) {
64    return nullptr;
65  }
66  return found->second.get();
67}
68
69Status FileSystemRegistryImpl::GetRegisteredFileSystemSchemes(
70    std::vector<string>* schemes) {
71  mutex_lock lock(mu_);
72  for (const auto& e : registry_) {
73    schemes->push_back(e.first);
74  }
75  return Status::OK();
76}
77
78Env::Env() : file_system_registry_(new FileSystemRegistryImpl) {}
79
80Status Env::GetFileSystemForFile(const string& fname, FileSystem** result) {
81  StringPiece scheme, host, path;
82  io::ParseURI(fname, &scheme, &host, &path);
83  FileSystem* file_system = file_system_registry_->Lookup(scheme.ToString());
84  if (!file_system) {
85    return errors::Unimplemented("File system scheme ", scheme,
86                                 " not implemented");
87  }
88  *result = file_system;
89  return Status::OK();
90}
91
92Status Env::GetRegisteredFileSystemSchemes(std::vector<string>* schemes) {
93  return file_system_registry_->GetRegisteredFileSystemSchemes(schemes);
94}
95
96Status Env::RegisterFileSystem(const string& scheme,
97                               FileSystemRegistry::Factory factory) {
98  return file_system_registry_->Register(scheme, factory);
99}
100
101Status Env::NewRandomAccessFile(const string& fname,
102                                std::unique_ptr<RandomAccessFile>* result) {
103  FileSystem* fs;
104  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
105  return fs->NewRandomAccessFile(fname, result);
106}
107
108Status Env::NewReadOnlyMemoryRegionFromFile(
109    const string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
110  FileSystem* fs;
111  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
112  return fs->NewReadOnlyMemoryRegionFromFile(fname, result);
113}
114
115Status Env::NewWritableFile(const string& fname,
116                            std::unique_ptr<WritableFile>* result) {
117  FileSystem* fs;
118  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
119  return fs->NewWritableFile(fname, result);
120}
121
122Status Env::NewAppendableFile(const string& fname,
123                              std::unique_ptr<WritableFile>* result) {
124  FileSystem* fs;
125  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
126  return fs->NewAppendableFile(fname, result);
127}
128
129Status Env::FileExists(const string& fname) {
130  FileSystem* fs;
131  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
132  return fs->FileExists(fname);
133}
134
135Status Env::GetChildren(const string& dir, std::vector<string>* result) {
136  FileSystem* fs;
137  TF_RETURN_IF_ERROR(GetFileSystemForFile(dir, &fs));
138  return fs->GetChildren(dir, result);
139}
140
141Status Env::GetMatchingPaths(const string& pattern,
142                             std::vector<string>* results) {
143  FileSystem* fs;
144  TF_RETURN_IF_ERROR(GetFileSystemForFile(pattern, &fs));
145  return fs->GetMatchingPaths(pattern, results);
146}
147
148Status Env::DeleteFile(const string& fname) {
149  FileSystem* fs;
150  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
151  return fs->DeleteFile(fname);
152}
153
154Status Env::RecursivelyCreateDir(const string& dirname) {
155  FileSystem* fs;
156  TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
157  return fs->RecursivelyCreateDir(dirname);
158}
159
160Status Env::CreateDir(const string& dirname) {
161  FileSystem* fs;
162  TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
163  return fs->CreateDir(dirname);
164}
165
166Status Env::DeleteDir(const string& dirname) {
167  FileSystem* fs;
168  TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
169  return fs->DeleteDir(dirname);
170}
171
172Status Env::Stat(const string& fname, FileStatistics* stat) {
173  FileSystem* fs;
174  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
175  return fs->Stat(fname, stat);
176}
177
178Status Env::IsDirectory(const string& fname) {
179  FileSystem* fs;
180  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
181  return fs->IsDirectory(fname);
182}
183
184Status Env::DeleteRecursively(const string& dirname, int64* undeleted_files,
185                              int64* undeleted_dirs) {
186  FileSystem* fs;
187  TF_RETURN_IF_ERROR(GetFileSystemForFile(dirname, &fs));
188  return fs->DeleteRecursively(dirname, undeleted_files, undeleted_dirs);
189}
190
191Status Env::GetFileSize(const string& fname, uint64* file_size) {
192  FileSystem* fs;
193  TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
194  return fs->GetFileSize(fname, file_size);
195}
196
197Status Env::RenameFile(const string& src, const string& target) {
198  FileSystem* src_fs;
199  FileSystem* target_fs;
200  TF_RETURN_IF_ERROR(GetFileSystemForFile(src, &src_fs));
201  TF_RETURN_IF_ERROR(GetFileSystemForFile(target, &target_fs));
202  if (src_fs != target_fs) {
203    return errors::Unimplemented("Renaming ", src, " to ", target,
204                                 " not implemented");
205  }
206  return src_fs->RenameFile(src, target);
207}
208
209string Env::GetExecutablePath() {
210  char exe_path[PATH_MAX] = {0};
211#ifdef __APPLE__
212  uint32_t buffer_size(0U);
213  _NSGetExecutablePath(nullptr, &buffer_size);
214  char unresolved_path[buffer_size];
215  _NSGetExecutablePath(unresolved_path, &buffer_size);
216  CHECK(realpath(unresolved_path, exe_path));
217#elif defined(PLATFORM_WINDOWS)
218  HMODULE hModule = GetModuleHandle(NULL);
219  GetModuleFileName(hModule, exe_path, MAX_PATH);
220#else
221  CHECK_NE(-1, readlink("/proc/self/exe", exe_path, sizeof(exe_path) - 1));
222#endif
223  // Make sure it's null-terminated:
224  exe_path[sizeof(exe_path) - 1] = 0;
225
226  return exe_path;
227}
228
229Thread::~Thread() {}
230
231EnvWrapper::~EnvWrapper() {}
232
233Status ReadFileToString(Env* env, const string& fname, string* data) {
234  uint64 file_size;
235  Status s = env->GetFileSize(fname, &file_size);
236  if (!s.ok()) {
237    return s;
238  }
239  std::unique_ptr<RandomAccessFile> file;
240  s = env->NewRandomAccessFile(fname, &file);
241  if (!s.ok()) {
242    return s;
243  }
244  gtl::STLStringResizeUninitialized(data, file_size);
245  char* p = gtl::string_as_array(data);
246  StringPiece result;
247  s = file->Read(0, file_size, &result, p);
248  if (!s.ok()) {
249    data->clear();
250  } else if (result.size() != file_size) {
251    s = errors::Aborted("File ", fname, " changed while reading: ", file_size,
252                        " vs. ", result.size());
253    data->clear();
254  } else if (result.data() == p) {
255    // Data is already in the correct location
256  } else {
257    memmove(p, result.data(), result.size());
258  }
259  return s;
260}
261
262Status WriteStringToFile(Env* env, const string& fname,
263                         const StringPiece& data) {
264  std::unique_ptr<WritableFile> file;
265  Status s = env->NewWritableFile(fname, &file);
266  if (!s.ok()) {
267    return s;
268  }
269  s = file->Append(data);
270  if (s.ok()) {
271    s = file->Close();
272  }
273  return s;
274}
275
276// A ZeroCopyInputStream on a RandomAccessFile.
277namespace {
278class FileStream : public ::tensorflow::protobuf::io::ZeroCopyInputStream {
279 public:
280  explicit FileStream(RandomAccessFile* file) : file_(file), pos_(0) {}
281
282  void BackUp(int count) override { pos_ -= count; }
283  bool Skip(int count) override {
284    pos_ += count;
285    return true;
286  }
287  protobuf_int64 ByteCount() const override { return pos_; }
288  Status status() const { return status_; }
289
290  bool Next(const void** data, int* size) override {
291    StringPiece result;
292    Status s = file_->Read(pos_, kBufSize, &result, scratch_);
293    if (result.empty()) {
294      status_ = s;
295      return false;
296    }
297    pos_ += result.size();
298    *data = result.data();
299    *size = result.size();
300    return true;
301  }
302
303 private:
304  static const int kBufSize = 512 << 10;
305
306  RandomAccessFile* file_;
307  int64 pos_;
308  Status status_;
309  char scratch_[kBufSize];
310};
311
312}  // namespace
313
314Status WriteBinaryProto(Env* env, const string& fname,
315                        const ::tensorflow::protobuf::MessageLite& proto) {
316  string serialized;
317  proto.AppendToString(&serialized);
318  return WriteStringToFile(env, fname, serialized);
319}
320
321Status ReadBinaryProto(Env* env, const string& fname,
322                       ::tensorflow::protobuf::MessageLite* proto) {
323  std::unique_ptr<RandomAccessFile> file;
324  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
325  std::unique_ptr<FileStream> stream(new FileStream(file.get()));
326
327  // TODO(jiayq): the following coded stream is for debugging purposes to allow
328  // one to parse arbitrarily large messages for MessageLite. One most likely
329  // doesn't want to put protobufs larger than 64MB on Android, so we should
330  // eventually remove this and quit loud when a large protobuf is passed in.
331  ::tensorflow::protobuf::io::CodedInputStream coded_stream(stream.get());
332  // Total bytes hard limit / warning limit are set to 1GB and 512MB
333  // respectively.
334  coded_stream.SetTotalBytesLimit(1024LL << 20, 512LL << 20);
335
336  if (!proto->ParseFromCodedStream(&coded_stream)) {
337    TF_RETURN_IF_ERROR(stream->status());
338    return errors::DataLoss("Can't parse ", fname, " as binary proto");
339  }
340  return Status::OK();
341}
342
343Status WriteTextProto(Env* env, const string& fname,
344                      const ::tensorflow::protobuf::Message& proto) {
345#if !defined(TENSORFLOW_LITE_PROTOS)
346  string serialized;
347  if (!::tensorflow::protobuf::TextFormat::PrintToString(proto, &serialized)) {
348    return errors::FailedPrecondition("Unable to convert proto to text.");
349  }
350  return WriteStringToFile(env, fname, serialized);
351#else
352  return errors::Unimplemented("Can't write text protos with protolite.");
353#endif
354}
355
356Status ReadTextProto(Env* env, const string& fname,
357                     ::tensorflow::protobuf::Message* proto) {
358#if !defined(TENSORFLOW_LITE_PROTOS)
359  std::unique_ptr<RandomAccessFile> file;
360  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(fname, &file));
361  std::unique_ptr<FileStream> stream(new FileStream(file.get()));
362
363  if (!::tensorflow::protobuf::TextFormat::Parse(stream.get(), proto)) {
364    TF_RETURN_IF_ERROR(stream->status());
365    return errors::DataLoss("Can't parse ", fname, " as text proto");
366  }
367  return Status::OK();
368#else
369  return errors::Unimplemented("Can't parse text protos with protolite.");
370#endif
371}
372
373}  // namespace tensorflow
374