Azure blob reads: add optional pooled buffers for streaming downloads

Introduce AzureBlobStorageOptions.UsePooledBufferForReads and switch ReadStateAsync to DownloadStreamingAsync with a rented buffer sized to ContentLength when enabled. This avoids large byte[] allocations from DownloadContentAsync, reducing LOH pressure and Gen2 GC churn for blob reads while keeping the behavior opt-in for safety with custom serializers.

Author: egil

src/Azure/Orleans.Persistence.AzureStorage/Providers/Storage/AzureBlobStorage.cs

@@ -1,8 +1,6 @@
 #nullable enable
-using System;
+using System.Buffers;
 using System.Diagnostics;
-using System.Threading;
-using System.Threading.Tasks;
 using Azure;
 using Azure.Storage.Blobs;
 using Azure.Storage.Blobs.Models;
@@ -11,7 +9,6 @@
 using Microsoft.Extensions.Options;
 using Orleans.Configuration;
 using Orleans.Providers.Azure;
-using Orleans.Runtime;
 using Orleans.Serialization.Serializers;
 using LogLevel = Microsoft.Extensions.Logging.LogLevel;
 
@@ -58,19 +55,46 @@ public async Task ReadStateAsync<T>(string grainType, GrainId grainId, IGrainSta
             {
                 var blob = container.GetBlobClient(blobName);
 
-                var response = await blob.DownloadContentAsync();
-                grainState.ETag = response.Value.Details.ETag.ToString();
-                var contents = response.Value.Content;
-                T? loadedState;
-                if (contents is null || contents.IsEmpty)
+                T? loadedState = default;
+                if (options.UsePooledBufferForReads)
                 {
-                    loadedState = default;
-                    LogTraceBlobEmptyReading(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    var response = await blob.DownloadStreamingAsync();
+                    grainState.ETag = response.Value.Details.ETag.ToString();
+
+                    var contentLength = response.Value.Details.ContentLength;
+                    if (contentLength <= 0)
+                    {
+                        LogTraceBlobEmptyReading(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    }
+                    else if (contentLength > int.MaxValue)
+                    {
+                        throw new InvalidOperationException($"Blob content length {contentLength} exceeds the maximum supported size.");
+                    }
+                    else
+                    {
+                        await using var content = response.Value.Content;
+                        var buffer = ArrayPool<byte>.Shared.Rent((int)contentLength);
+                        var memory = buffer.AsMemory(0, (int)contentLength);
+                        await content.ReadExactlyAsync(memory);
+                        loadedState = this.ConvertFromStorageFormat<T>(new BinaryData(memory));
+                        ArrayPool<byte>.Shared.Return(buffer);
+                        LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    }
                 }
                 else
                 {
-                    loadedState = this.ConvertFromStorageFormat<T>(contents);
-                    LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    var response = await blob.DownloadContentAsync();
+                    grainState.ETag = response.Value.Details.ETag.ToString();
+                    var contents = response.Value.Content;
+                    if (contents is null || contents.IsEmpty)
+                    {
+                        LogTraceBlobEmptyReading(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    }
+                    else
+                    {
+                        loadedState = this.ConvertFromStorageFormat<T>(contents);
+                        LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, container.Name);
+                    }
                 }
 
                 grainState.State = loadedState ?? CreateInstance<T>();

src/Azure/Orleans.Persistence.AzureStorage/Providers/Storage/AzureBlobStorageOptions.cs

@@ -57,6 +57,12 @@ public BlobServiceClient BlobServiceClient
         /// </summary>
         public bool DeleteStateOnClear { get; set; } = true;
 
+        /// <summary>
+        /// Gets or sets a value indicating whether to use pooled buffers when reading blob contents.
+        /// The deserializer must not retain the <see cref="BinaryData"/> or underlying buffer after deserialization.
+        /// </summary>
+        public bool UsePooledBufferForReads { get; set; }
+
         /// <summary>
         /// A function for building container factory instances
         /// </summary>

test/Benchmarks/GrainStorage/AzureBlobReadStateBenchmark.cs

@@ -0,0 +1,158 @@
+using System;
+using System.Threading.Tasks;
+using Azure.Storage.Blobs;
+using BenchmarkDotNet.Attributes;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
+using Orleans.Configuration;
+using Orleans.Runtime;
+using Orleans.Serialization;
+using Orleans.Serialization.Serializers;
+using Orleans.Storage;
+using TestExtensions;
+
+namespace Benchmarks.GrainStorage;
+
+[MemoryDiagnoser]
+[SimpleJob(launchCount: 1, warmupCount: 1, iterationCount: 6)]
+public class AzureBlobReadStateBenchmark
+{
+    private const int ReadIterations = 50;
+    private ServiceProvider _serviceProvider = null!;
+    private AzureBlobGrainStorage _nonPooledStorage = null!;
+    private AzureBlobGrainStorage _pooledStorage = null!;
+    private IGrainState<BenchmarkState> _nonPooledState = null!;
+    private IGrainState<BenchmarkState> _pooledState = null!;
+    private GrainId _grainId;
+    private string _grainType = null!;
+    private BlobContainerClient _containerClient = null!;
+
+    [Params(64 * 1024, 128 * 1024)]
+    public int PayloadSize { get; set; }
+
+    [GlobalSetup]
+    public async Task SetupAsync()
+    {
+        var client = CreateBlobServiceClient();
+        _serviceProvider = new ServiceCollection()
+            .AddLogging()
+            .AddSerializer()
+            .BuildServiceProvider();
+
+        var activatorProvider = _serviceProvider.GetRequiredService<IActivatorProvider>();
+        var serializer = new OrleansGrainStorageSerializer(_serviceProvider.GetRequiredService<Serializer>());
+        var containerName = $"bench-grainstate-{Guid.NewGuid():N}";
+        _grainType = "bench-grain";
+        _grainId = GrainId.Create("bench-grain", Guid.NewGuid().ToString("N"));
+
+        var nonPooledOptions = CreateOptions(client, serializer, containerName, usePooledReads: false);
+        var pooledOptions = CreateOptions(client, serializer, containerName, usePooledReads: true);
+
+        (_nonPooledStorage, var nonPooledFactory) = CreateStorage("bench-nonpooled", nonPooledOptions, activatorProvider);
+        (_pooledStorage, var pooledFactory) = CreateStorage("bench-pooled", pooledOptions, activatorProvider);
+
+        await nonPooledFactory.InitializeAsync(client);
+        await pooledFactory.InitializeAsync(client);
+
+        var writeState = new GrainState<BenchmarkState>
+        {
+            State = BenchmarkState.Create(PayloadSize)
+        };
+        await _nonPooledStorage.WriteStateAsync(_grainType, _grainId, writeState);
+
+        _nonPooledState = new GrainState<BenchmarkState>();
+        _pooledState = new GrainState<BenchmarkState>();
+        _containerClient = client.GetBlobContainerClient(containerName);
+    }
+
+    [GlobalCleanup]
+    public async Task CleanupAsync()
+    {
+        if (_containerClient is not null)
+        {
+            await _containerClient.DeleteIfExistsAsync();
+        }
+
+        _serviceProvider?.Dispose();
+    }
+
+    [Benchmark(Baseline = true, OperationsPerInvoke = ReadIterations)]
+    public async Task ReadStateNonPooledAsync()
+    {
+        for (var i = 0; i < ReadIterations; i++)
+        {
+            await _nonPooledStorage.ReadStateAsync(_grainType, _grainId, _nonPooledState);
+        }
+    }
+
+    [Benchmark(OperationsPerInvoke = ReadIterations)]
+    public async Task ReadStatePooledAsync()
+    {
+        for (var i = 0; i < ReadIterations; i++)
+        {
+            await _pooledStorage.ReadStateAsync(_grainType, _grainId, _pooledState);
+        }
+    }
+
+    private static BlobServiceClient CreateBlobServiceClient()
+    {
+        if (TestDefaultConfiguration.UseAadAuthentication)
+        {
+            if (!TestDefaultConfiguration.GetValue(nameof(TestDefaultConfiguration.DataBlobUri), out var blobUriValue) ||
+                string.IsNullOrWhiteSpace(blobUriValue))
+            {
+                throw new InvalidOperationException("DataBlobUri is required when UseAadAuthentication is true.");
+            }
+
+            return new BlobServiceClient(new Uri(blobUriValue), TestDefaultConfiguration.TokenCredential);
+        }
+
+        if (string.IsNullOrWhiteSpace(TestDefaultConfiguration.DataConnectionString))
+        {
+            throw new InvalidOperationException("OrleansDataConnectionString must be set for Azure Blob benchmarks.");
+        }
+
+        return new BlobServiceClient(TestDefaultConfiguration.DataConnectionString);
+    }
+
+    private static AzureBlobStorageOptions CreateOptions(
+        BlobServiceClient client,
+        IGrainStorageSerializer serializer,
+        string containerName,
+        bool usePooledReads)
+    {
+        return new AzureBlobStorageOptions
+        {
+            BlobServiceClient = client,
+            ContainerName = containerName,
+            GrainStorageSerializer = serializer,
+            UsePooledBufferForReads = usePooledReads
+        };
+    }
+
+    private (AzureBlobGrainStorage Storage, IBlobContainerFactory Factory) CreateStorage(
+        string name,
+        AzureBlobStorageOptions options,
+        IActivatorProvider activatorProvider)
+    {
+        var containerFactory = options.BuildContainerFactory(_serviceProvider, options);
+        var logger = NullLoggerFactory.Instance.CreateLogger<AzureBlobGrainStorage>();
+        var storage = new AzureBlobGrainStorage(name, options, containerFactory, activatorProvider, logger);
+        return (storage, containerFactory);
+    }
+
+    [GenerateSerializer]
+    public sealed class BenchmarkState
+    {
+        [Id(0)]
+        public byte[] Payload { get; set; } = Array.Empty<byte>();
+
+        public static BenchmarkState Create(int size)
+        {
+            var payload = new byte[size];
+            Random.Shared.NextBytes(payload);
+            return new BenchmarkState { Payload = payload };
+        }
+    }
+}

test/Benchmarks/Program.cs

@@ -260,6 +260,10 @@ internal class Program
             benchmark => benchmark.RunAsync().GetAwaiter().GetResult(),
             benchmark => benchmark.Teardown());
         },
+        ["GrainStorage.AzureBlob.ReadState"] = _ =>
+        {
+            BenchmarkRunner.Run<AzureBlobReadStateBenchmark>();
+        },
         ["GrainStorage.AdoNet"] = _ =>
         {
             RunBenchmark(

test/Benchmarks/Properties/launchSettings.json

@@ -1,5 +1,5 @@
 pushd %~dp0
 git log -n 1
 git --no-pager diff
-dotnet run -c Release -- ConcurrentPing
-popd
+dotnet run -c Release --framework net10.0 -- GrainStorage.AzureBlob.ReadState
+popd

test/Benchmarks/run_test.cmd

@@ -0,0 +1,45 @@
+using Orleans.Configuration;
+using Orleans.TestingHost;
+using Xunit;
+using Xunit.Abstractions;
+
+namespace Tester.AzureUtils.Persistence;
+
+/// <summary>
+/// PersistenceGrainTests using AzureStore with pooled read buffers enabled.
+/// </summary>
+[TestCategory("Persistence"), TestCategory("AzureStorage")]
+public class PersistenceGrainTests_AzureBlobStore_PooledReads : Base_PersistenceGrainTests_AzureStore, IClassFixture<PersistenceGrainTests_AzureBlobStore_PooledReads.Fixture>
+{
+    public class Fixture : BaseAzureTestClusterFixture
+    {
+        private class StorageSiloBuilderConfigurator : ISiloConfigurator
+        {
+            public void Configure(ISiloBuilder hostBuilder)
+            {
+                hostBuilder
+                    .AddAzureBlobGrainStorage("GrainStorageForTest", (AzureBlobStorageOptions options) =>
+                    {
+                        options.ConfigureTestDefaults();
+                        options.UsePooledBufferForReads = true;
+                    })
+                    .AddMemoryGrainStorage("MemoryStore")
+                    .AddMemoryGrainStorage("test1");
+            }
+        }
+
+        protected override void ConfigureTestCluster(TestClusterBuilder builder)
+        {
+            builder.Options.InitialSilosCount = 4;
+            builder.Options.UseTestClusterMembership = false;
+            builder.AddSiloBuilderConfigurator<SiloBuilderConfigurator>();
+            builder.AddSiloBuilderConfigurator<StorageSiloBuilderConfigurator>();
+            builder.AddClientBuilderConfigurator<ClientBuilderConfigurator>();
+        }
+    }
+
+    public PersistenceGrainTests_AzureBlobStore_PooledReads(ITestOutputHelper output, Fixture fixture) : base(output, fixture)
+    {
+        fixture.EnsurePreconditionsMet();
+    }
+}