addressed pr feedback

Author: bachuv

azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStore.java

@@ -6,8 +6,10 @@
 import com.azure.storage.blob.BlobContainerClient;
 import com.azure.storage.blob.BlobServiceClient;
 import com.azure.storage.blob.BlobServiceClientBuilder;
+import com.azure.storage.blob.models.BlobDownloadContentResponse;
 import com.azure.storage.blob.models.BlobHttpHeaders;
 import com.azure.storage.blob.models.BlobStorageException;
+import com.azure.storage.blob.options.BlobParallelUploadOptions;
 import com.microsoft.durabletask.PayloadStore;
 
 import java.io.ByteArrayInputStream;
@@ -43,6 +45,7 @@ public final class BlobPayloadStore implements PayloadStore {
     private final String blobPrefix;
     private final String containerName;
     private final boolean compressPayloads;
+    private volatile boolean containerEnsured;
 
     /**
      * Creates a new BlobPayloadStore with the given options.
@@ -72,8 +75,6 @@ public BlobPayloadStore(BlobPayloadStoreOptions options) {
         this.blobPrefix = options.getBlobPrefix();
         this.containerName = options.getContainerName();
         this.compressPayloads = options.isCompressPayloads();
-
-        ensureContainerExists();
     }
 
     @Override
@@ -82,6 +83,8 @@ public String upload(String payload) {
             throw new IllegalArgumentException("payload must not be null");
         }
 
+        ensureContainerExists();
+
         String blobName = this.blobPrefix + UUID.randomUUID().toString().replace("-", "") + BLOB_EXTENSION;
         BlobClient blobClient = this.containerClient.getBlobClient(blobName);
 
@@ -90,8 +93,11 @@ public String upload(String payload) {
 
         if (this.compressPayloads) {
             data = gzipCompress(rawData);
-            blobClient.upload(new ByteArrayInputStream(data), data.length, true);
-            blobClient.setHttpHeaders(new BlobHttpHeaders().setContentEncoding(GZIP_CONTENT_ENCODING));
+            BlobHttpHeaders headers = new BlobHttpHeaders().setContentEncoding(GZIP_CONTENT_ENCODING);
+            BlobParallelUploadOptions uploadOptions = new BlobParallelUploadOptions(
+                new ByteArrayInputStream(data), data.length)
+                .setHeaders(headers);
+            blobClient.uploadWithResponse(uploadOptions, null, null);
         } else {
             data = rawData;
             blobClient.upload(new ByteArrayInputStream(data), data.length, true);
@@ -112,13 +118,9 @@ public String download(String token) {
         String blobName = extractBlobName(token);
         BlobClient blobClient = this.containerClient.getBlobClient(blobName);
 
-        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
-        blobClient.downloadStream(outputStream);
-
-        byte[] rawBytes = outputStream.toByteArray();
-
-        // Check if the blob was compressed by inspecting content encoding
-        String contentEncoding = blobClient.getProperties().getContentEncoding();
+        BlobDownloadContentResponse response = blobClient.downloadContentWithResponse(null, null, null, null);
+        byte[] rawBytes = response.getValue().toBytes();
+        String contentEncoding = response.getDeserializedHeaders().getContentEncoding();
         if (GZIP_CONTENT_ENCODING.equalsIgnoreCase(contentEncoding)) {
             rawBytes = gzipDecompress(rawBytes);
         }
@@ -137,21 +139,27 @@ public boolean isKnownPayloadToken(String value) {
     }
 
     private void ensureContainerExists() {
+        if (this.containerEnsured) {
+            return;
+        }
         try {
             if (!this.containerClient.exists()) {
                 this.containerClient.create();
                 logger.info(() -> String.format("Created blob container: %s", this.containerClient.getBlobContainerName()));
             }
+            this.containerEnsured = true;
         } catch (BlobStorageException e) {
             // Container might have been created concurrently (409 Conflict)
             if (e.getStatusCode() != 409) {
                 throw e;
             }
+            this.containerEnsured = true;
         }
     }
 
     /**
-     * Extracts the blob name from a {@code blob:v1:<container>:<blobName>} token.
+     * Extracts the blob name from a {@code blob:v1:<container>:<blobName>} token
+     * and validates that the container matches the configured container.
      */
     private String extractBlobName(String token) {
         if (!token.startsWith(TOKEN_PREFIX)) {
@@ -165,6 +173,12 @@ private String extractBlobName(String token) {
             throw new IllegalArgumentException(
                 "Token does not have the expected format (blob:v1:<container>:<blobName>): " + token);
         }
+        String tokenContainer = remainder.substring(0, colonIndex);
+        if (!this.containerName.equals(tokenContainer)) {
+            throw new IllegalArgumentException(String.format(
+                "Token container '%s' does not match configured container '%s'",
+                tokenContainer, this.containerName));
+        }
         return remainder.substring(colonIndex + 1);
     }
 

azure-blob-payloads/src/main/java/com/microsoft/durabletask/azureblobpayloads/BlobPayloadStoreProvider.java

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.microsoft.durabletask.azureblobpayloads;
+
+import com.microsoft.durabletask.PayloadStore;
+import com.microsoft.durabletask.PayloadStoreProvider;
+
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * {@link PayloadStoreProvider} implementation that creates a {@link BlobPayloadStore}
+ * when the {@code DURABLETASK_LARGE_PAYLOADS_CONNECTION_STRING} environment variable is set.
+ * <p>
+ * This provider is discovered automatically via {@link java.util.ServiceLoader}.
+ */
+public final class BlobPayloadStoreProvider implements PayloadStoreProvider {
+
+    private static final Logger logger = Logger.getLogger(BlobPayloadStoreProvider.class.getName());
+    private static final String ENV_STORAGE_CONNECTION_STRING = "DURABLETASK_LARGE_PAYLOADS_CONNECTION_STRING";
+
+    @Override
+    public PayloadStore create() {
+        String connectionString = System.getenv(ENV_STORAGE_CONNECTION_STRING);
+        if (connectionString == null || connectionString.isEmpty()) {
+            return null;
+        }
+
+        try {
+            BlobPayloadStoreOptions options = new BlobPayloadStoreOptions.Builder()
+                .setConnectionString(connectionString)
+                .build();
+            logger.info("Large payload externalization enabled using Azure Blob Storage");
+            return new BlobPayloadStore(options);
+        } catch (Exception e) {
+            logger.log(Level.WARNING,
+                "Failed to initialize BlobPayloadStore; large payloads will not be externalized", e);
+            return null;
+        }
+    }
+}

azure-blob-payloads/src/main/resources/META-INF/services/com.microsoft.durabletask.PayloadStoreProvider

@@ -0,0 +1 @@
+com.microsoft.durabletask.azureblobpayloads.BlobPayloadStoreProvider

azurefunctions/build.gradle

@@ -35,7 +35,6 @@ repositories {
 
 dependencies {
     api project(':client')
-    implementation project(':azure-blob-payloads')
     implementation group: 'com.microsoft.azure.functions', name: 'azure-functions-java-library', version: '3.2.3'
     implementation "com.google.protobuf:protobuf-java:${protocVersion}"
     compileOnly "com.microsoft.azure.functions:azure-functions-java-spi:1.1.0"

azurefunctions/src/main/java/com/microsoft/durabletask/azurefunctions/internal/middleware/OrchestrationMiddleware.java

@@ -13,12 +13,12 @@
 import com.microsoft.durabletask.DataConverter;
 import com.microsoft.durabletask.OrchestrationRunner;
 import com.microsoft.durabletask.PayloadStore;
+import com.microsoft.durabletask.PayloadStoreProvider;
 import com.microsoft.durabletask.TaskFailedException;
-import com.microsoft.durabletask.azureblobpayloads.BlobPayloadStore;
-import com.microsoft.durabletask.azureblobpayloads.BlobPayloadStoreOptions;
 import com.microsoft.durabletask.interruption.ContinueAsNewInterruption;
 import com.microsoft.durabletask.interruption.OrchestratorBlockedException;
 
+import java.util.ServiceLoader;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -33,16 +33,6 @@ public class OrchestrationMiddleware implements Middleware {
     private static final String ORCHESTRATION_TRIGGER = "DurableOrchestrationTrigger";
     private static final Logger logger = Logger.getLogger(OrchestrationMiddleware.class.getName());
 
-    /**
-     * Environment variable for the Azure Storage connection string used for large payload externalization.
-     */
-    private static final String ENV_STORAGE_CONNECTION_STRING = "DURABLETASK_STORAGE_CONNECTION_STRING";
-
-    /**
-     * Fallback environment variable (standard Azure Functions storage connection).
-     */
-    private static final String ENV_AZURE_WEB_JOBS_STORAGE = "AzureWebJobsStorage";
-
     private final PayloadStore payloadStore;
 
     public OrchestrationMiddleware() {
@@ -98,27 +88,19 @@ public void invoke(MiddlewareContext context, MiddlewareChain chain) throws Exce
     }
 
     private static PayloadStore initializePayloadStore() {
-        // Check for explicit large-payload storage connection string
-        String connectionString = System.getenv(ENV_STORAGE_CONNECTION_STRING);
-        if (connectionString == null || connectionString.isEmpty()) {
-            // Fall back to standard Azure Functions storage connection
-            connectionString = System.getenv(ENV_AZURE_WEB_JOBS_STORAGE);
-        }
-
-        if (connectionString == null || connectionString.isEmpty()) {
-            logger.fine("No storage connection string configured for large payload externalization");
-            return null;
-        }
-
-        try {
-            BlobPayloadStoreOptions options = new BlobPayloadStoreOptions.Builder()
-                .setConnectionString(connectionString)
-                .build();
-            logger.info("Large payload externalization enabled using Azure Blob Storage");
-            return new BlobPayloadStore(options);
-        } catch (Exception e) {
-            logger.log(Level.WARNING, "Failed to initialize BlobPayloadStore; large payloads will not be externalized", e);
-            return null;
+        ServiceLoader<PayloadStoreProvider> loader = ServiceLoader.load(PayloadStoreProvider.class);
+        for (PayloadStoreProvider provider : loader) {
+            try {
+                PayloadStore store = provider.create();
+                if (store != null) {
+                    return store;
+                }
+            } catch (Exception e) {
+                logger.log(Level.WARNING,
+                    "PayloadStoreProvider " + provider.getClass().getName() + " failed to create store", e);
+            }
         }
+        logger.fine("No PayloadStoreProvider found or configured; large payload externalization is disabled");
+        return null;
     }
 }

client/src/main/java/com/microsoft/durabletask/DurableTaskGrpcWorker.java

@@ -311,7 +311,7 @@ public void startAndBlock() {
                                     response = externalizeOrchestratorResponsePayloads(response);
                                 }
                                 sendOrchestratorResponse(response);
-                            } catch (IllegalArgumentException | IllegalStateException e) {
+                            } catch (PayloadTooLargeException e) {
                                 logger.log(Level.WARNING,
                                     "Failed to send orchestrator response for instance '" +
                                     orchestratorRequest.getInstanceId() + "': " + e.getMessage(), e);
@@ -320,6 +320,7 @@ public void startAndBlock() {
                                     .setFailureDetails(TaskFailureDetails.newBuilder()
                                         .setErrorType(e.getClass().getName())
                                         .setErrorMessage(e.getMessage())
+                                        .setStackTrace(StringValue.of(FailureDetails.getFullStackTrace(e)))
                                         .build())
                                     .build();
                                 OrchestratorResponse failResponse = OrchestratorResponse.newBuilder()

client/src/main/java/com/microsoft/durabletask/PayloadHelper.java

@@ -48,6 +48,12 @@ String maybeExternalize(String value) {
             return value;
         }
 
+        // Fast path: if char count is below threshold, byte count is too
+        // (each Java char encodes to 1-3 UTF-8 bytes, so length() <= UTF-8 byte length)
+        if (value.length() <= this.options.getThresholdBytes()) {
+            return value;
+        }
+
         int byteSize = value.getBytes(StandardCharsets.UTF_8).length;
 
         // (2) below-threshold guard
@@ -57,7 +63,7 @@ String maybeExternalize(String value) {
 
         // (3) above-max-cap rejection
         if (byteSize > this.options.getMaxExternalizedPayloadBytes()) {
-            throw new IllegalArgumentException(String.format(
+            throw new PayloadTooLargeException(String.format(
                 "Payload size %d KB exceeds maximum of %d KB. " +
                 "Reduce the payload size or increase maxExternalizedPayloadBytes.",
                 byteSize / 1024,

client/src/main/java/com/microsoft/durabletask/PayloadInterceptionHelper.java

@@ -185,7 +185,7 @@ interface HistoryEventUpdater {
     private static HistoryEvent resolveStringValueField(HistoryEvent event, StringValue field,
                                                          PayloadHelper payloadHelper,
                                                          HistoryEventUpdater updater) {
-        if (!field.isInitialized() || field.getValue().isEmpty()) {
+        if (field.getValue().isEmpty()) {
             return event;
         }
         String resolved = payloadHelper.maybeResolve(field.getValue());
@@ -314,7 +314,7 @@ private static HistoryEvent externalizeHistoryEvent(HistoryEvent event, PayloadH
     private static HistoryEvent externalizeStringValueField(HistoryEvent event, StringValue field,
                                                              PayloadHelper payloadHelper,
                                                              HistoryEventUpdater updater) {
-        if (!field.isInitialized() || field.getValue().isEmpty()) {
+        if (field.getValue().isEmpty()) {
             return event;
         }
         String externalized = payloadHelper.maybeExternalize(field.getValue());

client/src/main/java/com/microsoft/durabletask/PayloadStore.java

@@ -13,6 +13,12 @@
  * The store implementation is solely responsible for generating blob names/keys and
  * defining the token format. The core framework treats tokens as opaque strings and
  * delegates token recognition to {@link #isKnownPayloadToken(String)}.
+ * <p>
+ * <b>Payload retention:</b> This interface does not define a deletion mechanism.
+ * Externalized payloads persist until removed by external means. When using Azure
+ * Blob Storage, configure
+ * <a href="https://learn.microsoft.com/azure/storage/blobs/lifecycle-management-overview">
+ * lifecycle management policies</a> to automatically expire old payloads.
  *
  * @see LargePayloadOptions
  */

client/src/main/java/com/microsoft/durabletask/PayloadStoreProvider.java

@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.microsoft.durabletask;
+
+/**
+ * Service provider interface for discovering {@link PayloadStore} implementations at runtime.
+ * <p>
+ * Implementations are discovered via {@link java.util.ServiceLoader}. To register a provider,
+ * create a file {@code META-INF/services/com.microsoft.durabletask.PayloadStoreProvider}
+ * containing the fully qualified class name of the implementation.
+ * <p>
+ * The provider is responsible for reading its own configuration (e.g., environment variables)
+ * and determining whether it can create a functional {@link PayloadStore}.
+ *
+ * @see PayloadStore
+ */
+public interface PayloadStoreProvider {
+
+    /**
+     * Attempts to create a {@link PayloadStore} based on available configuration.
+     * <p>
+     * Implementations should inspect environment variables or other configuration sources
+     * to determine if they can provide a store. If the required configuration is not present,
+     * this method should return {@code null} rather than throwing an exception.
+     *
+     * @return a configured {@link PayloadStore}, or {@code null} if the required configuration
+     *         is not available
+     */
+    PayloadStore create();
+}