wip: GrainStorage.AzureBlob.ReadState memory benchmark

Author: egil

test/Benchmarks/GrainStorage/AzureBlobReadStateBenchmark.cs

@@ -0,0 +1,157 @@
+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]
+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(4096, 256 * 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(