Repository: microsoft/autogen
Pull request: bugfix exception issue 5580

## Why are these changes needed?

fixing intermittent bug #5580 

was using the wrong scheduler. 

## Related issue number

closes #5580



Problem statement:
intermittent bug: Exception thrown: 'System.InvalidOperationException' in Orleans.Runtime.dll: 'Activation access violation. A non-activation thread attempted to access activation services.'
### What happened?

```Exception has occurred: CLR/System.InvalidOperationException
Exception thrown: 'System.InvalidOperationException' in Orleans.Runtime.dll: 'Activation access violation. A non-activation thread attempted to access activation services.'
   at Orleans.Runtime.GrainRuntime.<CheckRuntimeContext>g__ThrowMissingContext|22_0()
   at Orleans.Runtime.GrainRuntime.CheckRuntimeContext(IGrainContext context)
   at Orleans.Core.StateStorageBridge`1.get_State()
   at Microsoft.AutoGen.RuntimeGateway.Grpc.MessageRegistryGrain.<RemoveMessage>d__7.MoveNext() in /Users/ryan/0206/autogen/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs:line 52
```

the use of TaskScheduler.Default on line 41 is incorrect and will cause this error. Needs to be TaskScheduler.Current

### What did you expect to happen?

no error

### How can we reproduce it (as minimally and precisely as possible)?

happens intermittently on message removal. 

### AutoGen version

0.4.x

### Which package was this bug in

Core

### Model used

_No response_

### Python version

_No response_

### Operating system

_No response_

### Any additional info you think would be helpful for fixing this bug

_No response_

Hints from issue discussion:

Commit message: with help from Reuben making the message registry write more robust.

Patch to review:
diff --git a/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs b/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs
index 59dec750da41..cc00353dcf35 100644
--- a/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs
+++ b/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs
@@ -13,9 +13,13 @@ ILogger<MessageRegistryGrain> logger
 ) : Grain, IMessageRegistryGrain
 {
     // <summary>
-    // The number of times to retry writing the state before giving up.
+    // Helper class for managing state writes.
     // </summary>
+    private readonly StateManager _stateManager = new(state);
 
+    // <summary>
+    // The number of times to retry writing the state before giving up.
+    // </summary>
     private const int _retries = 5;
     /// <summary>
     /// The time to wait before removing a message from the event buffer.
@@ -38,7 +42,7 @@ public async Task AddMessageToEventBufferAsync(string topic, CloudEvent message)
         _ = Task.Delay(_bufferTime)
             .ContinueWith(
                 async _ => await RemoveMessage(topic, message),
-                TaskScheduler.Default
+                TaskScheduler.Current
             );
     }
     /// <summary>
@@ -54,7 +58,7 @@ private async ValueTask<bool> RemoveMessage(string topic, CloudEvent message)
             if (events != null && events.Remove(message))
             {
                 state.State.EventBuffer.AddOrUpdate(topic, events, (_, _) => events);
-                await state.WriteStateAsync().ConfigureAwait(true);
+                await _stateManager.WriteStateAsync().ConfigureAwait(true);
                 return true;
             }
         }
@@ -114,7 +118,7 @@ private async ValueTask<bool> WriteMessageAsync(string whichQueue, string topic,
             default:
                 throw new ArgumentException($"Invalid queue name: {whichQueue}");
         }
-        await state.WriteStateAsync().ConfigureAwait(true);
+        await _stateManager.WriteStateAsync().ConfigureAwait(true);
         return true;
     }
 
@@ -124,7 +128,7 @@ public async Task<List<CloudEvent>> RemoveMessagesAsync(string topic)
         var messages = new List<CloudEvent>();
         if (state.State.DeadLetterQueue != null && state.State.DeadLetterQueue.Remove(topic, out List<CloudEvent>? letters))
         {
-            await state.WriteStateAsync().ConfigureAwait(true);
+            await _stateManager.WriteStateAsync().ConfigureAwait(true);
             if (letters != null)
             {
                 messages.AddRange(letters);
@@ -132,7 +136,7 @@ public async Task<List<CloudEvent>> RemoveMessagesAsync(string topic)
         }
         if (state.State.EventBuffer != null && state.State.EventBuffer.Remove(topic, out List<CloudEvent>? events))
         {
-            await state.WriteStateAsync().ConfigureAwait(true);
+            await _stateManager.WriteStateAsync().ConfigureAwait(true);
             if (events != null)
             {
                 messages.AddRange(events);
diff --git a/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/StateManager.cs b/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/StateManager.cs
new file mode 100644
index 000000000000..7bce36845d13
--- /dev/null
+++ b/dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/StateManager.cs
@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// StateManager.cs
+
+using Orleans.Core;
+
+namespace Microsoft.AutoGen.RuntimeGateway.Grpc;
+
+/// <summary>
+/// A helper class which wraps a grain state instance and ensures that only a single write operation is outstanding at any moment in time.
+/// </summary>
+/// <param name="state">The grain state.</param>
+internal sealed class StateManager(IStorage state)
+{
+    /// <summary>
+    /// Allows state writing to happen in the background.
+    /// </summary>
+    private Task? _pendingOperation;
+
+    // When reentrant grain is doing WriteStateAsync, etag violations are possible due to concurrent writes.
+    // The solution is to serialize and batch writes, and make sure only a single write is outstanding at any moment in time.
+    public async ValueTask WriteStateAsync()
+    {
+        await PerformOperationAsync(static state => state.WriteStateAsync());
+    }
+
+    public async ValueTask ClearStateAsync()
+    {
+        await PerformOperationAsync(static state => state.ClearStateAsync());
+    }
+
+    public async ValueTask PerformOperationAsync(Func<IStorage, Task> performOperation)
+    {
+        if (_pendingOperation is Task currentWriteStateOperation)
+        {
+            // await the outstanding write, but ignore it since it doesn't include our changes
+            await currentWriteStateOperation.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing | ConfigureAwaitOptions.ContinueOnCapturedContext);
+            if (_pendingOperation == currentWriteStateOperation)
+            {
+                // only null out the outstanding operation if it's the same one as the one we awaited, otherwise
+                // another request might have already done so.
+                _pendingOperation = null;
+            }
+        }
+
+        Task operation;
+        if (_pendingOperation is null)
+        {
+            // If after the initial write is completed, no other request initiated a new write operation, do it now.
+            operation = performOperation(state);
+            _pendingOperation = operation;
+        }
+        else
+        {
+            // If there were many requests enqueued to persist state, there is no reason to enqueue a new write 
+            // operation for each, since any write (after the initial one that we already awaited) will have cumulative
+            // changes including the one requested by our caller. Just await the new outstanding write.
+            operation = _pendingOperation;
+        }
+
+        try
+        {
+            await operation;
+        }
+        finally
+        {
+            if (_pendingOperation == operation)
+            {
+                // only null out the outstanding operation if it's the same one as the one we awaited, otherwise
+                // another request might have already done so.
+                _pendingOperation = null;
+            }
+        }
+    }
+}
+ 
\ No newline at end of file


Code review comments:
diff_hunk: @@ -38,7 +42,7 @@ public async Task AddMessageToEventBufferAsync(string topic, CloudEvent message)
         _ = Task.Delay(_bufferTime)
             .ContinueWith(
                 async _ => await RemoveMessage(topic, message),
-                TaskScheduler.Default
+                TaskScheduler.Current
             );
path: dotnet/src/Microsoft.AutoGen/RuntimeGateway.Grpc/Services/Orleans/MessageRegistryGrain.cs
text: @user1:
It's best to `.Ignore()` tasks so that they don't trigger [`TaskScheduler.UnobservedTaskException`](https://learn.microsoft.com/en-us/dotnet/api/system.threading.tasks.taskscheduler.unobservedtaskexception?view=net-9.0) events

I would change this to remove `ContinueWith` and use `await` instead. You can use a local function:

```csharp
// Schedule removal of the message.
RemoveMessageAfterDelay().Ignore();

async Task RemoveMessageAfterDelay()
{
    await Task.Delay(_bufferTime)
    await RemoveMessage(topic, message);
}
```