Add streaming grain storage serializer and Azure Blob pooled IO

Introduce IGrainStorageStreamingSerializer and stream overloads for OrleansJsonSerializer plus Orleans/Json grain storage serializers. Azure blob storage now supports pooled read buffers, a buffered stream write mode, and logs a warning when large payloads force a fallback to DownloadContentAsync.

Add Azure Blob storage tests for pooled reads and streaming serializer behavior, and add focused grain storage benchmarks (binary vs streaming) across Orleans, Newtonsoft.Json, and STJ.

Alternatives considered: full streaming OpenWriteAsync with separate ETag readback (rejected due to concurrency race), and IBufferWriter/ReadOnlySequence buffer paths (explored, but didnt improve performance or allocation in a meaningful way, so but dropped for now). Buffered stream uploads and pooled reads were kept as the best allocation/compatibility tradeoff while keeping BinaryData writes as the default.

Author: egil

src/Azure/Orleans.Persistence.AzureStorage/Providers/AzureProviderErrorCode.cs

@@ -34,7 +34,7 @@ internal enum AzureProviderErrorCode
         AzureBlobProvider_ClearError                = AzureBlobProviderBase + 14,
         AzureBlobProvider_ClearingData              = AzureBlobProviderBase + 15,
         AzureBlobProvider_Cleared                   = AzureBlobProviderBase + 16,
-
+        AzureBlobProvider_LargePayloadFallback      = AzureBlobProviderBase + 17,
 
 
     }

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,7 @@
 using Microsoft.Extensions.Options;
 using Orleans.Configuration;
 using Orleans.Providers.Azure;
-using Orleans.Runtime;
+using Orleans.Serialization.Buffers.Adaptors;
 using Orleans.Serialization.Serializers;
 using LogLevel = Microsoft.Extensions.Logging.LogLevel;
 
@@ -28,6 +26,7 @@ public partial class AzureBlobGrainStorage : IGrainStorage, ILifecycleParticipan
         private readonly IActivatorProvider _activatorProvider;
         private readonly AzureBlobStorageOptions options;
         private readonly IGrainStorageSerializer grainStorageSerializer;
+        private readonly IGrainStorageStreamingSerializer? streamSerializer;
 
         /// <summary> Default constructor </summary>
         public AzureBlobGrainStorage(
@@ -42,6 +41,7 @@ public AzureBlobGrainStorage(
             this.blobContainerFactory = blobContainerFactory;
             _activatorProvider = activatorProvider;
             this.grainStorageSerializer = options.GrainStorageSerializer;
+            this.streamSerializer = options.GrainStorageSerializer as IGrainStorageStreamingSerializer;
             this.logger = logger;
         }
 
@@ -58,20 +58,12 @@ 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 = streamSerializer switch
                 {
-                    loadedState = default;
-                    LogTraceBlobEmptyReading(grainType, grainId, grainState.ETag, blobName, container.Name);
-                }
-                else
-                {
-                    loadedState = this.ConvertFromStorageFormat<T>(contents);
-                    LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, container.Name);
-                }
+                    not null => await ReadStateWithStreamAsync<T>(blob, grainType, grainId, grainState, blobName, container.Name),
+                    null when options.UsePooledBufferForReads => await ReadStateWithPooledBufferAsync<T>(blob, grainType, grainId, grainState, blobName, container.Name),
+                    _ => await ReadStateWithBinaryDataAsync<T>(blob, grainType, grainId, grainState, blobName, container.Name),
+                };
 
                 grainState.State = loadedState ?? CreateInstance<T>();
                 grainState.RecordExists = loadedState is not null;
@@ -115,11 +107,17 @@ public async Task WriteStateAsync<T>(string grainType, GrainId grainId, IGrainSt
             {
                 LogTraceWriting(grainType, grainId, grainState.ETag, blobName, container.Name);
 
-                var contents = ConvertToStorageFormat(grainState.State);
-
                 var blob = container.GetBlobClient(blobName);
 
-                await WriteStateAndCreateContainerIfNotExists(grainType, grainId, grainState, contents, "application/octet-stream", blob);
+                if (streamSerializer is null || options.WriteMode == AzureBlobStorageWriteMode.BinaryData)
+                {
+                    var contents = ConvertToStorageFormat(grainState.State);
+                    await WriteStateAndCreateContainerIfNotExists(grainType, grainId, grainState, contents, "application/octet-stream", blob);
+                }
+                else
+                {
+                    await WriteStateBufferedStreamAndCreateContainerIfNotExists(grainType, grainId, grainState, "application/octet-stream", blob);
+                }
 
                 LogTraceDataWritten(grainType, grainId, grainState.ETag, blobName, container.Name);
             }
@@ -200,8 +198,7 @@ private async Task WriteStateAndCreateContainerIfNotExists<T>(string grainType,
                     static state => state.blob.UploadAsync(state.contents, state.options),
                     (blob, contents, options),
                     blob,
-                    grainState.ETag)
-                        .ConfigureAwait(false);
+                    grainState.ETag).ConfigureAwait(false);
 
                 grainState.ETag = result.Value.ETag.ToString();
                 grainState.RecordExists = true;
@@ -211,11 +208,133 @@ private async Task WriteStateAndCreateContainerIfNotExists<T>(string grainType,
                 // if the container does not exist, create it, and make another attempt
                 LogTraceContainerNotFound(grainType, grainId, grainState.ETag, blob.Name, container.Name);
                 await container.CreateIfNotExistsAsync().ConfigureAwait(false);
-
                 await WriteStateAndCreateContainerIfNotExists(grainType, grainId, grainState, contents, mimeType, blob).ConfigureAwait(false);
             }
         }
 
+        private async Task WriteStateBufferedStreamAndCreateContainerIfNotExists<T>(string grainType, GrainId grainId, IGrainState<T> grainState, string mimeType, BlobClient blob)
+        {
+            var container = this.blobContainerFactory.GetBlobContainerClient(grainId);
+
+            try
+            {
+                var conditions = string.IsNullOrEmpty(grainState.ETag)
+                    ? new BlobRequestConditions { IfNoneMatch = ETag.All }
+                    : new BlobRequestConditions { IfMatch = new ETag(grainState.ETag) };
+
+                var options = new BlobUploadOptions
+                {
+                    HttpHeaders = new BlobHttpHeaders { ContentType = mimeType },
+                    Conditions = conditions,
+                };
+
+                var result = await DoOptimisticUpdate(
+                    static state => state.self.UploadSerializedStateBufferedAsync(state.blob, state.options, state.value),
+                    (self: this, blob, options, value: grainState.State),
+                    blob,
+                    grainState.ETag).ConfigureAwait(false);
+
+                grainState.ETag = result.Value.ETag.ToString();
+                grainState.RecordExists = true;
+            }
+            catch (RequestFailedException exception) when (exception.IsContainerNotFound())
+            {
+                // if the container does not exist, create it, and make another attempt
+                LogTraceContainerNotFound(grainType, grainId, grainState.ETag, blob.Name, container.Name);
+                await container.CreateIfNotExistsAsync().ConfigureAwait(false);
+                await WriteStateBufferedStreamAndCreateContainerIfNotExists(grainType, grainId, grainState, mimeType, blob).ConfigureAwait(false);
+            }
+        }
+
+        private async Task<Response<BlobContentInfo>> UploadSerializedStateBufferedAsync<T>(BlobClient blob, BlobUploadOptions options, T value)
+        {
+            if (streamSerializer is null)
+            {
+                throw new InvalidOperationException("Stream serializer is not configured.");
+            }
+
+            var bufferStream = PooledBufferStream.Rent();
+            try
+            {
+                await streamSerializer.SerializeAsync(value, bufferStream).ConfigureAwait(false);
+                bufferStream.Position = 0;
+                return await blob.UploadAsync(bufferStream, options).ConfigureAwait(false);
+            }
+            finally
+            {
+                PooledBufferStream.Return(bufferStream);
+            }
+        }
+
+        private async Task<T?> ReadStateWithStreamAsync<T>(BlobClient blob, string grainType, GrainId grainId, IGrainState<T> grainState, string blobName, string containerName)
+        {
+            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, containerName);
+                return default;
+            }
+
+            await using var content = response.Value.Content;
+            var loadedState = await streamSerializer!.DeserializeAsync<T>(content).ConfigureAwait(false);
+            LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, containerName);
+            return loadedState;
+        }
+
+        private async Task<T?> ReadStateWithBinaryDataAsync<T>(BlobClient blob, string grainType, GrainId grainId, IGrainState<T> grainState, string blobName, string containerName)
+        {
+            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, containerName);
+                return default;
+            }
+
+            var loadedState = this.ConvertFromStorageFormat<T>(contents);
+            LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, containerName);
+            return loadedState;
+        }
+
+        private async Task<T?> ReadStateWithPooledBufferAsync<T>(BlobClient blob, string grainType, GrainId grainId, IGrainState<T> grainState, string blobName, string containerName)
+        {
+            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, containerName);
+                return default;
+            }
+
+            if (contentLength > int.MaxValue)
+            {
+                LogWarningLargePayloadFallback(contentLength, grainType, grainId, grainState.ETag, blobName, containerName);
+                return await ReadStateWithBinaryDataAsync(blob, grainType, grainId, grainState, blobName, containerName);
+            }
+
+            await using var content = response.Value.Content;
+            var buffer = ArrayPool<byte>.Shared.Rent((int)contentLength);
+            try
+            {
+                var memory = buffer.AsMemory(0, (int)contentLength);
+                await content.ReadExactlyAsync(memory);
+                var loadedState = this.ConvertFromStorageFormat<T>(new BinaryData(memory));
+                LogTraceDataRead(grainType, grainId, grainState.ETag, blobName, containerName);
+                return loadedState;
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+
         private static async Task<TResult> DoOptimisticUpdate<TState, TResult>(Func<TState, Task<TResult>> updateOperation, TState state, BlobClient blob, string currentETag)
         {
             try
@@ -294,6 +413,13 @@ private async Task Init(CancellationToken ct)
         )]
         private partial void LogTraceDataRead(string grainType, GrainId grainId, string? eTag, string blobName, string containerName);
 
+        [LoggerMessage(
+            Level = LogLevel.Warning,
+            EventId = (int)AzureProviderErrorCode.AzureBlobProvider_LargePayloadFallback,
+            Message = "ContentLength={ContentLength} exceeds max array size; falling back to DownloadContentAsync. GrainType={GrainType} GrainId={GrainId} ETag={ETag} BlobName={BlobName} in Container={ContainerName}"
+        )]
+        private partial void LogWarningLargePayloadFallback(long contentLength, string grainType, GrainId grainId, string? eTag, string blobName, string containerName);
+
         [LoggerMessage(
             Level = LogLevel.Error,
             EventId = (int)AzureProviderErrorCode.AzureBlobProvider_ReadError,

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

@@ -11,6 +11,23 @@
 
 namespace Orleans.Configuration
 {
+    public enum AzureBlobStorageWriteMode
+    {
+        /// <summary>
+        /// Serialize to <see cref="BinaryData"/> and upload.
+        /// This uses the <see cref="IGrainStorageSerializer"/> binary path, materializing the full payload in memory.
+        /// It is typically the fastest path but can create large allocations (including LOH) for big payloads.
+        /// </summary>
+        BinaryData,
+
+        /// <summary>
+        /// Serialize using the stream serializer into a pooled in-memory stream and upload from that buffer.
+        /// This still buffers the full payload but avoids LOH churn by reusing pooled segments.
+        /// Requires <see cref="IGrainStorageStreamingSerializer"/>; otherwise the write falls back to <see cref="BinaryData"/>.
+        /// </summary>
+        BufferedStream,
+    }
+
     public class AzureBlobStorageOptions : IStorageProviderSerializerOptions
     {
         private BlobServiceClient _blobServiceClient;
@@ -57,6 +74,20 @@ 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.
+        /// When a stream serializer is configured, pooled buffers are used only if the content length fits in an <see cref="int"/>.
+        /// When pooled buffers are used, deserialization goes through the <see cref="IGrainStorageSerializer"/> binary path.
+        /// </summary>
+        public bool UsePooledBufferForReads { get; set; } = true;
+
+        /// <summary>
+        /// Gets or sets the write path to use when a stream serializer is available.
+        /// If the stream serializer is not configured, writes always use <see cref="BinaryData"/>.
+        /// </summary>
+        public AzureBlobStorageWriteMode WriteMode { get; set; } = AzureBlobStorageWriteMode.BinaryData;
+
         /// <summary>
         /// A function for building container factory instances
         /// </summary>
@@ -149,7 +180,7 @@ public void ValidateConfiguration()
                 AzureBlobUtils.ValidateContainerName(options.ContainerName);
                 AzureBlobUtils.ValidateBlobName(this.name);
             }
-            catch(ArgumentException e)
+            catch (ArgumentException e)
             {
                 throw new OrleansConfigurationException(
                     $"Configuration for AzureBlobStorageOptions {name} is invalid. {nameof(this.options.ContainerName)} is not valid", e);

src/Orleans.Core/Providers/IGrainStorageSerializer.cs

@@ -1,4 +1,7 @@
 using System;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Options;
 using Orleans.Runtime;
@@ -27,6 +30,30 @@ public interface IGrainStorageSerializer
         T Deserialize<T>(BinaryData input);
     }
 
+    /// <summary>
+    /// Optional stream-based serializer for grain state.
+    /// </summary>
+    public interface IGrainStorageStreamingSerializer
+    {
+        /// <summary>
+        /// Serializes the object input to a stream.
+        /// </summary>
+        /// <param name="input">The object to serialize.</param>
+        /// <param name="destination">The destination stream.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <typeparam name="T">The input type.</typeparam>
+        ValueTask SerializeAsync<T>(T input, Stream destination, CancellationToken cancellationToken = default);
+
+        /// <summary>
+        /// Deserializes the provided data from a stream.
+        /// </summary>
+        /// <param name="input">The input stream.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <typeparam name="T">The output type.</typeparam>
+        /// <returns>The deserialized object.</returns>
+        ValueTask<T> DeserializeAsync<T>(Stream input, CancellationToken cancellationToken = default);
+    }
+
     /// <summary>
     /// Extensions for <see cref="IGrainStorageSerializer"/>.
     /// </summary>

src/Orleans.Core/Providers/StorageSerializer/JsonGrainStorageSerializer.cs

@@ -1,12 +1,11 @@
-using System;
 using Orleans.Serialization;
 
 namespace Orleans.Storage
 {
     /// <summary>
     /// Grain storage serializer that uses Newtonsoft.Json
     /// </summary>
-    public class JsonGrainStorageSerializer : IGrainStorageSerializer
+    public class JsonGrainStorageSerializer : IGrainStorageSerializer, IGrainStorageStreamingSerializer
     {
         private readonly OrleansJsonSerializer _orleansJsonSerializer;
 
@@ -30,5 +29,20 @@ public T Deserialize<T>(BinaryData input)
         {
             return (T)_orleansJsonSerializer.Deserialize(typeof(T), input.ToString());
         }
+
+        /// <inheritdoc/>
+        public ValueTask SerializeAsync<T>(T value, Stream destination, CancellationToken cancellationToken = default)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            _orleansJsonSerializer.Serialize(value, typeof(T), destination);
+            return ValueTask.CompletedTask;
+        }
+
+        /// <inheritdoc/>
+        public ValueTask<T> DeserializeAsync<T>(Stream input, CancellationToken cancellationToken = default)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            return ValueTask.FromResult((T)_orleansJsonSerializer.Deserialize(typeof(T), input));
+        }
     }
 }