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

io.camunda.zeebe.backup.s3.FileSetManager Maven / Gradle / Ivy

There is a newer version: 8.7.0-alpha2
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.backup.s3;

import io.camunda.zeebe.backup.api.NamedFileSet;
import io.camunda.zeebe.backup.common.NamedFileSetImpl;
import io.camunda.zeebe.backup.s3.S3BackupStoreException.BackupCompressionFailed;
import io.camunda.zeebe.backup.s3.manifest.FileSet;
import io.camunda.zeebe.backup.s3.manifest.FileSet.FileMetadata;
import io.camunda.zeebe.backup.s3.util.CompletableFutureUtils;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Map.Entry;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Semaphore;
import org.apache.commons.compress.compressors.CompressorException;
import org.apache.commons.compress.compressors.CompressorStreamFactory;
import org.apache.commons.compress.utils.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.core.FileTransformerConfiguration.FailureBehavior;
import software.amazon.awssdk.core.FileTransformerConfiguration.FileWriteOption;
import software.amazon.awssdk.core.async.AsyncRequestBody;
import software.amazon.awssdk.core.async.AsyncResponseTransformer;
import software.amazon.awssdk.services.s3.S3AsyncClient;

/** Can save and restore {@link NamedFileSet NamedFileSets}. */
final class FileSetManager {

  private static final Logger LOG = LoggerFactory.getLogger(FileSetManager.class);
  private static final int COMPRESSION_SIZE_THRESHOLD = 8 * 1024 * 1024; // 8 MiB
  private static final String TMP_COMPRESSION_PREFIX = "zb-backup-compress-";
  private static final String TMP_DECOMPRESSION_PREFIX = "zb-backup-decompress-";

  private final S3AsyncClient client;
  private final S3BackupConfig config;
  private final Semaphore uploadLimit;

  public FileSetManager(final S3AsyncClient client, final S3BackupConfig config) {
    this.client = client;
    this.config = config;

    // We try not to exhaust the available connections by restricting the number of
    // concurrent uploads to half of the number of available connections.
    // This should prevent ConnectionAcquisitionTimeout for backups with many and/or large files
    // where we would otherwise occupy all connections, preventing some uploads from starting.
    uploadLimit = new Semaphore(Math.max(1, config.maxConcurrentConnections() / 2));
  }

  CompletableFuture save(final String prefix, final NamedFileSet files) {
    LOG.debug("Saving {} files to prefix {}", files.files().size(), prefix);
    return CompletableFutureUtils.mapAsync(
            files.namedFiles().entrySet(),
            Entry::getKey,
            namedFile -> saveFile(prefix, namedFile.getKey(), namedFile.getValue()))
        .thenApply(FileSet::new);
  }

  private CompletableFuture saveFile(
      final String prefix, final String fileName, final Path filePath) {

    if (shouldCompressFile(filePath)) {
      final var algorithm = config.compressionAlgorithm().orElseThrow();
      return CompletableFuture.runAsync(uploadLimit::acquireUninterruptibly)
          .thenApply((success) -> compressFile(filePath, algorithm))
          .thenCompose(
              (compressedFile) -> {
                LOG.trace(
                    "Saving compressed file {}({}) in prefix {}", fileName, compressedFile, prefix);
                return client
                    .putObject(
                        put -> put.bucket(config.bucketName()).key(prefix + fileName),
                        AsyncRequestBody.fromFile(compressedFile))
                    .thenRunAsync(() -> cleanupCompressedFile(compressedFile))
                    .thenApply(unused -> FileSet.FileMetadata.withCompression(algorithm));
              })
          .whenComplete((success, error) -> uploadLimit.release());
    }

    return CompletableFuture.runAsync(uploadLimit::acquireUninterruptibly)
        .thenCompose(
            (nothing) -> {
              LOG.trace("Saving file {}({}) in prefix {}", fileName, filePath, prefix);
              return client
                  .putObject(
                      put -> put.bucket(config.bucketName()).key(prefix + fileName),
                      AsyncRequestBody.fromFile(filePath))
                  .thenApply(unused -> FileSet.FileMetadata.none());
            })
        .whenComplete((success, error) -> uploadLimit.release());
  }

  private void cleanupCompressedFile(final Path compressedFile) {
    try {
      Files.delete(compressedFile);
    } catch (final IOException e) {
      LOG.warn(
          "Failed to clean up temporary file used for (de-)compression: {}", compressedFile, e);
    }
  }

  private boolean shouldCompressFile(final Path filePath) {
    try {
      return config.compressionAlgorithm().isPresent()
          && Files.size(filePath) > COMPRESSION_SIZE_THRESHOLD;
    } catch (final IOException e) {
      LOG.warn("Failed to determine if file should be compressed, assuming no: {}", filePath);
      return false;
    }
  }

  private Path compressFile(final Path file, final String algorithm) {
    try {
      final var compressedFile = Files.createTempFile(TMP_COMPRESSION_PREFIX, null);
      LOG.trace("Compressing file {} to {} using {}", file, compressedFile, algorithm);
      try (final var input = new BufferedInputStream(Files.newInputStream(file));
          final var output = new BufferedOutputStream(Files.newOutputStream(compressedFile));
          final var compressedOutput =
              new CompressorStreamFactory().createCompressorOutputStream(algorithm, output)) {
        IOUtils.copy(input, compressedOutput);
        if (LOG.isTraceEnabled()) {
          LOG.trace(
              "Compressed file {} to {}. Uncompressed: {} bytes, compressed: {} bytes",
              file,
              compressedFile,
              Files.size(file),
              Files.size(compressedFile));
        }
        return compressedFile;
      }
    } catch (final IOException | CompressorException e) {
      throw new BackupCompressionFailed(
          "Failed to compress file %s using %s".formatted(file, algorithm), e);
    }
  }

  CompletableFuture restore(
      final String sourcePrefix, final FileSet fileSet, final Path targetFolder) {
    LOG.debug(
        "Restoring {} files from prefix {} to {}",
        fileSet.files().size(),
        sourcePrefix,
        targetFolder);
    return CompletableFutureUtils.mapAsync(
            fileSet.files().entrySet(),
            Entry::getKey,
            namedFile ->
                restoreFile(sourcePrefix, targetFolder, namedFile.getKey(), namedFile.getValue()))
        .thenApply(NamedFileSetImpl::new);
  }

  private CompletableFuture restoreFile(
      final String sourcePrefix,
      final Path targetFolder,
      final String fileName,
      final FileMetadata metadata) {
    final var compressionAlgorithm = metadata.compressionAlgorithm();
    if (compressionAlgorithm.isPresent()) {
      final var decompressed = targetFolder.resolve(fileName);
      LOG.trace(
          "Restoring compressed file {} from prefix {} to {}",
          fileName,
          sourcePrefix,
          targetFolder);
      try {
        final var compressed = Files.createTempFile(TMP_DECOMPRESSION_PREFIX, null);
        return client
            .getObject(
                req -> req.bucket(config.bucketName()).key(sourcePrefix + fileName),
                AsyncResponseTransformer.toFile(
                    compressed,
                    cfg ->
                        cfg.fileWriteOption(FileWriteOption.CREATE_OR_REPLACE_EXISTING)
                            .failureBehavior(FailureBehavior.DELETE)))
            .thenApplyAsync(
                response -> decompressFile(compressed, decompressed, compressionAlgorithm.get()));

      } catch (final IOException e) {
        throw new UncheckedIOException(e);
      }
    }

    LOG.trace("Restoring file {} from prefix {} to {}", fileName, sourcePrefix, targetFolder);
    final var path = targetFolder.resolve(fileName);
    return client
        .getObject(req -> req.bucket(config.bucketName()).key(sourcePrefix + fileName), path)
        .thenApply(response -> path);
  }

  private Path decompressFile(
      final Path compressed, final Path decompressed, final String algorithm) {
    try (final var input = new BufferedInputStream(Files.newInputStream(compressed));
        final var output = new BufferedOutputStream(Files.newOutputStream(decompressed));
        final var decompressedOutput =
            new CompressorStreamFactory().createCompressorInputStream(algorithm, input)) {
      IOUtils.copy(decompressedOutput, output);
      if (LOG.isTraceEnabled()) {
        LOG.trace(
            "Decompressed file {} to {} using {}. Compressed: {} bytes, uncompressed: {} bytes",
            compressed,
            decompressed,
            algorithm,
            Files.size(compressed),
            Files.size(decompressed));
      }
      cleanupCompressedFile(compressed);
      return decompressed;
    } catch (final IOException | CompressorException e) {
      throw new BackupCompressionFailed(
          "Failed to decompress from %s to %s using %s"
              .formatted(compressed, decompressed, algorithm),
          e);
    }
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy