content: /*
   *
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   *
   */
  package org.apache.bookkeeper.statelib.impl.rocksdb.checkpoint;
  
  import com.google.common.hash.Hashing;
  import com.google.common.io.CountingInputStream;
  import com.google.common.io.Files;
  import java.io.File;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.file.Paths;
  import java.nio.file.StandardCopyOption;
  import java.time.Duration;
  import java.util.Arrays;
  import java.util.List;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionException;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.TimeoutException;
  import java.util.stream.Collectors;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.statelib.api.checkpoint.CheckpointStore;
  import org.apache.bookkeeper.statelib.impl.rocksdb.RocksUtils;
  import org.apache.bookkeeper.stream.proto.kv.store.CheckpointMetadata;
  import org.apache.bookkeeper.stream.proto.kv.store.FileInfo;
  
  /**
   * CheckpointFile encapsulates the attributes and operations for a file in checkpoint.
   */
  @Slf4j
  @lombok.Builder
  @lombok.EqualsAndHashCode
  public class CheckpointFile {
      private final File file;
      private final String checksum;
      private final boolean isSstFile;
  
      private CheckpointFile(File file, String checksum, boolean isSstFile) {
          this.file = file;
          this.checksum = checksum;
          this.isSstFile = isSstFile;
      }
  
      /**
       * CheckpointFileBuilder for building instances of CheckpointFile objects.
       */
      public static class CheckpointFileBuilder {
          CheckpointFileBuilder file(File file) {
              this.file = file;
              this.isSstFile = RocksUtils.isSstFile(this.file);
              return this;
          }
          CheckpointFileBuilder file(File checkpointDir, String filename) {
              file(new File(checkpointDir, filename));
              return this;
          }
  
          CheckpointFileBuilder computeChecksum() {
              return checksum(computeChecksum(this.file));
          }
  
          private static String computeChecksum(File file) {
              String ckSum = "invalid-" + System.currentTimeMillis();
              try {
                  ckSum = Files.asByteSource(file).hash(Hashing.sha256()).toString();
                  return ckSum;
              } catch (IOException e) {
                  log.error("Failed to get checksum for file {} {}", file.getName(), e.getMessage(), e);
              }
              return ckSum;
          }
      }
  
      public File getFile() {
          return file;
      }
  
      public String toString() {
          return String.format("CheckpointFile: %s", file.getName());
      }
  
      public boolean isSstFile() {
          return isSstFile;
      }
  
      public String getName() {
          return file.getName();
      }
  
      public String getNameWithChecksum() {
          if (checksum != null) {
              return file.getName() + "_" + checksum;
          }
          return getName();
      }
  
      public String getRemoteSstPath(String dbPrefix, boolean enableChecksum) {
          if (enableChecksum) {
              return RocksUtils.getDestSstPath(dbPrefix, getNameWithChecksum());
          } else {
              return RocksUtils.getDestSstPath(dbPrefix, getName());
          }
      }
  
      public String getRemotePath(String dbPrefix, String checkpointId, boolean enableChecksum) {
          if (isSstFile) {
              return getRemoteSstPath(dbPrefix, enableChecksum);
          }
          return RocksUtils.getDestPath(dbPrefix, checkpointId, file);
      }
  
      public static List<CheckpointFile> list(File checkpointedDir) {
          // List for files from checkpoint folder
          return Arrays.stream(checkpointedDir.listFiles())
              .map(f -> CheckpointFile.builder()
                  .file(f)
                  .computeChecksum()
                  .build())
              .collect(Collectors.toList());
      }
  
      public static List<CheckpointFile> list(File checkpointDir, CheckpointMetadata metadata) {
          // List for files from checkpoint metadata
          if (metadata.getFileInfosCount() != 0) {
              return metadata.getFileInfosList().stream()
                  .map(f -> CheckpointFile.builder()
                      .file(checkpointDir, f.getName())
                      .checksum(f.getChecksum())
                      .build())
                  .collect(Collectors.toList());
          } else {
              // Old checkpoint without checksums
              return metadata.getFilesList().stream()
                  .map(f -> CheckpointFile.builder()
                      .file(checkpointDir, f)
                      .build())
                  .collect(Collectors.toList());
          }
      }
  
      public boolean needCopy(CheckpointStore checkpointStore, String dbPrefix, boolean enableChecksum) {
          if (!isSstFile) {
              // Always copy the non SST files
              return true;
          }
          // sst files
          String destSstPath = getRemoteSstPath(dbPrefix, enableChecksum);
          try {
              if (checkpointStore.fileExists(destSstPath)) {
                  return false;
              }
          } catch (IOException e) {
              log.error("Failed fileExists {} {}", file.getName(), e.getMessage(), e);
          }
          return true;
      }
  
      public void copyToRemote(CheckpointStore checkpointStore, String dbPrefix, String checkpointId) throws IOException {
          String destPath = RocksUtils.getDestPath(dbPrefix, checkpointId, getName());
          try (OutputStream os = checkpointStore.openOutputStream(destPath)) {
              Files.copy(file, os);
          }
      }
  
      public void finalize(CheckpointStore checkpointStore,
                           String dbPrefix,
                           String checkpointId,
                           boolean enableChecksum,
                           boolean enableNonChecksumCompatibility) throws IOException {
          if (!RocksUtils.isSstFile(file)) {
              return;
          }
          // Move the SST file to common area where it can be shared among other checkpoints
          String destSstTempPath = RocksUtils.getDestPath(dbPrefix, checkpointId, getName());
          String destSstPath = getRemoteSstPath(dbPrefix, enableChecksum);
          checkpointStore.rename(destSstTempPath, destSstPath);
  
          if (enableChecksum && enableNonChecksumCompatibility) {
              copyToRemoteWithoutChecksum(checkpointStore, dbPrefix, checkpointId);
          }
      }
  
      public void copyToRemoteWithoutChecksum(CheckpointStore checkpointStore,
                                              String dbPrefix,
                                              String checkpointId) throws IOException {
          String destPath = getRemoteSstPath(dbPrefix, false);
  
          try (OutputStream os = checkpointStore.openOutputStream(destPath)) {
              Files.copy(file, os);
          }
      }
  
      public void copyFromRemote(CheckpointStore checkpointStore,
                                 String dbPrefix,
                                 String checkpointId,
                                 Duration idleWait) throws IOException, TimeoutException {
          String remoteFilePath = getRemotePath(dbPrefix, checkpointId, true);
          CountingInputStream cis = new CountingInputStream(
              checkpointStore.openInputStream(remoteFilePath)
          );
          CompletableFuture<Long> copyFuture = CompletableFuture.supplyAsync(() -> {
              try {
                  return java.nio.file.Files.copy(
                      cis, Paths.get(getFile().getAbsolutePath()),
                      StandardCopyOption.REPLACE_EXISTING);
              } catch (IOException e) {
                  throw new CompletionException(e);
              }
          });
  
          long startMs = System.currentTimeMillis();
  
          while (!copyFuture.isDone()) {
              long lastCount = cis.getCount();
              try {
                  // Wait of at most `idleWait` time for this to finish. If it is not done, we will get a
                  // TimeoutException. While handling the exception we will check if there was any progress. If there
                  // was some progress, we will try again.
                  copyFuture.get(idleWait.toMillis(), TimeUnit.MILLISECONDS);
              } catch (InterruptedException | ExecutionException e) {
                  throw new IOException("Failed to copy file from remote checkpoint: " + remoteFilePath, e);
              } catch (TimeoutException e) {
                  // Check if we made any progress
                  long endMs = System.currentTimeMillis();
                  long newCount = cis.getCount();
                  log.info("Timeout waiting for copy: {} last-read {} current-read {} runtime(ms) {} ",
                      remoteFilePath, lastCount, newCount, endMs - startMs);
                  if (lastCount == newCount) {
                      throw new TimeoutException("No progress reading: " + remoteFilePath
                          + " read " + lastCount + " runtime(ms) " + (endMs - startMs)
                      );
                  }
              }
          }
      }
  
  
      public FileInfo getFileInfo() {
          return FileInfo.newBuilder()
              .setName(file.getName())
              .setChecksum(checksum)
              .build();
      }
  }
  