fix opt 6 exit not working
also continue fix attempt on sheepit not starting
This commit is contained in:
@@ -5,9 +5,10 @@ param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$PayloadBase64,
|
||||
|
||||
[string]$PayloadBase64Path,
|
||||
|
||||
[string]$DataRoot = (Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'),
|
||||
|
||||
[int]$MaxRestarts = 5,
|
||||
@@ -16,6 +17,7 @@ param(
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
[System.Management.Automation.Runspaces.Runspace]::DefaultRunspace = $Host.Runspace
|
||||
|
||||
# region Path setup
|
||||
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
|
||||
@@ -33,11 +35,95 @@ $commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
|
||||
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload.ps1"
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
# region Logging
|
||||
$logStream = [System.IO.FileStream]::new(
|
||||
$logPath,
|
||||
[System.IO.FileMode]::Append,
|
||||
[System.IO.FileAccess]::Write,
|
||||
[System.IO.FileShare]::ReadWrite
|
||||
)
|
||||
$logWriter = [System.IO.StreamWriter]::new($logStream, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
if (-not ("UnifiedWorkers.WorkerLogSink" -as [type])) {
|
||||
Add-Type -Namespace UnifiedWorkers -Name WorkerLogSink -MemberDefinition @"
|
||||
using System;
|
||||
using System.Diagnostics;
|
||||
using System.IO;
|
||||
|
||||
public sealed class WorkerLogSink
|
||||
{
|
||||
private readonly TextWriter _writer;
|
||||
private readonly string _prefix;
|
||||
private readonly object _sync = new object();
|
||||
|
||||
public WorkerLogSink(TextWriter writer, string prefix)
|
||||
{
|
||||
if (writer == null)
|
||||
{
|
||||
throw new ArgumentNullException(\"writer\");
|
||||
}
|
||||
if (prefix == null)
|
||||
{
|
||||
throw new ArgumentNullException(\"prefix\");
|
||||
}
|
||||
|
||||
_writer = writer;
|
||||
_prefix = prefix;
|
||||
}
|
||||
|
||||
public void OnData(object sender, DataReceivedEventArgs e)
|
||||
{
|
||||
if (e == null || string.IsNullOrEmpty(e.Data))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
lock (_sync)
|
||||
{
|
||||
var timestamp = DateTime.UtcNow.ToString("u");
|
||||
var line = string.Format("[{0} {1}] {2}", _prefix, timestamp, e.Data);
|
||||
_writer.WriteLine(line);
|
||||
_writer.Flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
"@
|
||||
}
|
||||
|
||||
function Write-LogLine {
|
||||
param(
|
||||
[string]$Prefix,
|
||||
[string]$Message
|
||||
)
|
||||
|
||||
if (-not $logWriter) { return }
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[{0} {1}] {2}" -f $Prefix, $timestamp, $Message)
|
||||
}
|
||||
|
||||
function Write-ControllerLog {
|
||||
param([string]$Message)
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
Add-Content -Path $logPath -Value "[CTRL $timestamp] $Message" -Encoding UTF8
|
||||
Write-LogLine -Prefix 'CTRL' -Message $Message
|
||||
}
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
|
||||
function Resolve-PayloadBase64 {
|
||||
if ($PayloadBase64) {
|
||||
return $PayloadBase64
|
||||
}
|
||||
|
||||
if ($PayloadBase64Path) {
|
||||
if (-not (Test-Path $PayloadBase64Path)) {
|
||||
throw "Payload file '$PayloadBase64Path' not found."
|
||||
}
|
||||
|
||||
return (Get-Content -Path $PayloadBase64Path -Raw)
|
||||
}
|
||||
|
||||
throw "No payload data provided to controller."
|
||||
}
|
||||
|
||||
function Write-Metadata {
|
||||
@@ -85,9 +171,12 @@ function Get-PendingCommands {
|
||||
# record initial state before launching worker
|
||||
Write-Metadata -Status 'initializing' -WorkerPid $null -ControllerPid $PID -Restarts 0
|
||||
|
||||
$resolvedPayloadBase64 = Resolve-PayloadBase64
|
||||
$PayloadBase64 = $resolvedPayloadBase64
|
||||
|
||||
try {
|
||||
# Write payload script to disk
|
||||
$payloadBytes = [Convert]::FromBase64String($PayloadBase64)
|
||||
$payloadBytes = [Convert]::FromBase64String($resolvedPayloadBase64)
|
||||
[IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
|
||||
Write-ControllerLog "Payload written to $payloadPath"
|
||||
}
|
||||
@@ -128,23 +217,10 @@ while ($restartCount -le $MaxRestarts) {
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$logWriter = [System.IO.StreamWriter]::new($logPath, $true, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[OUT $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[ERR $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$stdoutSink = [UnifiedWorkers.WorkerLogSink]::new($logWriter, 'OUT')
|
||||
$stderrSink = [UnifiedWorkers.WorkerLogSink]::new($logWriter, 'ERR')
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]$stdoutSink.OnData
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]$stderrSink.OnData
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
@@ -199,3 +275,10 @@ while ($restartCount -le $MaxRestarts) {
|
||||
Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
|
||||
Write-ControllerLog "Controller exiting."
|
||||
|
||||
if ($logWriter) {
|
||||
$logWriter.Dispose()
|
||||
}
|
||||
if ($logStream) {
|
||||
$logStream.Dispose()
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user