Made key collection more efficient

rsksmart · May 30, 2023 · 33bbf49 · 33bbf49
1 parent 62c041b
commit 33bbf49
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 14 deletions.
diff --git a/rskj-core/src/main/java/co/rsk/RskContext.java b/rskj-core/src/main/java/co/rsk/RskContext.java
@@ -966,8 +966,6 @@ public synchronized List<InternalService> buildInternalServices() {
         boolean rpcWebSocketEnabled = getRskSystemProperties().isRpcWebSocketEnabled();
         boolean bloomServiceEnabled = getRskSystemProperties().bloomServiceEnabled();
 
-        logger.debug("State requester is {}", getRskSystemProperties().isSnapshotSyncEnabled());
-
         if (getRskSystemProperties().isSnapshotSyncEnabled()) {
             internalServices.add(getSnapshotProcessor());
         }

diff --git a/rskj-core/src/main/java/co/rsk/net/SnapshotProcessor.java b/rskj-core/src/main/java/co/rsk/net/SnapshotProcessor.java
@@ -3,11 +3,11 @@
 import co.rsk.config.InternalService;
 import co.rsk.net.messages.StateChunkRequestMessage;
 import co.rsk.net.messages.StateChunkResponseMessage;
+import co.rsk.trie.IterationElement;
 import co.rsk.trie.Trie;
 import co.rsk.trie.TrieStore;
 import org.ethereum.core.Block;
 import org.ethereum.core.Blockchain;
-import org.ethereum.db.ByteArrayWrapper;
 import org.ethereum.net.NodeHandler;
 import org.ethereum.net.NodeManager;
 import org.ethereum.net.client.PeerClient;
@@ -18,10 +18,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Optional;
+import java.util.*;
 
 public class SnapshotProcessor implements InternalService {
 
@@ -93,14 +90,16 @@ public void processStateChunkRequest(Peer sender, long requestId) {
         }
 
         Trie trie = retrieve.get();
-        List<ByteArrayWrapper> trieKeys = new ArrayList<>(trie.collectKeys(Integer.MAX_VALUE));
-        int nKeys = Math.min(chunkSize, trieKeys.size());
-        List<ByteArrayWrapper> sublistOfKeys = trieKeys.subList(0, nKeys);
         List<byte[]> trieEncoded = new ArrayList<>();
-
-        for (ByteArrayWrapper key : sublistOfKeys) {
-            byte[] value = trie.get(key.getData());
-            trieEncoded.add(RLP.encodeList(RLP.encodeElement(key.getData()), RLP.encodeElement(value)));
+        Iterator<IterationElement> it = trie.getInOrderIterator();
+        int i = 0;
+
+        while (it.hasNext() && i < chunkSize) {
+            IterationElement e = it.next();
+            byte[] key = e.getNodeKey().encode();
+            byte[] value = e.getNode().getValue();
+            trieEncoded.add(RLP.encodeList(RLP.encodeElement(key), RLP.encodeElement(value)));
+            i++;
         }
 
         byte[] chunkBytes = RLP.encodeList(trieEncoded.toArray(new byte[0][0]));