Coverage Summary for Class: MutableTrieCache (co.rsk.db)
Class |
Method, %
|
Line, %
|
MutableTrieCache |
52.2%
(12/23)
|
51.8%
(43/83)
|
MutableTrieCache$StorageKeysIterator |
0%
(0/5)
|
0%
(0/39)
|
Total |
42.9%
(12/28)
|
35.2%
(43/122)
|
1 /*
2 * This file is part of RskJ
3 * Copyright (C) 2019 RSK Labs Ltd.
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Lesser General Public License as published by
7 * the Free Software Foundation, either version 3 of the License, or
8 * (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 * GNU Lesser General Public License for more details.
14 *
15 * You should have received a copy of the GNU Lesser General Public License
16 * along with this program. If not, see <http://www.gnu.org/licenses/>.
17 */
18
19 package co.rsk.db;
20
21 import co.rsk.core.RskAddress;
22 import co.rsk.core.types.ints.Uint24;
23 import co.rsk.crypto.Keccak256;
24 import co.rsk.trie.MutableTrie;
25 import co.rsk.trie.Trie;
26 import co.rsk.trie.TrieKeySlice;
27 import org.ethereum.crypto.Keccak256Helper;
28 import org.ethereum.db.ByteArrayWrapper;
29 import org.ethereum.db.TrieKeyMapper;
30 import org.ethereum.vm.DataWord;
31
32 import java.nio.charset.StandardCharsets;
33 import java.util.*;
34 import java.util.function.Function;
35
36 public class MutableTrieCache implements MutableTrie {
37
38 private final TrieKeyMapper trieKeyMapper = new TrieKeyMapper();
39
40 private MutableTrie trie;
41 // We use a single cache to mark both changed elements and removed elements.
42 // null value means the element has been removed.
43 private final Map<ByteArrayWrapper, Map<ByteArrayWrapper, byte[]>> cache;
44
45 // this logs recursive delete operations to be performed at commit time
46 private final Set<ByteArrayWrapper> deleteRecursiveLog;
47
48 public MutableTrieCache(MutableTrie parentTrie) {
49 trie = parentTrie;
50 cache = new HashMap<>();
51 deleteRecursiveLog = new HashSet<>();
52 }
53
54 @Override
55 public Trie getTrie() {
56 assertNoCache();
57 return trie.getTrie();
58 }
59
60 @Override
61 public Keccak256 getHash() {
62 assertNoCache();
63 return trie.getHash();
64 }
65
66 @Override
67 public byte[] get(byte[] key) {
68 return internalGet(key, trie::get, Function.identity()).orElse(null);
69 }
70
71 private <T> Optional<T> internalGet(
72 byte[] key,
73 Function<byte[], T> trieRetriever,
74 Function<byte[], T> cacheTransformer) {
75 ByteArrayWrapper wrapper = new ByteArrayWrapper(key);
76 ByteArrayWrapper accountWrapper = getAccountWrapper(wrapper);
77
78 Map<ByteArrayWrapper, byte[]> accountItems = cache.get(accountWrapper);
79 boolean isDeletedAccount = deleteRecursiveLog.contains(accountWrapper);
80 if (accountItems == null || !accountItems.containsKey(wrapper)) {
81 if (isDeletedAccount) {
82 return Optional.empty();
83 }
84 // uncached account
85 return Optional.ofNullable(trieRetriever.apply(key));
86 }
87
88 byte[] cacheItem = accountItems.get(wrapper);
89 if (cacheItem == null) {
90 // deleted account key
91 return Optional.empty();
92 }
93
94 // cached account key
95 return Optional.ofNullable(cacheTransformer.apply(cacheItem));
96 }
97
98 public Iterator<DataWord> getStorageKeys(RskAddress addr) {
99 byte[] accountStoragePrefixKey = trieKeyMapper.getAccountStoragePrefixKey(addr);
100 ByteArrayWrapper accountWrapper = getAccountWrapper(new ByteArrayWrapper(accountStoragePrefixKey));
101
102 boolean isDeletedAccount = deleteRecursiveLog.contains(accountWrapper);
103 Map<ByteArrayWrapper, byte[]> accountItems = cache.get(accountWrapper);
104 if (accountItems == null && isDeletedAccount) {
105 return Collections.emptyIterator();
106 }
107
108 if (isDeletedAccount) {
109 // lower level is deleted, return cached items
110 return new StorageKeysIterator(Collections.emptyIterator(), accountItems, addr, trieKeyMapper);
111 }
112
113 Iterator<DataWord> storageKeys = trie.getStorageKeys(addr);
114 if (accountItems == null) {
115 // uncached account
116 return storageKeys;
117 }
118
119 return new StorageKeysIterator(storageKeys, accountItems, addr, trieKeyMapper);
120 }
121
122 // This method returns a wrapper with the same content and size expected for a account key
123 // when the key is from the same size than the original wrapper, it returns the same object
124 private ByteArrayWrapper getAccountWrapper(ByteArrayWrapper originalWrapper) {
125 byte[] key = originalWrapper.getData();
126 int size = TrieKeyMapper.domainPrefix().length + TrieKeyMapper.ACCOUNT_KEY_SIZE + TrieKeyMapper.SECURE_KEY_SIZE;
127 return key.length == size ? originalWrapper : new ByteArrayWrapper(Arrays.copyOf(key, size));
128 }
129
130 @Override
131 public void put(byte[] key, byte[] value) {
132 put(new ByteArrayWrapper(key), value);
133 }
134
135 // This method optimizes cache-to-cache transfers
136 @Override
137 public void put(ByteArrayWrapper wrapper, byte[] value) {
138 // If value==null, do we have the choice to either store it
139 // in cache with null or in deleteCache. Here we have the choice to
140 // to add it to cache with null value or to deleteCache.
141 ByteArrayWrapper accountWrapper = getAccountWrapper(wrapper);
142 Map<ByteArrayWrapper, byte[]> accountMap = cache.computeIfAbsent(accountWrapper, k -> new HashMap<>());
143 accountMap.put(wrapper, value);
144 }
145
146 @Override
147 public void put(String key, byte[] value) {
148 byte[] keybytes = key.getBytes(StandardCharsets.UTF_8);
149 put(keybytes, value);
150 }
151
152 ////////////////////////////////////////////////////////////////////////////////////
153 // The semantic of implementations is special, and not the same of the MutableTrie
154 // It is DELETE ON COMMIT, which means that changes are not applies until commit()
155 // is called, and changes are applied last.
156 ////////////////////////////////////////////////////////////////////////////////////
157 @Override
158 public void deleteRecursive(byte[] key) {
159 // Can there be wrongly unhandled interactions interactions between put() and deleteRecurse()
160 // In theory, yes. In practice, never.
161 // Suppose that a contract X calls a contract S.
162 // Contract S calls itself with CALL.
163 // Contract S suicides with SUICIDE opcode.
164 // This causes a return to prev contract.
165 // But the SUICIDE DOES NOT cause the storage keys to be removed YET.
166 // Now parent contract S is still running, and it then can create a new storage cell
167 // with SSTORE. This will be stored in the cache as a put(). The cache later receives a
168 // deleteRecursive, BUT NEVER IN THE OTHER order.
169 // See TransactionExecutor.finalization(), when it iterates the list with getDeleteAccounts().forEach()
170 ByteArrayWrapper wrap = new ByteArrayWrapper(key);
171 deleteRecursiveLog.add(wrap);
172 cache.remove(wrap);
173 }
174
175 @Override
176 public void commit() {
177 // in case something was deleted and then put again, we first have to delete all the previous data
178 deleteRecursiveLog.forEach(item -> trie.deleteRecursive(item.getData()));
179 cache.forEach((accountKey, accountData) -> {
180 if (accountData != null) {
181 // cached account
182 accountData.forEach((realKey, value) -> this.trie.put(realKey, value));
183 }
184 });
185
186 deleteRecursiveLog.clear();
187 cache.clear();
188 }
189
190 @Override
191 public void save() {
192 commit();
193 trie.save();
194 }
195
196 @Override
197 public void rollback() {
198 cache.clear();
199 deleteRecursiveLog.clear();
200 }
201
202 @Override
203 public Set<ByteArrayWrapper> collectKeys(int size) {
204 Set<ByteArrayWrapper> parentSet = trie.collectKeys(size);
205
206 // all cached items to be transferred to parent
207 cache.forEach((accountKey, account) ->
208 account.forEach((realKey, value) -> {
209 if (size == Integer.MAX_VALUE || realKey.getData().length == size) {
210 if (this.get(realKey.getData()) == null) {
211 parentSet.remove(realKey);
212 } else {
213 parentSet.add(realKey);
214 }
215 }
216 })
217 );
218 return parentSet;
219 }
220
221 private void assertNoCache() {
222 if (!cache.isEmpty()) {
223 throw new IllegalStateException();
224 }
225
226 if (!deleteRecursiveLog.isEmpty()) {
227 throw new IllegalStateException();
228 }
229 }
230
231 @Override
232 public Uint24 getValueLength(byte[] key) {
233 return internalGet(key, trie::getValueLength, cachedBytes -> new Uint24(cachedBytes.length)).orElse(Uint24.ZERO);
234 }
235
236 @Override
237 public Optional<Keccak256> getValueHash(byte[] key) {
238 return internalGet(key,
239 keyB -> trie.getValueHash(keyB).orElse(null),
240 cachedBytes -> new Keccak256(Keccak256Helper.keccak256(cachedBytes)));
241 }
242
243 private static class StorageKeysIterator implements Iterator<DataWord> {
244 private final Iterator<DataWord> keysIterator;
245 private final Map<ByteArrayWrapper, byte[]> accountItems;
246 private final RskAddress address;
247 private final int storageKeyOffset = (
248 TrieKeyMapper.domainPrefix().length +
249 TrieKeyMapper.SECURE_ACCOUNT_KEY_SIZE +
250 TrieKeyMapper.storagePrefix().length +
251 TrieKeyMapper.SECURE_KEY_SIZE)
252 * Byte.SIZE;
253 private final TrieKeyMapper trieKeyMapper;
254 private DataWord currentStorageKey;
255 private Iterator<Map.Entry<ByteArrayWrapper, byte[]>> accountIterator;
256
257 StorageKeysIterator(
258 Iterator<DataWord> keysIterator,
259 Map<ByteArrayWrapper, byte[]> accountItems,
260 RskAddress addr,
261 TrieKeyMapper trieKeyMapper) {
262 this.keysIterator = keysIterator;
263 this.accountItems = new HashMap<>(accountItems);
264 this.address = addr;
265 this.trieKeyMapper = trieKeyMapper;
266 }
267
268 @Override
269 public boolean hasNext() {
270 if (currentStorageKey != null) {
271 return true;
272 }
273
274 while (keysIterator.hasNext()) {
275 DataWord item = keysIterator.next();
276 ByteArrayWrapper fullKey = getCompleteKey(item);
277 if (accountItems.containsKey(fullKey)) {
278 byte[] value = accountItems.remove(fullKey);
279 if (value == null){
280 continue;
281 }
282 }
283 currentStorageKey = item;
284 return true;
285 }
286
287 if (accountIterator == null) {
288 accountIterator = accountItems.entrySet().iterator();
289 }
290
291 while (accountIterator.hasNext()) {
292 Map.Entry<ByteArrayWrapper, byte[]> entry = accountIterator.next();
293 byte[] key = entry.getKey().getData();
294 if (entry.getValue() != null && key.length * Byte.SIZE > storageKeyOffset) {
295 // cached account key
296 currentStorageKey = getPartialKey(key);
297 return true;
298 }
299 }
300
301 return false;
302 }
303
304 private DataWord getPartialKey(byte[] key) {
305 TrieKeySlice nodeKey = TrieKeySlice.fromKey(key);
306 byte[] storageExpandedKeySuffix = nodeKey.slice(storageKeyOffset, nodeKey.length()).encode();
307 return DataWord.valueOf(storageExpandedKeySuffix);
308 }
309
310 private ByteArrayWrapper getCompleteKey(DataWord subkey) {
311 byte[] secureKeyPrefix = trieKeyMapper.getAccountStorageKey(address, subkey);
312 return new ByteArrayWrapper(secureKeyPrefix);
313 }
314
315 @Override
316 public DataWord next() {
317 if (currentStorageKey == null && !hasNext()) {
318 throw new NoSuchElementException();
319 }
320
321 DataWord next = currentStorageKey;
322 currentStorageKey = null;
323 return next;
324 }
325 }
326 }