fix opt 6 exit not working

also continue fix attempt on sheepit not starting
This commit is contained in:
Nathan
2025-12-02 12:07:26 -07:00
parent 18683b4ad0
commit e568e02cdf
4 changed files with 5209 additions and 1532 deletions

View File

@@ -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()
}