Added electra validator index caching

beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorIndexCacheTracker.java

@@ -33,6 +33,6 @@ public void onNewFinalizedCheckpoint(
     final BeaconState finalizedState = recentChainData.getStore().getLatestFinalized().getState();
     BeaconStateCache.getTransitionCaches(finalizedState)
         .getValidatorIndexCache()
-        .updateLatestFinalizedIndex(finalizedState);
+        .updateLatestFinalizedState(finalizedState);
   }
 }

beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/ValidatorIndexCacheTrackerTest.java

@@ -34,7 +34,7 @@ class ValidatorIndexCacheTrackerTest {
 
   private final RecentChainData recentChainData = mock(RecentChainData.class);
 
-  private final Spec spec = TestSpecFactory.createDefault();
+  private final Spec spec = TestSpecFactory.createMinimalElectra();
 
   private final DataStructureUtil dataStructureUtil = new DataStructureUtil(spec);
 

eth-benchmark-tests/src/jmh/java/tech/pegasys/teku/benchmarks/ElectraValidatorIndexCacheBenchmark.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Consensys Software Inc., 2024
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package tech.pegasys.teku.benchmarks;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import tech.pegasys.teku.bls.BLSPublicKey;
+import tech.pegasys.teku.infrastructure.collections.cache.LRUCache;
+import tech.pegasys.teku.spec.Spec;
+import tech.pegasys.teku.spec.TestSpecFactory;
+import tech.pegasys.teku.spec.datastructures.state.beaconstate.BeaconState;
+import tech.pegasys.teku.spec.datastructures.state.beaconstate.common.ValidatorIndexCache;
+import tech.pegasys.teku.spec.util.DataStructureUtil;
+
+@Fork(1)
+@State(Scope.Thread)
+@BenchmarkMode(Mode.Throughput)
+@Warmup(iterations = 2)
+@Measurement(iterations = 5)
+public class ElectraValidatorIndexCacheBenchmark {
+  static final int VALIDATORS_MAX_IDX = 399_999;
+  private static final Spec SPEC = TestSpecFactory.createMinimalElectra();
+
+  private static final DataStructureUtil dataStructureUtil = new DataStructureUtil(0, SPEC);
+  private static final BeaconState STATE =
+      dataStructureUtil.randomBeaconState(VALIDATORS_MAX_IDX + 1);
+  private ValidatorIndexCache cache;
+  private static final BLSPublicKey RANDOM_KEY = dataStructureUtil.randomPublicKey();
+
+  @Setup(Level.Trial)
+  public void doSetup() {
+    cache =
+        new ValidatorIndexCache(
+            LRUCache.create(Integer.MAX_VALUE - 1), -1, -1, STATE.getSlot().minusMinZero(8));
+    cache.getValidatorIndex(STATE, STATE.getValidators().get(VALIDATORS_MAX_IDX).getPublicKey());
+    cache.updateLatestFinalizedState(STATE);
+  }
+
+  @Benchmark
+  public void cacheHit(Blackhole bh) {
+    bh.consume(
+        cache.getValidatorIndex(
+            STATE,
+            STATE
+                .getValidators()
+                .get(dataStructureUtil.randomPositiveInt(VALIDATORS_MAX_IDX))
+                .getPublicKey()));
+  }
+
+  @Benchmark
+  public void cacheMiss(Blackhole bh) {
+    bh.consume(cache.getValidatorIndex(STATE, RANDOM_KEY));
+  }
+}

ethereum/spec/src/main/java/tech/pegasys/teku/spec/datastructures/state/beaconstate/BeaconState.java

@@ -192,4 +192,8 @@ default Optional<BeaconStateDeneb> toVersionDeneb() {
   default Optional<BeaconStateElectra> toVersionElectra() {
     return Optional.empty();
   }
+
+  default boolean isVersionElectra() {
+    return false;
+  }
 }

ethereum/spec/src/main/java/tech/pegasys/teku/spec/datastructures/state/beaconstate/common/ValidatorIndexCache.java

@@ -21,6 +21,7 @@
 import tech.pegasys.teku.infrastructure.collections.cache.LRUCache;
 import tech.pegasys.teku.infrastructure.collections.cache.NoOpCache;
 import tech.pegasys.teku.infrastructure.ssz.SszList;
+import tech.pegasys.teku.infrastructure.unsigned.UInt64;
 import tech.pegasys.teku.spec.datastructures.state.Validator;
 import tech.pegasys.teku.spec.datastructures.state.beaconstate.BeaconState;
 
@@ -29,48 +30,98 @@ public class ValidatorIndexCache {
   private final AtomicInteger lastCachedIndex;
 
   private static final int INDEX_NONE = -1;
-  private final AtomicInteger latestFinalizedIndex;
+  private int latestFinalizedIndex;
+  private UInt64 latestFinalizedSlot;
   public static final ValidatorIndexCache NO_OP_INSTANCE =
-      new ValidatorIndexCache(NoOpCache.getNoOpCache(), INDEX_NONE, INDEX_NONE);
+      new ValidatorIndexCache(NoOpCache.getNoOpCache(), INDEX_NONE, INDEX_NONE, UInt64.ZERO);
 
   @VisibleForTesting
-  ValidatorIndexCache(
+  public ValidatorIndexCache(
       final Cache<BLSPublicKey, Integer> validatorIndices,
       final int latestFinalizedIndex,
-      final int lastCachedIndex) {
+      final int lastCachedIndex,
+      final UInt64 latestFinalizedSlot) {
     this.validatorIndices = validatorIndices;
-    this.latestFinalizedIndex = new AtomicInteger(latestFinalizedIndex);
+    this.latestFinalizedIndex = latestFinalizedIndex;
+    this.latestFinalizedSlot = latestFinalizedSlot;
     this.lastCachedIndex = new AtomicInteger(lastCachedIndex);
   }
 
   public ValidatorIndexCache() {
     this.validatorIndices = LRUCache.create(Integer.MAX_VALUE - 1);
     this.lastCachedIndex = new AtomicInteger(INDEX_NONE);
-    latestFinalizedIndex = new AtomicInteger(INDEX_NONE);
+    latestFinalizedIndex = INDEX_NONE;
+    this.latestFinalizedSlot = UInt64.ZERO;
   }
 
   public Optional<Integer> getValidatorIndex(
       final BeaconState state, final BLSPublicKey publicKey) {
-    final Optional<Integer> validatorIndex = validatorIndices.getCached(publicKey);
-    if (validatorIndex.isPresent()) {
-      return validatorIndex.filter(index -> index < state.getValidators().size());
+    if (!state.isVersionElectra() || state.getSlot().isLessThan(latestFinalizedSlot)) {
+      return getStableValidatorIndex(state, publicKey);
     }
-
-    return findIndexFromState(state.getValidators(), publicKey);
+    return getUnstableValidatorIndex(state, publicKey);
   }
 
   public void invalidateWithNewValue(final BLSPublicKey pubKey, final int updatedIndex) {
     validatorIndices.invalidateWithNewValue(pubKey, updatedIndex);
   }
 
-  public void updateLatestFinalizedIndex(final BeaconState finalizedState) {
-    latestFinalizedIndex.updateAndGet(
-        curr -> Math.max(curr, finalizedState.getValidators().size() - 1));
+  public synchronized void updateLatestFinalizedState(final BeaconState finalizedState) {
+    if (finalizedState.isVersionElectra()) {
+      // fix cache entries that may be incorrect
+      for (int i = latestFinalizedIndex + 1; i < finalizedState.getValidators().size(); i++) {
+        final BLSPublicKey publicKey = finalizedState.getValidators().get(i).getPublicKey();
+        invalidateWithNewValue(publicKey, i);
+      }
+      latestFinalizedIndex =
+          Math.max(latestFinalizedIndex, finalizedState.getValidators().size() - 1);
+      lastCachedIndex.updateAndGet(curr -> Math.max(curr, latestFinalizedIndex));
+      latestFinalizedSlot = latestFinalizedSlot.max(finalizedState.getSlot());
+    }
+  }
+
+  /*
+   * Pre Electra logic
+   *   - validator indices always stable
+   *   - because of stability, no need to track finalized context
+   *
+   * Post Electra
+   *   - non-finalized indices are not guaranteed
+   *   - finalization will update any indices that are now 'stable'
+   *   - this requires tracking and updating on finalization boundary.
+   *   - Even non final indices go into the LRU, its just they're not
+   *     guaranteed, and need validating after fetching.
+   */
+
+  Optional<Integer> getStableValidatorIndex(final BeaconState state, final BLSPublicKey publicKey) {
+    final Optional<Integer> validatorIndex = validatorIndices.getCached(publicKey);
+    if (validatorIndex.isPresent()) {
+      return validatorIndex.filter(index -> index < state.getValidators().size());
+    }
+    return findIndexFromState(state.getValidators(), publicKey, lastCachedIndex.get() + 1);
+  }
+
+  Optional<Integer> getUnstableValidatorIndex(
+      final BeaconState state, final BLSPublicKey publicKey) {
+    final Optional<Integer> validatorIndex = validatorIndices.getCached(publicKey);
+    if (validatorIndex.isPresent()
+        && state.getValidators().get(validatorIndex.get()).getPublicKey().equals(publicKey)) {
+      return validatorIndex.filter(index -> index < state.getValidators().size());
+    }
+    return findIndexFromState(
+        state.getValidators(),
+        publicKey,
+        Math.min(getLastCachedIndex() + 1, latestFinalizedIndex + 1));
   }
 
   @VisibleForTesting
   public int getLatestFinalizedIndex() {
-    return latestFinalizedIndex.get();
+    return latestFinalizedIndex;
+  }
+
+  @VisibleForTesting
+  UInt64 getLatestFinalizedSlot() {
+    return latestFinalizedSlot;
   }
 
   @VisibleForTesting
@@ -88,9 +139,11 @@ private void updateLastIndex(final int i) {
   }
 
   private Optional<Integer> findIndexFromState(
-      final SszList<Validator> validatorList, final BLSPublicKey publicKey) {
+      final SszList<Validator> validatorList,
+      final BLSPublicKey publicKey,
+      final int searchStartIndex) {
     final int initialCacheSize = getCacheSize();
-    for (int i = Math.max(lastCachedIndex.get() + 1, 0); i < validatorList.size(); i++) {
+    for (int i = searchStartIndex; i < validatorList.size(); i++) {
       final BLSPublicKey pubKey = validatorList.get(i).getPublicKey();
       validatorIndices.invalidateWithNewValue(pubKey, i);
       if (pubKey.equals(publicKey)) {

ethereum/spec/src/main/java/tech/pegasys/teku/spec/datastructures/state/beaconstate/versions/electra/BeaconStateElectra.java

@@ -84,6 +84,11 @@ default Optional<BeaconStateElectra> toVersionElectra() {
     return Optional.of(this);
   }
 
+  @Override
+  default boolean isVersionElectra() {
+    return true;
+  }
+
   default UInt64 getDepositReceiptsStartIndex() {
     final int index = getSchema().getFieldIndex(DEPOSIT_RECEIPTS_START_INDEX);
     return ((SszUInt64) get(index)).get();