Add tests for large payload externalization and fix review issues

bachuv · bachuv · commit 5eb3ac7e92e1 · 2026-04-10T14:44:31.000-07:00
diff --git a/azure-blob-payloads/build.gradle b/azure-blob-payloads/build.gradle
@@ -72,7 +72,7 @@ test {
     }
 }
 
-// Integration tests require DTS emulator on localhost:4001 and Azurite on localhost:10000.
+// Integration tests require DTS emulator (default localhost:8080) and Azurite on localhost:10000.
 task integrationTest(type: Test) {
     useJUnitPlatform {
         includeTags 'integration'
diff --git a/azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStore.java b/azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStore.java
@@ -2,23 +2,25 @@
 // Licensed under the MIT License.
 package com.microsoft.durabletask.azureblobpayloads;
 
-import com.azure.core.credential.TokenCredential;
-import com.azure.core.http.rest.Response;
 import com.azure.core.util.Context;
 import com.azure.storage.blob.BlobClient;
 import com.azure.storage.blob.BlobContainerClient;
 import com.azure.storage.blob.BlobServiceClient;
 import com.azure.storage.blob.BlobServiceClientBuilder;
+import com.azure.storage.blob.models.BlobDownloadResponse;
 import com.azure.storage.blob.models.BlobHttpHeaders;
+import com.azure.storage.blob.models.BlobRange;
 import com.azure.storage.blob.models.BlobStorageException;
-import com.azure.storage.blob.models.PublicAccessType;
+import com.azure.storage.blob.options.BlobDownloadToFileOptions;
+import com.azure.storage.common.policy.RequestRetryOptions;
+import com.azure.storage.common.policy.RetryPolicyType;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.UUID;
 import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
@@ -62,15 +64,27 @@ public BlobPayloadStore(LargePayloadStorageOptions options) {
                 "Either ConnectionString or AccountUri and Credential must be provided.");
         }
 
+        // Retry policy: exponential (8 retries, 250ms base, 10s max, 2min network timeout)
+        // Matches the .NET BlobPayloadStore retry configuration.
+        RequestRetryOptions retryOptions = new RequestRetryOptions(
+            RetryPolicyType.EXPONENTIAL,
+            8,           // maxTries
+            120,         // tryTimeoutInSeconds (2 min network timeout)
+            250L,        // retryDelayInMs (250ms base)
+            10_000L,     // maxRetryDelayInMs (10s max)
+            null);       // secondaryHost
+
         BlobServiceClient serviceClient;
         if (hasIdentityAuth) {
             serviceClient = new BlobServiceClientBuilder()
                 .endpoint(options.getAccountUri().toString())
                 .credential(options.getCredential())
+                .retryOptions(retryOptions)
                 .buildClient();
         } else {
             serviceClient = new BlobServiceClientBuilder()
                 .connectionString(options.getConnectionString())
+                .retryOptions(retryOptions)
                 .buildClient();
         }
 
@@ -150,11 +164,20 @@ public String download(String token) {
 
         try {
             ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
-            blob.downloadStream(outputStream);
+            // Use downloadStreamWithResponse to get content-encoding header in the same call,
+            // avoiding a separate getProperties() round-trip.
+            BlobDownloadResponse downloadResponse = blob.downloadStreamWithResponse(
+                outputStream,
+                null,  // range (full blob)
+                null,  // options
+                null,  // requestConditions
+                false, // getMD5
+                null,  // timeout
+                Context.NONE);
             byte[] rawBytes = outputStream.toByteArray();
 
-            // Check if the content is gzip-compressed by trying to read the gzip header
-            String contentEncoding = blob.getProperties().getContentEncoding();
+            // Check if the content is gzip-compressed via the response header
+            String contentEncoding = downloadResponse.getDeserializedHeaders().getContentEncoding();
             boolean isGzip = CONTENT_ENCODING_GZIP.equalsIgnoreCase(contentEncoding);
 
             if (isGzip) {
diff --git a/azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/LargePayloadInterceptor.java b/azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/LargePayloadInterceptor.java
@@ -8,7 +8,6 @@
 
 import io.grpc.*;
 
-import java.nio.charset.StandardCharsets;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -712,7 +711,7 @@ private StringValue maybeExternalize(StringValue value) {
         }
 
         String strValue = value.getValue();
-        int size = strValue.getBytes(StandardCharsets.UTF_8).length;
+        int size = utf8ByteLength(strValue);
 
         if (size < this.options.getThresholdBytes()) {
             return value;
@@ -766,4 +765,25 @@ private static String getStackTraceString(Throwable t) {
         t.printStackTrace(new java.io.PrintWriter(sw));
         return sw.toString();
     }
+
+    /**
+     * Computes the UTF-8 encoded byte length of a string without allocating a byte array.
+     */
+    private static int utf8ByteLength(String s) {
+        int count = 0;
+        for (int i = 0; i < s.length(); i++) {
+            char ch = s.charAt(i);
+            if (ch <= 0x7F) {
+                count++;
+            } else if (ch <= 0x7FF) {
+                count += 2;
+            } else if (Character.isHighSurrogate(ch)) {
+                count += 4;
+                i++; // skip low surrogate
+            } else {
+                count += 3;
+            }
+        }
+        return count;
+    }
 }
diff --git a/azure-blob-payloads/src/test/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStoreTest.java b/azure-blob-payloads/src/test/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStoreTest.java
@@ -4,8 +4,9 @@
 
 import com.azure.storage.blob.BlobClient;
 import com.azure.storage.blob.BlobContainerClient;
+import com.azure.storage.blob.models.BlobDownloadHeaders;
+import com.azure.storage.blob.models.BlobDownloadResponse;
 import com.azure.storage.blob.models.BlobHttpHeaders;
-import com.azure.storage.blob.models.BlobProperties;
 import com.azure.storage.blob.models.BlobStorageException;
 
 import org.junit.jupiter.api.BeforeEach;
@@ -121,15 +122,7 @@ void download_uncompressedBlob_returnsOriginalPayload() {
         String expectedPayload = "Hello, uncompressed world!";
         byte[] payloadBytes = expectedPayload.getBytes(StandardCharsets.UTF_8);
 
-        doAnswer(inv -> {
-            OutputStream out = inv.getArgument(0);
-            out.write(payloadBytes);
-            return null;
-        }).when(mockBlobClient).downloadStream(any(OutputStream.class));
-
-        BlobProperties properties = mock(BlobProperties.class);
-        when(properties.getContentEncoding()).thenReturn(null);
-        when(mockBlobClient.getProperties()).thenReturn(properties);
+        mockDownloadResponse(payloadBytes, null);
 
         BlobPayloadStore store = new BlobPayloadStore(mockContainerClient, options);
         String token = BlobPayloadStore.encodeToken("durabletask-payloads", "testblob");
@@ -150,15 +143,7 @@ void download_gzipCompressedBlob_decompressesCorrectly() throws Exception {
         }
         byte[] compressedBytes = compressedBuffer.toByteArray();
 
-        doAnswer(inv -> {
-            OutputStream out = inv.getArgument(0);
-            out.write(compressedBytes);
-            return null;
-        }).when(mockBlobClient).downloadStream(any(OutputStream.class));
-
-        BlobProperties properties = mock(BlobProperties.class);
-        when(properties.getContentEncoding()).thenReturn("gzip");
-        when(mockBlobClient.getProperties()).thenReturn(properties);
+        mockDownloadResponse(compressedBytes, "gzip");
 
         BlobPayloadStore store = new BlobPayloadStore(mockContainerClient, options);
         String token = BlobPayloadStore.encodeToken("durabletask-payloads", "compressed-blob");
@@ -172,7 +157,8 @@ void download_gzipCompressedBlob_decompressesCorrectly() throws Exception {
     void download_blobNotFound_throwsPayloadStorageException() {
         BlobStorageException notFound = mock(BlobStorageException.class);
         when(notFound.getStatusCode()).thenReturn(404);
-        doThrow(notFound).when(mockBlobClient).downloadStream(any(OutputStream.class));
+        doThrow(notFound).when(mockBlobClient).downloadStreamWithResponse(
+            any(OutputStream.class), isNull(), isNull(), isNull(), eq(false), isNull(), any());
 
         BlobPayloadStore store = new BlobPayloadStore(mockContainerClient, options);
         String token = BlobPayloadStore.encodeToken("durabletask-payloads", "missing-blob");
@@ -194,7 +180,8 @@ void download_nonMatchingContainer_throwsIllegalArgument() {
     void download_blobStorageError_throwsPayloadStorageException() {
         BlobStorageException serverError = mock(BlobStorageException.class);
         when(serverError.getStatusCode()).thenReturn(500);
-        doThrow(serverError).when(mockBlobClient).downloadStream(any(OutputStream.class));
+        doThrow(serverError).when(mockBlobClient).downloadStreamWithResponse(
+            any(OutputStream.class), isNull(), isNull(), isNull(), eq(false), isNull(), any());
 
         BlobPayloadStore store = new BlobPayloadStore(mockContainerClient, options);
         String token = BlobPayloadStore.encodeToken("durabletask-payloads", "error-blob");
@@ -228,15 +215,7 @@ void upload_download_roundTrip_withCompression() throws Exception {
         String token = store.upload(originalPayload);
 
         // Now set up download to return the captured bytes
-        doAnswer(inv -> {
-            OutputStream out = inv.getArgument(0);
-            out.write(capturedBytes[0]);
-            return null;
-        }).when(mockBlobClient).downloadStream(any(OutputStream.class));
-
-        BlobProperties properties = mock(BlobProperties.class);
-        when(properties.getContentEncoding()).thenReturn("gzip");
-        when(mockBlobClient.getProperties()).thenReturn(properties);
+        mockDownloadResponse(capturedBytes[0], "gzip");
 
         // Need to map the token's blob name back to our mock
         String[] decoded = BlobPayloadStore.decodeToken(token);
@@ -269,15 +248,7 @@ void upload_download_roundTrip_withoutCompression() throws Exception {
         String token = store.upload(originalPayload);
 
         // Set up download
-        doAnswer(inv -> {
-            OutputStream out = inv.getArgument(0);
-            out.write(capturedBytes[0]);
-            return null;
-        }).when(mockBlobClient).downloadStream(any(OutputStream.class));
-
-        BlobProperties properties = mock(BlobProperties.class);
-        when(properties.getContentEncoding()).thenReturn(null);
-        when(mockBlobClient.getProperties()).thenReturn(properties);
+        mockDownloadResponse(capturedBytes[0], null);
 
         String[] decoded = BlobPayloadStore.decodeToken(token);
         when(mockContainerClient.getBlobClient(decoded[1])).thenReturn(mockBlobClient);
@@ -340,4 +311,25 @@ void upload_compressionProducesGzipHeaders() {
         BlobHttpHeaders capturedHeaders = headersCaptor.getValue();
         assertEquals("gzip", capturedHeaders.getContentEncoding());
     }
+
+    // ==================== Test helpers ====================
+
+    /**
+     * Sets up the mockBlobClient to return a downloadStreamWithResponse that writes
+     * the given bytes and returns headers with the specified content-encoding.
+     */
+    private void mockDownloadResponse(byte[] content, String contentEncoding) {
+        BlobDownloadHeaders downloadHeaders = mock(BlobDownloadHeaders.class);
+        when(downloadHeaders.getContentEncoding()).thenReturn(contentEncoding);
+
+        BlobDownloadResponse downloadResponse = mock(BlobDownloadResponse.class);
+        when(downloadResponse.getDeserializedHeaders()).thenReturn(downloadHeaders);
+
+        doAnswer(inv -> {
+            OutputStream out = inv.getArgument(0);
+            out.write(content);
+            return downloadResponse;
+        }).when(mockBlobClient).downloadStreamWithResponse(
+            any(OutputStream.class), isNull(), isNull(), isNull(), eq(false), isNull(), any());
+    }
 }
diff --git a/azure-blob-payloads/src/test/java/com/microsoft/durabletask/azureblobpayloads/LargePayloadIntegrationTest.java b/azure-blob-payloads/src/test/java/com/microsoft/durabletask/azureblobpayloads/LargePayloadIntegrationTest.java
@@ -24,8 +24,8 @@
  * <p>
  * These tests require:
  * <ul>
- *   <li>DTS emulator running on localhost:4001:
- *       {@code docker run --name durabletask-emulator -p 4001:8080 -d mcr.microsoft.com/dts/dts-emulator:latest}</li>
+ *   <li>DTS emulator running on localhost:8080:
+ *       {@code docker run --name durabletask-emulator -p 8080:8080 -d mcr.microsoft.com/dts/dts-emulator:latest}</li>
  *   <li>Azurite running on localhost:10000:
  *       {@code docker run --name azurite -p 10000:10000 -p 10001:10001 -p 10002:10002 -d mcr.microsoft.com/azure-storage/azurite}</li>
  * </ul>

Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ test {`
`72`	`72`	`}`
`73`	`73`	`}`
`74`	`74`
`75`		`-// Integration tests require DTS emulator on localhost:4001 and Azurite on localhost:10000.`
	`75`	`+// Integration tests require DTS emulator (default localhost:8080) and Azurite on localhost:10000.`
`76`	`76`	`task integrationTest(type: Test) {`
`77`	`77`	`useJUnitPlatform {`
`78`	`78`	`includeTags 'integration'`