From b6e902daace48a18cfb7dfcbaac2eecc684884de Mon Sep 17 00:00:00 2001
From: Raghuvansh Raj <raghraaj@amazon.com>
Date: Thu, 8 Jun 2023 14:10:24 +0530
Subject: [PATCH] Addressing PR comments

Signed-off-by: Raghuvansh Raj <raghraaj@amazon.com>
---
 .../blobstore/stream/write/WriteContext.java  |  9 +++++----
 .../transfer/RemoteTransferContainer.java     | 20 +++++++++----------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/server/src/main/java/org/opensearch/common/blobstore/stream/write/WriteContext.java b/server/src/main/java/org/opensearch/common/blobstore/stream/write/WriteContext.java
index 3a6c3951890d2..ef5e3d1e8c26c 100644
--- a/server/src/main/java/org/opensearch/common/blobstore/stream/write/WriteContext.java
+++ b/server/src/main/java/org/opensearch/common/blobstore/stream/write/WriteContext.java
@@ -8,10 +8,11 @@
 
 package org.opensearch.common.blobstore.stream.write;
 
+import org.opensearch.common.CheckedConsumer;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.StreamContext;
 
-import java.util.function.Consumer;
+import java.io.IOException;
 
 /**
  * WriteContext is used to encapsulate all data needed by <code>BlobContainer#writeStreams</code>
@@ -25,7 +26,7 @@ public class WriteContext {
     private final long fileSize;
     private final boolean failIfAlreadyExists;
     private final WritePriority writePriority;
-    private final Consumer<Boolean> uploadFinalizer;
+    private final CheckedConsumer<Boolean, IOException> uploadFinalizer;
     private final boolean doRemoteDataIntegrityCheck;
     private final Long expectedChecksum;
 
@@ -46,7 +47,7 @@ public WriteContext(
         long fileSize,
         boolean failIfAlreadyExists,
         WritePriority writePriority,
-        Consumer<Boolean> uploadFinalizer,
+        CheckedConsumer<Boolean, IOException> uploadFinalizer,
         boolean doRemoteDataIntegrityCheck,
         @Nullable Long expectedChecksum
     ) {
@@ -99,7 +100,7 @@ public WritePriority getWritePriority() {
     /**
      * @return The <code>UploadFinalizer</code> for this upload
      */
-    public Consumer<Boolean> getUploadFinalizer() {
+    public CheckedConsumer<Boolean, IOException> getUploadFinalizer() {
         return uploadFinalizer;
     }
 
diff --git a/server/src/main/java/org/opensearch/common/blobstore/transfer/RemoteTransferContainer.java b/server/src/main/java/org/opensearch/common/blobstore/transfer/RemoteTransferContainer.java
index 11e820147cf0e..7864c3ab5c794 100644
--- a/server/src/main/java/org/opensearch/common/blobstore/transfer/RemoteTransferContainer.java
+++ b/server/src/main/java/org/opensearch/common/blobstore/transfer/RemoteTransferContainer.java
@@ -172,7 +172,7 @@ private boolean isRemoteDataIntegrityCheckPossible() {
         return isRemoteDataIntegritySupported;
     }
 
-    private void finalizeUpload(boolean uploadSuccessful) {
+    private void finalizeUpload(boolean uploadSuccessful) throws IOException {
         if (isRemoteDataIntegrityCheckPossible()) {
             return;
         }
@@ -180,16 +180,14 @@ private void finalizeUpload(boolean uploadSuccessful) {
         if (uploadSuccessful) {
             long actualChecksum = getActualChecksum();
             if (actualChecksum != expectedChecksum) {
-                throw new RuntimeException(
-                    new CorruptIndexException(
-                        "Data integrity check done after upload for file "
-                            + fileName
-                            + " failed, actual checksum: "
-                            + actualChecksum
-                            + ", expected checksum: "
-                            + expectedChecksum,
-                        fileName
-                    )
+                throw new CorruptIndexException(
+                    "Data integrity check done after upload for file "
+                        + fileName
+                        + " failed, actual checksum: "
+                        + actualChecksum
+                        + ", expected checksum: "
+                        + expectedChecksum,
+                    fileName
                 );
             }
         }