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 }