diff --git a/Services/OrchestratorService.cs b/Services/OrchestratorService.cs index 4293524..a7830b6 100644 --- a/Services/OrchestratorService.cs +++ b/Services/OrchestratorService.cs @@ -607,8 +607,7 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task, CheckRunCompletion(run); } } - await _store.UpsertTaskAsync(run.Name, task); - await _store.UpsertRunAsync(run); + PersistTaskAndRunAsync(run, task); return; } @@ -616,6 +615,7 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task, { task.Status = "Running"; } + // Pre-script "Running" write is awaited — it's the durability marker for crash recovery. await _store.UpsertTaskAsync(run.Name, task); try @@ -646,8 +646,11 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task, task.Parameters = null!; CheckRunCompletion(run); } - await _store.UpsertTaskAsync(run.Name, task); - await _store.UpsertRunAsync(run); + // Post-script writes are fire-and-forget so the JobManager slot releases + // immediately and the dispatch loop can hand the worker to the next task. + // Crash recovery still works: the next startup re-reads task state from the + // table and re-runs anything not marked Completed (idempotent). + PersistTaskAndRunAsync(run, task); _logger.LogDebug("[Scheduler] Task completed: {TaskId}", task.Id); } @@ -668,8 +671,7 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task, task.Parameters = null!; CheckRunCompletion(run); } - await _store.UpsertTaskAsync(run.Name, task); - await _store.UpsertRunAsync(run); + PersistTaskAndRunAsync(run, task); _logger.LogError(ex, "[Scheduler] Task failed: {TaskId}", task.Id); throw; // Let JobManager also track the failure } @@ -677,6 +679,24 @@ private void DispatchSingleTask(OrchestratorRun run, OrchestratorTaskItem task, ); } + /// + /// Fire-and-forget persistence of task + run state. Callers do not await this — it lets the + /// JobManager slot release immediately so the dispatch loop can hand the worker to the next + /// task. Errors are logged; on host crash, ResumeInterruptedRunsAsync re-derives state from + /// whatever made it to the table (writes are idempotent). + /// + private void PersistTaskAndRunAsync(OrchestratorRun run, OrchestratorTaskItem task) + { + _ = Task.Run(async () => + { + try { await _store.UpsertTaskAsync(run.Name, task); } + catch (Exception ex) { _logger.LogWarning(ex, "[Scheduler] Background UpsertTask failed for {Run}/{Task}", run.Name, task.Id); } + + try { await _store.UpsertRunAsync(run); } + catch (Exception ex) { _logger.LogWarning(ex, "[Scheduler] Background UpsertRun failed for {Run}", run.Name); } + }); + } + private void LogRunStatus(OrchestratorRun run) { var elapsed = DateTime.UtcNow - run.StartedUtc; diff --git a/Services/PowerShellRunnerService.cs b/Services/PowerShellRunnerService.cs index 146ef09..b75463b 100644 --- a/Services/PowerShellRunnerService.cs +++ b/Services/PowerShellRunnerService.cs @@ -293,6 +293,26 @@ private async Task ExecuteHttpScriptInternal(string route, Hashtab } } + /// + /// Drain pending orchestrator/queue triggers off the calling thread. Bridges are thread-safe + /// concurrent queues, so multiple in-flight drain calls are fine — each TryDequeue serialises. + /// + private void DrainBridgesInBackground() + { + _ = Task.Run(async () => + { + try + { + await OrchestratorBridge.DrainPendingAsync(); + QueueBridge.DrainPending(); + } + catch (Exception ex) + { + _logger.LogError(ex, "[Scheduler] Background bridge drain failed"); + } + }); + } + /// /// Execute a script by function name (for scheduler / orchestrator). No HTTP context needed. /// Runs on the background pool. @@ -380,10 +400,6 @@ public async Task ExecuteScript(string functionName, Dictionary? sw.Stop(); _logger.LogInformation("[Scheduler] {InvocationId} {Function} completed {Ms}ms", invocation.Id, functionName, sw.ElapsedMilliseconds); - - // Process any orchestrator/queue triggers queued during execution - await OrchestratorBridge.DrainPendingAsync(); - QueueBridge.DrainPending(); } catch (OperationCanceledException) when (sw.ElapsedMilliseconds > 0) { @@ -410,6 +426,11 @@ public async Task ExecuteScript(string functionName, Dictionary? if (onVerbose != null) worker.Streams.Verbose.DataAdded -= onVerbose; _pool.Reclaim(worker, isHttp: false, faulted: exceptionOccurred); } + + // Worker has been returned to the pool — drain any orchestrator/queue triggers + // the script enqueued in the background so the next job can grab the worker now + // instead of waiting for child-run table writes. + DrainBridgesInBackground(); } /// @@ -483,14 +504,14 @@ public async Task ExecuteScriptWithOutput(string functionName, Dictionar : null; var results = await worker.InvokeAsync(resolvedName, psParams, cts?.Token ?? default); - // Process any orchestrator/queue triggers queued during execution - await OrchestratorBridge.DrainPendingAsync(); - QueueBridge.DrainPending(); - sw.Stop(); _logger.LogInformation("[Planner] {InvocationId} {Function} completed {Ms}ms", invocation.Id, functionName, sw.ElapsedMilliseconds); - return string.Join("\n", (results ?? new Collection()).Select(r => r?.ToString() ?? "")); + var output = string.Join("\n", (results ?? new Collection()).Select(r => r?.ToString() ?? "")); + // Drain triggered child orchestrators/queue commands after returning — they should + // not block the planner's caller. (Note: Reclaim still happens in finally below.) + DrainBridgesInBackground(); + return output; } catch (OperationCanceledException) when (sw.ElapsedMilliseconds > 0) { diff --git a/Services/PowerShellWorker.cs b/Services/PowerShellWorker.cs index c17da72..6bc6df5 100644 --- a/Services/PowerShellWorker.cs +++ b/Services/PowerShellWorker.cs @@ -105,16 +105,39 @@ class HttpResponseContext { "); } - // Load shared assemblies from config + // Load shared assemblies from config. + // ISS-level registration happens in PowerShellWorkerPool.RegisterSharedAssemblies(); this + // runtime LoadFile is a defence-in-depth fallback that surfaces any load failure to the log + // (RunScript silently swallows streams, so we run a labelled invocation here instead). foreach (var asmRelPath in settings.Worker.SharedAssemblies) { + if (string.IsNullOrWhiteSpace(asmRelPath)) continue; var asmPath = Path.Combine(apiBasePath, asmRelPath).Replace("\\", "/"); - RunScript($@" -if (Test-Path '{asmPath}') {{ - if (-not ([System.AppDomain]::CurrentDomain.GetAssemblies().Location -contains '{asmPath}')) {{ + var asmLabel = Path.GetFileNameWithoutExtension(asmPath); + try + { + _pwsh.AddScript($@" +if (Test-Path -LiteralPath '{asmPath}') {{ + if (-not ([System.AppDomain]::CurrentDomain.GetAssemblies() | Where-Object {{ $_.Location -ieq '{asmPath}' }})) {{ [void][Reflection.Assembly]::LoadFile('{asmPath}') }} -}}"); + [System.AppDomain]::CurrentDomain.GetAssemblies() | Where-Object {{ $_.Location -ieq '{asmPath}' }} | Select-Object -First 1 -ExpandProperty FullName +}} else {{ + Write-Error ""SharedAssembly not found: {asmPath}"" +}}").Invoke(); + + foreach (var err in _pwsh.Streams.Error) + _logger.LogError("Worker{Id}: SharedAssembly '{Label}' load error: {Error}", Id, asmLabel, err.ToString()); + foreach (var warn in _pwsh.Streams.Warning) + _logger.LogWarning("Worker{Id}: SharedAssembly '{Label}' warning: {Message}", Id, asmLabel, warn.Message); + + _logger.LogDebug("Worker{Id}: SharedAssembly '{Label}' available at {Path}", Id, asmLabel, asmPath); + } + catch (Exception ex) + { + _logger.LogError(ex, "Worker{Id}: SharedAssembly '{Label}' load threw", Id, asmLabel); + } + finally { _pwsh.Commands.Clear(); _pwsh.Streams.ClearStreams(); } } // Deploy background scripts as Function:\ items. diff --git a/Services/PowerShellWorkerPool.cs b/Services/PowerShellWorkerPool.cs index 891a282..0eeb560 100644 --- a/Services/PowerShellWorkerPool.cs +++ b/Services/PowerShellWorkerPool.cs @@ -440,13 +440,32 @@ public void Reclaim(PowerShellWorker worker, bool isHttp, bool faulted = false) var oldId = worker.Id; WorkerMetricsBridge.DeregisterWorker(oldId); worker.Dispose(); - var cloned = isHttp ? _httpClonedState : _bgClonedState; - var iss = cloned != null ? BuildClonedISS(cloned) : BuildISS(isHttp: isHttp); - worker = new PowerShellWorker(Interlocked.Increment(ref _nextId), iss, _logger); - worker.Initialize(_repo, _apiBasePath, _settings); - WorkerMetricsBridge.RegisterWorker(worker.Id, isHttp); - _logger.LogInformation("[Pool] Replaced W{OldId} → W{NewId} ({Type})", - oldId, worker.Id, isHttp ? "HTTP" : "BG"); + + // Build the replacement off the calling thread so the dispatch loop is not held + // for the 6-13s ISS rebuild + Initialize() cost. Capacity briefly drops by 1 (the + // pool is short one worker until this Task completes), but the thread that just + // finished a task returns to the dispatch loop immediately. + var ish = isHttp; + _ = Task.Run(() => + { + try + { + var cloned = ish ? _httpClonedState : _bgClonedState; + var iss = cloned != null ? BuildClonedISS(cloned) : BuildISS(isHttp: ish); + var fresh = new PowerShellWorker(Interlocked.Increment(ref _nextId), iss, _logger); + fresh.Initialize(_repo, _apiBasePath, _settings); + WorkerMetricsBridge.RegisterWorker(fresh.Id, ish); + if (ish) _httpPool.Add(fresh); else _bgPool.Add(fresh); + _logger.LogInformation("[Pool] Replaced W{OldId} → W{NewId} ({Type}) (background recycle)", + oldId, fresh.Id, ish ? "HTTP" : "BG"); + } + catch (Exception ex) + { + _logger.LogError(ex, "[Pool] Background recycle failed for W{OldId} ({Type}); pool capacity reduced by 1", + oldId, ish ? "HTTP" : "BG"); + } + }); + return; } if (isHttp) _httpPool.Add(worker); else _bgPool.Add(worker); @@ -458,6 +477,43 @@ private InitialSessionState BuildISS(bool isHttp) return BuildISSForModules(allowList.Count > 0 ? allowList : null); } + /// + /// Register configured SharedAssemblies on the InitialSessionState so the runspace's + /// type resolver knows about them from creation. This complements the runtime + /// [Reflection.Assembly]::LoadFile script in PowerShellWorker.Initialize() and ensures + /// type literals (e.g. [CIPP.TestDataCache]) resolve in cloned/recycled runspaces. + /// Logs each path so silent load failures during cold-start become observable. + /// + private void RegisterSharedAssemblies(InitialSessionState iss, string buildContext) + { + if (_settings.Worker.SharedAssemblies.Count == 0) return; + + foreach (var asmRelPath in _settings.Worker.SharedAssemblies) + { + if (string.IsNullOrWhiteSpace(asmRelPath)) continue; + + var asmPath = Path.GetFullPath(Path.Combine(_apiBasePath, asmRelPath)); + if (!File.Exists(asmPath)) + { + _logger.LogError("[Pool] SharedAssembly missing for {Context}: {Path}", buildContext, asmPath); + continue; + } + + try + { + var asmName = Path.GetFileNameWithoutExtension(asmPath); + iss.Assemblies.Add(new SessionStateAssemblyEntry(asmName, asmPath)); + _logger.LogDebug("[Pool] SharedAssembly registered for {Context}: {Name} ({Path})", + buildContext, asmName, asmPath); + } + catch (Exception ex) + { + _logger.LogError(ex, "[Pool] SharedAssembly registration failed for {Context}: {Path}", + buildContext, asmPath); + } + } + } + /// /// Build an ISS that imports only the specified modules. /// If moduleList is null, imports all modules (minus SkipModules). @@ -468,6 +524,8 @@ private InitialSessionState BuildISSForModules(List? moduleList) if (OperatingSystem.IsWindows()) iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass; + RegisterSharedAssemblies(iss, "BuildISSForModules"); + // Copy environment variables into the runspace foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables()) iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null)); @@ -515,6 +573,8 @@ private InitialSessionState BuildClonedISSWithModules(ExportedModuleState baseSt if (OperatingSystem.IsWindows()) iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass; + RegisterSharedAssemblies(iss, "BuildClonedISSWithModules"); + foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables()) iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null)); @@ -581,6 +641,8 @@ private InitialSessionState BuildClonedISS(ExportedModuleState state) if (OperatingSystem.IsWindows()) iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Bypass; + RegisterSharedAssemblies(iss, "BuildClonedISS"); + // Copy environment variables foreach (System.Collections.DictionaryEntry env in Environment.GetEnvironmentVariables()) iss.EnvironmentVariables.Add(new SessionStateVariableEntry((string)env.Key, env.Value, null)); diff --git a/Services/SetupPages.cs b/Services/SetupPages.cs index dccca98..5121b33 100644 --- a/Services/SetupPages.cs +++ b/Services/SetupPages.cs @@ -216,8 +216,7 @@ function setTenantMode(section, isMulti) { } document.getElementById('manual-redirect-uri').textContent = window.location.origin + '/.auth/login/aad/callback'; if (setupState.isEasyAuthConfigured) { - showBanner('Authentication is already configured. Redirecting...'); - setTimeout(() => window.location.href = '/', 2000); + showBanner('Authentication is already configured. This page will redirect shortly.'); return; } if (!setupState.isRunningInAppService || !setupState.hasManagedIdentity) { @@ -459,57 +458,52 @@ async function submitManual() { } } function showRestartScreen() { - // Stop background status polling — restart screen has its own polling + // Stop background status polling clearInterval(statusPollTimer); - // Hide setup sections, show restart polling UI + // Hide setup sections, show restart status document.getElementById('user-section').classList.add('hidden'); document.getElementById('auto-section').classList.add('hidden'); document.getElementById('manual-section').classList.add('hidden'); document.getElementById('divider-user').classList.add('hidden'); document.getElementById('divider-auth').classList.add('hidden'); document.querySelector('.subtitle').textContent = ''; - document.getElementById('page-title').textContent = 'Restarting...'; + document.getElementById('page-title').textContent = 'Setup Complete'; const banner = document.getElementById('status-banner'); banner.innerHTML = '
' + '
Authentication has been configured. The application is restarting to apply changes.
' + '
' + - '
Waiting for the application to come back online...
' + + '
Waiting for the new container to come online...
' + '
'; banner.classList.remove('hidden'); - let attempts = 0; - const maxAttempts = 120; // 10 min at 5s intervals + // Poll /api/setup/status — once the new container comes online with EasyAuth + // configured, isEasyAuthConfigured will be true and setup mode will no longer + // be active. At that point redirect to / and let the app take over. const pollInterval = 5000; - const pollRestart = async () => { - attempts++; - const statusEl = document.getElementById('restart-status'); try { - const res = await fetch('/api/setup/health', { cache: 'no-store' }); + const res = await fetch('/api/setup/status', { cache: 'no-store' }); if (res.ok) { const data = await res.json(); - if (data.ready) { - statusEl.textContent = 'Application is ready! Redirecting...'; + if (data.isEasyAuthConfigured && !data.isSetupCompleted) { + // New container is online with EasyAuth active — redirect to app + document.getElementById('restart-status').textContent = 'Application is ready!'; document.getElementById('restart-spinner').textContent = '\u2705'; - setTimeout(() => window.location.href = '/', 1500); + setTimeout(() => window.location.href = '/', 1000); return; } + // Still on the old container (isSetupCompleted=true) — keep waiting + document.getElementById('restart-status').textContent = 'Auth configured, waiting for container restart...'; } } catch (e) { - // App still restarting — expected - } - if (attempts >= maxAttempts) { - statusEl.textContent = 'The application is taking longer than expected. Try refreshing the page manually.'; - document.getElementById('restart-spinner').textContent = '\u26A0\uFE0F'; - return; + // Container is cycling — expected during restart + document.getElementById('restart-status').textContent = 'Container restarting...'; } - statusEl.textContent = 'Waiting for the application to come back online... (' + attempts + ')'; setTimeout(pollRestart, pollInterval); }; - // Start polling after a short delay to let the restart begin - setTimeout(pollRestart, 8000); + setTimeout(pollRestart, 5000); } diff --git a/Services/WorkerMetricsBridge.cs b/Services/WorkerMetricsBridge.cs index 054a9c8..ce38def 100644 --- a/Services/WorkerMetricsBridge.cs +++ b/Services/WorkerMetricsBridge.cs @@ -28,6 +28,17 @@ public static class WorkerMetricsBridge private static readonly DateTime s_startTimeUtc = DateTime.UtcNow; private static long s_globalInvocations; + // Retired-worker totals: when a worker is recycled, its accumulated counters are + // moved here so pool aggregates remain monotonic across recycles. Without these, + // pool.TotalInvocations would drop on recycle and StatsHistoryService deltas + // (BgInvocations, BgBusyMs) would go negative. + private static long s_retiredHttpInvocations; + private static long s_retiredHttpBusyMs; + private static long s_retiredHttpFaults; + private static long s_retiredBgInvocations; + private static long s_retiredBgBusyMs; + private static long s_retiredBgFaults; + public static void Initialize(PowerShellWorkerPool pool, BackgroundTaskLimiter limiter, JobManager jobManager, ILogger? logger = null) { s_pool = pool; @@ -48,7 +59,26 @@ public static void RegisterWorker(int workerId, bool isHttp) /// Remove a worker's stats when it is recycled/replaced. public static void DeregisterWorker(int workerId) { - s_workerStats.TryRemove(workerId, out _); + // Accumulate the retiring worker's totals into the per-pool "retired" buckets + // so pool-level sums (and delta-based history) stay monotonic across recycles. + if (s_workerStats.TryRemove(workerId, out var stats)) + { + var inv = Interlocked.Read(ref stats._totalInvocations); + var busy = Interlocked.Read(ref stats._totalBusyMs); + var faults = Interlocked.Read(ref stats._totalFaults); + if (stats.IsHttp) + { + Interlocked.Add(ref s_retiredHttpInvocations, inv); + Interlocked.Add(ref s_retiredHttpBusyMs, busy); + Interlocked.Add(ref s_retiredHttpFaults, faults); + } + else + { + Interlocked.Add(ref s_retiredBgInvocations, inv); + Interlocked.Add(ref s_retiredBgBusyMs, busy); + Interlocked.Add(ref s_retiredBgFaults, faults); + } + } } /// Record that a worker was checked out (started processing). @@ -135,9 +165,15 @@ public static WorkerMetricsSnapshot GetSnapshot() Workers = bgWorkers, }; - // Aggregate pool-level stats - AggregatePoolStats(snapshot.HttpPool, httpWorkers); - AggregatePoolStats(snapshot.BgPool, bgWorkers); + // Aggregate pool-level stats (live workers + retired buckets so totals are monotonic) + AggregatePoolStats(snapshot.HttpPool, httpWorkers, + Interlocked.Read(ref s_retiredHttpInvocations), + Interlocked.Read(ref s_retiredHttpBusyMs), + Interlocked.Read(ref s_retiredHttpFaults)); + AggregatePoolStats(snapshot.BgPool, bgWorkers, + Interlocked.Read(ref s_retiredBgInvocations), + Interlocked.Read(ref s_retiredBgBusyMs), + Interlocked.Read(ref s_retiredBgFaults)); } if (s_limiter != null) @@ -490,13 +526,22 @@ private static WorkerDetail BuildWorkerDetail(WorkerStats stats) }; } - private static void AggregatePoolStats(PoolMetrics pool, List workers) + private static void AggregatePoolStats(PoolMetrics pool, List workers, + long retiredInvocations, long retiredBusyMs, long retiredFaults) { - if (workers.Count == 0) return; - pool.TotalInvocations = workers.Sum(w => w.TotalInvocations); - pool.TotalBusyMs = workers.Sum(w => w.TotalBusyMs); - pool.TotalFaults = workers.Sum(w => w.TotalFaults); - pool.AvgUtilizationPct = Math.Round(workers.Average(w => w.UtilizationPct), 1); + // Sum live workers and add retired-worker totals (workers that were recycled out + // of the pool). Without the retired buckets, recycling a worker would decrease + // pool.TotalInvocations and break delta-based history (StatsHistoryService). + var liveInv = workers.Sum(w => w.TotalInvocations); + var liveBusy = workers.Sum(w => w.TotalBusyMs); + var liveFaults = workers.Sum(w => w.TotalFaults); + + pool.TotalInvocations = liveInv + retiredInvocations; + pool.TotalBusyMs = liveBusy + retiredBusyMs; + pool.TotalFaults = liveFaults + retiredFaults; + pool.AvgUtilizationPct = workers.Count > 0 + ? Math.Round(workers.Average(w => w.UtilizationPct), 1) + : 0; pool.AvgDurationMs = pool.TotalInvocations > 0 ? pool.TotalBusyMs / pool.TotalInvocations : 0;