Cache Azure Blob container initialization to avoid per-upload CreateIfNotExistsAsync

Microsoft.DurableTask.sln

@@ -145,6 +145,10 @@ Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "AzureManaged", "AzureManage
 EndProject
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Grpc", "Grpc", "{3B8F957E-7773-4C0C-ACD7-91A1591D9312}"
 EndProject
+Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "Extensions", "Extensions", "{00205C88-F000-28F2-A910-C6FA00E065EE}"
+EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AzureBlobPayloads.Tests", "test\Extensions\AzureBlobPayloads.Tests\AzureBlobPayloads.Tests.csproj", "{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -839,6 +843,18 @@ Global
 		{C1995163-1DCE-405D-BE82-8B4B2584893E}.Release|x64.Build.0 = Release|Any CPU
 		{C1995163-1DCE-405D-BE82-8B4B2584893E}.Release|x86.ActiveCfg = Release|Any CPU
 		{C1995163-1DCE-405D-BE82-8B4B2584893E}.Release|x86.Build.0 = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|x64.ActiveCfg = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|x64.Build.0 = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|x86.ActiveCfg = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Debug|x86.Build.0 = Debug|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|Any CPU.Build.0 = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|x64.ActiveCfg = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|x64.Build.0 = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|x86.ActiveCfg = Release|Any CPU
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1}.Release|x86.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -911,6 +927,8 @@ Global
 		{C1995163-1DCE-405D-BE82-8B4B2584893E} = {9686B8F9-2644-6C9B-E567-55B0471E4584}
 		{53193780-CD18-2643-6953-C26F59EAEDF5} = {5B448FF6-EC42-491D-A22E-1DC8B618E6D5}
 		{3B8F957E-7773-4C0C-ACD7-91A1591D9312} = {5B448FF6-EC42-491D-A22E-1DC8B618E6D5}
+		{00205C88-F000-28F2-A910-C6FA00E065EE} = {E5637F81-2FB9-4CD7-900D-455363B142A7}
+		{3E509481-3CCC-4006-BCB2-9E8FA7C275F1} = {00205C88-F000-28F2-A910-C6FA00E065EE}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {AB41CB55-35EA-4986-A522-387AB3402E71}

src/Extensions/AzureBlobPayloads/PayloadStore/BlobPayloadStore.cs

@@ -17,6 +17,10 @@ namespace Microsoft.DurableTask;
 /// <c>blob:v2:{fullBlobUrl}</c>, where the URL is the blob's absolute URI including the storage account.
 /// Legacy <c>blob:v1:{container}:{blobName}</c> tokens are still recognized for read back-compatibility.
 /// </summary>
+[System.Diagnostics.CodeAnalysis.SuppressMessage(
+    "Design",
+    "CA1001:Types that own disposable fields should be disposable",
+    Justification = "SemaphoreSlim does not allocate a disposable resource unless AvailableWaitHandle is accessed.")]
 public sealed class BlobPayloadStore : PayloadStore
 {
     const string TokenPrefixV1 = "blob:v1:";
@@ -26,9 +30,15 @@ public sealed class BlobPayloadStore : PayloadStore
     const int BaseDelayMs = 250;
     const int MaxDelayMs = 10_000;
     const int NetworkTimeoutMinutes = 2;
+
     readonly BlobContainerClient containerClient;
     readonly LargePayloadStorageOptions options;
     readonly BlobClientOptions clientOptions;
+    readonly SemaphoreSlim containerInitializationLock = new(initialCount: 1, maxCount: 1);
+
+    // Each successful initialization publishes a unique token. The token lets a stale
+    // ContainerNotFound failure invalidate only the initialization generation it used.
+    object? containerGeneration;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="BlobPayloadStore"/> class.
@@ -71,6 +81,19 @@ public BlobPayloadStore(LargePayloadStorageOptions options)
         this.containerClient = serviceClient.GetBlobContainerClient(options.ContainerName);
     }
 
+    /// <summary>
+    /// Initializes a new instance of the <see cref="BlobPayloadStore"/> class using an existing
+    /// container client. Intended for unit testing only.
+    /// </summary>
+    /// <param name="options">The options for the blob payload store.</param>
+    /// <param name="containerClient">The blob container client to use.</param>
+    internal BlobPayloadStore(LargePayloadStorageOptions options, BlobContainerClient containerClient)
+    {
+        this.options = options ?? throw new ArgumentNullException(nameof(options));
+        this.containerClient = containerClient ?? throw new ArgumentNullException(nameof(containerClient));
+        this.clientOptions = new BlobClientOptions();
+    }
+
     /// <inheritdoc/>
     public override async Task<string> UploadAsync(string payLoad, CancellationToken cancellationToken)
     {
@@ -80,36 +103,58 @@ public override async Task<string> UploadAsync(string payLoad, CancellationToken
 
         byte[] payloadBuffer = Encoding.UTF8.GetBytes(payLoad);
 
-        // Ensure container exists (idempotent)
-        await this.containerClient.CreateIfNotExistsAsync(PublicAccessType.None, default, default, cancellationToken);
-
-        if (this.options.CompressionEnabled)
+        // Ensure container exists. Cached/single-flight after the first successful call so we
+        // don't pay for an extra CreateIfNotExistsAsync request/transaction on every upload.
+        // Retry one write after an out-of-band container deletion so the cached path preserves
+        // the recovery behavior that an unconditional CreateIfNotExistsAsync provided before
+        // initialization was cached.
+        bool retryAfterContainerNotFound = true;
+        while (true)
         {
-            BlobOpenWriteOptions writeOptions = new()
-            {
-                HttpHeaders = new BlobHttpHeaders { ContentEncoding = ContentEncodingGzip },
-            };
-            using Stream blobStream = await blob.OpenWriteAsync(true, writeOptions, cancellationToken);
-            using GZipStream compressedBlobStream = new(blobStream, System.IO.Compression.CompressionLevel.Optimal, leaveOpen: true);
-
-            // using MemoryStream payloadStream = new(payloadBuffer, writable: false);
+            object generation = await this.EnsureContainerExistsAsync(cancellationToken).ConfigureAwait(false);
 
-            // await payloadStream.CopyToAsync(compressedBlobStream, bufferSize: DefaultCopyBufferSize, cancellationToken);
-            await WritePayloadAsync(payloadBuffer, compressedBlobStream, cancellationToken);
-            await compressedBlobStream.FlushAsync(cancellationToken);
-            await blobStream.FlushAsync(cancellationToken);
-        }
-        else
-        {
-            using Stream blobStream = await blob.OpenWriteAsync(true, default, cancellationToken);
+            try
+            {
+                if (this.options.CompressionEnabled)
+                {
+                    BlobOpenWriteOptions writeOptions = new()
+                    {
+                        HttpHeaders = new BlobHttpHeaders { ContentEncoding = ContentEncodingGzip },
+                    };
+                    using Stream blobStream = await blob.OpenWriteAsync(true, writeOptions, cancellationToken);
+                    using GZipStream compressedBlobStream = new(blobStream, System.IO.Compression.CompressionLevel.Optimal, leaveOpen: true);
+
+                    // using MemoryStream payloadStream = new(payloadBuffer, writable: false);
+
+                    // await payloadStream.CopyToAsync(compressedBlobStream, bufferSize: DefaultCopyBufferSize, cancellationToken);
+                    await WritePayloadAsync(payloadBuffer, compressedBlobStream, cancellationToken);
+                    await compressedBlobStream.FlushAsync(cancellationToken);
+                    await blobStream.FlushAsync(cancellationToken);
+                }
+                else
+                {
+                    using Stream blobStream = await blob.OpenWriteAsync(true, default, cancellationToken);
+
+                    // using MemoryStream payloadStream = new(payloadBuffer, writable: false);
+                    // await payloadStream.CopyToAsync(blobStream, bufferSize: DefaultCopyBufferSize, cancellationToken);
+                    await WritePayloadAsync(payloadBuffer, blobStream, cancellationToken);
+                    await blobStream.FlushAsync(cancellationToken);
+                }
+            }
+            catch (RequestFailedException ex) when (
+                retryAfterContainerNotFound &&
+                ex.ErrorCode == BlobErrorCode.ContainerNotFound)
+            {
+                // Clear only the generation used by this upload, then retry once. If Azure is still
+                // deleting the container, CreateIfNotExistsAsync can surface ContainerBeingDeleted,
+                // matching the behavior before initialization was cached.
+                _ = Interlocked.CompareExchange(ref this.containerGeneration, null, generation);
+                retryAfterContainerNotFound = false;
+                continue;
+            }
 
-            // using MemoryStream payloadStream = new(payloadBuffer, writable: false);
-            // await payloadStream.CopyToAsync(blobStream, bufferSize: DefaultCopyBufferSize, cancellationToken);
-            await WritePayloadAsync(payloadBuffer, blobStream, cancellationToken);
-            await blobStream.FlushAsync(cancellationToken);
+            return EncodeToken(blob.Uri);
         }
-
-        return EncodeToken(blob.Uri);
     }
 
     /// <inheritdoc/>
@@ -277,6 +322,39 @@ bool IsConfiguredContainer(Uri tokenContainerUri)
                 StringComparison.Ordinal);
     }
 
+    /// <summary>
+    /// Ensures the container exists and returns the initialization generation used by the caller.
+    /// </summary>
+    async Task<object> EnsureContainerExistsAsync(CancellationToken cancellationToken)
+    {
+        object? generation = Volatile.Read(ref this.containerGeneration);
+        if (generation is not null)
+        {
+            return generation;
+        }
+
+        await this.containerInitializationLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+        try
+        {
+            generation = Volatile.Read(ref this.containerGeneration);
+            if (generation is not null)
+            {
+                return generation;
+            }
+
+            await this.containerClient.CreateIfNotExistsAsync(PublicAccessType.None, cancellationToken: cancellationToken)
+                .ConfigureAwait(false);
+
+            generation = new object();
+            Volatile.Write(ref this.containerGeneration, generation);
+            return generation;
+        }
+        finally
+        {
+            this.containerInitializationLock.Release();
+        }
+    }
+
     /// <summary>
     /// The result of decoding an externalized payload token.
     /// </summary>

test/Extensions/AzureBlobPayloads.Tests/AzureBlobPayloads.Tests.csproj

@@ -0,0 +1,14 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>net10.0</TargetFramework>
+    <!-- Must match the $(AssemblyName).Tests InternalsVisibleTo generated for the product assembly. -->
+    <AssemblyName>Microsoft.DurableTask.Extensions.AzureBlobPayloads.Tests</AssemblyName>
+    <RootNamespace>$(AssemblyName)</RootNamespace>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="$(SrcRoot)Extensions/AzureBlobPayloads/AzureBlobPayloads.csproj" />
+  </ItemGroup>
+
+</Project>