Fix bulk load storage slots

web3/src/main/java/org/hiero/mirror/web3/service/ContractStateServiceImpl.java

@@ -83,35 +83,35 @@
      * @return slotKey-value pairs for contractId
      */
     private Optional<byte[]> findStorageBatch(final EntityId contractId, final byte[] key) {
-        final var contractSlotsCache = this.contractSlotsCache.get(
-                contractId, () -> cacheManagerSlotsPerContract.getCache(contractId.toString()));
-        contractSlotsCache.putIfAbsent(ByteBuffer.wrap(key), EMPTY_VALUE);
-
+        final var contractSlotsCache = ((CaffeineCache) this.contractSlotsCache.get(
+                contractId, () -> cacheManagerSlotsPerContract.getCache(contractId.toString())));
         // Cached slot keys for contract, whose slot values are not present in the contractStateCache
         final var cachedSlots = new ArrayList<byte[]>();
-        ((CaffeineCache) contractSlotsCache).getNativeCache().asMap().keySet().forEach(slot -> {
+        contractSlotsCache.putIfAbsent(ByteBuffer.wrap(key), EMPTY_VALUE);
+        (contractSlotsCache).getNativeCache().asMap().keySet().forEach(slot -> {
             final var slotBytes = ((ByteBuffer) slot).array();
-            final var value = contractStateCache.putIfAbsent(generateCacheKey(contractId, slotBytes), EMPTY_VALUE);
-
-            if (value == null) {
-                cachedSlots.add(slotBytes);
-            }
+            cachedSlots.add(slotBytes);
         });
 
         final var contractSlotValues = contractStateRepository.findStorageBatch(contractId.getId(), cachedSlots);
+        boolean evicted = !cachedSlots.contains(ByteBuffer.wrap(key).array());
         byte[] cachedValue = null;
 
         for (final var contractSlotValue : contractSlotValues) {
             final byte[] slotKey = contractSlotValue.getSlot();
             final byte[] slotValue = contractSlotValue.getValue();
             contractStateCache.put(generateCacheKey(contractId, slotKey), slotValue);
-            contractSlotsCache.put(ByteBuffer.wrap(slotKey), EMPTY_VALUE);
 
             if (Arrays.equals(slotKey, key)) {
                 cachedValue = slotValue;
             }
         }
 
+        // If the cache key was evicted and hasn't been requested since, the cached value will be null.
+        // In that case, fall back to the original query.
+        if (cachedValue == null && evicted) {
+            return contractStateRepository.findStorage(contractId.getId(), key);
+        }
         return Optional.ofNullable(cachedValue);
     }
 

web3/src/test/java/org/hiero/mirror/web3/service/ContractStateServiceTest.java

@@ -2,6 +2,7 @@
 
 package org.hiero.mirror.web3.service;
 
+import static net.i2p.crypto.eddsa.Utils.hexToBytes;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.assertj.core.api.InstanceOfAssertFactories.LIST;
 import static org.awaitility.Awaitility.await;
@@ -12,6 +13,10 @@
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
 import org.awaitility.Durations;
@@ -277,6 +282,131 @@
                 .isEmpty();
     }
 
+    @Test
+    void verifyConcurrentBatchSlotLoadingReturnsCorrectValues() throws Exception {
+        // Given
+        final var contract = persistContract();
+        final var slot1 = generateSlotKey(1);
+        final var slot2 = generateSlotKey(2);
+
+        final byte[] value1 = hexToBytes("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
+        final byte[] value2 = hexToBytes("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb");
+
+        persistContractState(contract.getId(), slot1, value1);
+        persistContractState(contract.getId(), slot2, value2);
+
+        final var contractId = contract.toEntityId();
+
+        // When: run two parallel lookups
+        final var executor = Executors.newFixedThreadPool(2);
+        final var future1 = executor.submit(() -> contractStateService.findStorage(contractId, slot1));
+        final var future2 = executor.submit(() -> contractStateService.findStorage(contractId, slot2));
+
+        final var result1 = future1.get(2, TimeUnit.SECONDS);
+        final var result2 = future2.get(2, TimeUnit.SECONDS);
+
+        // Then
+        assertThat(result1).get().isEqualTo(value1);
+        assertThat(result2).get().isEqualTo(value2);
+
+        executor.shutdown();
+    }
+
+    @Test
+    void verifyConcurrentBatchSlotLoadingReturnsCorrectValuesWithFourConcurrentValues() throws Exception {
+        // Given
+        try {
+            final int maxCacheSize = 3;
+            cacheProperties.setSlotsPerContract("expireAfterAccess=10s,maximumSize=" + maxCacheSize);
+            cacheManagerSlotsPerContract.setCacheSpecification(cacheProperties.getSlotsPerContract());
+            final var contract = persistContract();
+
+            final var slots = List.of(generateSlotKey(1), generateSlotKey(2), generateSlotKey(3), generateSlotKey(4));
+
+            final var values = List.of(
+                    hexToBytes("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
+                    hexToBytes("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"),
+                    hexToBytes("cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"),
+                    hexToBytes("dddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddddd"));
+
+            final var contractId = contract.toEntityId();
+            for (int i = 0; i < 4; i++) {
+                persistContractState(contract.getId(), slots.get(i), values.get(i));
+            }
+
+            final var executor = Executors.newFixedThreadPool(4);
+
+            // First parallel lookup
+            final List<Future<Optional<byte[]>>> firstFutures = new ArrayList<>();
+            for (var slot : slots) {
+                firstFutures.add(executor.submit(() -> contractStateService.findStorage(contractId, slot)));
+            }
+
+            final List<Optional<byte[]>> firstResults = new ArrayList<>();
+            for (var future : firstFutures) {
+                firstResults.add(future.get(2, TimeUnit.SECONDS));
+            }
+
+            for (int i = 0; i < 4; i++) {
+                assertThat(firstResults.get(i)).get().isEqualTo(values.get(i));
+            }
+
+            // Wait for contract state cache to expire
+            Thread.sleep(6000);
+
+            final List<Future<Optional<byte[]>>> secondFutures = new ArrayList<>();
+            for (var slot : slots) {
+                secondFutures.add(executor.submit(() -> contractStateService.findStorage(contractId, slot)));
+            }
+
+            final List<Optional<byte[]>> secondResults = new ArrayList<>();
+            for (var future : secondFutures) {
+                secondResults.add(future.get(2, TimeUnit.SECONDS));
+            }
+
+            for (int i = 0; i < 4; i++) {
+                assertThat(secondResults.get(i)).get().isEqualTo(values.get(i));
+            }
+            executor.shutdown();
+        } finally {
+            // reset cache
+            final int initialSize = 10;
+            cacheProperties.setSlotsPerContract("expireAfterAccess=2s,maximumSize=" + initialSize);
+            cacheManagerSlotsPerContract.setCacheSpecification(cacheProperties.getSlotsPerContract());
+        }
+    }
+
+    @Test
+    void verifyBatchSlotLoadingReturnsCorrectValuesSequentially() throws InterruptedException {
+        // Given
+        final var contract = persistContract();
+        final var slot1 = generateSlotKey(1);
+        final var slot2 = generateSlotKey(2);
+
+        final byte[] value1 = hexToBytes("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
+        final byte[] value2 = hexToBytes("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb");
+
+        persistContractState(contract.getId(), slot1, value1);
+        persistContractState(contract.getId(), slot2, value2);
+
+        final var contractId = contract.toEntityId();
+
+        // When: read both slots one after the other
+        final var result1 = contractStateService.findStorage(contractId, slot1);
+        final var result2 = contractStateService.findStorage(contractId, slot2);
+
+        Thread.sleep(6000);
+
+        final var result1Again = contractStateService.findStorage(contractId, slot1);
+        final var result2Again = contractStateService.findStorage(contractId, slot2);
+
+        // Then: both should return the correct values
+        assertThat(result1).get().isEqualTo(value1);
+        assertThat(result2).get().isEqualTo(value2);
+        assertThat(result1Again).get().isEqualTo(value1);
+        assertThat(result2Again).get().isEqualTo(value2);
+    }
+
     @Test
     void verifyDeletedHistoricalContractSlotIsNotReturned() {
         // Given
@@ -350,6 +480,13 @@
                 .persist();
     }
 
+    private void persistContractState(final long contractId, final byte[] slotKey, final byte[] value) {
+        domainBuilder
+                .contractState()
+                .customize(cs -> cs.contractId(contractId).slot(slotKey).value(value))
+                .persist();
+    }
+
     private byte[] generateSlotKey(final int index) {
         final byte[] slotKey = new byte[32];
         final byte[] indexBytes = ByteBuffer.allocate(4).putInt(index).array();