perf(grain): move runtime access to grain context

playground/ActivationRepartitioning/ActivationRepartitioning.Frontend/Program.cs

@@ -97,7 +97,7 @@ public async ValueTask Reset()
     {
         ++_resetCount;
         _numForests = 0;
-        await ServiceProvider.GetRequiredService<ClusterDiagnosticsService>().ResetAsync();
+        await ServiceProvider!.GetRequiredService<ClusterDiagnosticsService>().ResetAsync();
         await GrainFactory.GetGrain<IManagementGrain>(0).ResetGrainCallFrequencies();
     }
 

src/Orleans.Core.Abstractions/Core/Grain.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection;
@@ -14,6 +15,8 @@ namespace Orleans;
 /// </summary>
 public abstract partial class Grain : IGrainBase, IAddressable
 {
+    private static readonly ConditionalWeakTable<Grain, IGrainRuntime> StandaloneRuntimes = new();
+
     // Do not use this directly because we currently don't provide a way to inject it;
     // any interaction with it will result in non unit-testable code. Any behavior that can be accessed
     // from within client code (including subclasses of this class), should be exposed through IGrainRuntime.
@@ -23,18 +26,24 @@ public abstract partial class Grain : IGrainBase, IAddressable
 
     public GrainReference GrainReference { get { return GrainContext.GrainReference; } }
 
-    internal IGrainRuntime Runtime { get; }
+    private IGrainRuntime? RuntimeOrDefault => GrainContext?.GrainRuntime
+        ?? (StandaloneRuntimes.TryGetValue(this, out var runtime) ? runtime : null);
+
+    internal IGrainRuntime Runtime => RuntimeOrDefault
+        ?? throw new InvalidOperationException("Grain was created outside of the Orleans creation process and no runtime was specified.");
 
     /// <summary>
-    /// Gets an object which can be used to access other grains. Null if this grain is not associated with a Runtime, such as when created directly for unit testing.
+    /// Gets an object which can be used to access other grains.
     /// </summary>
+    /// <exception cref="InvalidOperationException">
+    /// The grain was created outside of the Orleans activation process and no runtime was provided.
+    /// </exception>
     protected IGrainFactory GrainFactory => Runtime.GrainFactory;
 
     /// <summary>
     /// Gets the IServiceProvider managed by the runtime. Null if this grain is not associated with a Runtime, such as when created directly for unit testing.
     /// </summary>
-    // ! The runtime ensures that this is not null and Unit testing frameworks must make sure that this is not null.
-    protected internal IServiceProvider ServiceProvider => GrainContext?.ActivationServices ?? Runtime?.ServiceProvider!;
+    protected internal IServiceProvider? ServiceProvider => GrainContext?.ActivationServices ?? RuntimeOrDefault?.ServiceProvider;
 
     internal GrainId GrainId => GrainContext.GrainId;
 
@@ -51,12 +60,24 @@ protected Grain() : this(RuntimeContext.Current!, grainRuntime: null)
     /// This constructor is particularly useful for unit testing where test code can create a Grain and replace
     /// the IGrainIdentity and IGrainRuntime with test doubles (mocks/stubs).
     /// </summary>
+    /// <remarks>
+    /// When <paramref name="grainRuntime"/> is provided, it is associated with this grain and registered as an
+    /// <see cref="IGrainRuntime"/> component when <paramref name="grainContext"/> is available.
+    /// </remarks>
     protected Grain(IGrainContext grainContext, IGrainRuntime? grainRuntime = null)
     {
         GrainContext = grainContext;
-
-        // ! The runtime ensures that this is not null and Unit testing frameworks must make sure that this is not null.
-        Runtime = grainRuntime ?? grainContext?.ActivationServices.GetService<IGrainRuntime>()!;
+        if (grainRuntime is not null)
+        {
+            if (grainContext is null)
+            {
+                StandaloneRuntimes.Add(this, grainRuntime);
+            }
+            else
+            {
+                grainContext.GrainRuntime = grainRuntime;
+            }
+        }
     }
 
     /// <summary>
@@ -68,7 +89,7 @@ protected Grain(IGrainContext grainContext, IGrainRuntime? grainRuntime = null)
     /// A unique identifier for the current silo.
     /// There is no semantic content to this string, but it may be useful for logging.
     /// </summary>
-    public string RuntimeIdentity => Runtime?.SiloIdentity ?? string.Empty;
+    public string RuntimeIdentity => RuntimeOrDefault?.SiloIdentity ?? string.Empty;
 
     /// <summary>
     /// Registers a timer to send periodic callbacks to this grain.
@@ -163,13 +184,7 @@ protected void DelayDeactivation(TimeSpan timeSpan)
     /// <param name="cancellationToken">A cancellation token which signals when deactivation should complete promptly.</param>
     public virtual Task OnDeactivateAsync(DeactivationReason reason, CancellationToken cancellationToken) => Task.CompletedTask;
 
-    internal void EnsureRuntime()
-    {
-        if (Runtime == null)
-        {
-            throw new InvalidOperationException("Grain was created outside of the Orleans creation process and no runtime was specified.");
-        }
-    }
+    internal void EnsureRuntime() => _ = Runtime;
 }
 
 /// <summary>

src/Orleans.Core.Abstractions/Core/IGrainContext.cs

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Orleans.Serialization.Invocation;
 
 namespace Orleans.Runtime
@@ -41,6 +42,20 @@ public interface IGrainContext : ITargetHolder, IEquatable<IGrainContext>
         /// </summary>
         IServiceProvider ActivationServices { get; }
 
+        /// <summary>
+        /// Gets or sets the grain runtime associated with this context.
+        /// </summary>
+        /// <remarks>
+        /// Runtime contexts provide the activation runtime directly. Custom contexts resolve an
+        /// <see cref="IGrainRuntime"/> component before resolving an activation service.
+        /// </remarks>
+        IGrainRuntime? GrainRuntime
+        {
+            get => GetComponent(typeof(IGrainRuntime)) as IGrainRuntime
+                ?? ActivationServices?.GetService<IGrainRuntime>();
+            set => SetComponent(value);
+        }
+
         /// <summary>
         /// Gets the observable <see cref="Grain"/> lifecycle, which can be used to add lifecycle hooks.
         /// </summary>

src/Orleans.EventSourcing/LogConsistency/LogConsistentGrain.cs

@@ -55,11 +55,12 @@ private async Task OnDeactivateState(CancellationToken ct)
         private Task OnSetupState(CancellationToken ct)
         {
             if (ct.IsCancellationRequested) return Task.CompletedTask;
-            IGrainContextAccessor grainContextAccessor = this.ServiceProvider.GetRequiredService<IGrainContextAccessor>();
-            Factory<IGrainContext, ILogConsistencyProtocolServices> protocolServicesFactory = this.ServiceProvider.GetRequiredService<Factory<IGrainContext, ILogConsistencyProtocolServices>>();
+            var serviceProvider = ServiceProvider!;
+            IGrainContextAccessor grainContextAccessor = serviceProvider.GetRequiredService<IGrainContextAccessor>();
+            Factory<IGrainContext, ILogConsistencyProtocolServices> protocolServicesFactory = serviceProvider.GetRequiredService<Factory<IGrainContext, ILogConsistencyProtocolServices>>();
             IGrainContext grainContext = grainContextAccessor.GrainContext!;
             ILogViewAdaptorFactory consistencyProvider = SetupLogConsistencyProvider(grainContext);
-            IGrainStorage? grainStorage = consistencyProvider.UsesStorageProvider ? GrainStorageHelpers.GetGrainStorage((grainContext?.GrainInstance!.GetType())!, this.ServiceProvider) : null;
+            IGrainStorage? grainStorage = consistencyProvider.UsesStorageProvider ? GrainStorageHelpers.GetGrainStorage((grainContext?.GrainInstance!.GetType())!, serviceProvider) : null;
             InstallLogViewAdaptor(grainContext!, protocolServicesFactory, consistencyProvider, grainStorage);
             return Task.CompletedTask;
         }
@@ -91,10 +92,11 @@ private void InstallLogViewAdaptor(
         private ILogViewAdaptorFactory SetupLogConsistencyProvider(IGrainContext activationContext)
         {
             var attr = this.GetType().GetCustomAttributes<LogConsistencyProviderAttribute>(true).FirstOrDefault();
+            var serviceProvider = ServiceProvider!;
 
             ILogViewAdaptorFactory? defaultFactory = attr != null
-                ? this.ServiceProvider.GetKeyedService<ILogViewAdaptorFactory>(attr.ProviderName)
-                : this.ServiceProvider.GetService<ILogViewAdaptorFactory>();
+                ? serviceProvider.GetKeyedService<ILogViewAdaptorFactory>(attr.ProviderName)
+                : serviceProvider.GetService<ILogViewAdaptorFactory>();
             if (attr != null && defaultFactory == null)
             {
                 var errMsg = $"Cannot find consistency provider with Name={attr.ProviderName} for grain type {this.GetType().FullName}";

src/Orleans.Journaling/DurableGrain.cs

@@ -6,7 +6,7 @@ public abstract class DurableGrain : Grain, IGrainBase
 {
     protected DurableGrain()
     {
-        StateManager = ServiceProvider.GetRequiredService<IJournaledStateManager>();
+        StateManager = ServiceProvider!.GetRequiredService<IJournaledStateManager>();
         if (StateManager is ILifecycleParticipant<IGrainLifecycle> participant)
         {
             participant.Participate(((IGrainBase)this).GrainContext.ObservableLifecycle);
@@ -16,7 +16,7 @@ protected DurableGrain()
     protected IJournaledStateManager StateManager { get; }
 
     protected TState GetOrCreateState<TState>(string name) where TState : class, IJournaledState
-        => GetOrCreateState(name, static sp => sp.GetRequiredService<TState>(), ServiceProvider);
+        => GetOrCreateState(name, static sp => sp.GetRequiredService<TState>(), ServiceProvider!);
 
     protected TState GetOrCreateState<TArg, TState>(string name, Func<TArg, TState> createState, TArg arg) where TState : class, IJournaledState
     {

src/Orleans.Runtime/Catalog/ActivationData.cs

@@ -142,7 +142,20 @@ public void Start(IGrainActivator grainActivator)
     }
 
     public ActivationTaskScheduler ActivationTaskScheduler => _workItemGroup.TaskScheduler;
-    public IGrainRuntime GrainRuntime => _shared.Runtime;
+    public IGrainRuntime GrainRuntime
+    {
+        get
+        {
+            var extras = Volatile.Read(ref _extras);
+            return extras is null ? _shared.Runtime : Volatile.Read(ref extras.GrainRuntime) ?? _shared.Runtime;
+        }
+    }
+
+    IGrainRuntime? IGrainContext.GrainRuntime
+    {
+        get => GrainRuntime;
+        set => SetComponent(value);
+    }
     public object? GrainInstance { get; private set; }
     public GrainAddress Address { get; private set; }
     public GrainReference GrainReference => _selfReference ??= _shared.GrainReferenceActivator.CreateReference(GrainId, default);
@@ -327,6 +340,10 @@ private DehydrationContextHolder? DehydrationContext
         {
             result = this;
         }
+        else if (componentType == typeof(IGrainRuntime))
+        {
+            result = GrainRuntime;
+        }
         else if (_extras is { } components && components.TryGetValue(componentType, out var resultObj))
         {
             result = resultObj;
@@ -363,6 +380,25 @@ public void SetComponent(Type componentType, object? instance)
 
         lock (this)
         {
+            if (componentType == typeof(IGrainRuntime))
+            {
+                var extras = Volatile.Read(ref _extras);
+                if (extras is null)
+                {
+                    if (instance is null)
+                    {
+                        return;
+                    }
+
+                    extras = new() { GrainRuntime = (IGrainRuntime)instance };
+                    Volatile.Write(ref _extras, extras);
+                    return;
+                }
+
+                Volatile.Write(ref extras.GrainRuntime, (IGrainRuntime?)instance);
+                return;
+            }
+
             if (instance == null)
             {
                 _extras?.Remove(componentType);
@@ -2318,6 +2354,8 @@ private class ActivationDataExtra : Dictionary<object, object>
         private const int IsDisposingFlag = 1 << 2;
         private byte _flags;
 
+        public IGrainRuntime? GrainRuntime;
+
         public HashSet<IGrainTimer>? Timers { get => GetValueOrDefault<HashSet<IGrainTimer>>(nameof(Timers)); set => SetOrRemoveValue(nameof(Timers), value); }
 
         /// <summary>

src/Orleans.Runtime/Catalog/StatelessWorkerGrainContext.cs

@@ -77,6 +77,20 @@ public StatelessWorkerGrainContext(
 
     public GrainAddress Address { get; }
 
+    public IGrainRuntime GrainRuntime => _shared.Shared.Runtime;
+
+    IGrainRuntime? IGrainContext.GrainRuntime
+    {
+        get => GrainRuntime;
+        set
+        {
+            if (!ReferenceEquals(value, GrainRuntime))
+            {
+                throw new ArgumentException("The runtime for a stateless worker context is provided by its shared grain type context.", nameof(value));
+            }
+        }
+    }
+
     public IServiceProvider ActivationServices => throw new NotImplementedException();
 
     public IGrainLifecycle ObservableLifecycle => throw new NotImplementedException();

src/Orleans.TestingHost/TestStorageProviders/StorageFaultGrain.cs

@@ -23,7 +23,7 @@ public partial class StorageFaultGrain : Grain, IStorageFaultGrain
         public override async Task OnActivateAsync(CancellationToken cancellationToken)
         {
             await base.OnActivateAsync(cancellationToken);
-            logger = this.ServiceProvider.GetService<ILoggerFactory>()!.CreateLogger($"{typeof (StorageFaultGrain).FullName}-{IdentityString}-{RuntimeIdentity}");
+            logger = ServiceProvider!.GetService<ILoggerFactory>()!.CreateLogger($"{typeof (StorageFaultGrain).FullName}-{IdentityString}-{RuntimeIdentity}");
             readFaults = new();
             writeFaults = new();
             clearfaults = new();

src/api/Orleans.Core.Abstractions/Orleans.Core.Abstractions.cs

@@ -1071,7 +1071,7 @@ protected Grain(Runtime.IGrainContext grainContext, Runtime.IGrainRuntime? grain
 
         public string RuntimeIdentity { get { throw null; } }
 
-        protected internal System.IServiceProvider ServiceProvider { get { throw null; } }
+        protected internal System.IServiceProvider? ServiceProvider { get { throw null; } }
 
         protected void DeactivateOnIdle() { }
 
@@ -2555,6 +2555,8 @@ public partial interface IGrainContext : Orleans.Serialization.Invocation.ITarge
 
         GrainReference GrainReference { get; }
 
+        IGrainRuntime? GrainRuntime { get; set; }
+
         IGrainLifecycle ObservableLifecycle { get; }
 
         IWorkItemScheduler Scheduler { get; }

test/Grains/TestGrains/MessageSerializationGrain.cs

@@ -59,8 +59,9 @@ private async Task<IMessageSerializationGrain> GetGrainOnOtherSilo()
             IMessageSerializationGrain otherGrain;
             var id = this.GetPrimaryKeyLong();
             var currentSiloIdentity = await this.GetSiloIdentity();
-            var silos = ServiceProvider.GetRequiredService<IClusterMembershipService>().CurrentSnapshot.Members.Where(kv => kv.Value.Status == SiloStatus.Active).Select(kv => kv.Key).ToHashSet();
-            silos.Remove(ServiceProvider.GetRequiredService<ILocalSiloDetails>().SiloAddress);
+            var serviceProvider = ServiceProvider!;
+            var silos = serviceProvider.GetRequiredService<IClusterMembershipService>().CurrentSnapshot.Members.Where(kv => kv.Value.Status == SiloStatus.Active).Select(kv => kv.Key).ToHashSet();
+            silos.Remove(serviceProvider.GetRequiredService<ILocalSiloDetails>().SiloAddress);
             while (true)
             {
                 RequestContext.Set(IPlacementDirector.PlacementHintKey, silos.First());

test/Grains/TestGrains/ProgrammaticSubscribe/SubscribeGrain.cs

@@ -14,7 +14,7 @@ public class SubscribeGrain : Grain, ISubscribeGrain
     {
         public Task<bool> CanGetSubscriptionManager(string providerName)
         {
-            return Task.FromResult(this.ServiceProvider.GetKeyedService<IStreamProvider>(providerName)!.TryGetStreamSubscriptionManager(out _)); // The test configures the named provider before invoking this grain.
+            return Task.FromResult(ServiceProvider!.GetKeyedService<IStreamProvider>(providerName)!.TryGetStreamSubscriptionManager(out _)); // The test configures the named provider before invoking this grain.
         }
     }
 

test/Grains/TestGrains/SimpleDIGrain.cs

@@ -65,9 +65,10 @@ public Task DoDeactivate()
 
         public Task AssertCanResolveSameServiceInstances()
         {
-            if (!ReferenceEquals(this.ServiceProvider.GetRequiredService<IInjectedService>(), this.injectedService)) throw new Exception("singleton not equal");
-            if (!ReferenceEquals(this.ServiceProvider.GetRequiredService<IInjectedScopedService>(), this.injectedScopedService)) throw new Exception("scoped not equal");
-            if (!ReferenceEquals(this.ServiceProvider.GetRequiredService<IGrainContextAccessor>().GrainContext, this.originalGrainContext)) throw new Exception("scoped grain activation context not equal");
+            var serviceProvider = ServiceProvider!;
+            if (!ReferenceEquals(serviceProvider.GetRequiredService<IInjectedService>(), this.injectedService)) throw new Exception("singleton not equal");
+            if (!ReferenceEquals(serviceProvider.GetRequiredService<IInjectedScopedService>(), this.injectedScopedService)) throw new Exception("scoped not equal");
+            if (!ReferenceEquals(serviceProvider.GetRequiredService<IGrainContextAccessor>().GrainContext, this.originalGrainContext)) throw new Exception("scoped grain activation context not equal");
 
             return Task.CompletedTask;
         }

test/Grains/TestInternalGrains/CollectionTestGrain.cs

@@ -29,7 +29,7 @@ protected virtual ILogger Logger()
 
         public override Task OnActivateAsync(CancellationToken cancellationToken)
         {
-            logger = this.ServiceProvider.GetRequiredService<ILoggerFactory>()
+            logger = ServiceProvider!.GetRequiredService<ILoggerFactory>()
                 .CreateLogger(string.Format("CollectionTestGrain {0} {1} on {2}.", GrainId, _grainContext.ActivationId, RuntimeIdentity));
             logger.LogInformation("OnActivateAsync.");
             activated = DateTime.UtcNow;
@@ -127,7 +127,7 @@ protected override ILogger Logger()
 
         public override Task OnActivateAsync(CancellationToken cancellationToken)
         {
-            logger = this.ServiceProvider.GetRequiredService<ILoggerFactory>()
+            logger = ServiceProvider!.GetRequiredService<ILoggerFactory>()
                 .CreateLogger($"CollectionTestGrain {GrainId} {_grainContext.ActivationId} on {RuntimeIdentity}.");
             logger.LogInformation("OnActivateAsync.");
             counter = 0;

test/Grains/TestInternalGrains/PersistenceTestGrains.cs

@@ -67,7 +67,7 @@ public Task<bool> CheckStateInit()
 
         public Task<string> CheckProviderType()
         {
-            IGrainStorage grainStorage = GrainStorageHelpers.GetGrainStorage(GetType(), this.ServiceProvider);
+            IGrainStorage grainStorage = GrainStorageHelpers.GetGrainStorage(GetType(), ServiceProvider!);
             Assert.NotNull(grainStorage);
             return Task.FromResult(grainStorage.GetType().FullName!);
         }

test/Grains/TestInternalGrains/PlacementTestGrain.cs

@@ -252,7 +252,7 @@ internal class DefaultPlacementGrain : Grain, IDefaultPlacementGrain
     {
         public Task<PlacementStrategy> GetDefaultPlacement()
         {
-            var defaultStrategy = this.ServiceProvider.GetRequiredService<PlacementStrategy>();
+            var defaultStrategy = ServiceProvider!.GetRequiredService<PlacementStrategy>();
             return Task.FromResult(defaultStrategy);
         }
     }

test/Grains/TestInternalGrains/SerializerPresenceTestGrain.cs

@@ -8,7 +8,7 @@ internal class SerializerPresenceTestGrain : Grain, ISerializerPresenceTest
     {
         public Task<bool> SerializerExistsForType(Type t)
         {
-            return Task.FromResult(this.ServiceProvider.GetRequiredService<Serializer>().CanSerialize(t));
+            return Task.FromResult(ServiceProvider!.GetRequiredService<Serializer>().CanSerialize(t));
         }
 
         public Task TakeSerializedData(object data)