1010#include < memory>
1111#include < type_traits>
1212
13+ #include < boost/core/noncopyable.hpp>
14+ #include < boost/optional.hpp>
1315#include " interfaces/iroha_internal/transaction_batch.hpp"
1416#include " multi_sig_transactions/mst_types.hpp"
1517
@@ -81,6 +83,39 @@ namespace iroha {
8183 std::atomic<size_t > txs_quantity_{0 };
8284 };
8385
86+ // / RAII batch wrapper for transfers between limited storages
87+ class MovedBatchPtr : public boost ::noncopyable {
88+ public:
89+ ~MovedBatchPtr () {
90+ if (not is_extracted_.test_and_set ()) {
91+ limit_->remove (batch_);
92+ }
93+ }
94+
95+ BatchPtr get () const {
96+ return batch_;
97+ }
98+
99+ BatchPtr extract () {
100+ if (not is_extracted_.test_and_set ()) {
101+ limit_->remove (batch_);
102+ }
103+ return batch_;
104+ }
105+
106+ protected:
107+ template <typename T>
108+ friend class LimitedStorage ;
109+
110+ MovedBatchPtr (BatchPtr batch, std::shared_ptr<StorageLimit> limit)
111+ : batch_(std::move(batch)), limit_(std::move(limit)) {}
112+
113+ private:
114+ std::atomic_flag is_extracted_ = ATOMIC_FLAG_INIT;
115+ BatchPtr batch_;
116+ std::shared_ptr<StorageLimit> limit_;
117+ };
118+
84119 template <class StorageImpl >
85120 class LimitedStorage {
86121 static_assert (
@@ -108,9 +143,7 @@ namespace iroha {
108143 if (not limit_->addIfAllowed (batch)) {
109144 return false ;
110145 }
111-
112- txs_quantity_ += batch->transactions ().size ();
113- ++batches_quantity_;
146+ updateCountersOnInsertedBatch (batch);
114147 return storage_->insert (std::move (batch));
115148 }
116149
@@ -120,11 +153,7 @@ namespace iroha {
120153 auto extracted = extractor (static_cast <StorageImpl &>(*storage_));
121154 for (const auto &batch : extracted) {
122155 limit_->remove (batch);
123- const size_t extracted_txs = batch->transactions ().size ();
124- assert (txs_quantity_ >= extracted_txs);
125- txs_quantity_ -= extracted_txs;
126- assert (batches_quantity_ > 0 );
127- --batches_quantity_;
156+ updateCountersOnRemovedBatch (batch);
128157 }
129158 return extracted;
130159 }
@@ -135,7 +164,50 @@ namespace iroha {
135164 return func (static_cast <const StorageImpl &>(*storage_));
136165 }
137166
167+ template <typename Lambda>
168+ std::vector<std::shared_ptr<MovedBatchPtr>> move (Lambda extractor) {
169+ auto moved_batches = extractor (static_cast <StorageImpl &>(*storage_));
170+ std::vector<std::shared_ptr<MovedBatchPtr>> wrapped_moved_batches;
171+ // wrapped_moved_batches.reserve(moved_batches.size());
172+ std::transform (moved_batches.begin (),
173+ moved_batches.end (),
174+ std::back_inserter (wrapped_moved_batches),
175+ [this ](auto &&moved_batch) {
176+ this ->updateCountersOnRemovedBatch (moved_batch);
177+ return std::shared_ptr<MovedBatchPtr>(
178+ new MovedBatchPtr (std::move (moved_batch), limit_));
179+ });
180+ return wrapped_moved_batches;
181+ }
182+
183+ bool insert (std::shared_ptr<MovedBatchPtr> moved) {
184+ if (moved->limit_ == limit_) {
185+ moved->is_extracted_ .test_and_set ();
186+ return insertUnsafe (moved->batch_ );
187+ } else {
188+ return insert (moved->extract ());
189+ }
190+ }
191+
138192 private:
193+ bool insertUnsafe (BatchPtr batch) {
194+ updateCountersOnInsertedBatch (batch);
195+ return storage_->insert (std::move (batch));
196+ }
197+
198+ void updateCountersOnInsertedBatch (const BatchPtr &batch) {
199+ txs_quantity_ += batch->transactions ().size ();
200+ ++batches_quantity_;
201+ }
202+
203+ void updateCountersOnRemovedBatch (const BatchPtr &batch) {
204+ const size_t extracted_txs = batch->transactions ().size ();
205+ assert (txs_quantity_ >= extracted_txs);
206+ txs_quantity_ -= extracted_txs;
207+ assert (batches_quantity_ > 0 );
208+ --batches_quantity_;
209+ }
210+
139211 const std::shared_ptr<StorageLimit> limit_;
140212 const std::shared_ptr<LimitableStorage> storage_;
141213 size_t txs_quantity_{0 };
0 commit comments