Adding additional checks during upload and download of chunk

astra/pom.xml

@@ -29,7 +29,7 @@
         <opensearch.version>2.11.1</opensearch.version>
         <curator.version>5.8.0</curator.version>
         <log4j.version>2.24.3</log4j.version>
-        <aws.sdk.version>2.31.63</aws.sdk.version>
+        <aws.sdk.version>2.35.6</aws.sdk.version>
         <error.prone.version>2.38.0</error.prone.version>
         <junit.jupiter.version>5.13.1</junit.jupiter.version>
     </properties>
@@ -357,7 +357,7 @@
         <dependency>
             <groupId>software.amazon.awssdk.crt</groupId>
             <artifactId>aws-crt</artifactId>
-            <version>0.38.5</version>
+            <version>0.39.0</version>
         </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>

astra/src/main/java/com/slack/astra/blobfs/BlobStore.java

@@ -7,12 +7,18 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
 import org.slf4j.Logger;
@@ -21,6 +27,7 @@
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.internal.async.ByteArrayAsyncResponseTransformer;
+import software.amazon.awssdk.crt.s3.ChecksumAlgorithm;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
 import software.amazon.awssdk.services.s3.model.Delete;
@@ -41,6 +48,7 @@
 import software.amazon.awssdk.transfer.s3.model.DownloadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.model.DownloadRequest;
 import software.amazon.awssdk.transfer.s3.model.UploadDirectoryRequest;
+import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
 
 /**
  * Blob store abstraction for basic operations on chunk/snapshots on remote storage. All operations
@@ -68,6 +76,69 @@ public BlobStore(S3AsyncClient s3AsyncClient, String bucketName) {
    * @throws IllegalStateException Thrown if any files fail to upload
    * @throws RuntimeException Thrown when error is considered generally non-retryable
    */
+  public void uploadSequentially(String prefix, Path directoryToUpload) {
+    assert prefix != null && !prefix.isEmpty();
+    assert directoryToUpload.toFile().isDirectory();
+    assert Objects.requireNonNull(directoryToUpload.toFile().listFiles()).length > 0;
+
+    try {
+      // Get all files and sort by size (largest first for better progress visibility)
+      List<Path> allFiles;
+      try (Stream<Path> pathStream = Files.walk(directoryToUpload)) {
+        allFiles =
+            pathStream
+                .filter(Files::isRegularFile)
+                .sorted((a, b) -> Long.compare(b.toFile().length(), a.toFile().length()))
+                .collect(Collectors.toList());
+      }
+
+      List<String> failedUploads = new ArrayList<>();
+      int completedCount = 0;
+
+      for (Path file : allFiles) {
+        String relativePath = directoryToUpload.relativize(file).toString();
+        String s3Key = prefix + "/" + relativePath;
+
+        try {
+          // Upload single file with checksum validation
+          transferManager
+              .uploadFile(
+                  UploadFileRequest.builder()
+                      .putObjectRequest(
+                          req ->
+                              req.bucket(bucketName)
+                                  .key(s3Key)
+                                  .checksumAlgorithm(
+                                      String.valueOf(
+                                          ChecksumAlgorithm.SHA256)) // Critical for integrity
+                          )
+                      .source(file)
+                      .build())
+              .completionFuture()
+              .get(); // Wait for completion
+
+          completedCount++;
+        } catch (Exception e) {
+          LOG.error("Failed to upload file: {} for prefix: {}", relativePath, prefix, e);
+          failedUploads.add(relativePath);
+        }
+      }
+
+      // Handle failures same as original code
+      if (!failedUploads.isEmpty()) {
+        throw new IllegalStateException(
+            String.format(
+                "Some files failed to upload for prefix %s - attempted to upload %s files, failed %s.",
+                prefix, allFiles.size(), failedUploads.size()));
+      }
+      LOG.info(
+          "Successfully uploaded all {} files sequentially for prefix {}", completedCount, prefix);
+
+    } catch (IOException e) {
+      throw new RuntimeException("Failed to walk directory", e);
+    }
+  }
+
   public void upload(String prefix, Path directoryToUpload) {
     assert prefix != null && !prefix.isEmpty();
     assert directoryToUpload.toFile().isDirectory();
@@ -132,14 +203,17 @@ public void download(String prefix, Path destinationDirectory) {
         download
             .failedTransfers()
             .forEach(
-                failedFileUpload ->
-                    LOG.error(
-                        "Error attempting to download file from S3", failedFileUpload.exception()));
+                failedFileDownload -> {
+                  LOG.error(
+                      "Error attempting to download file from S3 for prefix {}",
+                      prefix,
+                      failedFileDownload.exception());
+                });
 
         throw new IllegalStateException(
             String.format(
-                "Some files failed to download - failed to download %s files.",
-                download.failedTransfers().size()));
+                "Some files failed to download for prefix %s - failed to download %s files.",
+                prefix, download.failedTransfers().size()));
       }
     } catch (ExecutionException | InterruptedException e) {
       throw new RuntimeException(e);
@@ -195,6 +269,28 @@ public List<String> listFiles(String prefix) {
     return filesList;
   }
 
+  public Map<String, Long> listFilesWithSize(String prefix) {
+    assert prefix != null && !prefix.isEmpty();
+
+    ListObjectsV2Request listRequest = builder().bucket(bucketName).prefix(prefix).build();
+    ListObjectsV2Publisher asyncPaginatedListResponse =
+        s3AsyncClient.listObjectsV2Paginator(listRequest);
+
+    Map<String, Long> filesListWithSize = new HashMap<>();
+    try {
+      asyncPaginatedListResponse
+          .subscribe(
+              listResponse ->
+                  listResponse
+                      .contents()
+                      .forEach(s3Object -> filesListWithSize.put(s3Object.key(), s3Object.size())))
+          .get();
+    } catch (InterruptedException | ExecutionException e) {
+      throw new RuntimeException(e);
+    }
+    return filesListWithSize;
+  }
+
   /**
    * Deletes a chunk off of object storage by chunk id. If object was not found returns false.
    *

astra/src/main/java/com/slack/astra/blobfs/S3AsyncUtil.java

@@ -1,5 +1,7 @@
 package com.slack.astra.blobfs;
 
+import static software.amazon.awssdk.core.checksums.ResponseChecksumValidation.WHEN_SUPPORTED;
+
 import com.google.common.base.Preconditions;
 import com.slack.astra.proto.config.AstraConfigs;
 import java.net.URI;
@@ -58,7 +60,9 @@ public static S3AsyncClient initS3Client(AstraConfigs.S3Config config) {
               .targetThroughputInGbps(config.getS3TargetThroughputGbps())
               .region(Region.of(region))
               .maxNativeMemoryLimitInBytes(maxNativeMemoryLimitBytes)
-              .credentialsProvider(awsCredentialsProvider);
+              .credentialsProvider(awsCredentialsProvider)
+              .responseChecksumValidation(WHEN_SUPPORTED)
+              .minimumPartSizeInBytes(64 * 1024 * 1024L); // 64MiB minimum part size
 
       // We add a healthcheck to prevent an error with the CRT client, where it will
       // continue to attempt to read data from a socket that is no longer returning data

astra/src/main/java/com/slack/astra/chunk/ChunkValidationUtils.java

@@ -0,0 +1,17 @@
+package com.slack.astra.chunk;
+
+import java.nio.file.Path;
+import org.apache.lucene.index.CheckIndex;
+import org.apache.lucene.store.FSDirectory;
+import org.apache.lucene.store.NoLockFactory;
+
+public class ChunkValidationUtils {
+
+  public static boolean isChunkClean(Path path) throws Exception {
+    FSDirectory existingDir = FSDirectory.open(path, NoLockFactory.INSTANCE);
+    try (CheckIndex checker = new CheckIndex(existingDir)) {
+      CheckIndex.Status status = checker.checkIndex();
+      return status.clean;
+    }
+  }
+}

astra/src/main/java/com/slack/astra/chunk/ReadOnlyChunkImpl.java

@@ -1,5 +1,6 @@
 package com.slack.astra.chunk;
 
+import static com.slack.astra.chunk.ChunkValidationUtils.isChunkClean;
 import static com.slack.astra.chunkManager.CachingChunkManager.ASTRA_NG_DYNAMIC_CHUNK_SIZES_FLAG;
 import static com.slack.astra.server.AstraConfig.DEFAULT_ZK_TIMEOUT_SECS;
 
@@ -28,9 +29,11 @@
 import com.slack.astra.proto.metadata.Metadata;
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.Timer;
+import java.io.File;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.time.Instant;
 import java.util.EnumSet;
 import java.util.List;
@@ -228,6 +231,62 @@ public CacheNodeAssignment getCacheNodeAssignment() {
     return assignment;
   }
 
+  private boolean validateS3vsLocalDownLoad() {
+    // check if the number of files in S3 matches the local directory
+    Map<String, Long> filesWithSizeInS3 = blobStore.listFilesWithSize(snapshotMetadata.snapshotId);
+
+    Map<String, Long> localFilesSizeMap;
+    List<String> mismatchFiles = new java.util.ArrayList<>();
+    try (Stream<Path> fileList = Files.list(dataDirectory)) {
+      localFilesSizeMap =
+          fileList
+              .filter(Files::isRegularFile)
+              .collect(
+                  Collectors.toMap(
+                      path ->
+                          dataDirectory.relativize(path).toString().replace(File.separator, "/"),
+                      path -> path.toFile().length()));
+    } catch (IOException e) {
+      throw new RuntimeException(
+          String.format("Error reading local files in directory %s", dataDirectory), e);
+    }
+    if (localFilesSizeMap.size() != filesWithSizeInS3.size()) {
+      LOG.error(
+          "Mismatch in number of files in S3 ({}) and local directory ({}) for snapshot {}",
+          filesWithSizeInS3.size(),
+          localFilesSizeMap.size(),
+          snapshotMetadata.toString());
+      return false;
+    }
+
+    for (Map.Entry<String, Long> entry : filesWithSizeInS3.entrySet()) {
+      String s3Path = entry.getKey();
+      long s3Size = entry.getValue();
+      String fileName = Paths.get(s3Path).getFileName().toString();
+
+      if (!localFilesSizeMap.containsKey(fileName)
+          || !localFilesSizeMap.get(fileName).equals(s3Size)) {
+        mismatchFiles.add(fileName);
+      }
+    }
+    if (!mismatchFiles.isEmpty()) {
+      String mismatchFilesAndSize =
+          mismatchFiles.stream()
+              .map(
+                  e ->
+                      String.format(
+                          "%s (S3Size: %s, LocalSize: %s)",
+                          e, filesWithSizeInS3.get(e), localFilesSizeMap.get(e)))
+              .collect(Collectors.joining(", "));
+      LOG.error(
+          "Mismatch in file sizes between S3 and local directory for snapshot {}. Mismatch files: {}",
+          snapshotMetadata.toString(),
+          mismatchFilesAndSize);
+      return false;
+    }
+    return true;
+  }
+
   public void downloadChunkData() {
     Timer.Sample assignmentTimer = Timer.start(meterRegistry);
     // lock
@@ -265,7 +324,24 @@ public void downloadChunkData() {
                 "No files found on blob storage, released slot for re-assignment");
           }
         }
+        // validate if the number of files in S3 matches the local directory
+        if (!validateS3vsLocalDownLoad()) {
+          String errorString =
+              String.format(
+                  "Mismatch in number or size of files in S3 and local directory for snapshot %s",
+                  snapshotMetadata);
+          throw new RuntimeException(errorString);
+        }
+
+        // check if lucene index is valid and not corrupted
+        boolean luceneStatus = isChunkClean(dataDirectory);
+        if (!luceneStatus) {
+          throw new IOException(
+              String.format(
+                  "Lucene index is not clean. Found issues for snapshot: %s.", snapshotMetadata));
+        }
 
+        // check if schema file exists
         Path schemaPath = Path.of(dataDirectory.toString(), ReadWriteChunk.SCHEMA_FILE_NAME);
         if (!Files.exists(schemaPath)) {
           throw new RuntimeException("We expect a schema.json file to exist within the index");
@@ -305,7 +381,7 @@ public void downloadChunkData() {
       // disregarding any errors
       setAssignmentState(
           getCacheNodeAssignment(), Metadata.CacheNodeAssignment.CacheNodeAssignmentState.EVICT);
-      LOG.error("Error handling chunk assignment", e);
+      LOG.error("Error handling chunk assignment for snapshot: {}", snapshotMetadata, e);
       assignmentTimer.stop(chunkAssignmentTimerFailure);
     } finally {
       chunkAssignmentLock.unlock();
@@ -538,6 +614,7 @@ private void cleanDirectory() {
     if (dataDirectory != null) {
       try {
         FileUtils.cleanDirectory(dataDirectory.toFile());
+        FileUtils.deleteDirectory(dataDirectory.toFile());
       } catch (Exception e) {
         LOG.error("Error removing files {}", dataDirectory.toString(), e);
       }