refactor(runtime): start grain activation synchronously

src/Orleans.Runtime/Activation/ActivationDataActivatorProvider.cs

@@ -49,14 +49,13 @@ public bool TryGet(GrainType grainType, [NotNullWhen(true)] out IGrainContextAct
         return true;
     }
 
-    private partial class ActivationDataActivator : IGrainContextActivator
+    private partial class ActivationDataActivator : IPreparedGrainContextActivator
     {
         private readonly IOptions<SchedulingOptions> _schedulingOptions;
         private readonly IGrainActivator _grainActivator;
         private readonly IServiceProvider _serviceProvider;
         private readonly GrainTypeSharedContext _sharedComponents;
         private readonly Func<IGrainContext, WorkItemGroup> _createWorkItemGroup;
-        private readonly Action<object?> _startActivation;
 
         public ActivationDataActivator(
             IGrainActivator grainActivator,
@@ -73,30 +72,40 @@ public ActivationDataActivator(
                 context,
                 _schedulingOptions,
                 schedulerInstruments);
-            _startActivation = state => ((ActivationData)state!).Start(_grainActivator);
         }
 
         public IGrainContext CreateContext(GrainAddress activationAddress, IConfigureGrainContext[] configureActions)
+        {
+            var preparedContext = CreatePreparedContext(activationAddress, configureActions);
+            using var startup = preparedContext.Start();
+            return preparedContext.Context;
+        }
+
+        public PreparedGrainContext CreatePreparedContext(
+            GrainAddress activationAddress,
+            IConfigureGrainContext[] configureActions)
         {
             var context = new ActivationData(
                 activationAddress,
                 _createWorkItemGroup,
                 _serviceProvider,
-                _sharedComponents);
+                _sharedComponents,
+                _grainActivator);
 
-            foreach (var configure in configureActions)
+            try
             {
-                configure.Configure(context);
-            }
+                foreach (var configure in configureActions)
+                {
+                    configure.Configure(context);
+                }
 
-            using var ecSuppressor = ExecutionContext.SuppressFlow();
-            _ = Task.Factory.StartNew(
-                _startActivation,
-                context,
-                CancellationToken.None,
-                TaskCreationOptions.DenyChildAttach,
-                context.ActivationTaskScheduler);
-            return context;
+                return new(context, (IGrainContextStartup)context);
+            }
+            catch
+            {
+                ((IGrainContextStartup)context).Abort();
+                throw;
+            }
         }
     }
 }

src/Orleans.Runtime/Activation/IGrainContextActivator.cs

@@ -57,14 +57,21 @@ public GrainContextActivator(
         /// <param name="address">The grain address.</param>
         /// <returns>The grain context.</returns>
         public IGrainContext CreateInstance(GrainAddress address)
+        {
+            var preparedContext = CreatePreparedContext(address);
+            using var startup = preparedContext.Start();
+            return preparedContext.Context;
+        }
+
+        internal PreparedGrainContext CreatePreparedContext(GrainAddress address)
         {
             var grainId = address.GrainId;
             if (!_activators.TryGetValue(grainId.Type, out var activator))
             {
                 activator = this.CreateActivator(grainId.Type);
             }
 
-            return activator.Activator.CreateContext(address, activator.ConfigureActions);
+            return PreparedGrainContext.Create(activator.Activator, address, activator.ConfigureActions);
         }
 
         private (IGrainContextActivator, IConfigureGrainContext[]) CreateActivator(GrainType grainType)
@@ -134,6 +141,79 @@ public interface IGrainContextActivator
         public IGrainContext CreateContext(GrainAddress address, IConfigureGrainContext[] configureActions);
     }
 
+    internal interface IPreparedGrainContextActivator : IGrainContextActivator
+    {
+        PreparedGrainContext CreatePreparedContext(
+            GrainAddress address,
+            IConfigureGrainContext[] configureActions);
+    }
+
+    internal readonly struct PreparedGrainContext(IGrainContext context, IGrainContextStartup? startup)
+    {
+        private readonly IGrainContext? _context = context;
+        private readonly IGrainContextStartup? _startup = startup;
+
+        public IGrainContext Context
+            => _context ?? throw new InvalidOperationException("The grain context activation is not initialized.");
+
+        public bool HasStartup => _startup is not null;
+
+        public static PreparedGrainContext Create(
+            IGrainContextActivator activator,
+            GrainAddress address,
+            IConfigureGrainContext[] configureActions) =>
+            activator is IPreparedGrainContextActivator preparedActivator
+                ? preparedActivator.CreatePreparedContext(address, configureActions)
+                : new(activator.CreateContext(address, configureActions), startup: null);
+
+        public IDisposable Start()
+        {
+            if (_startup is not { } startup)
+            {
+                return NoopDisposable.Instance;
+            }
+
+            try
+            {
+                return startup.Start();
+            }
+            catch (Exception startupException)
+            {
+                try
+                {
+                    startup.Abort();
+                }
+                catch (Exception abortException)
+                {
+                    throw new AggregateException(
+                        "Grain context startup failed and aborting the startup also failed.",
+                        startupException,
+                        abortException);
+                }
+
+                throw;
+            }
+        }
+
+        public void Abort() => _startup?.Abort();
+
+        private sealed class NoopDisposable : IDisposable
+        {
+            public static NoopDisposable Instance { get; } = new();
+
+            public void Dispose()
+            {
+            }
+        }
+    }
+
+    internal interface IGrainContextStartup
+    {
+        IDisposable Start();
+
+        void Abort();
+    }
+
     /// <summary>
     /// Provides a <see cref="IConfigureGrainContext"/> instance for the provided grain type.
     /// </summary>