Add support for MayInterleave in StatelessWorker (#9050)

Co-authored-by: Ledjon Behluli <Ledjon@notiphy.io>

src/Orleans.Runtime/Catalog/StatelessWorkerGrainContext.cs

@@ -103,7 +103,15 @@ public async ValueTask DisposeAsync()
             _ => _shared.GetComponent<TComponent>()
         };
 
-        public void SetComponent<TComponent>(TComponent instance) where TComponent : class => throw new ArgumentException($"Cannot set a component on a {nameof(StatelessWorkerGrainContext)}");
+        public void SetComponent<TComponent>(TComponent instance) where TComponent : class
+        {
+            if (typeof(TComponent) != typeof(GrainCanInterleave))
+            {
+                throw new ArgumentException($"Cannot set a component of type '{instance.GetType().FullName}' on a {nameof(StatelessWorkerGrainContext)}");
+            }
+
+            _shared.SetComponent(instance);
+        }
 
         public TTarget GetTarget<TTarget>() where TTarget : class => throw new NotImplementedException();
 

test/DefaultCluster.Tests/StatelessWorkerTests.cs

@@ -130,5 +130,44 @@ public async Task StatelessWorkerFastActivationsDontFailInMultiSiloDeployment()
 
             // Calls should not have thrown ForwardingFailed exceptions.
         }
+
+        [Fact, TestCategory("SlowBVT"), TestCategory("StatelessWorker")]
+        public async Task StatelessWorker_DoesNotThrow_IfMarkedWithMayInterleave()
+        {
+            var grain = GrainFactory.GetGrain<IStatelessWorkerWithMayInterleaveGrain>(0);
+            var exception = await Record.ExceptionAsync(grain.GoFast);
+
+            Assert.Null(exception);
+        }
+
+        [Fact, TestCategory("SlowBVT"), TestCategory("StatelessWorker")]
+        public async Task StatelessWorker_ShouldNotInterleaveCalls_IfMayInterleavePredicatedDoesntMatch()
+        {
+            var grain = GrainFactory.GetGrain<IStatelessWorkerWithMayInterleaveGrain>(0);
+
+            var delay = TimeSpan.FromSeconds(1);
+            await grain.SetDelay(delay);
+
+            var stopwatch = Stopwatch.StartNew();
+            await Task.WhenAll(grain.GoSlow(), grain.GoSlow(), grain.GoSlow());
+            stopwatch.Stop();
+
+            Assert.InRange(stopwatch.Elapsed.TotalSeconds, 2.85, 3.15); // theoretically it should be 3.0
+        }
+
+        [Fact, TestCategory("SlowBVT"), TestCategory("StatelessWorker")]
+        public async Task StatelessWorker_ShouldInterleaveCalls_IfMayInterleavePredicatedMatches()
+        {
+            var grain = GrainFactory.GetGrain<IStatelessWorkerWithMayInterleaveGrain>(0);
+
+            var delay = TimeSpan.FromSeconds(1);
+            await grain.SetDelay(delay);
+
+            var stopwatch = Stopwatch.StartNew();
+            await Task.WhenAll(grain.GoFast(), grain.GoFast(), grain.GoFast());
+            stopwatch.Stop();
+
+            Assert.InRange(stopwatch.Elapsed.TotalSeconds, 0.85, 1.15); // theoretically it should be 1.0
+        }
     }
-}
+}

test/Grains/TestGrainInterfaces/IStatelessWorkerWithMayInterleaveGrain.cs

@@ -0,0 +1,8 @@
+namespace UnitTests.GrainInterfaces;
+
+public interface IStatelessWorkerWithMayInterleaveGrain : IGrainWithIntegerKey
+{
+    Task SetDelay(TimeSpan delay);
+    Task GoSlow();
+    Task GoFast();
+}

test/Grains/TestGrains/StatelessWorkerWithMayInterleaveGrain.cs

@@ -0,0 +1,23 @@
+using Orleans.Concurrency;
+using Orleans.Serialization.Invocation;
+using UnitTests.GrainInterfaces;
+
+namespace UnitTests.Grains;
+
+[StatelessWorker(1)] // '1' to force interleaving, otherwise it just creates a new worker
+[MayInterleave(nameof(MayInterleaveMethod))]
+public class StatelessWorkerWithMayInterleaveGrain : Grain, IStatelessWorkerWithMayInterleaveGrain
+{
+    private TimeSpan _delay = TimeSpan.FromMilliseconds(1);
+
+    public static bool MayInterleaveMethod(IInvokable req) => req.GetMethodName() == nameof(GoFast);
+
+    public Task SetDelay(TimeSpan delay)
+    {
+        _delay = delay;
+        return Task.CompletedTask;
+    }
+
+    public Task GoSlow() => Task.Delay(_delay);
+    public Task GoFast() => Task.Delay(_delay);
+}