Skip to content

Commit

Permalink
Merge pull request #1581 from rjmholt/pt-eventing
Browse files Browse the repository at this point in the history
Implement OnIdle engine events
  • Loading branch information
rjmholt authored Oct 13, 2021
2 parents 79a586b + 08df557 commit fa3413a
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ public BlockingConcurrentDeque()
};
}

public int Count => _queues[0].Count + _queues[1].Count;
public bool IsEmpty => _queues[0].Count == 0 && _queues[1].Count == 0;

public void Prepend(T item)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ internal class PsesInternalHost : PSHost, IHostSupportsInteractiveSession, IRuns

private readonly IdempotentLatch _isRunningLatch = new();

private EngineIntrinsics _mainRunspaceEngineIntrinsics;

private bool _shouldExit = false;

private string _localComputerName;
Expand Down Expand Up @@ -347,17 +349,18 @@ public Task SetInitialWorkingDirectoryAsync(string path, CancellationToken cance

private void Run()
{
(PowerShell pwsh, RunspaceInfo localRunspaceInfo) = CreateInitialPowerShellSession();
(PowerShell pwsh, RunspaceInfo localRunspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
_mainRunspaceEngineIntrinsics = engineIntrinsics;
_localComputerName = localRunspaceInfo.SessionDetails.ComputerName;
_runspaceStack.Push(new RunspaceFrame(pwsh.Runspace, localRunspaceInfo));
PushPowerShellAndRunLoop(pwsh, PowerShellFrameType.Normal, localRunspaceInfo);
}

private (PowerShell, RunspaceInfo) CreateInitialPowerShellSession()
private (PowerShell, RunspaceInfo, EngineIntrinsics) CreateInitialPowerShellSession()
{
PowerShell pwsh = CreateInitialPowerShell(_hostInfo, _readLineProvider);
(PowerShell pwsh, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShell(_hostInfo, _readLineProvider);
RunspaceInfo localRunspaceInfo = RunspaceInfo.CreateFromLocalPowerShell(_logger, pwsh);
return (pwsh, localRunspaceInfo);
return (pwsh, localRunspaceInfo, engineIntrinsics);
}

private void PushPowerShellAndRunLoop(PowerShell pwsh, PowerShellFrameType frameType, RunspaceInfo newRunspaceInfo = null)
Expand Down Expand Up @@ -613,7 +616,7 @@ private static PowerShell CreatePowerShellForRunspace(Runspace runspace)
return pwsh;
}

public PowerShell CreateInitialPowerShell(
public (PowerShell, EngineIntrinsics) CreateInitialPowerShell(
HostStartupInfo hostStartupInfo,
ReadLineProvider readLineProvider)
{
Expand Down Expand Up @@ -649,7 +652,7 @@ public PowerShell CreateInitialPowerShell(
}
}

return pwsh;
return (pwsh, engineIntrinsics);
}

private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)
Expand All @@ -668,7 +671,27 @@ private Runspace CreateInitialRunspace(InitialSessionState initialSessionState)

private void OnPowerShellIdle()
{
if (_taskQueue.Count == 0)
IReadOnlyList<PSEventSubscriber> eventSubscribers = _mainRunspaceEngineIntrinsics.Events.Subscribers;

// Go through pending event subscribers and:
// - if we have any subscribers, ensure we process any events
// - if we have any idle events, generate an idle event and process that
bool runPipelineForEventProcessing = false;
foreach (PSEventSubscriber subscriber in eventSubscribers)
{
runPipelineForEventProcessing = true;

if (string.Equals(subscriber.SourceIdentifier, PSEngineEvent.OnIdle, StringComparison.OrdinalIgnoreCase))
{
// We control the pipeline thread, so it's not possible for PowerShell to generate events while we're here.
// But we know we're sitting waiting for the prompt, so we generate the idle event ourselves
// and that will flush idle event subscribers in PowerShell so we can service them
_mainRunspaceEngineIntrinsics.Events.GenerateEvent(PSEngineEvent.OnIdle, sender: null, args: null, extraData: null);
break;
}
}

if (!runPipelineForEventProcessing && _taskQueue.IsEmpty)
{
return;
}
Expand All @@ -688,9 +711,21 @@ private void OnPowerShellIdle()
return;
}

// If we're executing a task, we don't need to run an extra pipeline later for events
// TODO: This may not be a PowerShell task, so ideally we can differentiate that here.
// For now it's mostly true and an easy assumption to make.
runPipelineForEventProcessing = false;
task.ExecuteSynchronously(cancellationScope.CancellationToken);
}
}

// We didn't end up executing anything in the background,
// so we need to run a small artificial pipeline instead
// to force event processing
if (runPipelineForEventProcessing)
{
InvokePSCommand(new PSCommand().AddScript("0", useLocalScope: true), PowerShellExecutionOptions.Default, CancellationToken.None);
}
}

private void OnCancelKeyPress(object sender, ConsoleCancelEventArgs args)
Expand Down Expand Up @@ -771,7 +806,8 @@ private Task PopOrReinitializeRunspaceAsync()
// If our main runspace was corrupted,
// we must re-initialize our state.
// TODO: Use runspace.ResetRunspaceState() here instead
(PowerShell pwsh, RunspaceInfo runspaceInfo) = CreateInitialPowerShellSession();
(PowerShell pwsh, RunspaceInfo runspaceInfo, EngineIntrinsics engineIntrinsics) = CreateInitialPowerShellSession();
_mainRunspaceEngineIntrinsics = engineIntrinsics;
PushPowerShell(new PowerShellContextFrame(pwsh, runspaceInfo, PowerShellFrameType.Normal));

_logger.LogError($"Top level runspace entered state '{oldRunspaceState.State}' for reason '{oldRunspaceState.Reason}' and was reinitialized."
Expand Down

0 comments on commit fa3413a

Please sign in to comment.