better way

StefanBratanov · Oct 20, 2023 · e9d0726 · e9d0726
1 parent 45901f3
commit e9d0726
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 9 deletions.
diff --git a/...erenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java b/...erenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java
@@ -42,10 +42,13 @@
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
 import tech.pegasys.teku.infrastructure.async.eventthread.InlineEventThread;
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
+import tech.pegasys.teku.kzg.KZG;
 import tech.pegasys.teku.kzg.KZGProof;
 import tech.pegasys.teku.reference.TestDataUtils;
 import tech.pegasys.teku.reference.TestExecutor;
 import tech.pegasys.teku.spec.Spec;
+import tech.pegasys.teku.spec.SpecMilestone;
+import tech.pegasys.teku.spec.SpecVersion;
 import tech.pegasys.teku.spec.datastructures.attestation.ValidatableAttestation;
 import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.Blob;
 import tech.pegasys.teku.spec.datastructures.blocks.BeaconBlock;
@@ -63,7 +66,9 @@
 import tech.pegasys.teku.spec.executionlayer.ExecutionLayerChannelStub;
 import tech.pegasys.teku.spec.executionlayer.ExecutionPayloadStatus;
 import tech.pegasys.teku.spec.executionlayer.PayloadStatus;
+import tech.pegasys.teku.spec.logic.common.helpers.MiscHelpers;
 import tech.pegasys.teku.spec.logic.common.statetransition.results.BlockImportResult;
+import tech.pegasys.teku.spec.logic.versions.deneb.helpers.MiscHelpersDeneb;
 import tech.pegasys.teku.statetransition.forkchoice.ForkChoice;
 import tech.pegasys.teku.statetransition.forkchoice.ForkChoiceStateProvider;
 import tech.pegasys.teku.statetransition.forkchoice.MergeTransitionBlockValidator;
@@ -156,6 +161,8 @@ spec, new SignedBlockAndState(anchorBlock, anchorState)),
               + "\nProtoarray data:\n"
               + protoArrayData,
           e);
+    } finally {
+      freeTrustedSetupIfRequired(spec);
     }
   }
 
@@ -482,6 +489,14 @@ private void assertCheckpoint(
         .isEqualTo(new Checkpoint(expectedEpoch, expectedRoot));
   }
 
+  private void freeTrustedSetupIfRequired(final Spec spec) {
+    Optional.ofNullable(spec.forMilestone(SpecMilestone.DENEB))
+        .map(SpecVersion::miscHelpers)
+        .flatMap(MiscHelpers::toVersionDeneb)
+        .map(MiscHelpersDeneb::getKzg)
+        .ifPresent(KZG::freeTrustedSetup);
+  }
+
   @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"})
   private static <T> T get(final Map<String, Object> yamlData, final String key) {
     return (T) yamlData.get(key);

diff --git a/...c/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java b/...c/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java
@@ -60,6 +60,10 @@ private static KZG initKZG(final SpecConfigDeneb config) {
     return kzg;
   }
 
+  public KZG getKzg() {
+    return kzg;
+  }
+
   /**
    * Performs complete data availability check <a
    * href="https://github.com/ethereum/consensus-specs/blob/dev/specs/deneb/fork-choice.md#is_data_available">is_data_available</a>

diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java
@@ -63,15 +63,6 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws
     }
 
     try {
-      // allow loading different trusted setup
-      loadedTrustedSetupFile.ifPresent(
-          currentTrustedSetupFile -> {
-            LOG.debug(
-                "Replacing current trusted setup {} with {}",
-                currentTrustedSetupFile,
-                trustedSetupFile);
-            freeTrustedSetup();
-          });
       final TrustedSetup trustedSetup = CKZG4844Utils.parseTrustedSetupFile(trustedSetupFile);
       final List<Bytes> g1Points = trustedSetup.g1Points();
       final List<Bytes> g2Points = trustedSetup.g2Points();