From a1be3a140eaba16303889236aab84d5604706a86 Mon Sep 17 00:00:00 2001 From: Scott Beddall <45376673+scbedd@users.noreply.github.com> Date: Fri, 13 Dec 2024 16:43:15 -0800 Subject: [PATCH] [TestProxy] Add auto-shutdown to the proxy (#9540) * if `auto-shutdown-in-seconds X` is provided to the StartInvocation, enable the proxy to shutdown after X seconds of no traffic over its routes --- .../CommandOptions/OptionsGenerator.cs | 8 +++- .../CommandOptions/StartOptions.cs | 7 +++- .../AutoShutdown/ShutdownConfiguration.cs | 8 ++++ .../Common/AutoShutdown/ShutdownTimer.cs | 42 +++++++++++++++++++ .../AutoShutdown/ShutdownTimerMiddleware.cs | 28 +++++++++++++ .../Azure.Sdk.Tools.TestProxy/Startup.cs | 16 ++++++- 6 files changed, 106 insertions(+), 3 deletions(-) create mode 100644 tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownConfiguration.cs create mode 100644 tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimer.cs create mode 100644 tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimerMiddleware.cs diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/OptionsGenerator.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/OptionsGenerator.cs index 3ad5748f5ca..fde07504ea3 100644 --- a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/OptionsGenerator.cs +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/OptionsGenerator.cs @@ -55,6 +55,11 @@ public static RootCommand GenerateCommandLineOptions(Func getDefaultValue: () => false); universalOption.AddAlias("-u"); + var autoShutdownOption = new Option( + name: "--auto-shutdown-in-seconds", + description: "Integer argument; When provided, the proxy will auto-shutdown after not being contacted over any HTTP route for this number of seconds.", + getDefaultValue: () => -1); + var breakGlassOption = new Option( name: "--break-glass", description: "Flag; Ignore secret push protection results when pushing.", @@ -88,10 +93,11 @@ public static RootCommand GenerateCommandLineOptions(Func startCommand.AddOption(insecureOption); startCommand.AddOption(dumpOption); startCommand.AddOption(universalOption); + startCommand.AddOption(autoShutdownOption); startCommand.AddArgument(collectedArgs); startCommand.SetHandler(async (startOpts) => await callback(startOpts), - new StartOptionsBinder(storageLocationOption, storagePluginOption, insecureOption, dumpOption, universalOption, collectedArgs) + new StartOptionsBinder(storageLocationOption, storagePluginOption, insecureOption, dumpOption, universalOption, autoShutdownOption, collectedArgs) ); root.Add(startCommand); diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/StartOptions.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/StartOptions.cs index 6014568e1bf..db8a44eeb5c 100644 --- a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/StartOptions.cs +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/CommandOptions/StartOptions.cs @@ -10,6 +10,8 @@ public class StartOptions : DefaultOptions public bool Dump { get; set; } public bool UniversalOutput { get; set; } + public int AutoShutdownTime { get; set; } + // On the command line, use -- and everything after that becomes arguments to Host.CreateDefaultBuilder // For example Test-Proxy start -i -d -- --urls https://localhost:8002 would set AdditionaArgs to a list containing // --urls and https://localhost:8002 as individual entries. This is converted to a string[] before being @@ -23,15 +25,17 @@ public class StartOptionsBinder : BinderBase private readonly Option _insecureOption; private readonly Option _dumpOption; private readonly Option _univeralOutputOption; + private readonly Option _autoShutdownTime; private readonly Argument _additionalArgs; - public StartOptionsBinder(Option storageLocationOption, Option storagePluginOption, Option insecureOption, Option dumpOption, Option universalOutput, Argument additionalArgs) + public StartOptionsBinder(Option storageLocationOption, Option storagePluginOption, Option insecureOption, Option dumpOption, Option universalOutput, Option autoShutdownTime, Argument additionalArgs) { _storageLocationOption = storageLocationOption; _storagePluginOption = storagePluginOption; _insecureOption = insecureOption; _dumpOption = dumpOption; _univeralOutputOption = universalOutput; + _autoShutdownTime = autoShutdownTime; _additionalArgs = additionalArgs; } @@ -43,6 +47,7 @@ protected override StartOptions GetBoundValue(BindingContext bindingContext) => Insecure = bindingContext.ParseResult.GetValueForOption(_insecureOption), Dump = bindingContext.ParseResult.GetValueForOption(_dumpOption), UniversalOutput = bindingContext.ParseResult.GetValueForOption(_univeralOutputOption), + AutoShutdownTime = bindingContext.ParseResult.GetValueForOption(_autoShutdownTime), AdditionalArgs = bindingContext.ParseResult.GetValueForArgument(_additionalArgs) }; } diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownConfiguration.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownConfiguration.cs new file mode 100644 index 00000000000..18f11245bea --- /dev/null +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownConfiguration.cs @@ -0,0 +1,8 @@ +namespace Azure.Sdk.Tools.TestProxy.Common.AutoShutdown +{ + public class ShutdownConfiguration + { + public bool EnableAutoShutdown { get; set; } = false; + public int TimeoutInSeconds { get; set; } = 300; + } +} diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimer.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimer.cs new file mode 100644 index 00000000000..cf8ac998d96 --- /dev/null +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimer.cs @@ -0,0 +1,42 @@ +using System.Threading.Tasks; +using System.Threading; +using System; + +namespace Azure.Sdk.Tools.TestProxy.Common.AutoShutdown +{ + public class ShutdownTimer + { + private readonly ShutdownConfiguration _shutdownConfig; + private CancellationTokenSource _cts; + private Task _shutdownTask; + + public ShutdownTimer(ShutdownConfiguration shutdownConfiguration) + { + _shutdownConfig = shutdownConfiguration; + } + + public void ResetTimer() + { + _cts?.Cancel(); + _cts = new CancellationTokenSource(); + + _shutdownTask = Task.Run(async () => + { + try + { + await Task.Delay(_shutdownConfig.TimeoutInSeconds * 1000, _cts.Token); + + if (_shutdownConfig.EnableAutoShutdown) + { + System.Console.WriteLine("Server idle timeout reached. Shutting down..."); + Environment.Exit(0); + } + } + catch (TaskCanceledException) + { + // Timer was reset or canceled + } + }); + } + } +} diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimerMiddleware.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimerMiddleware.cs new file mode 100644 index 00000000000..369a04b798a --- /dev/null +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/AutoShutdown/ShutdownTimerMiddleware.cs @@ -0,0 +1,28 @@ +using Microsoft.AspNetCore.Http; +using System.Threading.Tasks; + +namespace Azure.Sdk.Tools.TestProxy.Common.AutoShutdown +{ + public class ShutdownTimerMiddleware + { + private readonly RequestDelegate _next; + private readonly ShutdownTimer _shutdownTimer; + private readonly ShutdownConfiguration _shutdownConfig; + + public ShutdownTimerMiddleware(RequestDelegate next, ShutdownTimer shutdownTimer, ShutdownConfiguration shutdownConfig) + { + _next = next; + _shutdownTimer = shutdownTimer; + _shutdownConfig = shutdownConfig; + } + + public async Task InvokeAsync(HttpContext context) + { + if (_shutdownConfig.EnableAutoShutdown) + { + _shutdownTimer.ResetTimer(); + } + await _next(context); + } + } +} diff --git a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Startup.cs b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Startup.cs index eecdedab1c7..d703dd7559e 100644 --- a/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Startup.cs +++ b/tools/test-proxy/Azure.Sdk.Tools.TestProxy/Startup.cs @@ -24,6 +24,7 @@ using System.Text.Json; using Microsoft.Extensions.Logging.Console; using Microsoft.Extensions.Options; +using Azure.Sdk.Tools.TestProxy.Common.AutoShutdown; namespace Azure.Sdk.Tools.TestProxy { @@ -125,6 +126,7 @@ private static async Task Run(object commandObj) StorageLocation = defaultOpts.StorageLocation, StoragePlugin = defaultOpts.StoragePlugin, Insecure = false, + AutoShutdownTime = -1, Dump = false }); break; @@ -147,7 +149,7 @@ private static void StartServer(StartOptions startOptions) newLine: true, statusThreadCts.Token); var host = Host.CreateDefaultBuilder((startOptions.AdditionalArgs??new string[] { }).ToArray()); - + host.ConfigureWebHostDefaults( builder => builder.UseStartup() @@ -180,6 +182,15 @@ private static void StartServer(StartOptions startOptions) var app = host.Build(); + var shutdownService = app.Services.GetRequiredService(); + if (startOptions.AutoShutdownTime > -1) + { + shutdownService.EnableAutoShutdown = true; + shutdownService.TimeoutInSeconds = startOptions.AutoShutdownTime; + // start the first iteration of the shutdown timer + app.Services.GetRequiredService().ResetTimer(); + } + if (startOptions.Dump) { var config = app.Services?.GetService(); @@ -227,6 +238,8 @@ public void ConfigureServices(IServiceCollection services) ); services.AddSingleton(singletonRecordingHandler); + services.AddSingleton(); + services.AddSingleton(); } public void Configure(IApplicationBuilder app, IWebHostEnvironment env, ILoggerFactory loggerFactory) @@ -237,6 +250,7 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env, ILoggerF } app.UseCors("DefaultPolicy"); app.UseMiddleware(); + app.UseMiddleware(); DebugLogger.ConfigureLogger(loggerFactory);