memenv.cc raw

   1  // Copyright (c) 2011 The LevelDB Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style license that can be
   3  // found in the LICENSE file. See the AUTHORS file for names of contributors.
   4  
   5  #include "helpers/memenv/memenv.h"
   6  
   7  #include <string.h>
   8  
   9  #include <limits>
  10  #include <map>
  11  #include <string>
  12  #include <vector>
  13  
  14  #include "leveldb/env.h"
  15  #include "leveldb/status.h"
  16  #include "port/port.h"
  17  #include "port/thread_annotations.h"
  18  #include "util/mutexlock.h"
  19  
  20  namespace leveldb {
  21  
  22  namespace {
  23  
  24  class FileState {
  25   public:
  26    // FileStates are reference counted. The initial reference count is zero
  27    // and the caller must call Ref() at least once.
  28    FileState() : refs_(0), size_(0) {}
  29  
  30    // No copying allowed.
  31    FileState(const FileState&) = delete;
  32    FileState& operator=(const FileState&) = delete;
  33  
  34    // Increase the reference count.
  35    void Ref() {
  36      MutexLock lock(&refs_mutex_);
  37      ++refs_;
  38    }
  39  
  40    // Decrease the reference count. Delete if this is the last reference.
  41    void Unref() {
  42      bool do_delete = false;
  43  
  44      {
  45        MutexLock lock(&refs_mutex_);
  46        --refs_;
  47        assert(refs_ >= 0);
  48        if (refs_ <= 0) {
  49          do_delete = true;
  50        }
  51      }
  52  
  53      if (do_delete) {
  54        delete this;
  55      }
  56    }
  57  
  58    uint64_t Size() const {
  59      MutexLock lock(&blocks_mutex_);
  60      return size_;
  61    }
  62  
  63    void Truncate() {
  64      MutexLock lock(&blocks_mutex_);
  65      for (char*& block : blocks_) {
  66        delete[] block;
  67      }
  68      blocks_.clear();
  69      size_ = 0;
  70    }
  71  
  72    Status Read(uint64_t offset, size_t n, Slice* result, char* scratch) const {
  73      MutexLock lock(&blocks_mutex_);
  74      if (offset > size_) {
  75        return Status::IOError("Offset greater than file size.");
  76      }
  77      const uint64_t available = size_ - offset;
  78      if (n > available) {
  79        n = static_cast<size_t>(available);
  80      }
  81      if (n == 0) {
  82        *result = Slice();
  83        return Status::OK();
  84      }
  85  
  86      assert(offset / kBlockSize <= std::numeric_limits<size_t>::max());
  87      size_t block = static_cast<size_t>(offset / kBlockSize);
  88      size_t block_offset = offset % kBlockSize;
  89      size_t bytes_to_copy = n;
  90      char* dst = scratch;
  91  
  92      while (bytes_to_copy > 0) {
  93        size_t avail = kBlockSize - block_offset;
  94        if (avail > bytes_to_copy) {
  95          avail = bytes_to_copy;
  96        }
  97        memcpy(dst, blocks_[block] + block_offset, avail);
  98  
  99        bytes_to_copy -= avail;
 100        dst += avail;
 101        block++;
 102        block_offset = 0;
 103      }
 104  
 105      *result = Slice(scratch, n);
 106      return Status::OK();
 107    }
 108  
 109    Status Append(const Slice& data) {
 110      const char* src = data.data();
 111      size_t src_len = data.size();
 112  
 113      MutexLock lock(&blocks_mutex_);
 114      while (src_len > 0) {
 115        size_t avail;
 116        size_t offset = size_ % kBlockSize;
 117  
 118        if (offset != 0) {
 119          // There is some room in the last block.
 120          avail = kBlockSize - offset;
 121        } else {
 122          // No room in the last block; push new one.
 123          blocks_.push_back(new char[kBlockSize]);
 124          avail = kBlockSize;
 125        }
 126  
 127        if (avail > src_len) {
 128          avail = src_len;
 129        }
 130        memcpy(blocks_.back() + offset, src, avail);
 131        src_len -= avail;
 132        src += avail;
 133        size_ += avail;
 134      }
 135  
 136      return Status::OK();
 137    }
 138  
 139   private:
 140    enum { kBlockSize = 8 * 1024 };
 141  
 142    // Private since only Unref() should be used to delete it.
 143    ~FileState() { Truncate(); }
 144  
 145    port::Mutex refs_mutex_;
 146    int refs_ GUARDED_BY(refs_mutex_);
 147  
 148    mutable port::Mutex blocks_mutex_;
 149    std::vector<char*> blocks_ GUARDED_BY(blocks_mutex_);
 150    uint64_t size_ GUARDED_BY(blocks_mutex_);
 151  };
 152  
 153  class SequentialFileImpl : public SequentialFile {
 154   public:
 155    explicit SequentialFileImpl(FileState* file) : file_(file), pos_(0) {
 156      file_->Ref();
 157    }
 158  
 159    ~SequentialFileImpl() override { file_->Unref(); }
 160  
 161    Status Read(size_t n, Slice* result, char* scratch) override {
 162      Status s = file_->Read(pos_, n, result, scratch);
 163      if (s.ok()) {
 164        pos_ += result->size();
 165      }
 166      return s;
 167    }
 168  
 169    Status Skip(uint64_t n) override {
 170      if (pos_ > file_->Size()) {
 171        return Status::IOError("pos_ > file_->Size()");
 172      }
 173      const uint64_t available = file_->Size() - pos_;
 174      if (n > available) {
 175        n = available;
 176      }
 177      pos_ += n;
 178      return Status::OK();
 179    }
 180  
 181    virtual std::string GetName() const override { return "[memenv]"; }
 182   private:
 183    FileState* file_;
 184    uint64_t pos_;
 185  };
 186  
 187  class RandomAccessFileImpl : public RandomAccessFile {
 188   public:
 189    explicit RandomAccessFileImpl(FileState* file) : file_(file) { file_->Ref(); }
 190  
 191    ~RandomAccessFileImpl() override { file_->Unref(); }
 192  
 193    Status Read(uint64_t offset, size_t n, Slice* result,
 194                char* scratch) const override {
 195      return file_->Read(offset, n, result, scratch);
 196    }
 197  
 198    virtual std::string GetName() const override { return "[memenv]"; }
 199   private:
 200    FileState* file_;
 201  };
 202  
 203  class WritableFileImpl : public WritableFile {
 204   public:
 205    WritableFileImpl(FileState* file) : file_(file) { file_->Ref(); }
 206  
 207    ~WritableFileImpl() override { file_->Unref(); }
 208  
 209    Status Append(const Slice& data) override { return file_->Append(data); }
 210  
 211    Status Close() override { return Status::OK(); }
 212    Status Flush() override { return Status::OK(); }
 213    Status Sync() override { return Status::OK(); }
 214  
 215    virtual std::string GetName() const override { return "[memenv]"; }
 216   private:
 217    FileState* file_;
 218  };
 219  
 220  class NoOpLogger : public Logger {
 221   public:
 222    void Logv(const char* format, va_list ap) override {}
 223  };
 224  
 225  class InMemoryEnv : public EnvWrapper {
 226   public:
 227    explicit InMemoryEnv(Env* base_env) : EnvWrapper(base_env) {}
 228  
 229    ~InMemoryEnv() override {
 230      for (const auto& kvp : file_map_) {
 231        kvp.second->Unref();
 232      }
 233    }
 234  
 235    // Partial implementation of the Env interface.
 236    Status NewSequentialFile(const std::string& fname,
 237                             SequentialFile** result) override {
 238      MutexLock lock(&mutex_);
 239      if (file_map_.find(fname) == file_map_.end()) {
 240        *result = nullptr;
 241        return Status::IOError(fname, "File not found");
 242      }
 243  
 244      *result = new SequentialFileImpl(file_map_[fname]);
 245      return Status::OK();
 246    }
 247  
 248    Status NewRandomAccessFile(const std::string& fname,
 249                               RandomAccessFile** result) override {
 250      MutexLock lock(&mutex_);
 251      if (file_map_.find(fname) == file_map_.end()) {
 252        *result = nullptr;
 253        return Status::IOError(fname, "File not found");
 254      }
 255  
 256      *result = new RandomAccessFileImpl(file_map_[fname]);
 257      return Status::OK();
 258    }
 259  
 260    Status NewWritableFile(const std::string& fname,
 261                           WritableFile** result) override {
 262      MutexLock lock(&mutex_);
 263      FileSystem::iterator it = file_map_.find(fname);
 264  
 265      FileState* file;
 266      if (it == file_map_.end()) {
 267        // File is not currently open.
 268        file = new FileState();
 269        file->Ref();
 270        file_map_[fname] = file;
 271      } else {
 272        file = it->second;
 273        file->Truncate();
 274      }
 275  
 276      *result = new WritableFileImpl(file);
 277      return Status::OK();
 278    }
 279  
 280    Status NewAppendableFile(const std::string& fname,
 281                             WritableFile** result) override {
 282      MutexLock lock(&mutex_);
 283      FileState** sptr = &file_map_[fname];
 284      FileState* file = *sptr;
 285      if (file == nullptr) {
 286        file = new FileState();
 287        file->Ref();
 288      }
 289      *result = new WritableFileImpl(file);
 290      return Status::OK();
 291    }
 292  
 293    bool FileExists(const std::string& fname) override {
 294      MutexLock lock(&mutex_);
 295      return file_map_.find(fname) != file_map_.end();
 296    }
 297  
 298    Status GetChildren(const std::string& dir,
 299                       std::vector<std::string>* result) override {
 300      MutexLock lock(&mutex_);
 301      result->clear();
 302  
 303      for (const auto& kvp : file_map_) {
 304        const std::string& filename = kvp.first;
 305  
 306        if (filename.size() >= dir.size() + 1 && filename[dir.size()] == '/' &&
 307            Slice(filename).starts_with(Slice(dir))) {
 308          result->push_back(filename.substr(dir.size() + 1));
 309        }
 310      }
 311  
 312      return Status::OK();
 313    }
 314  
 315    void DeleteFileInternal(const std::string& fname)
 316        EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
 317      if (file_map_.find(fname) == file_map_.end()) {
 318        return;
 319      }
 320  
 321      file_map_[fname]->Unref();
 322      file_map_.erase(fname);
 323    }
 324  
 325    Status DeleteFile(const std::string& fname) override {
 326      MutexLock lock(&mutex_);
 327      if (file_map_.find(fname) == file_map_.end()) {
 328        return Status::IOError(fname, "File not found");
 329      }
 330  
 331      DeleteFileInternal(fname);
 332      return Status::OK();
 333    }
 334  
 335    Status CreateDir(const std::string& dirname) override { return Status::OK(); }
 336  
 337    Status DeleteDir(const std::string& dirname) override { return Status::OK(); }
 338  
 339    Status GetFileSize(const std::string& fname, uint64_t* file_size) override {
 340      MutexLock lock(&mutex_);
 341      if (file_map_.find(fname) == file_map_.end()) {
 342        return Status::IOError(fname, "File not found");
 343      }
 344  
 345      *file_size = file_map_[fname]->Size();
 346      return Status::OK();
 347    }
 348  
 349    Status RenameFile(const std::string& src,
 350                      const std::string& target) override {
 351      MutexLock lock(&mutex_);
 352      if (file_map_.find(src) == file_map_.end()) {
 353        return Status::IOError(src, "File not found");
 354      }
 355  
 356      DeleteFileInternal(target);
 357      file_map_[target] = file_map_[src];
 358      file_map_.erase(src);
 359      return Status::OK();
 360    }
 361  
 362    Status LockFile(const std::string& fname, FileLock** lock) override {
 363      *lock = new FileLock;
 364      return Status::OK();
 365    }
 366  
 367    Status UnlockFile(FileLock* lock) override {
 368      delete lock;
 369      return Status::OK();
 370    }
 371  
 372    Status GetTestDirectory(std::string* path) override {
 373      *path = "/test";
 374      return Status::OK();
 375    }
 376  
 377    Status NewLogger(const std::string& fname, Logger** result) override {
 378      *result = new NoOpLogger;
 379      return Status::OK();
 380    }
 381  
 382   private:
 383    // Map from filenames to FileState objects, representing a simple file system.
 384    typedef std::map<std::string, FileState*> FileSystem;
 385  
 386    port::Mutex mutex_;
 387    FileSystem file_map_ GUARDED_BY(mutex_);
 388  };
 389  
 390  }  // namespace
 391  
 392  Env* NewMemEnv(Env* base_env) { return new InMemoryEnv(base_env); }
 393  
 394  }  // namespace leveldb
 395