All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.camunda.zeebe.snapshots.impl.FileBasedReceivedSnapshot Maven / Gradle / Ivy

There is a newer version: 8.7.0-alpha2-rc1
Show newest version
/*
 * Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH under
 * one or more contributor license agreements. See the NOTICE file distributed
 * with this work for additional information regarding copyright ownership.
 * Licensed under the Camunda License 1.0. You may not use this file
 * except in compliance with the Camunda License 1.0.
 */
package io.camunda.zeebe.snapshots.impl;

import io.camunda.zeebe.scheduler.ConcurrencyControl;
import io.camunda.zeebe.scheduler.future.ActorFuture;
import io.camunda.zeebe.scheduler.future.CompletableActorFuture;
import io.camunda.zeebe.snapshots.PersistedSnapshot;
import io.camunda.zeebe.snapshots.ReceivedSnapshot;
import io.camunda.zeebe.snapshots.SnapshotChunk;
import io.camunda.zeebe.snapshots.SnapshotId;
import io.camunda.zeebe.util.FileUtil;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.Arrays;
import java.util.Objects;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class FileBasedReceivedSnapshot implements ReceivedSnapshot {

  private static final Logger LOGGER = LoggerFactory.getLogger(FileBasedReceivedSnapshot.class);
  private static final int BLOCK_SIZE = 512 * 1024;

  private final Path directory;
  private final ConcurrencyControl actor;
  private final FileBasedSnapshotStoreImpl snapshotStore;

  private final FileBasedSnapshotId snapshotId;
  private int expectedTotalCount;
  private FileBasedSnapshotMetadata metadata;
  private ByteBuffer metadataBuffer;
  private long writtenMetadataBytes;
  private SfvChecksumImpl checksumCollection;

  FileBasedReceivedSnapshot(
      final FileBasedSnapshotId snapshotId,
      final Path directory,
      final FileBasedSnapshotStoreImpl snapshotStore,
      final ConcurrencyControl actor) {
    this.snapshotId = snapshotId;
    this.snapshotStore = snapshotStore;
    this.directory = directory;
    this.actor = actor;
    expectedTotalCount = Integer.MIN_VALUE;
    writtenMetadataBytes = 0;
  }

  @Override
  public long index() {
    return snapshotId.getIndex();
  }

  @Override
  public ActorFuture apply(final SnapshotChunk snapshotChunk) {
    return actor.call(
        () -> {
          applyInternal(snapshotChunk);
          return null;
        });
  }

  private void applyInternal(final SnapshotChunk snapshotChunk) throws SnapshotWriteException {
    checkSnapshotIdIsValid(snapshotChunk.getSnapshotId());

    final var currentTotalCount = snapshotChunk.getTotalCount();
    checkTotalCountIsValid(currentTotalCount);

    final String snapshotId = snapshotChunk.getSnapshotId();
    final String chunkName = snapshotChunk.getChunkName();

    if (snapshotStore.hasSnapshotId(snapshotId)) {
      LOGGER.debug(
          "Ignore snapshot snapshotChunk {}, because snapshot {} already exists.",
          chunkName,
          snapshotId);
      return;
    }

    checkChunkChecksumIsValid(snapshotChunk, snapshotId, chunkName);

    final var tmpSnapshotDirectory = directory;
    try {
      FileUtil.ensureDirectoryExists(tmpSnapshotDirectory);
    } catch (final IOException e) {
      throw new SnapshotWriteException(
          String.format("Failed to ensure that directory %s exists.", tmpSnapshotDirectory), e);
    }

    final var snapshotFile = tmpSnapshotDirectory.resolve(chunkName);

    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(FileBasedSnapshotStoreImpl.METADATA_FILE_NAME)) {
      try {
        collectMetadata(snapshotChunk);
      } catch (final IOException e) {
        throw new SnapshotWriteException("Cannot decode snapshot metadata");
      }
    }
  }

  private void collectMetadata(final SnapshotChunk chunk) throws IOException {
    if (metadataBuffer == null) {
      metadataBuffer = ByteBuffer.allocate(Math.toIntExact(chunk.getTotalFileSize()));
    }

    metadataBuffer.put(Math.toIntExact(chunk.getFileBlockPosition()), chunk.getContent());
    writtenMetadataBytes += chunk.getContent().length;

    if (writtenMetadataBytes == chunk.getTotalFileSize()) {
      metadata = FileBasedSnapshotMetadata.decode(metadataBuffer.array());
    }
  }

  private void checkChunkChecksumIsValid(
      final SnapshotChunk snapshotChunk, final String snapshotId, final String chunkName)
      throws SnapshotWriteException {
    final long expectedChecksum = snapshotChunk.getChecksum();
    final long actualChecksum = SnapshotChunkUtil.createChecksum(snapshotChunk.getContent());

    if (expectedChecksum != actualChecksum) {
      throw new SnapshotWriteException(
          String.format(
              "Expected to have checksum %d for snapshot chunk %s (%s), but calculated %d",
              expectedChecksum, chunkName, snapshotId, actualChecksum));
    }
  }

  private void checkTotalCountIsValid(final int currentTotalCount) throws SnapshotWriteException {
    if (expectedTotalCount == Integer.MIN_VALUE) {
      expectedTotalCount = currentTotalCount;
    }

    if (expectedTotalCount != currentTotalCount) {
      throw new SnapshotWriteException(
          String.format(
              "Expected snapshot chunk with equal snapshot total count %d, but got chunk with total count %d.",
              expectedTotalCount, currentTotalCount));
    }
  }

  private void checkSnapshotIdIsValid(final String snapshotId) throws SnapshotWriteException {
    final var receivedSnapshotId = FileBasedSnapshotId.ofFileName(snapshotId);
    if (receivedSnapshotId.isEmpty()) {
      throw new SnapshotWriteException(
          String.format("Snapshot file name '%s' has unexpected format", snapshotId));
    }

    final FileBasedSnapshotId chunkSnapshotId = receivedSnapshotId.get();
    if (this.snapshotId.compareTo(chunkSnapshotId) != 0) {
      throw new SnapshotWriteException(
          String.format(
              "Expected snapshot id in chunk to be '%s' but was '%s' instead",
              this.snapshotId, chunkSnapshotId));
    }
  }

  private void writeReceivedSnapshotChunk(
      final SnapshotChunk snapshotChunk, final Path snapshotFile) throws SnapshotWriteException {

    try (final var channel =
        FileChannel.open(snapshotFile, StandardOpenOption.WRITE, StandardOpenOption.CREATE)) {
      final ByteBuffer buffer = ByteBuffer.wrap(snapshotChunk.getContent());

      while (buffer.hasRemaining()) {
        final int newLimit = Math.min(buffer.capacity(), buffer.position() + BLOCK_SIZE);
        channel.position(snapshotChunk.getFileBlockPosition() + buffer.position());
        channel.write(buffer.limit(newLimit));
        buffer.limit(buffer.capacity());
      }

      channel.force(true);

    } catch (final IOException e) {
      throw new SnapshotWriteException(
          String.format("Failed to write snapshot chunk %s", snapshotChunk), e);
    }

    LOGGER.trace("Wrote replicated snapshot chunk to file {}", snapshotFile);
  }

  @Override
  public ActorFuture abort() {
    final CompletableActorFuture abortFuture = new CompletableActorFuture<>();
    actor.run(
        () -> {
          abortInternal();
          abortFuture.complete(null);
        });
    return abortFuture;
  }

  @Override
  public ActorFuture persist() {
    final CompletableActorFuture future = new CompletableActorFuture<>();
    actor.call(
        () -> {
          persistInternal(future);
          return null;
        });
    return future;
  }

  @Override
  public SnapshotId snapshotId() {
    return snapshotId;
  }

  @Override
  public Path getPath() {
    return directory;
  }

  private void abortInternal() {
    try {
      LOGGER.debug("Aborting received snapshot in dir {}", directory);
      FileUtil.deleteFolderIfExists(directory);
    } catch (final IOException e) {
      LOGGER.warn("Failed to delete pending snapshot {}", this, e);
    } finally {
      snapshotStore.removePendingSnapshot(this);
    }
  }

  private void persistInternal(final CompletableActorFuture future) {
    if (snapshotStore.hasSnapshotId(snapshotId.getSnapshotIdAsString())) {
      abortInternal();
      future.complete(snapshotStore.getLatestSnapshot().orElseThrow());
      return;
    }

    final var files = directory.toFile().listFiles();
    try {
      Objects.requireNonNull(files, "No chunks have been applied yet");
    } catch (final Exception e) {
      future.completeExceptionally(e);
      return;
    }

    if (files.length != expectedTotalCount) {
      future.completeExceptionally(
          new IllegalStateException(
              String.format(
                  "Expected '%d' chunk files for this snapshot, but found '%d'. Files are: %s.",
                  expectedTotalCount, files.length, Arrays.toString(files))));
      return;
    }

    try {
      if (metadata == null) {
        // backward compatibility
        metadata =
            new FileBasedSnapshotMetadata(
                FileBasedSnapshotStoreImpl.VERSION,
                snapshotId.getProcessedPosition(),
                snapshotId.getExportedPosition(),
                Long.MAX_VALUE);
      }
      final PersistedSnapshot value =
          snapshotStore.persistNewSnapshot(snapshotId, checksumCollection, metadata);
      future.complete(value);
    } catch (final Exception e) {
      future.completeExceptionally(e);
    }

    snapshotStore.removePendingSnapshot(this);
  }

  @Override
  public String toString() {
    return "FileBasedReceivedSnapshot{"
        + "directory="
        + directory
        + ", metadata="
        + snapshotId
        + '}';
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy