add feature flag

clohfink · clohfink · commit de34805c72e7 · 2025-08-18T08:18:51.000-05:00
diff --git a/priam/src/main/java/com/netflix/priam/aws/DataPart.java b/priam/src/main/java/com/netflix/priam/aws/DataPart.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright 2013 Netflix, Inc.
+ *
+ * Licensed 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 com.netflix.priam.aws;
+
+import com.netflix.priam.utils.SystemUtils;
+
+/** Class for holding part data of a backup file, which will be used for multi-part uploading */
+public class DataPart {
+    private final String bucketName;
+    private final String uploadID;
+    private final String s3key;
+    private int partNo;
+    private byte[] partData;
+    private byte[] md5;
+
+    public DataPart(String bucket, String s3key, String mUploadId) {
+        this.bucketName = bucket;
+        this.uploadID = mUploadId;
+        this.s3key = s3key;
+    }
+
+    public DataPart(int partNumber, byte[] data, String bucket, String s3key, String mUploadId) {
+        this(bucket, s3key, mUploadId);
+        this.partNo = partNumber;
+        this.partData = data;
+        this.md5 = SystemUtils.md5(data);
+    }
+
+    public String getBucketName() {
+        return bucketName;
+    }
+
+    public String getUploadID() {
+        return uploadID;
+    }
+
+    public String getS3key() {
+        return s3key;
+    }
+
+    public int getPartNo() {
+        return partNo;
+    }
+
+    public byte[] getPartData() {
+        return partData;
+    }
+
+    public byte[] getMd5() {
+        return md5;
+    }
+}
diff --git a/priam/src/main/java/com/netflix/priam/aws/S3FileSystem.java b/priam/src/main/java/com/netflix/priam/aws/S3FileSystem.java
@@ -131,12 +131,21 @@ private ObjectMetadata getObjectMetadata(File file) {
 
     private long uploadMultipart(AbstractBackupPath path, Instant target)
             throws BackupRestoreException {
+        if (config.useReusableBufferForMultipartUploads()) {
+            return uploadMultipartWithBuffers(path, target);
+        } else {
+            return uploadMultipartLegacy(path, target);
+        }
+    }
+
+    private long uploadMultipartWithBuffers(AbstractBackupPath path, Instant target)
+            throws BackupRestoreException {
         Path localPath = Paths.get(path.getBackupFile().getAbsolutePath());
         String remotePath = path.getRemotePath();
         long chunkSize = getChunkSize(localPath);
         String prefix = config.getBackupPrefix();
         if (logger.isDebugEnabled())
-            logger.debug("Uploading to {}/{} with chunk size {}", prefix, remotePath, chunkSize);
+            logger.debug("Uploading to {}/{} with chunk size {} (using ByteBuffer implementation)", prefix, remotePath, chunkSize);
         File localFile = localPath.toFile();
         long fileSize = localFile.length();
         
@@ -244,13 +253,74 @@ public UploadPartResult retriableCall() {
         }
     }
 
+    private long uploadMultipartLegacy(AbstractBackupPath path, Instant target)
+            throws BackupRestoreException {
+        Path localPath = Paths.get(path.getBackupFile().getAbsolutePath());
+        String remotePath = path.getRemotePath();
+        long chunkSize = getChunkSize(localPath);
+        String prefix = config.getBackupPrefix();
+        if (logger.isDebugEnabled())
+            logger.debug("Uploading to {}/{} with chunk size {} (using legacy implementation)", prefix, remotePath, chunkSize);
+        File localFile = localPath.toFile();
+        InitiateMultipartUploadRequest initRequest =
+                new InitiateMultipartUploadRequest(prefix, remotePath)
+                        .withObjectMetadata(getObjectMetadata(localFile));
+        String uploadId = s3Client.initiateMultipartUpload(initRequest).getUploadId();
+        DataPart part = new DataPart(prefix, remotePath, uploadId);
+        List<PartETag> partETags = Collections.synchronizedList(new ArrayList<>());
+
+        try (InputStream in = new FileInputStream(localFile)) {
+            Iterator<byte[]> chunks = new ChunkedStream(in, chunkSize, path.getCompression());
+            int partNum = 0;
+            AtomicInteger partsPut = new AtomicInteger(0);
+            long compressedFileSize = 0;
+
+            while (chunks.hasNext()) {
+                byte[] chunk = chunks.next();
+                rateLimiter.acquire(chunk.length);
+                dynamicRateLimiter.acquire(path, target, chunk.length);
+                DataPart dp = new DataPart(++partNum, chunk, prefix, remotePath, uploadId);
+                S3PartUploader partUploader = new S3PartUploader(s3Client, dp, partETags, partsPut);
+                compressedFileSize += chunk.length;
+                executor.submit(partUploader);
+            }
+
+            executor.sleepTillEmpty();
+            logger.info("{} done. part count: {} expected: {}", localFile, partsPut.get(), partNum);
+            Preconditions.checkState(partNum == partETags.size(), "part count mismatch");
+            CompleteMultipartUploadResult resultS3MultiPartUploadComplete =
+                    new S3PartUploader(s3Client, part, partETags).completeUpload();
+            checkSuccessfulUpload(resultS3MultiPartUploadComplete, localPath);
+
+            if (logger.isDebugEnabled()) {
+                final S3ResponseMetadata info = s3Client.getCachedResponseMetadata(initRequest);
+                logger.debug("Request Id: {}, Host Id: {}", info.getRequestId(), info.getHostId());
+            }
+
+            return compressedFileSize;
+        } catch (Exception e) {
+            new S3PartUploader(s3Client, part, partETags).abortUpload();
+            throw new BackupRestoreException("Error uploading file: " + localPath.toString(), e);
+        }
+    }
+
     protected long uploadFileImpl(AbstractBackupPath path, Instant target)
             throws BackupRestoreException {
         File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
 
         if (localFile.length() >= config.getBackupChunkSize())
             return uploadMultipart(path, target);
 
+        if (config.useReusableBufferForMultipartUploads()) {
+            return uploadFileWithByteBuffer(path, target);
+        } else {
+            return uploadFileWithByteArray(path, target);
+        }
+    }
+
+    private long uploadFileWithByteBuffer(AbstractBackupPath path, Instant target)
+            throws BackupRestoreException {
+        File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
         ByteBuffer buffer = getFileByteBuffer(path);
         // C* snapshots may have empty files. That is probably unintentional.
         if (buffer.remaining() > 0) {
@@ -261,7 +331,7 @@ protected long uploadFileImpl(AbstractBackupPath path, Instant target)
             new BoundedExponentialRetryCallable<PutObjectResult>(1000, 10000, 5) {
                 @Override
                 public PutObjectResult retriableCall() {
-                    return s3Client.putObject(generatePut(path, buffer));
+                    return s3Client.putObject(generatePutFromBuffer(path, buffer));
                 }
             }.call();
         } catch (Exception e) {
@@ -270,7 +340,29 @@ public PutObjectResult retriableCall() {
         return buffer.remaining();
     }
 
-    private PutObjectRequest generatePut(AbstractBackupPath path, ByteBuffer buffer) {
+    private long uploadFileWithByteArray(AbstractBackupPath path, Instant target)
+            throws BackupRestoreException {
+        File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
+        byte[] chunk = getFileContents(path);
+        // C* snapshots may have empty files. That is probably unintentional.
+        if (chunk.length > 0) {
+            rateLimiter.acquire(chunk.length);
+            dynamicRateLimiter.acquire(path, target, chunk.length);
+        }
+        try {
+            new BoundedExponentialRetryCallable<PutObjectResult>(1000, 10000, 5) {
+                @Override
+                public PutObjectResult retriableCall() {
+                    return s3Client.putObject(generatePutFromByteArray(path, chunk));
+                }
+            }.call();
+        } catch (Exception e) {
+            throw new BackupRestoreException("Error uploading file: " + localFile.getName(), e);
+        }
+        return chunk.length;
+    }
+
+    private PutObjectRequest generatePutFromBuffer(AbstractBackupPath path, ByteBuffer buffer) {
         File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
         ObjectMetadata metadata = getObjectMetadata(localFile);
         metadata.setContentLength(buffer.remaining());
@@ -286,6 +378,22 @@ private PutObjectRequest generatePut(AbstractBackupPath path, ByteBuffer buffer)
         return put;
     }
 
+    private PutObjectRequest generatePutFromByteArray(AbstractBackupPath path, byte[] chunk) {
+        File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
+        ObjectMetadata metadata = getObjectMetadata(localFile);
+        metadata.setContentLength(chunk.length);
+        PutObjectRequest put =
+                new PutObjectRequest(
+                        config.getBackupPrefix(),
+                        path.getRemotePath(),
+                        new ByteArrayInputStream(chunk),
+                        metadata);
+        if (config.addMD5ToBackupUploads()) {
+            put.getMetadata().setContentMD5(SystemUtils.toBase64(SystemUtils.md5(chunk)));
+        }
+        return put;
+    }
+
     @VisibleForTesting
     public ByteBuffer getFileByteBuffer(AbstractBackupPath path) throws BackupRestoreException {
         File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
@@ -314,4 +422,19 @@ public ByteBuffer getFileByteBuffer(AbstractBackupPath path) throws BackupRestor
             throw new BackupRestoreException("Error reading file: " + localFile.getName(), e);
         }
     }
+
+    private byte[] getFileContents(AbstractBackupPath path) throws BackupRestoreException {
+        File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
+        try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+                InputStream in = new BufferedInputStream(new FileInputStream(localFile))) {
+            Iterator<byte[]> chunks =
+                    new ChunkedStream(in, config.getBackupChunkSize(), path.getCompression());
+            while (chunks.hasNext()) {
+                byteArrayOutputStream.write(chunks.next());
+            }
+            return byteArrayOutputStream.toByteArray();
+        } catch (Exception e) {
+            throw new BackupRestoreException("Error reading file: " + localFile.getName(), e);
+        }
+    }
 }
diff --git a/priam/src/main/java/com/netflix/priam/aws/S3PartUploader.java b/priam/src/main/java/com/netflix/priam/aws/S3PartUploader.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2013 Netflix, Inc.
+ *
+ * Licensed 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 com.netflix.priam.aws;
+
+import com.amazonaws.AmazonClientException;
+import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.s3.model.*;
+import com.netflix.priam.backup.BackupRestoreException;
+import com.netflix.priam.utils.BoundedExponentialRetryCallable;
+import com.netflix.priam.utils.SystemUtils;
+import java.io.ByteArrayInputStream;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class S3PartUploader extends BoundedExponentialRetryCallable<Void> {
+    private final AmazonS3 client;
+    private final DataPart dataPart;
+    private final List<PartETag> partETags;
+    private AtomicInteger partsUploaded = null; // num of data parts successfully uploaded
+
+    private static final Logger logger = LoggerFactory.getLogger(S3PartUploader.class);
+    private static final int MAX_RETRIES = 5;
+    private static final int DEFAULT_MIN_SLEEP_MS = 200;
+
+    public S3PartUploader(AmazonS3 client, DataPart dp, List<PartETag> partETags) {
+        super(DEFAULT_MIN_SLEEP_MS, BoundedExponentialRetryCallable.MAX_SLEEP, MAX_RETRIES);
+        this.client = client;
+        this.dataPart = dp;
+        this.partETags = partETags;
+    }
+
+    public S3PartUploader(
+            AmazonS3 client, DataPart dp, List<PartETag> partETags, AtomicInteger partsUploaded) {
+        super(DEFAULT_MIN_SLEEP_MS, BoundedExponentialRetryCallable.MAX_SLEEP, MAX_RETRIES);
+        this.client = client;
+        this.dataPart = dp;
+        this.partETags = partETags;
+        this.partsUploaded = partsUploaded;
+    }
+
+    private Void uploadPart() throws AmazonClientException, BackupRestoreException {
+        UploadPartRequest req = new UploadPartRequest();
+        req.setBucketName(dataPart.getBucketName());
+        req.setKey(dataPart.getS3key());
+        req.setUploadId(dataPart.getUploadID());
+        req.setPartNumber(dataPart.getPartNo());
+        req.setPartSize(dataPart.getPartData().length);
+        req.setMd5Digest(SystemUtils.toBase64(dataPart.getMd5()));
+        req.setInputStream(new ByteArrayInputStream(dataPart.getPartData()));
+        UploadPartResult res = client.uploadPart(req);
+        PartETag partETag = res.getPartETag();
+        if (!partETag.getETag().equals(SystemUtils.toHex(dataPart.getMd5())))
+            throw new BackupRestoreException(
+                    "Unable to match MD5 for part " + dataPart.getPartNo());
+        partETags.add(partETag);
+        if (this.partsUploaded != null) this.partsUploaded.incrementAndGet();
+        return null;
+    }
+
+    public CompleteMultipartUploadResult completeUpload() throws BackupRestoreException {
+        CompleteMultipartUploadRequest compRequest =
+                new CompleteMultipartUploadRequest(
+                        dataPart.getBucketName(),
+                        dataPart.getS3key(),
+                        dataPart.getUploadID(),
+                        partETags);
+        return client.completeMultipartUpload(compRequest);
+    }
+
+    // Abort
+    public void abortUpload() {
+        AbortMultipartUploadRequest abortRequest =
+                new AbortMultipartUploadRequest(
+                        dataPart.getBucketName(), dataPart.getS3key(), dataPart.getUploadID());
+        client.abortMultipartUpload(abortRequest);
+    }
+
+    @Override
+    public Void retriableCall() throws AmazonClientException, BackupRestoreException {
+        logger.debug(
+                "Picked up part {} size {}", dataPart.getPartNo(), dataPart.getPartData().length);
+        return uploadPart();
+    }
+}
diff --git a/priam/src/main/java/com/netflix/priam/config/IConfiguration.java b/priam/src/main/java/com/netflix/priam/config/IConfiguration.java
@@ -1138,4 +1138,15 @@ default int getBlockForPeersTimeoutInSecs() {
     default boolean skipMetaFileValidationOnRestore() {
         return false;
     }
+
+    /**
+     * Use reusable ByteBuffer implementation for S3 multipart uploads.
+     * When true, uses thread-local ByteBuffers to reduce memory allocation overhead.
+     * When false, uses the legacy byte array based implementation.
+     * 
+     * @return true to use reusable ByteBuffer implementation, false for legacy implementation
+     */
+    default boolean useReusableBufferForMultipartUploads() {
+        return false;
+    }
 }
