Merge pull request #14184 from camunda/ck-reuse-sfv-checksum

Reuse already created checksum collections on persist

snapshot/src/main/java/io/camunda/zeebe/snapshots/MutableChecksumsSFV.java

@@ -21,6 +21,18 @@ public interface MutableChecksumsSFV extends ImmutableChecksumsSFV {
    */
   void updateFromFile(final Path filePath) throws IOException;
 
+  /**
+   * Update the checksum collection, and add a new checksum from given bytes, likely a read file or
+   * soon to be written file.
+   *
+   * <p>Useful, if we want to avoid re-reading files etc.
+   *
+   * @param fileName the name of the file (which relates to the given bytes), that is used in the
+   *     checksum collection in SFV file format
+   * @param bytes the bytes for which the checksum should be created
+   */
+  void updateFromBytes(final String fileName, final byte[] bytes);
+
   /**
    * Build the checksum collection from a SFV format string array.
    *

snapshot/src/main/java/io/camunda/zeebe/snapshots/impl/FileBasedReceivedSnapshot.java

@@ -39,6 +39,7 @@ public class FileBasedReceivedSnapshot implements ReceivedSnapshot {
   private long expectedSnapshotChecksum;
   private int expectedTotalCount;
   private FileBasedSnapshotMetadata metadata;
+  private SfvChecksumImpl checksumCollection;
 
   FileBasedReceivedSnapshot(
       final FileBasedSnapshotId snapshotId,
@@ -115,6 +116,12 @@ private void applyInternal(final SnapshotChunk snapshotChunk) throws SnapshotWri
     LOGGER.trace("Consume snapshot snapshotChunk {} of snapshot {}", chunkName, snapshotId);
     writeReceivedSnapshotChunk(snapshotChunk, snapshotFile);
 
+    if (checksumCollection == null) {
+      checksumCollection = new SfvChecksumImpl();
+    }
+    checksumCollection.updateFromBytes(
+        snapshotFile.getFileName().toString(), snapshotChunk.getContent());
+
     if (snapshotChunk.getChunkName().equals(FileBasedSnapshotStore.METADATA_FILE_NAME)) {
       try {
         collectMetadata(snapshotChunk.getContent());
@@ -269,6 +276,13 @@ private void persistInternal(final CompletableActorFuture<PersistedSnapshot> fut
       return;
     }
 
+    if (expectedSnapshotChecksum != checksumCollection.getCombinedValue()) {
+      future.completeExceptionally(
+          new InvalidSnapshotChecksum(
+              directory, expectedSnapshotChecksum, checksumCollection.getCombinedValue()));
+      return;
+    }
+
     try {
       if (metadata == null) {
         // backward compatibility
@@ -280,8 +294,7 @@ private void persistInternal(final CompletableActorFuture<PersistedSnapshot> fut
                 Long.MAX_VALUE);
       }
       final PersistedSnapshot value =
-          snapshotStore.persistNewSnapshot(
-              snapshotId, directory, expectedSnapshotChecksum, metadata);
+          snapshotStore.persistNewSnapshot(snapshotId, directory, checksumCollection, metadata);
       future.complete(value);
     } catch (final Exception e) {
       future.completeExceptionally(e);

snapshot/src/main/java/io/camunda/zeebe/snapshots/impl/FileBasedSnapshotStore.java

@@ -14,7 +14,7 @@
 import io.camunda.zeebe.scheduler.future.ActorFuture;
 import io.camunda.zeebe.scheduler.future.CompletableActorFuture;
 import io.camunda.zeebe.snapshots.ConstructableSnapshotStore;
-import io.camunda.zeebe.snapshots.MutableChecksumsSFV;
+import io.camunda.zeebe.snapshots.ImmutableChecksumsSFV;
 import io.camunda.zeebe.snapshots.PersistableSnapshot;
 import io.camunda.zeebe.snapshots.PersistedSnapshot;
 import io.camunda.zeebe.snapshots.PersistedSnapshotListener;
@@ -484,12 +484,10 @@ private boolean isCurrentSnapshotNewer(final FileBasedSnapshotId snapshotId) {
         && persistedSnapshot.getSnapshotId().compareTo(snapshotId) >= 0);
   }
 
-  // TODO(npepinpe): using Either here would allow easy rollback regardless of when or where an
-  // exception is thrown, without having to catch and rollback for every possible case
   FileBasedSnapshot persistNewSnapshot(
       final FileBasedSnapshotId snapshotId,
       final Path directory,
-      final long expectedChecksum,
+      final ImmutableChecksumsSFV immutableChecksumsSFV,
       final FileBasedSnapshotMetadata metadata) {
     final var currentPersistedSnapshot = currentPersistedSnapshotRef.get();
 
@@ -513,19 +511,8 @@ FileBasedSnapshot persistNewSnapshot(
       moveToSnapshotDirectory(directory, destination);
 
       final var checksumPath = buildSnapshotsChecksumPath(snapshotId);
-      final MutableChecksumsSFV actualChecksum;
       try {
-        // computing the checksum on the final destination also lets us detect any failures during
-        // the
-        // copy/move that could occur
-        actualChecksum = SnapshotChecksum.calculate(destination);
-        if (actualChecksum.getCombinedValue() != expectedChecksum) {
-          rollbackPartialSnapshot(destination);
-          throw new InvalidSnapshotChecksum(
-              directory, expectedChecksum, actualChecksum.getCombinedValue());
-        }
-
-        SnapshotChecksum.persist(checksumPath, actualChecksum);
+        SnapshotChecksum.persist(checksumPath, immutableChecksumsSFV);
       } catch (final IOException e) {
         rollbackPartialSnapshot(destination);
         throw new UncheckedIOException(e);
@@ -535,7 +522,7 @@ FileBasedSnapshot persistNewSnapshot(
           new FileBasedSnapshot(
               destination,
               checksumPath,
-              actualChecksum.getCombinedValue(),
+              immutableChecksumsSFV.getCombinedValue(),
               snapshotId,
               metadata,
               this::onSnapshotDeleted,

snapshot/src/main/java/io/camunda/zeebe/snapshots/impl/FileBasedTransientSnapshot.java

@@ -149,9 +149,7 @@ private void persistInternal(final CompletableActorFuture<PersistedSnapshot> fut
               snapshotId.getExportedPosition(),
               lastFollowupEventPosition);
       writeMetadataAndUpdateChecksum(metadata);
-      snapshot =
-          snapshotStore.persistNewSnapshot(
-              snapshotId, directory, checksum.getCombinedValue(), metadata);
+      snapshot = snapshotStore.persistNewSnapshot(snapshotId, directory, checksum, metadata);
       future.complete(snapshot);
     } catch (final Exception e) {
       future.completeExceptionally(e);

snapshot/src/main/java/io/camunda/zeebe/snapshots/impl/SfvChecksumImpl.java

@@ -122,6 +122,15 @@ public void updateFromFile(final Path filePath) throws IOException {
     checksums.put(fileName, checksum.getValue());
   }
 
+  @Override
+  public void updateFromBytes(final String fileName, final byte[] bytes) {
+    combinedChecksum.update(fileName.getBytes(UTF_8));
+    final Checksum checksum = new CRC32C();
+    checksum.update(bytes);
+    combinedChecksum.update(bytes);
+    checksums.put(fileName, checksum.getValue());
+  }
+
   @Override
   public void updateFromSfvFile(final String... lines) {
     for (String line : lines) {

snapshot/src/main/java/io/camunda/zeebe/snapshots/impl/SnapshotChecksum.java

@@ -62,7 +62,7 @@ private static boolean isNotMetadataFile(final Path file) {
 
   public static void persist(final Path checksumPath, final ImmutableChecksumsSFV checksum)
       throws IOException {
-    try(final var stream = new FileOutputStream(checksumPath.toFile())) {
+    try (final var stream = new FileOutputStream(checksumPath.toFile())) {
       checksum.write(stream);
     }
   }