diff --git a/helpers/memenv/memenv.cc b/helpers/memenv/memenv.cc index d44627b..b78a998 100644 --- a/helpers/memenv/memenv.cc +++ b/helpers/memenv/memenv.cc @@ -51,9 +51,22 @@ class FileState { } } - uint64_t Size() const { return size_; } + uint64_t Size() const { + MutexLock lock(&blocks_mutex_); + return size_; + } + + void Truncate() { + MutexLock lock(&blocks_mutex_); + for (char*& block : blocks_) { + delete[] block; + } + blocks_.clear(); + size_ = 0; + } Status Read(uint64_t offset, size_t n, Slice* result, char* scratch) const { + MutexLock lock(&blocks_mutex_); if (offset > size_) { return Status::IOError("Offset greater than file size."); } @@ -100,6 +113,7 @@ class FileState { const char* src = data.data(); size_t src_len = data.size(); + MutexLock lock(&blocks_mutex_); while (src_len > 0) { size_t avail; size_t offset = size_ % kBlockSize; @@ -128,10 +142,7 @@ class FileState { private: // Private since only Unref() should be used to delete it. ~FileState() { - for (std::vector::iterator i = blocks_.begin(); i != blocks_.end(); - ++i) { - delete [] *i; - } + Truncate(); } // No copying allowed. @@ -141,11 +152,9 @@ class FileState { port::Mutex refs_mutex_; int refs_ GUARDED_BY(refs_mutex_); - // The following fields are not protected by any mutex. They are only mutable - // while the file is being written, and concurrent access is not allowed - // to writable files. - std::vector blocks_; - uint64_t size_; + mutable port::Mutex blocks_mutex_; + std::vector blocks_ GUARDED_BY(blocks_mutex_); + uint64_t size_ GUARDED_BY(blocks_mutex_); enum { kBlockSize = 8 * 1024 }; }; @@ -269,13 +278,18 @@ class InMemoryEnv : public EnvWrapper { virtual Status NewWritableFile(const std::string& fname, WritableFile** result) { MutexLock lock(&mutex_); - if (file_map_.find(fname) != file_map_.end()) { - DeleteFileInternal(fname); - } + FileSystem::iterator it = file_map_.find(fname); - FileState* file = new FileState(); - file->Ref(); - file_map_[fname] = file; + FileState* file; + if (it == file_map_.end()) { + // File is not currently open. + file = new FileState(); + file->Ref(); + file_map_[fname] = file; + } else { + file = it->second; + file->Truncate(); + } *result = new WritableFileImpl(file); return Status::OK(); diff --git a/helpers/memenv/memenv_test.cc b/helpers/memenv/memenv_test.cc index 5cff776..4664795 100644 --- a/helpers/memenv/memenv_test.cc +++ b/helpers/memenv/memenv_test.cc @@ -191,6 +191,29 @@ TEST(MemEnvTest, LargeWrite) { delete [] scratch; } +TEST(MemEnvTest, OverwriteOpenFile) { + const char kWrite1Data[] = "Write #1 data"; + const size_t kFileDataLen = sizeof(kWrite1Data) - 1; + const std::string kTestFileName = test::TmpDir() + "/leveldb-TestFile.dat"; + + ASSERT_OK(WriteStringToFile(env_, kWrite1Data, kTestFileName)); + + RandomAccessFile* rand_file; + ASSERT_OK(env_->NewRandomAccessFile(kTestFileName, &rand_file)); + + const char kWrite2Data[] = "Write #2 data"; + ASSERT_OK(WriteStringToFile(env_, kWrite2Data, kTestFileName)); + + // Verify that overwriting an open file will result in the new file data + // being read from files opened before the write. + Slice result; + char scratch[kFileDataLen]; + ASSERT_OK(rand_file->Read(0, kFileDataLen, &result, scratch)); + ASSERT_EQ(0, result.compare(kWrite2Data)); + + delete rand_file; +} + TEST(MemEnvTest, DBTest) { Options options; options.create_if_missing = true;