Simplify loading/freeing trusted setup

StefanBratanov · Oct 23, 2023 · 4b530df · 4b530df
1 parent 98724d3
commit 4b530df
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 70 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,13 +42,10 @@
 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;
@@ -66,9 +63,7 @@
 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;
@@ -161,8 +156,6 @@ spec, new SignedBlockAndState(anchorBlock, anchorState)),
               + "\nProtoarray data:\n"
               + protoArrayData,
           e);
-    } finally {
-      freeTrustedSetupIfRequired(spec);
     }
   }
 
@@ -489,14 +482,6 @@ 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/...-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java b/...-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/kzg/KzgTestExecutor.java
@@ -40,11 +40,9 @@ public final void runTest(final TestDefinition testDefinition) throws Throwable
     final SpecConfigDeneb specConfigDeneb =
         SpecConfigDeneb.required(networkConfig.getSpec().getGenesisSpecConfig());
 
-    try {
+    try (kzg) {
       kzg.loadTrustedSetup(specConfigDeneb.getTrustedSetupPath().orElseThrow());
       runTestImpl(testDefinition);
-    } finally {
-      kzg.freeTrustedSetup();
     }
   }
 

diff --git a/...c/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java b/...c/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java
@@ -56,15 +56,16 @@ private KZG getKzgWithTrustedSetup() {
   private static class KzgAutoLoadFree implements Store.CloseOnReset {
     private static final String TRUSTED_SETUP =
         Resources.getResource(TrustedSetups.class, "trusted_setup.txt").toExternalForm();
+
     private final KZG kzg = CKZG4844.createInstance();
 
     private KzgAutoLoadFree() {
       kzg.loadTrustedSetup(TRUSTED_SETUP);
     }
 
     @Override
-    public void close() {
-      kzg.freeTrustedSetup();
+    public void close() throws Exception {
+      kzg.close();
     }
   }
 }
diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java
@@ -20,15 +20,13 @@
  * This interface specifies all the KZG functions needed for the Deneb specification and is the
  * entry-point for all KZG operations in Teku.
  */
-public interface KZG {
+public interface KZG extends AutoCloseable {
 
   KZG NOOP =
       new KZG() {
-        @Override
-        public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {}
 
         @Override
-        public void freeTrustedSetup() throws KZGException {}
+        public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {}
 
         @Override
         public boolean verifyBlobKzgProofBatch(
@@ -49,12 +47,13 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
             throws KZGException {
           return KZGProof.INFINITY;
         }
+
+        @Override
+        public void close() {}
       };
 
   void loadTrustedSetup(String trustedSetupFile) throws KZGException;
 
-  void freeTrustedSetup() throws KZGException;
-
   boolean verifyBlobKzgProofBatch(
       List<Bytes> blobs, List<KZGCommitment> kzgCommitments, List<KZGProof> kzgProofs)
       throws KZGException;

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
@@ -54,6 +54,7 @@ private CKZG4844() {
     }
   }
 
+  /** Only one trusted setup at a time can be loaded. */
   @Override
   public synchronized void loadTrustedSetup(final String trustedSetupFile) throws KZGException {
     if (loadedTrustedSetupFile.isPresent()
@@ -62,6 +63,14 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws
       return;
     }
     try {
+      loadedTrustedSetupFile.ifPresent(
+          currentTrustedSetupFile -> {
+            LOG.debug(
+                "Freeing current trusted setup {} in order to load trusted setup from {}",
+                currentTrustedSetupFile,
+                trustedSetupFile);
+            freeTrustedSetup();
+          });
       final TrustedSetup trustedSetup = CKZG4844Utils.parseTrustedSetupFile(trustedSetupFile);
       final List<Bytes> g1Points = trustedSetup.g1Points();
       final List<Bytes> g2Points = trustedSetup.g2Points();
@@ -77,17 +86,6 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws
     }
   }
 
-  @Override
-  public synchronized void freeTrustedSetup() throws KZGException {
-    try {
-      CKZG4844JNI.freeTrustedSetup();
-      loadedTrustedSetupFile = Optional.empty();
-      LOG.debug("Trusted setup was freed");
-    } catch (final Exception ex) {
-      throw new KZGException("Failed to free trusted setup", ex);
-    }
-  }
-
   @Override
   public boolean verifyBlobKzgProofBatch(
       final List<Bytes> blobs,
@@ -128,4 +126,20 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
           "Failed to compute KZG proof for blob with commitment " + kzgCommitment, ex);
     }
   }
+
+  /** Frees the current trusted setup if any is loaded */
+  @Override
+  public void close() {
+    loadedTrustedSetupFile.ifPresent(__ -> freeTrustedSetup());
+  }
+
+  private void freeTrustedSetup() throws KZGException {
+    try {
+      CKZG4844JNI.freeTrustedSetup();
+      loadedTrustedSetupFile = Optional.empty();
+      LOG.debug("Trusted setup was freed");
+    } catch (final Exception ex) {
+      throw new KZGException("Failed to free trusted setup", ex);
+    }
+  }
 }
diff --git a/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java b/infrastructure/kzg/src/test/java/tech/pegasys/teku/kzg/KZGTest.java
@@ -34,7 +34,7 @@
 import java.util.stream.Stream;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.units.bigints.UInt256;
-import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -53,38 +53,28 @@ public final class KZGTest {
   @BeforeAll
   public static void setUp() {
     kzg = CKZG4844.createInstance();
+    loadTrustedSetup();
   }
 
-  @AfterEach
-  public void cleanUpIfNeeded() {
-    try {
-      kzg.freeTrustedSetup();
-    } catch (final KZGException ex) {
-      // NOOP
-    }
-  }
-
-  @Test
-  public void testKzgLoadSameTrustedSetupTwice_shouldNotThrowException() {
-    loadTrustedSetup();
-    loadTrustedSetup();
+  @AfterAll
+  public static void cleanUp() throws Exception {
+    kzg.close();
   }
 
-  @Test
-  public void testKzLoadDifferentTrustedSetupTwice_shouldThrowException() {
-    loadTrustedSetup();
-    assertThrows(KZGException.class, () -> kzg.loadTrustedSetup("trusted_setup-not-existing.txt"));
+  private static void loadTrustedSetup() {
+    final String trustedSetup =
+        Resources.getResource(TrustedSetups.class, TRUSTED_SETUP_PATH).toExternalForm();
+    kzg.loadTrustedSetup(trustedSetup);
   }
 
   @Test
-  public void testKzgFreeTrustedSetupTwice_shouldThrowException() {
+  public void testKzgLoadSameTrustedSetupTwice_shouldNotThrowException() {
     loadTrustedSetup();
-    kzg.freeTrustedSetup();
-    assertThrows(KZGException.class, kzg::freeTrustedSetup);
   }
 
   @Test
-  public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() {
+  public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() throws Exception {
+    kzg.close();
     final List<KZGException> exceptions =
         List.of(
             assertThrows(
@@ -102,11 +92,12 @@ public void testUsageWithoutLoadedTrustedSetup_shouldThrowException() {
     assertThat(exceptions)
         .allSatisfy(
             exception -> assertThat(exception).cause().hasMessage("Trusted Setup is not loaded."));
+    // load trusted setup again for other tests
+    loadTrustedSetup();
   }
 
   @Test
   public void testComputingAndVerifyingBatchProofs() {
-    loadTrustedSetup();
     final int numberOfBlobs = 4;
     final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
     final List<KZGCommitment> kzgCommitments =
@@ -133,13 +124,11 @@ public void testComputingAndVerifyingBatchProofs() {
 
   @Test
   public void testVerifyingEmptyBatch() {
-    loadTrustedSetup();
     assertThat(kzg.verifyBlobKzgProofBatch(List.of(), List.of(), List.of())).isTrue();
   }
 
   @Test
   public void testComputingAndVerifyingBatchSingleProof() {
-    loadTrustedSetup();
     final int numberOfBlobs = 1;
     final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
     final List<KZGCommitment> kzgCommitments =
@@ -167,7 +156,6 @@ public void testComputingAndVerifyingBatchSingleProof() {
 
   @Test
   public void testVerifyingBatchProofsThrowsIfSizesDoesntMatch() {
-    loadTrustedSetup();
     final int numberOfBlobs = 4;
     final List<Bytes> blobs = getSampleBlobs(numberOfBlobs);
     final List<KZGCommitment> kzgCommitments =
@@ -209,7 +197,6 @@ public void testVerifyingBatchProofsThrowsIfSizesDoesntMatch() {
         "0x925668a49d06f4"
       })
   public void testComputingProofWithIncorrectLengthBlobDoesNotCauseSegfault(final String blobHex) {
-    loadTrustedSetup();
     final Bytes blob = Bytes.fromHexString(blobHex);
 
     final KZGException kzgException =
@@ -241,6 +228,8 @@ public void incorrectTrustedSetupFilesShouldThrow(final String path) {
     final Throwable cause =
         assertThrows(KZGException.class, () -> kzg.loadTrustedSetup(trustedSetup)).getCause();
     assertThat(cause.getMessage()).contains("Failed to parse trusted setup file");
+    // reload real trusted setup for other tests
+    loadTrustedSetup();
   }
 
   @Test
@@ -252,6 +241,8 @@ public void monomialTrustedSetupFilesShouldThrow() {
     assertThat(kzgException.getMessage()).contains("Failed to load trusted setup");
     assertThat(kzgException.getCause().getMessage())
         .contains("There was an error while loading the Trusted Setup. (C_KZG_BADARGS)");
+    // reload real trusted setup for other tests
+    loadTrustedSetup();
   }
 
   @Test
@@ -261,12 +252,6 @@ public void testInvalidLengthG2PointInNewTrustedSetup() {
         .hasMessage("Expected G2 point to be 96 bytes");
   }
 
-  private void loadTrustedSetup() {
-    final String trustedSetup =
-        Resources.getResource(TrustedSetups.class, TRUSTED_SETUP_PATH).toExternalForm();
-    kzg.loadTrustedSetup(trustedSetup);
-  }
-
   private List<Bytes> getSampleBlobs(final int count) {
     return IntStream.range(0, count).mapToObj(__ -> getSampleBlob()).collect(Collectors.toList());
   }