HDDS-9526. Two S3G instances writing the same key may cause data loss in case of an exception.

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/BlockDataStreamOutputEntryPool.java

@@ -289,6 +289,10 @@ long computeBufferData() {
     return totalDataLen;
   }
 
+  public long getDataSize() {
+    return keyArgs.getDataSize();
+  }
+
   @Override
   public Map<String, String> getMetadata() {
     return this.keyArgs.getMetadata();

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/BlockOutputStreamEntryPool.java

@@ -438,4 +438,7 @@ public Map<String, String> getMetadata() {
     return null;
   }
 
+  long getDataSize() {
+    return keyArgs.getDataSize();
+  }
 }

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/ECKeyOutputStream.java

@@ -78,6 +78,14 @@ public final class ECKeyOutputStream extends KeyOutputStream
   private final Future<Boolean> flushFuture;
   private final AtomicLong flushCheckpoint;
 
+  /**
+   * Indicates if an atomic write is required. When set to true,
+   * the amount of data written must match the declared size during the commit.
+   * A mismatch will prevent the commit from succeeding.
+   * This is essential for operations like S3 put to ensure atomicity.
+   */
+  private boolean atomicKeyCreation;
+
   private enum StripeWriteStatus {
     SUCCESS,
     FAILED
@@ -155,6 +163,7 @@ private ECKeyOutputStream(Builder builder) {
     flushExecutor.submit(() -> s3CredentialsProvider.set(s3Auth));
     this.flushFuture = this.flushExecutor.submit(this::flushStripeFromQueue);
     this.flushCheckpoint = new AtomicLong(0);
+    this.atomicKeyCreation = builder.getAtomicKeyCreation();
   }
 
   /**
@@ -512,6 +521,12 @@ public void close() throws IOException {
         Preconditions.checkArgument(writeOffset == offset,
             "Expected writeOffset= " + writeOffset
                 + " Expected offset=" + offset);
+        if (atomicKeyCreation) {
+          long expectedSize = blockOutputStreamEntryPool.getDataSize();
+          Preconditions.checkState(expectedSize == offset, String.format(
+              "Expected: %d and actual %d write sizes do not match",
+                  expectedSize, offset));
+        }
         blockOutputStreamEntryPool.commitKey(offset);
       }
     } catch (ExecutionException e) {

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyDataStreamOutput.java

@@ -81,6 +81,14 @@ enum StreamAction {
 
   private long clientID;
 
+  /**
+   * Indicates if an atomic write is required. When set to true,
+   * the amount of data written must match the declared size during the commit.
+   * A mismatch will prevent the commit from succeeding.
+   * This is essential for operations like S3 put to ensure atomicity.
+   */
+  private boolean atomicKeyCreation;
+
   @VisibleForTesting
   public List<BlockDataStreamOutputEntry> getStreamEntries() {
     return blockDataStreamOutputEntryPool.getStreamEntries();
@@ -109,7 +117,8 @@ public KeyDataStreamOutput(
       OzoneManagerProtocol omClient, int chunkSize,
       String requestId, ReplicationConfig replicationConfig,
       String uploadID, int partNumber, boolean isMultipart,
-      boolean unsafeByteBufferConversion
+      boolean unsafeByteBufferConversion,
+      boolean atomicKeyCreation
   ) {
     super(HddsClientUtils.getRetryPolicyByException(
         config.getMaxRetryCount(), config.getRetryInterval()));
@@ -130,6 +139,7 @@ public KeyDataStreamOutput(
     // encrypted bucket.
     this.writeOffset = 0;
     this.clientID = handler.getId();
+    this.atomicKeyCreation = atomicKeyCreation;
   }
 
   /**
@@ -387,6 +397,12 @@ public void close() throws IOException {
       if (!isException()) {
         Preconditions.checkArgument(writeOffset == offset);
       }
+      if (atomicKeyCreation) {
+        long expectedSize = blockDataStreamOutputEntryPool.getDataSize();
+        Preconditions.checkArgument(expectedSize == offset,
+            String.format("Expected: %d and actual %d write sizes do not match",
+                expectedSize, offset));
+      }
       blockDataStreamOutputEntryPool.commitKey(offset);
     } finally {
       blockDataStreamOutputEntryPool.cleanup();
@@ -422,6 +438,7 @@ public static class Builder {
     private boolean unsafeByteBufferConversion;
     private OzoneClientConfig clientConfig;
     private ReplicationConfig replicationConfig;
+    private boolean atomicKeyCreation = false;
 
     public Builder setMultipartUploadID(String uploadID) {
       this.multipartUploadID = uploadID;
@@ -474,6 +491,11 @@ public Builder setReplicationConfig(ReplicationConfig replConfig) {
       return this;
     }
 
+    public Builder setAtomicKeyCreation(boolean atomicKey) {
+      this.atomicKeyCreation = atomicKey;
+      return this;
+    }
+
     public KeyDataStreamOutput build() {
       return new KeyDataStreamOutput(
           clientConfig,
@@ -486,7 +508,8 @@ public KeyDataStreamOutput build() {
           multipartUploadID,
           multipartNumber,
           isMultipartKey,
-          unsafeByteBufferConversion);
+          unsafeByteBufferConversion,
+          atomicKeyCreation);
     }
 
   }

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/KeyOutputStream.java

@@ -95,6 +95,14 @@ enum StreamAction {
 
   private long clientID;
 
+  /**
+   * Indicates if an atomic write is required. When set to true,
+   * the amount of data written must match the declared size during the commit.
+   * A mismatch will prevent the commit from succeeding.
+   * This is essential for operations like S3 put to ensure atomicity.
+   */
+  private boolean atomicKeyCreation;
+
   public KeyOutputStream(ReplicationConfig replicationConfig,
       ContainerClientMetrics clientMetrics) {
     this.replication = replicationConfig;
@@ -142,7 +150,8 @@ public KeyOutputStream(
       String requestId, ReplicationConfig replicationConfig,
       String uploadID, int partNumber, boolean isMultipart,
       boolean unsafeByteBufferConversion,
-      ContainerClientMetrics clientMetrics
+      ContainerClientMetrics clientMetrics,
+      boolean atomicKeyCreation
   ) {
     this.config = config;
     this.replication = replicationConfig;
@@ -163,6 +172,7 @@ public KeyOutputStream(
     this.isException = false;
     this.writeOffset = 0;
     this.clientID = handler.getId();
+    this.atomicKeyCreation = atomicKeyCreation;
   }
 
   /**
@@ -555,6 +565,12 @@ public synchronized void close() throws IOException {
       if (!isException) {
         Preconditions.checkArgument(writeOffset == offset);
       }
+      if (atomicKeyCreation) {
+        long expectedSize = blockOutputStreamEntryPool.getDataSize();
+        Preconditions.checkState(expectedSize == offset,
+            String.format("Expected: %d and actual %d write sizes do not match",
+                expectedSize, offset));
+      }
       blockOutputStreamEntryPool.commitKey(offset);
     } finally {
       blockOutputStreamEntryPool.cleanup();
@@ -591,6 +607,7 @@ public static class Builder {
     private OzoneClientConfig clientConfig;
     private ReplicationConfig replicationConfig;
     private ContainerClientMetrics clientMetrics;
+    private boolean atomicKeyCreation = false;
 
     public String getMultipartUploadID() {
       return multipartUploadID;
@@ -677,6 +694,11 @@ public Builder setReplicationConfig(ReplicationConfig replConfig) {
       return this;
     }
 
+    public Builder setAtomicKeyCreation(boolean atomicKey) {
+      this.atomicKeyCreation = atomicKey;
+      return this;
+    }
+
     public Builder setClientMetrics(ContainerClientMetrics clientMetrics) {
       this.clientMetrics = clientMetrics;
       return this;
@@ -686,6 +708,10 @@ public ContainerClientMetrics getClientMetrics() {
       return clientMetrics;
     }
 
+    public boolean getAtomicKeyCreation() {
+      return atomicKeyCreation;
+    }
+
     public KeyOutputStream build() {
       return new KeyOutputStream(
           clientConfig,
@@ -698,7 +724,8 @@ public KeyOutputStream build() {
           multipartNumber,
           isMultipartKey,
           unsafeByteBufferConversion,
-          clientMetrics);
+          clientMetrics,
+          atomicKeyCreation);
     }
 
   }

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/OzoneDataStreamOutput.java

@@ -59,30 +59,35 @@ public synchronized void close() throws IOException {
   }
 
   public OmMultipartCommitUploadPartInfo getCommitUploadPartInfo() {
+    KeyDataStreamOutput keyDataStreamOutput = getKeyDataStreamOutput();
+    if (keyDataStreamOutput != null) {
+      return keyDataStreamOutput.getCommitUploadPartInfo();
+    }
+    // Otherwise return null.
+    return null;
+  }
+
+  public KeyDataStreamOutput getKeyDataStreamOutput() {
     if (byteBufferStreamOutput instanceof OzoneOutputStream) {
       OutputStream outputStream =
           ((OzoneOutputStream) byteBufferStreamOutput).getOutputStream();
       if (outputStream instanceof KeyDataStreamOutput) {
-        return ((KeyDataStreamOutput)
-            outputStream).getCommitUploadPartInfo();
+        return ((KeyDataStreamOutput) outputStream);
       } else if (outputStream instanceof CryptoOutputStream) {
         OutputStream wrappedStream =
             ((CryptoOutputStream) outputStream).getWrappedStream();
         if (wrappedStream instanceof KeyDataStreamOutput) {
-          return ((KeyDataStreamOutput) wrappedStream)
-              .getCommitUploadPartInfo();
+          return ((KeyDataStreamOutput) wrappedStream);
         }
       } else if (outputStream instanceof CipherOutputStreamOzone) {
         OutputStream wrappedStream =
             ((CipherOutputStreamOzone) outputStream).getWrappedStream();
         if (wrappedStream instanceof KeyDataStreamOutput) {
-          return ((KeyDataStreamOutput) wrappedStream)
-              .getCommitUploadPartInfo();
+          return ((KeyDataStreamOutput) wrappedStream);
         }
       }
     } else if (byteBufferStreamOutput instanceof KeyDataStreamOutput) {
-      return ((KeyDataStreamOutput)
-          byteBufferStreamOutput).getCommitUploadPartInfo();
+      return ((KeyDataStreamOutput) byteBufferStreamOutput);
     }
     // Otherwise return null.
     return null;

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/io/OzoneOutputStream.java

@@ -123,29 +123,38 @@ public void hsync() throws IOException {
   }
 
   public OmMultipartCommitUploadPartInfo getCommitUploadPartInfo() {
+    KeyOutputStream keyOutputStream = getKeyOutputStream();
+    if (keyOutputStream != null) {
+      return keyOutputStream.getCommitUploadPartInfo();
+    }
+    // Otherwise return null.
+    return null;
+  }
+
+  public OutputStream getOutputStream() {
+    return outputStream;
+  }
+
+  public KeyOutputStream getKeyOutputStream() {
     if (outputStream instanceof KeyOutputStream) {
-      return ((KeyOutputStream) outputStream).getCommitUploadPartInfo();
+      return ((KeyOutputStream) outputStream);
     } else  if (outputStream instanceof CryptoOutputStream) {
       OutputStream wrappedStream =
           ((CryptoOutputStream) outputStream).getWrappedStream();
       if (wrappedStream instanceof KeyOutputStream) {
-        return ((KeyOutputStream) wrappedStream).getCommitUploadPartInfo();
+        return ((KeyOutputStream) wrappedStream);
       }
     } else if (outputStream instanceof CipherOutputStreamOzone) {
       OutputStream wrappedStream =
           ((CipherOutputStreamOzone) outputStream).getWrappedStream();
       if (wrappedStream instanceof KeyOutputStream) {
-        return ((KeyOutputStream)wrappedStream).getCommitUploadPartInfo();
+        return ((KeyOutputStream)wrappedStream);
       }
     }
     // Otherwise return null.
     return null;
   }
 
-  public OutputStream getOutputStream() {
-    return outputStream;
-  }
-
   @Override
   public Map<String, String> getMetadata() {
     if (outputStream instanceof CryptoOutputStream) {

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/protocol/ClientProtocol.java

@@ -1008,6 +1008,9 @@ OzoneKey headObject(String volumeName, String bucketName,
    */
   void setThreadLocalS3Auth(S3Auth s3Auth);
 
+
+  void setIsS3Request(boolean isS3Request);
+
   /**
    * Gets the S3 Authentication information that is attached to the thread.
    * @return S3 Authentication information.