shared storage limit

libs/storage_shared_limit/batch_storage_limit_by_txs.hpp

@@ -0,0 +1,59 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_MST_STORAGE_GUARD_HPP
+#define IROHA_MST_STORAGE_GUARD_HPP
+
+#include <atomic>
+
+#include "interfaces/iroha_internal/transaction_batch.hpp"
+#include "multi_sig_transactions/mst_types.hpp"
+#include "storage_shared_limit/storage_limit.hpp"
+
+namespace iroha {
+
+  class BatchStorageLimitByTxs : public StorageLimit<BatchPtr> {
+   public:
+    explicit BatchStorageLimitByTxs(size_t txs_limit) : txs_limit_(txs_limit) {}
+
+    bool addIfAllowed(const BatchPtr &batch) override {
+      const auto added_txs_quantity = batch->transactions().size();
+
+      size_t current_txs_quantity;
+      size_t new_txs_quantity;
+      do {
+        current_txs_quantity = txs_quantity_.load(std::memory_order_relaxed);
+        new_txs_quantity = current_txs_quantity + added_txs_quantity;
+        if (new_txs_quantity > txs_limit_) {
+          return false;
+        }
+      } while (not std::atomic_compare_exchange_weak_explicit(
+          &txs_quantity_,
+          &current_txs_quantity,
+          new_txs_quantity,
+          std::memory_order_relaxed,
+          std::memory_order_relaxed));
+      assert(txs_limit_ >= txs_quantity_.load(std::memory_order_relaxed));
+      return true;
+    }
+
+    void remove(const BatchPtr &batch) override {
+      const size_t extracted_txs = batch->transactions().size();
+      assert(txs_quantity_.load(std::memory_order_relaxed) >= extracted_txs);
+      txs_quantity_.fetch_sub(extracted_txs, std::memory_order_relaxed);
+    }
+
+    size_t transactionsQuantity() const {
+      return txs_quantity_;
+    }
+
+   private:
+    const size_t txs_limit_;
+    std::atomic<size_t> txs_quantity_{0};
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_MST_STORAGE_GUARD_HPP

libs/storage_shared_limit/limitable_storage.hpp

@@ -0,0 +1,21 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_LIBS_LIMITABLE_STORAGE_HPP
+#define IROHA_LIBS_LIMITABLE_STORAGE_HPP
+
+namespace iroha {
+
+  template <typename Item>
+  struct LimitableStorage {
+    using ItemType = Item;
+
+    virtual ~LimitableStorage() = default;
+    virtual bool insert(Item item) = 0;
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_LIBS_LIMITABLE_STORAGE_HPP

libs/storage_shared_limit/limited_storage.hpp

@@ -0,0 +1,118 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_LIBS_LIMITED_STORAGE_HPP
+#define IROHA_LIBS_LIMITED_STORAGE_HPP
+
+#include <memory>
+#include <type_traits>
+
+#include <boost/core/noncopyable.hpp>
+#include "storage_shared_limit/limitable_storage.hpp"
+#include "storage_shared_limit/moved_item.hpp"
+#include "storage_shared_limit/storage_limit.hpp"
+
+namespace iroha {
+
+  template <class StorageImpl>
+  class LimitedStorage : boost::noncopyable {
+   public:
+    using ItemType = typename StorageImpl::ItemType;
+    using LimitType = StorageLimit<ItemType>;
+    using MovedItemType = MovedItem<ItemType>;
+
+    static_assert(
+        std::is_base_of<LimitableStorage<ItemType>, StorageImpl>::value,
+        "The storage implementation must be derived from LimitableStorage!");
+
+    LimitedStorage(std::shared_ptr<LimitType> limit,
+                   std::unique_ptr<StorageImpl> storage)
+        : limit_(std::move(limit)), storage_(std::move(storage)) {}
+
+    LimitedStorage(LimitedStorage &&other)
+        : limit_(std::move(other.limit_)),
+          storage_(std::move(other.storage_)),
+          items_quantity_(other.items_quantity_) {}
+
+    std::shared_ptr<LimitType> sharedLimit() const {
+      return limit_;
+    }
+
+    size_t itemsQuantity() const {
+      return items_quantity_;
+    }
+
+    bool insert(ItemType item) {
+      if (not limit_->addIfAllowed(item)) {
+        return false;
+      }
+      updateCountersOnInsertedItem(item);
+      return storage_->insert(std::move(item));
+    }
+
+    template <typename Lambda>
+    auto extract(Lambda extractor) ->
+        typename std::result_of<Lambda(StorageImpl &)>::type {
+      auto extracted = extractor(*storage_);
+      for (const auto &item : extracted) {
+        limit_->remove(item);
+        updateCountersOnRemovedItem(item);
+      }
+      return extracted;
+    }
+
+    template <typename Lambda>
+    auto access(Lambda func) const ->
+        typename std::result_of<Lambda(const StorageImpl &)>::type {
+      return func(static_cast<const StorageImpl &>(*storage_));
+    }
+
+    template <typename Lambda>
+    std::vector<std::shared_ptr<MovedItemType>> move(Lambda extractor) {
+      auto moved_items = extractor(*storage_);
+      std::vector<std::shared_ptr<MovedItemType>> wrapped_moved_items;
+      std::transform(moved_items.begin(),
+                     moved_items.end(),
+                     std::back_inserter(wrapped_moved_items),
+                     [this](auto &&moved_item) {
+                       this->updateCountersOnRemovedItem(moved_item);
+                       return std::shared_ptr<MovedItemType>(
+                           new MovedItemType(std::move(moved_item), limit_));
+                     });
+      return wrapped_moved_items;
+    }
+
+    bool insert(std::shared_ptr<MovedItemType> moved) {
+      if (moved->limit_ == limit_) {
+        moved->is_extracted_.test_and_set();
+        return insertUnsafe(moved->item_);
+      } else {
+        return insert(moved->extract());
+      }
+    }
+
+   private:
+    bool insertUnsafe(ItemType item) {
+      updateCountersOnInsertedItem(item);
+      return storage_->insert(std::move(item));
+    }
+
+    void updateCountersOnInsertedItem(const ItemType &item) {
+      ++items_quantity_;
+    }
+
+    void updateCountersOnRemovedItem(const ItemType &item) {
+      assert(items_quantity_ > 0);
+      --items_quantity_;
+    }
+
+    std::shared_ptr<LimitType> limit_;
+    std::unique_ptr<StorageImpl> storage_;
+    size_t items_quantity_{0};
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_LIBS_LIMITED_STORAGE_HPP

libs/storage_shared_limit/moved_item.hpp

@@ -0,0 +1,53 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_LIBS_LIMITED_STORAGE_MOVED_ITEM_HPP
+#define IROHA_LIBS_LIMITED_STORAGE_MOVED_ITEM_HPP
+
+#include <atomic>
+#include <memory>
+
+#include <boost/core/noncopyable.hpp>
+#include "storage_shared_limit/storage_limit.hpp"
+
+namespace iroha {
+
+  /// RAII item wrapper for transfers between limited storages
+  template <typename Item>
+  class MovedItem : public boost::noncopyable {
+   public:
+    ~MovedItem() {
+      if (not is_extracted_.test_and_set()) {
+        limit_->remove(item_);
+      }
+    }
+
+    Item get() const {
+      return item_;
+    }
+
+    Item extract() {
+      if (not is_extracted_.test_and_set()) {
+        limit_->remove(item_);
+      }
+      return item_;
+    }
+
+   protected:
+    template <typename T>
+    friend class LimitedStorage;
+
+    MovedItem(Item item, std::shared_ptr<StorageLimit<Item>> limit)
+        : item_(std::move(item)), limit_(std::move(limit)) {}
+
+   private:
+    std::atomic_flag is_extracted_ = ATOMIC_FLAG_INIT;
+    Item item_;
+    std::shared_ptr<StorageLimit<Item>> limit_;
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_LIBS_LIMITED_STORAGE_MOVED_ITEM_HPP

libs/storage_shared_limit/storage_limit.hpp

@@ -0,0 +1,22 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_LIBS_STORAGE_LIMIT_HPP
+#define IROHA_LIBS_STORAGE_LIMIT_HPP
+
+namespace iroha {
+
+  template <typename Item>
+  struct StorageLimit {
+    virtual ~StorageLimit() = default;
+
+    virtual bool addIfAllowed(const Item &item) = 0;
+
+    virtual void remove(const Item &item) = 0;
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_LIBS_STORAGE_LIMIT_HPP

libs/storage_shared_limit/storage_limit_none.hpp

@@ -0,0 +1,23 @@
+/**
+ * Copyright Soramitsu Co., Ltd. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#ifndef IROHA_LIBS_STORAGE_LIMIT_NONE_HPP
+#define IROHA_LIBS_STORAGE_LIMIT_NONE_HPP
+
+#include "storage_shared_limit/storage_limit.hpp"
+
+namespace iroha {
+
+  template <typename Item>
+  struct StorageLimitNone : public StorageLimit<Item> {
+    bool addIfAllowed(const Item &item) override {
+      return true;
+    }
+    void remove(const Item &item) override {}
+  };
+
+}  // namespace iroha
+
+#endif  // IROHA_LIBS_STORAGE_LIMIT_NONE_HPP

test/module/libs/CMakeLists.txt

@@ -13,3 +13,4 @@ add_subdirectory(crypto)
 add_subdirectory(datetime)
 add_subdirectory(converter)
 add_subdirectory(common)
+add_subdirectory(storage_shared_limit)

test/module/libs/storage_shared_limit/CMakeLists.txt

@@ -0,0 +1,8 @@
+AddTest(storage_shared_limit_test storage_shared_limit_test.cpp)
+target_link_libraries(storage_shared_limit_test
+    schema
+    shared_model_interfaces_factories
+    shared_model_proto_backend
+    test_logger
+    )
+