Implement BlockStore::forEachBlock

src/blockstore/implementations/caching/CachingBlockStore.cpp

@@ -18,7 +18,7 @@ namespace blockstore {
 namespace caching {
 
 CachingBlockStore::CachingBlockStore(cpputils::unique_ref<BlockStore> baseBlockStore)
-  :_baseBlockStore(std::move(baseBlockStore)), _cache(), _numNewBlocks(0) {
+  :_baseBlockStore(std::move(baseBlockStore)), _newBlocks(), _cache() {
 }
 
 Key CachingBlockStore::createKey() {
@@ -29,7 +29,6 @@ optional<unique_ref<Block>> CachingBlockStore::tryCreate(const Key &key, Data da
   ASSERT(_cache.pop(key) == none, "Key already exists in cache");
   //TODO Shouldn't we return boost::none if the key already exists?
   //TODO Key can also already exist but not be in the cache right now.
-  ++_numNewBlocks;
   return unique_ref<Block>(make_unique_ref<CachedBlock>(make_unique_ref<NewBlock>(key, std::move(data), this), this));
 }
 
@@ -54,9 +53,6 @@ void CachingBlockStore::remove(cpputils::unique_ref<Block> block) {
   auto baseBlock = (*cached_block)->releaseBlock();
   auto baseNewBlock = dynamic_pointer_move<NewBlock>(baseBlock);
   if (baseNewBlock != none) {
-	if(!(*baseNewBlock)->alreadyExistsInBaseStore()) {
-	  --_numNewBlocks;
-	}
     (*baseNewBlock)->remove();
   } else {
     _baseBlockStore->remove(std::move(baseBlock));
@@ -64,7 +60,7 @@ void CachingBlockStore::remove(cpputils::unique_ref<Block> block) {
 }
 
 uint64_t CachingBlockStore::numBlocks() const {
-  return _baseBlockStore->numBlocks() + _numNewBlocks;
+  return _baseBlockStore->numBlocks() + _newBlocks.size();
 }
 
 uint64_t CachingBlockStore::estimateNumFreeBytes() const {
@@ -77,11 +73,7 @@ void CachingBlockStore::release(unique_ref<Block> block) {
 }
 
 optional<unique_ref<Block>> CachingBlockStore::tryCreateInBaseStore(const Key &key, Data data) {
-  auto block = _baseBlockStore->tryCreate(key, std::move(data));
-  if (block != none) {
-	--_numNewBlocks;
-  }
-  return block;
+  return _baseBlockStore->tryCreate(key, std::move(data));
 }
 
 void CachingBlockStore::removeFromBaseStore(cpputils::unique_ref<Block> block) {
@@ -96,5 +88,20 @@ uint64_t CachingBlockStore::blockSizeFromPhysicalBlockSize(uint64_t blockSize) c
   return _baseBlockStore->blockSizeFromPhysicalBlockSize(blockSize);
 }
 
+void CachingBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  _baseBlockStore->forEachBlock(callback);
+  for (NewBlock *newBlock : _newBlocks) {
+    callback(newBlock->key());
+  }
+}
+
+void CachingBlockStore::registerNewBlock(NewBlock *newBlock) {
+  _newBlocks.insert(newBlock);
+}
+
+void CachingBlockStore::unregisterNewBlock(NewBlock *newBlock) {
+  _newBlocks.erase(newBlock);
+}
+
 }
 }

src/blockstore/implementations/caching/CachingBlockStore.h

@@ -4,10 +4,13 @@
 
 #include "cache/Cache.h"
 #include "../../interface/BlockStore.h"
+#include <unordered_set>
 
 namespace blockstore {
 namespace caching {
 
+class NewBlock;
+
 //TODO Check that this blockstore allows parallel destructing of blocks (otherwise we won't encrypt blocks in parallel)
 class CachingBlockStore final: public BlockStore {
 public:
@@ -20,18 +23,21 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
   void release(cpputils::unique_ref<Block> block);
 
   boost::optional<cpputils::unique_ref<Block>> tryCreateInBaseStore(const Key &key, cpputils::Data data);
   void removeFromBaseStore(cpputils::unique_ref<Block> block);
+  void registerNewBlock(NewBlock *newBlock);
+  void unregisterNewBlock(NewBlock *newBlock);
 
   void flush();
 
 private:
   cpputils::unique_ref<BlockStore> _baseBlockStore;
+  std::unordered_set<NewBlock*> _newBlocks; // List of all new blocks that aren't in the base store yet.
   Cache<Key, cpputils::unique_ref<Block>, 1000> _cache;
-  uint32_t _numNewBlocks;
 
   DISALLOW_COPY_AND_ASSIGN(CachingBlockStore);
 };

src/blockstore/implementations/caching/NewBlock.cpp

@@ -15,6 +15,7 @@ NewBlock::NewBlock(const Key &key, Data data, CachingBlockStore *blockStore)
      _data(std::move(data)),
      _baseBlock(none),
      _dataChanged(true) {
+  _blockStore->registerNewBlock(this);
 }
 
 NewBlock::~NewBlock() {
@@ -37,6 +38,7 @@ void NewBlock::writeToBaseBlockIfChanged() {
       //TODO _data.copy() necessary?
       auto newBase = _blockStore->tryCreateInBaseStore(key(), _data.copy());
       ASSERT(newBase != boost::none, "Couldn't create base block"); //TODO What if tryCreate fails due to a duplicate key? We should ensure we don't use duplicate keys.
+      _blockStore->unregisterNewBlock(this); // block now exists in base store, we have to remove it from the list of blocks not in the base store.
       _baseBlock = std::move(*newBase);
     } else {
         (*_baseBlock)->write(_data.data(), 0, _data.size());
@@ -48,6 +50,9 @@ void NewBlock::writeToBaseBlockIfChanged() {
 void NewBlock::remove() {
   if (_baseBlock != none) {
 	_blockStore->removeFromBaseStore(std::move(*_baseBlock));
+  } else {
+    // block doesn't exists in base store, we have to remove it from the list of blocks not in the base store.
+    _blockStore->unregisterNewBlock(this);
   }
   _dataChanged = false;
 }
@@ -67,9 +72,5 @@ void NewBlock::resize(size_t newSize) {
     _dataChanged = true;
 }
 
-bool NewBlock::alreadyExistsInBaseStore() const {
-  return _baseBlock != none;
-}
-
 }
 }

src/blockstore/implementations/caching/NewBlock.h

@@ -33,8 +33,6 @@ public:
 
   void remove();
 
-  bool alreadyExistsInBaseStore() const;
-
 private:
   CachingBlockStore *_blockStore;
   cpputils::Data _data;

src/blockstore/implementations/compressing/CompressingBlockStore.h

@@ -21,6 +21,7 @@ public:
     uint64_t numBlocks() const override;
     uint64_t estimateNumFreeBytes() const override;
     uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+    void forEachBlock(std::function<void (const Key &)> callback) const override;
 
 private:
     cpputils::unique_ref<BlockStore> _baseBlockStore;
@@ -78,6 +79,11 @@ uint64_t CompressingBlockStore<Compressor>::estimateNumFreeBytes() const {
     return _baseBlockStore->estimateNumFreeBytes();
 }
 
+template<class Compressor>
+void CompressingBlockStore<Compressor>::forEachBlock(std::function<void (const Key &)> callback) const {
+    return _baseBlockStore->forEachBlock(callback);
+}
+
 template<class Compressor>
 uint64_t CompressingBlockStore<Compressor>::blockSizeFromPhysicalBlockSize(uint64_t blockSize) const {
     //We probably have more since we're compressing, but we don't know exactly how much.

src/blockstore/implementations/encrypted/EncryptedBlockStore.h

@@ -24,6 +24,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
   //This function should only be used by test cases
   void __setKey(const typename Cipher::EncryptionKey &encKey);
@@ -96,6 +97,11 @@ uint64_t EncryptedBlockStore<Cipher>::blockSizeFromPhysicalBlockSize(uint64_t bl
   return EncryptedBlock<Cipher>::blockSizeFromPhysicalBlockSize(_baseBlockStore->blockSizeFromPhysicalBlockSize(blockSize));
 }
 
+template<class Cipher>
+void EncryptedBlockStore<Cipher>::forEachBlock(std::function<void (const Key &)> callback) const {
+  return _baseBlockStore->forEachBlock(callback);
+}
+
 }
 }
 

src/blockstore/implementations/inmemory/InMemoryBlockStore.cpp

@@ -61,5 +61,11 @@ uint64_t InMemoryBlockStore::blockSizeFromPhysicalBlockSize(uint64_t blockSize)
   return blockSize;
 }
 
+void InMemoryBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  for (const auto &entry : _blocks) {
+    callback(Key::FromString(entry.first));
+  }
+}
+
 }
 }

src/blockstore/implementations/inmemory/InMemoryBlockStore.h

@@ -22,6 +22,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
 private:
   std::unordered_map<Key, InMemoryBlock> _blocks;

src/blockstore/implementations/ondisk/OnDiskBlockStore.cpp

@@ -77,9 +77,9 @@ void OnDiskBlockStore::remove(unique_ref<Block> block) {
 
 uint64_t OnDiskBlockStore::numBlocks() const {
   uint64_t count = 0;
-  for (auto entry = bf::directory_iterator(_rootdir); entry != bf::directory_iterator(); ++entry) {
-    if (bf::is_directory(entry->path())) {
-      count += std::distance(bf::directory_iterator(entry->path()), bf::directory_iterator());
+  for (auto prefixDir = bf::directory_iterator(_rootdir); prefixDir != bf::directory_iterator(); ++prefixDir) {
+    if (bf::is_directory(prefixDir->path())) {
+      count += std::distance(bf::directory_iterator(prefixDir->path()), bf::directory_iterator());
     }
   }
   return count;
@@ -95,5 +95,17 @@ uint64_t OnDiskBlockStore::blockSizeFromPhysicalBlockSize(uint64_t blockSize) co
   return OnDiskBlock::blockSizeFromPhysicalBlockSize(blockSize);
 }
 
+void OnDiskBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  for (auto prefixDir = bf::directory_iterator(_rootdir); prefixDir != bf::directory_iterator(); ++prefixDir) {
+    if (bf::is_directory(prefixDir->path())) {
+      string blockKeyPrefix = prefixDir->path().filename().native();
+      for (auto block = bf::directory_iterator(prefixDir->path()); block != bf::directory_iterator(); ++block) {
+        string blockKeyPostfix = block->path().filename().native();
+        callback(Key::FromString(blockKeyPrefix + blockKeyPostfix));
+      }
+    }
+  }
+}
+
 }
 }

src/blockstore/implementations/ondisk/OnDiskBlockStore.h

@@ -21,6 +21,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
 private:
   const boost::filesystem::path _rootdir;

src/blockstore/implementations/parallelaccess/ParallelAccessBlockStore.cpp

@@ -64,5 +64,9 @@ uint64_t ParallelAccessBlockStore::blockSizeFromPhysicalBlockSize(uint64_t block
   return _baseBlockStore->blockSizeFromPhysicalBlockSize(blockSize);
 }
 
+void ParallelAccessBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  return _baseBlockStore->forEachBlock(callback);
+}
+
 }
 }

src/blockstore/implementations/parallelaccess/ParallelAccessBlockStore.h

@@ -22,6 +22,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
 private:
   cpputils::unique_ref<BlockStore> _baseBlockStore;

src/blockstore/implementations/testfake/FakeBlockStore.cpp

@@ -84,5 +84,11 @@ uint64_t FakeBlockStore::blockSizeFromPhysicalBlockSize(uint64_t blockSize) cons
   return blockSize;
 }
 
+void FakeBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  for (const auto &entry : _blocks) {
+    callback(Key::FromString(entry.first));
+  }
+}
+
 }
 }

src/blockstore/implementations/testfake/FakeBlockStore.h

@@ -37,6 +37,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
   void updateData(const Key &key, const cpputils::Data &data);
 

src/blockstore/implementations/versioncounting/VersionCountingBlockStore.h

@@ -23,6 +23,7 @@ public:
   uint64_t numBlocks() const override;
   uint64_t estimateNumFreeBytes() const override;
   uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const override;
+  void forEachBlock(std::function<void (const Key &)> callback) const override;
 
 private:
   cpputils::unique_ref<BlockStore> _baseBlockStore;
@@ -76,6 +77,10 @@ inline uint64_t VersionCountingBlockStore::blockSizeFromPhysicalBlockSize(uint64
   return VersionCountingBlock::blockSizeFromPhysicalBlockSize(_baseBlockStore->blockSizeFromPhysicalBlockSize(blockSize));
 }
 
+inline void VersionCountingBlockStore::forEachBlock(std::function<void (const Key &)> callback) const {
+  return _baseBlockStore->forEachBlock(callback);
+}
+
 }
 }
 

src/blockstore/interface/BlockStore.h

@@ -29,6 +29,8 @@ public:
   // This can be used to create blocks with a certain physical block size.
   virtual uint64_t blockSizeFromPhysicalBlockSize(uint64_t blockSize) const = 0;
 
+  virtual void forEachBlock(std::function<void (const Key &)> callback) const = 0;
+
   cpputils::unique_ref<Block> create(const cpputils::Data &data) {
     while(true) {
       //TODO Copy (data.copy()) necessary?

test/blockstore/interface/helpers/BlockStoreWithRandomKeysTest.cpp

@@ -32,6 +32,7 @@ public:
   MOCK_CONST_METHOD0(numBlocks, uint64_t());
   MOCK_CONST_METHOD0(estimateNumFreeBytes, uint64_t());
   MOCK_CONST_METHOD1(blockSizeFromPhysicalBlockSize, uint64_t(uint64_t));
+  MOCK_CONST_METHOD1(forEachBlock, void(std::function<void (const blockstore::Key &)>));
 };
 
 class BlockMock: public Block {

test/blockstore/testutils/BlockStoreTest.h

@@ -3,9 +3,22 @@
 #define MESSMER_BLOCKSTORE_TEST_IMPLEMENTATIONS_TESTUTILS_BLOCKSTORETEST_H_
 
 #include <gtest/gtest.h>
+#include <gmock/gmock.h>
 #include <cpp-utils/data/DataFixture.h>
+
 #include "blockstore/interface/BlockStore.h"
 
+class MockForEachBlockCallback final {
+public:
+    std::function<void (const blockstore::Key &)> callback() {
+      return [this] (const blockstore::Key &key) {
+          called_with.push_back(key);
+      };
+    }
+
+    std::vector<blockstore::Key> called_with;
+};
+
 class BlockStoreTestFixture {
 public:
   virtual ~BlockStoreTestFixture() {}
@@ -36,6 +49,24 @@ public:
     block = blockStore->load(key).value();
     EXPECT_EQ(0, std::memcmp(fixture.data(), block->data(), fixture.size()));
   }
+
+  template<class Entry>
+  void EXPECT_UNORDERED_EQ(const std::vector<Entry> &expected, std::vector<Entry> actual) {
+    EXPECT_EQ(expected.size(), actual.size());
+    for (const Entry &expectedEntry : expected) {
+      removeOne(&actual, expectedEntry);
+    }
+  }
+
+  template<class Entry>
+  void removeOne(std::vector<Entry> *entries, const Entry &toRemove) {
+    auto found = std::find(entries->begin(), entries->end(), toRemove);
+    if (found != entries->end()) {
+      entries->erase(found);
+      return;
+    }
+    EXPECT_TRUE(false);
+  }
 };
 
 TYPED_TEST_CASE_P(BlockStoreTest);
@@ -117,11 +148,64 @@ TYPED_TEST_P(BlockStoreTest, NumBlocksIsCorrectAfterRemovingABlock) {
 }
 
 TYPED_TEST_P(BlockStoreTest, CanRemoveModifiedBlock) {
-    auto blockStore = this->fixture.createBlockStore();
-    auto block = blockStore->create(cpputils::Data(5));
-    block->write("data", 0, 4);
-    blockStore->remove(std::move(block));
-    EXPECT_EQ(0u, blockStore->numBlocks());
+  auto blockStore = this->fixture.createBlockStore();
+  auto block = blockStore->create(cpputils::Data(5));
+  block->write("data", 0, 4);
+  blockStore->remove(std::move(block));
+  EXPECT_EQ(0u, blockStore->numBlocks());
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_zeroblocks) {
+  auto blockStore = this->fixture.createBlockStore();
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({}, mockForEachBlockCallback.called_with);
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_oneblock) {
+  auto blockStore = this->fixture.createBlockStore();
+  auto block = blockStore->create(cpputils::Data(1));
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({block->key()}, mockForEachBlockCallback.called_with);
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_twoblocks) {
+  auto blockStore = this->fixture.createBlockStore();
+  auto block1 = blockStore->create(cpputils::Data(1));
+  auto block2 = blockStore->create(cpputils::Data(1));
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({block1->key(), block2->key()}, mockForEachBlockCallback.called_with);
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_threeblocks) {
+  auto blockStore = this->fixture.createBlockStore();
+  auto block1 = blockStore->create(cpputils::Data(1));
+  auto block2 = blockStore->create(cpputils::Data(1));
+  auto block3 = blockStore->create(cpputils::Data(1));
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({block1->key(), block2->key(), block3->key()}, mockForEachBlockCallback.called_with);
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_doesntListRemovedBlocks_oneblock) {
+  auto blockStore = this->fixture.createBlockStore();
+  auto block1 = blockStore->create(cpputils::Data(1));
+  blockStore->remove(std::move(block1));
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({}, mockForEachBlockCallback.called_with);
+}
+
+TYPED_TEST_P(BlockStoreTest, ForEachBlock_doesntListRemovedBlocks_twoblocks) {
+  auto blockStore = this->fixture.createBlockStore();
+  auto block1 = blockStore->create(cpputils::Data(1));
+  auto block2 = blockStore->create(cpputils::Data(1));
+  blockStore->remove(std::move(block1));
+  MockForEachBlockCallback mockForEachBlockCallback;
+  blockStore->forEachBlock(mockForEachBlockCallback.callback());
+  this->EXPECT_UNORDERED_EQ({block2->key()}, mockForEachBlockCallback.called_with);
 }
 
 TYPED_TEST_P(BlockStoreTest, Resize_Larger_FromZero) {
@@ -211,6 +295,12 @@ REGISTER_TYPED_TEST_CASE_P(BlockStoreTest,
     WriteAndReadAfterLoading,
     OverwriteAndRead,
     CanRemoveModifiedBlock,
+    ForEachBlock_zeroblocks,
+    ForEachBlock_oneblock,
+    ForEachBlock_twoblocks,
+    ForEachBlock_threeblocks,
+    ForEachBlock_doesntListRemovedBlocks_oneblock,
+    ForEachBlock_doesntListRemovedBlocks_twoblocks,
     Resize_Larger_FromZero,
     Resize_Larger_FromZero_BlockIsStillUsable,
     Resize_Larger,