Load content hash as the etag of the object when the UFS is S3 in 2.8

core/client/fs/src/main/java/alluxio/client/block/stream/BlockWorkerDataWriter.java

@@ -30,6 +30,7 @@
 import io.netty.buffer.ByteBuf;
 
 import java.io.IOException;
+import java.util.Optional;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -90,6 +91,11 @@ public int chunkSize() {
     return mChunkSize;
   }
 
+  @Override
+  public Optional<String> getUfsContentHash() {
+    return Optional.empty();
+  }
+
   @Override
   public void writeChunk(final ByteBuf buf) throws IOException {
     try {

core/client/fs/src/main/java/alluxio/client/block/stream/DataWriter.java

@@ -29,6 +29,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.Optional;
 import javax.annotation.concurrent.ThreadSafe;
 
 /**
@@ -97,6 +98,12 @@ public static DataWriter create(FileSystemContext context, long blockId, long bl
     }
   }
 
+  /**
+   * @return the content hash of the file if it is written to the UFS. Will only
+   * return a non-empty value after the data writer has been closed.
+   */
+  Optional<String> getUfsContentHash();
+
   /**
    * Writes a chunk. This method takes the ownership of this chunk even if it fails to write
    * the chunk.

core/client/fs/src/main/java/alluxio/client/block/stream/GrpcBlockingStream.java

@@ -28,6 +28,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.Optional;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -250,14 +251,20 @@ public void cancel() {
    * Wait for server to complete the inbound stream.
    *
    * @param timeoutMs maximum time to wait for server response
+   * @return the last response of the stream
    */
-  public void waitForComplete(long timeoutMs) throws IOException {
+  public Optional<ResT> waitForComplete(long timeoutMs) throws IOException {
     if (mCompleted || mCanceled) {
-      return;
+      return Optional.empty();
     }
-    while (receive(timeoutMs) != null) {
+    ResT prevResponse;
+    ResT response = null;
+    do {
       // wait until inbound stream is closed from server.
-    }
+      prevResponse = response;
+      response = receive(timeoutMs);
+    } while (response != null);
+    return Optional.ofNullable(prevResponse);
   }
 
   /**

core/client/fs/src/main/java/alluxio/client/block/stream/GrpcDataMessageBlockingStream.java

@@ -20,6 +20,7 @@
 import io.grpc.stub.StreamObserver;
 
 import java.io.IOException;
+import java.util.Optional;
 import java.util.function.Function;
 import javax.annotation.concurrent.NotThreadSafe;
 
@@ -105,17 +106,21 @@ public void sendDataMessage(DataMessage<ReqT, DataBuffer> message, long timeoutM
   }
 
   @Override
-  public void waitForComplete(long timeoutMs) throws IOException {
+  public Optional<ResT> waitForComplete(long timeoutMs) throws IOException {
     if (mResponseMarshaller == null) {
-      super.waitForComplete(timeoutMs);
-      return;
+      return super.waitForComplete(timeoutMs);
     }
+    // loop until the last response is received, whose result will be returned
     DataMessage<ResT, DataBuffer> message;
+    DataMessage<ResT, DataBuffer> prevMessage = null;
     while (!isCanceled() && (message = receiveDataMessage(timeoutMs)) != null) {
-      if (message.getBuffer() != null) {
-        message.getBuffer().release();
+      if (prevMessage != null && prevMessage.getBuffer() != null) {
+        prevMessage.getBuffer().release();
       }
+      prevMessage = message;
     }
-    super.waitForComplete(timeoutMs);
+    // note that the combineData call is responsible for releasing the buffer of prevMessage
+    ResT result = mResponseMarshaller.combineData(prevMessage);
+    return Optional.ofNullable(super.waitForComplete(timeoutMs).orElse(result));
   }
 }

core/client/fs/src/main/java/alluxio/client/block/stream/GrpcDataWriter.java

@@ -38,6 +38,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.Optional;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -75,6 +76,9 @@ public final class GrpcDataWriter implements DataWriter {
   private final GrpcBlockingStream<WriteRequest, WriteResponse> mStream;
   private final WriteRequestMarshaller mMarshaller;
 
+  /** The content hash resulting from the write operation if one is available. */
+  private String mContentHash = null;
+
   /**
    * The next pos to queue to the buffer.
    */
@@ -186,6 +190,11 @@ public long pos() {
     return mPosToQueue;
   }
 
+  @Override
+  public Optional<String> getUfsContentHash() {
+    return Optional.ofNullable(mContentHash);
+  }
+
   @Override
   public void writeChunk(final ByteBuf buf) throws IOException {
     mPosToQueue += buf.readableBytes();
@@ -247,6 +256,9 @@ public void flush() throws IOException {
             "Flush request %s to worker %s is not acked before complete.", writeRequest, mAddress));
       }
       posWritten = response.getOffset();
+      if (response.hasContentHash()) {
+        mContentHash = response.getContentHash();
+      }
     } while (mPosToQueue != posWritten);
   }
 
@@ -257,7 +269,9 @@ public void close() throws IOException {
         return;
       }
       mStream.close();
-      mStream.waitForComplete(mWriterCloseTimeoutMs);
+      mStream.waitForComplete(mWriterCloseTimeoutMs)
+          .ifPresent(writeResponse -> mContentHash = writeResponse.hasContentHash()
+              ? writeResponse.getContentHash() : null);
     } finally {
       mClient.close();
     }

core/client/fs/src/main/java/alluxio/client/block/stream/LocalFileDataWriter.java

@@ -33,6 +33,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.Optional;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -122,6 +123,11 @@ public int chunkSize() {
     return (int) mChunkSize;
   }
 
+  @Override
+  public Optional<String> getUfsContentHash() {
+    return Optional.empty();
+  }
+
   @Override
   public void writeChunk(final ByteBuf buf) throws IOException {
     try {

core/client/fs/src/main/java/alluxio/client/block/stream/UfsFallbackLocalFileDataWriter.java

@@ -22,6 +22,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.Optional;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -82,6 +83,11 @@ public static UfsFallbackLocalFileDataWriter create(FileSystemContext context,
     mIsWritingToLocal = mLocalFileDataWriter != null;
   }
 
+  @Override
+  public Optional<String> getUfsContentHash() {
+    return mGrpcDataWriter.getUfsContentHash();
+  }
+
   @Override
   public void writeChunk(ByteBuf chunk) throws IOException {
     if (mIsWritingToLocal) {

core/client/fs/src/main/java/alluxio/client/block/stream/UnderFileSystemFileOutStream.java

@@ -30,6 +30,7 @@
 @NotThreadSafe
 public class UnderFileSystemFileOutStream extends BlockOutStream {
   private static final int ID_UNUSED = -1;
+  private final DataWriter mDataWriter;
 
   /**
    * Creates an instance of {@link UnderFileSystemFileOutStream} that writes to a UFS file.
@@ -52,6 +53,14 @@ public static UnderFileSystemFileOutStream create(FileSystemContext context,
    */
   protected UnderFileSystemFileOutStream(DataWriter dataWriter, WorkerNetAddress address) {
     super(dataWriter, Long.MAX_VALUE, address);
+    mDataWriter = dataWriter;
+  }
+
+  /**
+   * @return the data writer for the stream
+   */
+  public DataWriter getDataWriter() {
+    return mDataWriter;
   }
 
   @Override

core/client/fs/src/main/java/alluxio/client/file/AlluxioFileOutStream.java

@@ -167,6 +167,8 @@ public void close() throws IOException {
         } else {
           mUnderStorageOutputStream.close();
           optionsBuilder.setUfsLength(mBytesWritten);
+          mUnderStorageOutputStream.getDataWriter().getUfsContentHash().ifPresent(
+              optionsBuilder::setContentHash);
         }
       }
 

core/client/fs/src/test/java/alluxio/client/block/stream/TestDataWriter.java

@@ -15,6 +15,7 @@
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Optional;
 
 /**
  * A {@link DataWriter} which writes data to a bytebuffer.
@@ -26,6 +27,11 @@ public TestDataWriter(ByteBuffer buffer) {
     mBuffer = buffer;
   }
 
+  @Override
+  public Optional<String> getUfsContentHash() {
+    return Optional.empty();
+  }
+
   @Override
   public void writeChunk(ByteBuf chunk) throws IOException {
     try {

core/common/src/main/java/alluxio/Constants.java

@@ -228,5 +228,8 @@ public final class Constants {
   public static final String MEDIUM_HDD = "HDD";
   public static final String MEDIUM_SSD = "SSD";
 
+  /* xAttr keys */
+  public static final String ETAG_XATTR_KEY = "s3_etag";
+
   private Constants() {} // prevent instantiation
 }

core/common/src/main/java/alluxio/concurrent/ManagedBlockingUfsForwarder.java

@@ -37,6 +37,7 @@
 import java.io.OutputStream;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 /**
  * Forwarder for {@link UnderFileSystem} objects that works through with ForkJoinPool's
@@ -288,6 +289,11 @@ public Fingerprint getParsedFingerprint(String path) {
     return mUfs.getParsedFingerprint(path);
   }
 
+  @Override
+  public Fingerprint getParsedFingerprint(String path, @Nullable String contentHash) {
+    return mUfs.getParsedFingerprint(path, contentHash);
+  }
+
   @Override
   public UfsMode getOperationMode(Map<String, UfsMode> physicalUfsState) {
     return mUfs.getOperationMode(physicalUfsState);