Ensure Snappy compression matches snappy-java framing when uploading in parts and reusing upload buffers. Consolidate this logic into single new class BufferIterator.

mattl-netflix · mattl-netflix · commit 2d0cddff02df · 2026-05-21T20:48:02.000-07:00
diff --git a/priam/src/main/java/com/netflix/priam/aws/BufferIterator.java b/priam/src/main/java/com/netflix/priam/aws/BufferIterator.java
@@ -0,0 +1,99 @@
+package com.netflix.priam.aws;
+
+import com.netflix.priam.backup.AbstractBackupPath;
+import com.netflix.priam.compress.CompressionType;
+import com.netflix.priam.utils.ThreadLocalByteBuffer;
+import org.xerial.snappy.Snappy;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+
+public class BufferIterator implements Closeable {
+    private static final int SNAPPY_CHUNK_SIZE = 4;
+    static final byte[] SNAPPY_HEADER = new byte[] {
+            (byte) 0x82, 'S', 'N', 'A', 'P', 'P', 'Y', (byte) 0x00,
+            (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x01,
+            (byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x01
+    };
+    private final ByteBuffer readBuffer;
+    private final ByteBuffer compressBuffer;
+    private final FileChannel channel;
+    private final CompressionType compressionType;
+    private final long fileSize;
+    private final long chunkSize;
+    private boolean headerApplied;
+    private long position;
+
+    BufferIterator(ThreadLocal<ByteBuffer> inputBuffer,
+                   ThreadLocal<ByteBuffer> compressBuffer,
+                   AbstractBackupPath path,
+                   long minimumChunkSize) throws IOException {
+        this(inputBuffer, compressBuffer, path, minimumChunkSize,
+                FileChannel.open(Paths.get(path.getBackupFile().getAbsolutePath()), StandardOpenOption.READ));
+    }
+
+    @com.google.common.annotations.VisibleForTesting
+    BufferIterator(ThreadLocal<ByteBuffer> inputBuffer,
+                   ThreadLocal<ByteBuffer> compressBuffer,
+                   AbstractBackupPath path,
+                   long minimumChunkSize,
+                   FileChannel fileChannel) throws IOException {
+        Path localPath = Paths.get(path.getBackupFile().getAbsolutePath());
+        fileSize = localPath.toFile().length();
+        chunkSize = S3FileSystem.getChunkSize(localPath, minimumChunkSize);
+        compressionType = path.getCompression();
+        // Direct buffer use is obviated when compressing.
+        boolean useDirectBuffer = compressionType != CompressionType.SNAPPY;
+        readBuffer = ThreadLocalByteBuffer.get(inputBuffer, (int) Math.min(chunkSize, fileSize), useDirectBuffer);
+        if (compressionType == CompressionType.SNAPPY) {
+            int maxCompressedLength = Snappy.maxCompressedLength((int) chunkSize) + SNAPPY_HEADER.length + SNAPPY_CHUNK_SIZE;
+            this.compressBuffer = ThreadLocalByteBuffer.get(compressBuffer, maxCompressedLength, false);
+        } else {
+            this.compressBuffer = null;
+        }
+        this.channel = fileChannel;
+    }
+
+    public boolean hasNext() {
+        return position < fileSize;
+    }
+
+    public ByteBuffer next() throws IOException {
+        long bytesToRead = Math.min(chunkSize, fileSize - position);
+        readBuffer.clear();
+        readBuffer.limit((int) bytesToRead);
+        int bytesRead = 0;
+        while (bytesRead < bytesToRead) {
+            int read = channel.read(readBuffer, position + bytesRead);
+            if (read == -1) break;
+            bytesRead += read;
+        }
+        readBuffer.flip();
+        position += bytesRead;
+        if (compressionType != CompressionType.SNAPPY) {
+            return readBuffer;
+        }
+        compressBuffer.clear();
+        int offset = headerApplied ? SNAPPY_CHUNK_SIZE : SNAPPY_CHUNK_SIZE + SNAPPY_HEADER.length;
+        int compressedSize = Snappy.compress(readBuffer.array(), 0, bytesRead, compressBuffer.array(), offset);
+        compressBuffer.position(0);
+        if (!headerApplied) {
+            compressBuffer.put(SNAPPY_HEADER);
+            headerApplied = true;
+        }
+        compressBuffer.putInt(compressedSize);
+        compressBuffer.limit(compressedSize + offset);
+        compressBuffer.position(0);
+        return compressBuffer;
+    }
+
+    @Override
+    public void close() throws IOException {
+        channel.close();
+    }
+}
diff --git a/priam/src/main/java/com/netflix/priam/aws/S3FileSystem.java b/priam/src/main/java/com/netflix/priam/aws/S3FileSystem.java
@@ -33,11 +33,9 @@
 import com.netflix.priam.utils.BoundedExponentialRetryCallable;
 import com.netflix.priam.utils.ByteBufferInputStream;
 import com.netflix.priam.utils.SystemUtils;
-import com.netflix.priam.utils.ThreadLocalByteBuffer;
 import org.apache.commons.io.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.xerial.snappy.Snappy;
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.sync.RequestBody;
 import software.amazon.awssdk.regions.Region;
@@ -50,10 +48,8 @@
 import javax.inject.Singleton;
 import java.io.*;
 import java.nio.ByteBuffer;
-import java.nio.channels.FileChannel;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.nio.file.StandardOpenOption;
 import java.time.Instant;
 import java.util.*;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -62,11 +58,12 @@
 @Singleton
 public class S3FileSystem extends S3FileSystemBase {
     private static final Logger logger = LoggerFactory.getLogger(S3FileSystem.class);
+    private static final int MAX_CHUNKS = 9995; // 10K is AWS limit, minus a small buffer
     private static final long MAX_BUFFER_SIZE = 5L * 1024L * 1024L;
+
     private final DynamicRateLimiter dynamicRateLimiter;
     private final ThreadLocal<ByteBuffer> inputBufferThreadLocal = new ThreadLocal<>();
     private final ThreadLocal<ByteBuffer> compressBufferThreadLocal = new ThreadLocal<>();
-    private final ThreadLocal<ByteBuffer> chunkBufferThreadLocal = new ThreadLocal<>();
 
     @Inject
     public S3FileSystem(
@@ -87,6 +84,10 @@ public S3FileSystem(
         this.dynamicRateLimiter = dynamicRateLimiter;
     }
 
+    static long getChunkSize(Path path, long minimumChunkSize) {
+        return Math.max(path.toFile().length() / MAX_CHUNKS, minimumChunkSize);
+    }
+
     @Override
     protected void downloadFileImpl(AbstractBackupPath path, String suffix)
             throws BackupRestoreException {
@@ -149,12 +150,11 @@ private long uploadMultipartWithBuffers(AbstractBackupPath path, Instant target)
             throws BackupRestoreException {
         Path localPath = Paths.get(path.getBackupFile().getAbsolutePath());
         String remotePath = path.getRemotePath();
-        long chunkSize = getChunkSize(localPath);
+        long chunkSize = getChunkSize(localPath, config.getBackupChunkSize());
         String prefix = config.getBackupPrefix();
         if (logger.isDebugEnabled())
             logger.debug("Uploading to {}/{} with chunk size {} (using ByteBuffer implementation)", prefix, remotePath, chunkSize);
         File localFile = localPath.toFile();
-        long fileSize = localFile.length();
 
         CreateMultipartUploadRequest.Builder initRequestBuilder = CreateMultipartUploadRequest.builder()
                 .bucket(prefix)
@@ -169,66 +169,25 @@ private long uploadMultipartWithBuffers(AbstractBackupPath path, Instant target)
         String uploadId = s3Client.createMultipartUpload(initRequest).uploadId();
         List<CompletedPart> completedParts = new ArrayList<>();
 
-        try (FileChannel channel = FileChannel.open(localPath, StandardOpenOption.READ)) {
+        try (BufferIterator bufferIterator = new BufferIterator(inputBufferThreadLocal, compressBufferThreadLocal, path, config.getBackupChunkSize())) {
             int partNum = 0;
             long compressedFileSize = 0;
-            long position = 0;
-            
-            // Allocate thread-local buffers for reading and compression
-            ByteBuffer readBuffer = ThreadLocalByteBuffer.get(inputBufferThreadLocal, (int) Math.min(chunkSize, fileSize));
-            ByteBuffer compressBuffer = null;
-            if (path.getCompression() == CompressionType.SNAPPY) {
-                int maxCompressedLength = Snappy.maxCompressedLength((int) chunkSize);
-                compressBuffer = ThreadLocalByteBuffer.get(compressBufferThreadLocal, maxCompressedLength);
-            }
-            
-            while (position < fileSize) {
-                // Read chunk from file
-                long bytesToRead = Math.min(chunkSize, fileSize - position);
-                readBuffer.clear();
-                readBuffer.limit((int) bytesToRead);
-                
-                int bytesRead = 0;
-                while (bytesRead < bytesToRead) {
-                    int read = channel.read(readBuffer, position + bytesRead);
-                    if (read == -1) break;
-                    bytesRead += read;
-                }
-                readBuffer.flip();
-                position += bytesRead;
-                
-                // Compress if needed
-                ByteBuffer uploadBuffer;
-                int uploadSize;
-                if (path.getCompression() == CompressionType.SNAPPY) {
-                    compressBuffer.clear();
-                    uploadSize = Snappy.compress(readBuffer, compressBuffer);
-                    compressBuffer.limit(uploadSize);
-                    compressBuffer.position(0);
-                    uploadBuffer = compressBuffer;
-                } else {
-                    uploadBuffer = readBuffer;
-                    uploadSize = bytesRead;
-                }
-                
-                // Apply rate limiting
+            while (bufferIterator.hasNext()) {
+                ByteBuffer uploadBuffer = bufferIterator.next();
+                int uploadSize = uploadBuffer.limit();
                 rateLimiter.acquire(uploadSize);
                 dynamicRateLimiter.acquire(path, target, uploadSize);
-                
-                // Upload part directly from ByteBuffer
                 partNum++;
                 UploadPartRequest.Builder req = UploadPartRequest.builder()
                         .bucket(prefix)
                         .key(remotePath)
                         .uploadId(uploadId)
                         .partNumber(partNum)
                         .contentLength((long) uploadSize);
-
                 byte[] md5 = SystemUtils.md5(uploadBuffer);
                 if (config.addMD5ToBackupUploads()) {
                     req.contentMD5(SystemUtils.toBase64(md5));
                 }
-
                 UploadPartResponse uploadResult = new BoundedExponentialRetryCallable<UploadPartResponse>(200, 10000, 5) {
                     @Override
                     public UploadPartResponse retriableCall() {
@@ -246,15 +205,12 @@ public UploadPartResponse retriableCall() {
                 validateUpload(eTag, md5, partNum);
                 completedParts.add(CompletedPart.builder().partNumber(partNum).eTag(eTag).build());
                 compressedFileSize += uploadSize;
-                
+
                 if (logger.isDebugEnabled()) {
                     logger.debug("Uploaded part {} of size {}", partNum, uploadSize);
                 }
             }
-            
             logger.info("{} done. part count: {}", localFile, partNum);
-            
-            // Complete the multipart upload
             CompleteMultipartUploadResponse multipartUploadResponse =
                     s3Client.completeMultipartUpload(
                             CompleteMultipartUploadRequest.builder()
@@ -298,7 +254,7 @@ private long uploadMultipartLegacy(AbstractBackupPath path, Instant target)
             throws BackupRestoreException {
         Path localPath = Paths.get(path.getBackupFile().getAbsolutePath());
         String remotePath = path.getRemotePath();
-        long chunkSize = getChunkSize(localPath);
+        long chunkSize = getChunkSize(localPath, config.getBackupChunkSize());
         String prefix = config.getBackupPrefix();
         if (logger.isDebugEnabled())
             logger.debug("Uploading to {}/{} with chunk size {} (using legacy implementation)", prefix, remotePath, chunkSize);
@@ -450,29 +406,14 @@ private PutObjectRequest generatePutFromByteArray(AbstractBackupPath path, byte[
 
     @VisibleForTesting
     public ByteBuffer getFileByteBuffer(AbstractBackupPath path) throws BackupRestoreException {
-        File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
-        long fileSize = localFile.length();
-        
-        try {
-            // Get input buffer and read file data directly into it
-            ByteBuffer inputBuffer = ThreadLocalByteBuffer.get(inputBufferThreadLocal, (int) fileSize);
-            try (FileChannel channel = FileChannel.open(localFile.toPath(), StandardOpenOption.READ)) {
-                channel.read(inputBuffer);
-                inputBuffer.flip();
-            }
-            
-            // Compress if needed, otherwise return input buffer
-            if (path.getCompression() == CompressionType.SNAPPY) {
-                int maxCompressedLength = Snappy.maxCompressedLength((int) fileSize);
-                ByteBuffer compressBuffer = ThreadLocalByteBuffer.get(compressBufferThreadLocal, maxCompressedLength);
-                int compressedSize = Snappy.compress(inputBuffer, compressBuffer);
-                compressBuffer.limit(compressedSize);
-                compressBuffer.position(0);
-                return compressBuffer;
-            } else {
-                return inputBuffer;
-            }
+        try (BufferIterator bufferIterator = new BufferIterator(
+                inputBufferThreadLocal,
+                compressBufferThreadLocal,
+                path,
+                config.getBackupChunkSize())) {
+            return bufferIterator.next();
         } catch (Exception e) {
+            File localFile = Paths.get(path.getBackupFile().getAbsolutePath()).toFile();
             throw new BackupRestoreException("Error reading file: " + localFile.getName(), e);
         }
     }
diff --git a/priam/src/main/java/com/netflix/priam/aws/S3FileSystemBase.java b/priam/src/main/java/com/netflix/priam/aws/S3FileSystemBase.java
@@ -37,7 +37,6 @@
 import software.amazon.awssdk.services.s3.model.*;
 
 public abstract class S3FileSystemBase extends AbstractFileSystem {
-    private static final int MAX_CHUNKS = 9995; // 10K is AWS limit, minus a small buffer
     private static final Logger logger = LoggerFactory.getLogger(S3FileSystemBase.class);
     S3Client s3Client;
     final IConfiguration config;
@@ -189,8 +188,4 @@ public void deleteFiles(List<Path> remotePaths) throws BackupRestoreException {
             throw new BackupRestoreException(e + " while trying to delete the objects");
         }
     }
-
-    final long getChunkSize(Path path) {
-        return Math.max(path.toFile().length() / MAX_CHUNKS, config.getBackupChunkSize());
-    }
 }
diff --git a/priam/src/main/java/com/netflix/priam/utils/ThreadLocalByteBuffer.java b/priam/src/main/java/com/netflix/priam/utils/ThreadLocalByteBuffer.java
@@ -21,13 +21,17 @@
 public class ThreadLocalByteBuffer {
     
     public static ByteBuffer get(ThreadLocal<ByteBuffer> threadLocalBuffer, int bytes) {
+        return get(threadLocalBuffer, bytes, true);
+    }
+
+    public static ByteBuffer get(ThreadLocal<ByteBuffer> threadLocalBuffer, int bytes, boolean direct) {
         ByteBuffer buffer = threadLocalBuffer.get();
-        
-        if (buffer == null || buffer.capacity() < bytes) {
-            buffer = ByteBuffer.allocateDirect(bytes);
+
+        if (buffer == null || buffer.capacity() < bytes || buffer.isDirect() != direct) {
+            buffer = direct ? ByteBuffer.allocateDirect(bytes) : ByteBuffer.allocate(bytes);
             threadLocalBuffer.set(buffer);
         }
-        
+
         buffer.clear();
         buffer.limit(bytes);
         return buffer;
diff --git a/priam/src/test/java/com/netflix/priam/aws/TestBufferIterator.java b/priam/src/test/java/com/netflix/priam/aws/TestBufferIterator.java

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,6 @@`
`37`	`37`	`import software.amazon.awssdk.services.s3.model.*;`
`38`	`38`
`39`	`39`	`public abstract class S3FileSystemBase extends AbstractFileSystem {`
`40`		`- private static final int MAX_CHUNKS = 9995; // 10K is AWS limit, minus a small buffer`
`41`	`40`	`private static final Logger logger = LoggerFactory.getLogger(S3FileSystemBase.class);`
`42`	`41`	`S3Client s3Client;`
`43`	`42`	`final IConfiguration config;`
`@@ -189,8 +188,4 @@ public void deleteFiles(List<Path> remotePaths) throws BackupRestoreException {`
`189`	`188`	`throw new BackupRestoreException(e + " while trying to delete the objects");`
`190`	`189`	`}`
`191`	`190`	`}`
`192`		`-`
`193`		`- final long getChunkSize(Path path) {`
`194`		`- return Math.max(path.toFile().length() / MAX_CHUNKS, config.getBackupChunkSize());`
`195`		`- }`
`196`	`191`	`}`