wip

src/Microsoft.DotNet.Interactive.Tests/KernelCommandSchedulerTests.cs → src/Microsoft.DotNet.Interactive.Tests/KernelSchedulerTests.cs

@@ -18,11 +18,11 @@
 
 namespace Microsoft.DotNet.Interactive.Tests
 {
-    public class KernelCommandSchedulerTests : IDisposable
+    public class KernelSchedulerTests : IDisposable
     {
         private readonly CompositeDisposable _disposables = new();
 
-        public KernelCommandSchedulerTests(ITestOutputHelper output)
+        public KernelSchedulerTests(ITestOutputHelper output)
         {
             DisposeAfterTest(output.SubscribeToPocketLogger());
         }
@@ -35,7 +35,7 @@ public void Dispose()
             }
             catch (Exception ex)
             {
-                Logger<KernelCommandSchedulerTests>.Log.Error(exception: ex);
+                Logger<KernelSchedulerTests>.Log.Error(exception: ex);
             }
         }
 

src/Microsoft.DotNet.Interactive/CompositeKernel.cs

@@ -73,6 +73,7 @@ public void Add(Kernel kernel, IReadOnlyCollection<string> aliases = null)
             }
 
             kernel.ParentKernel = this;
+            kernel.SetScheduler(Scheduler);
             kernel.AddMiddleware(LoadExtensions);
 
             AddChooseKernelDirective(kernel, aliases);
@@ -202,38 +203,6 @@ private Kernel GetHandlingKernel(
             return kernel ?? this;
         }
 
-        internal override async Task HandleAsync(
-            KernelCommand command,
-            KernelInvocationContext context)
-        {
-            var kernel = context.HandlingKernel;
-
-            if (kernel is null)
-            {
-                throw new NoSuitableKernelException(command);
-            }
-
-            switch (command)
-            {
-                case Cancel _:
-                    CancelCommands();
-                    kernel.CancelCommands();
-                    break;
-            }
-
-            await kernel.RunDeferredCommandsAsync();
-
-            if (kernel != this)
-            {
-                // route to a subkernel
-                await kernel.Pipeline.SendAsync(command, context);
-            }
-            else
-            {
-                await base.HandleAsync(command, context);
-            }
-        }
-
         private protected override IEnumerable<Parser> GetDirectiveParsersForCompletion(
             DirectiveNode directiveNode, 
             int requestPosition)

src/Microsoft.DotNet.Interactive/Kernel.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.CommandLine;
 using System.CommandLine.Parsing;
@@ -24,11 +25,12 @@ public abstract partial class Kernel : IDisposable
     {
         private readonly Subject<KernelEvent> _kernelEvents = new();
         private readonly CompositeDisposable _disposables;
-
+        private readonly ConcurrentQueue<KernelCommand> _deferredCommands = new();
         private readonly Dictionary<Type, KernelCommandInvocation> _dynamicHandlers = new();
         private FrontendEnvironment _frontendEnvironment;
         private ChooseKernelDirective _chooseKernelDirective;
-        private readonly KernelCommandScheduler _scheduler;
+
+        private KernelScheduler<KernelCommand, KernelCommandResult> _commandScheduler = null;
 
         protected Kernel(string name)
         {
@@ -45,8 +47,7 @@ protected Kernel(string name)
 
             Pipeline = new KernelCommandPipeline(this);
 
-            _scheduler = new KernelCommandScheduler();
-
+
             AddSetKernelMiddleware();
 
             AddDirectiveMiddlewareAndCommonCommandHandlers();
@@ -56,7 +57,17 @@ protected Kernel(string name)
                 ));
         }
 
-        internal KernelCommandScheduler Scheduler => _scheduler;
+        internal KernelScheduler<KernelCommand, KernelCommandResult> Scheduler
+        {
+            get
+            {
+                if(_commandScheduler is null)
+                {
+                    SetScheduler(new ());
+                }
+                return _commandScheduler;
+            }
+        }
 
         internal KernelCommandPipeline Pipeline { get; }
 
@@ -76,7 +87,13 @@ public void DeferCommand(KernelCommand command)
             }
 
             command.SetToken($"deferredCommand::{Guid.NewGuid():N}");
-            Scheduler.DeferCommand(command, this);
+
+            if(string.IsNullOrWhiteSpace(command.TargetKernelName))
+            {
+                command.TargetKernelName = Name;
+            }
+
+            _deferredCommands.Enqueue(command);
         }
 
         private void AddSetKernelMiddleware()
@@ -235,35 +252,20 @@ internal virtual async Task HandleAsync(
         public Task<KernelCommandResult> SendAsync(
             KernelCommand command,
             CancellationToken cancellationToken)
-        {
-            return SendAsync(command, cancellationToken, null);
-        }
-
-        internal Task<KernelCommandResult> SendAsync(
-            KernelCommand command,
-            CancellationToken cancellationToken,
-            Action onDone)
         {
             if (command == null)
             {
                 throw new ArgumentNullException(nameof(command));
             }
 
-            return Scheduler.Schedule(command, this, cancellationToken, onDone);
-
+            return Scheduler.Schedule(command, OnExecuteAsync, Name);
+            
         }
 
         protected internal void CancelCommands()
         {
-            Scheduler.CancelCommands();
-        }
-
-        internal Task RunDeferredCommandsAsync()
-        {
-            return Scheduler.RunDeferredCommandsAsync(this);
-
+            Scheduler.Cancel();
         }
-
         protected internal void PublishEvent(KernelEvent kernelEvent)
         {
             if (kernelEvent == null)
@@ -436,5 +438,61 @@ protected virtual ChooseKernelDirective CreateChooseKernelDirective()
         }
 
         internal ChooseKernelDirective ChooseKernelDirective => _chooseKernelDirective ??= CreateChooseKernelDirective();
+
+        internal void SetScheduler(KernelScheduler<KernelCommand, KernelCommandResult> scheduler)
+        {
+
+            _commandScheduler = scheduler;
+
+            IEnumerable<KernelCommand> GetDeferredOperations(KernelCommand command, string scope)
+            {
+                if (scope != Name)
+                {
+                    yield break;
+                }
+
+                while (_deferredCommands.TryDequeue(out var kernelCommand))
+                {
+                    yield return kernelCommand;
+                }
+            }
+
+            Scheduler.RegisterDeferredOperationSource(GetDeferredOperations, OnExecuteAsync);
+        }
+
+        internal  async Task<KernelCommandResult> OnExecuteAsync(KernelCommand command)
+        {
+            var context = KernelInvocationContext.Establish(command);
+
+            // only subscribe for the root command 
+            using var _ = context.Command == command
+                ? context.KernelEvents.Subscribe(PublishEvent)
+                : Disposable.Empty;
+
+            try
+            {
+                await Pipeline.SendAsync(command, context);
+
+                if (command == context.Command)
+                {
+                    await context.DisposeAsync();
+                }
+                else
+                {
+                    context.Complete(command);
+                }
+
+                return context.Result;
+            }
+            catch (Exception exception)
+            {
+                if (!context.IsComplete)
+                {
+                    context.Fail(exception);
+                }
+
+                throw;
+            }
+        }
     }
 }

src/Microsoft.DotNet.Interactive/KernelCommandScheduler.cs

@@ -20,16 +20,11 @@ public class KernelCommandScheduler
 
         private readonly ConcurrentQueue<KernelOperation> _commandQueue = new();
 
-        public Task<KernelCommandResult> Schedule(KernelCommand command, Kernel kernel)
-        {
-            return Schedule(command, kernel, CancellationToken.None, () => { });
-        }
 
 
-        public Task<KernelCommandResult> Schedule(KernelCommand command, Kernel kernel, CancellationToken cancellationToken, Action onDone)
+        public Task<KernelCommandResult> Schedule(KernelCommand command, Kernel kernel, CancellationToken cancellationToken)
         {
 
-
             switch (command)
             {
                 case Cancel _:
@@ -46,15 +41,14 @@ public Task<KernelCommandResult> Schedule(KernelCommand command, Kernel kernel,
 
             _commandQueue.Enqueue(operation);
 
-            ProcessCommandQueue(_commandQueue, cancellationToken, onDone);
+            ProcessCommandQueue(_commandQueue, cancellationToken);
 
             return kernelCommandResultSource.Task;
         }
 
         private void ProcessCommandQueue(
             ConcurrentQueue<KernelOperation> commandQueue,
-            CancellationToken cancellationToken,
-            Action onDone)
+            CancellationToken cancellationToken)
         {
             if (commandQueue.TryDequeue(out var currentOperation))
             {
@@ -64,13 +58,9 @@ private void ProcessCommandQueue(
 
                     await ExecuteCommand(currentOperation);
 
-                    ProcessCommandQueue(commandQueue, cancellationToken, onDone);
+                    ProcessCommandQueue(commandQueue, cancellationToken);
                 }, cancellationToken).ConfigureAwait(false);
             }
-            else
-            {
-                onDone?.Invoke();
-            }
         }
 
         private async Task ExecuteCommand(KernelOperation operation)
@@ -142,8 +132,7 @@ internal Task RunDeferredCommandsAsync(Kernel kernel)
             UndeferCommandsFor(kernel);
             ProcessCommandQueue(
                 _commandQueue,
-                CancellationToken.None,
-                () => tcs.SetResult(Unit.Default));
+                CancellationToken.None);
             return tcs.Task;
         }
 

src/Microsoft.DotNet.Interactive/KernelScheduler.cs

@@ -161,7 +161,7 @@ public void Dispose()
 
         public delegate Task<U> OnExecuteDelegate(T value);
 
-        public delegate IEnumerable<T> GetDeferredOperationsDelegate(T operationToExecute, string queueName);
+        public delegate IEnumerable<T> GetDeferredOperationsDelegate(T state, string scope);
 
         private class ScheduledOperation
         {

src/Microsoft.DotNet.Interactive/KernelSupportsNugetExtensions.cs

@@ -263,8 +263,6 @@ internal static KernelCommandInvocation DoNugetRestore()
                 };
 
                 await invocationContext.QueueAction(restore);
-                var kernel = invocationContext.HandlingKernel;
-                await kernel.RunDeferredCommandsAsync();
             };
 
             static string InstallingPackageMessage(PackageReference package)