[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
Azure · Dec 14, 2024 · a1be3a1 · a1be3a1
1 parent f6efa0a
commit a1be3a1
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 3 deletions.
diff --git 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<DefaultOptions, Task>
                 getDefaultValue: () => false);
             universalOption.AddAlias("-u");
 
+            var autoShutdownOption = new Option<int>(
+                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<bool>(
                 name: "--break-glass",
                 description: "Flag; Ignore secret push protection results when pushing.",
@@ -88,10 +93,11 @@ public static RootCommand GenerateCommandLineOptions(Func<DefaultOptions, Task>
             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
@@ -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<StartOptions>
         private readonly Option<bool> _insecureOption;
         private readonly Option<bool> _dumpOption;
         private readonly Option<bool> _univeralOutputOption;
+        private readonly Option<int> _autoShutdownTime;
         private readonly Argument<string[]> _additionalArgs;
 
-        public StartOptionsBinder(Option<string> storageLocationOption, Option<string> storagePluginOption, Option<bool> insecureOption, Option<bool> dumpOption, Option<bool> universalOutput, Argument<string[]> additionalArgs)
+        public StartOptionsBinder(Option<string> storageLocationOption, Option<string> storagePluginOption, Option<bool> insecureOption, Option<bool> dumpOption, Option<bool> universalOutput, Option<int> autoShutdownTime, Argument<string[]> 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
@@ -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
@@ -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
@@ -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
@@ -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<int> 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<Startup>()
@@ -180,6 +182,15 @@ private static void StartServer(StartOptions startOptions)
 
             var app = host.Build();
 
+            var shutdownService = app.Services.GetRequiredService<ShutdownConfiguration>();
+            if (startOptions.AutoShutdownTime > -1)
+            {
+                shutdownService.EnableAutoShutdown = true;
+                shutdownService.TimeoutInSeconds = startOptions.AutoShutdownTime;
+                // start the first iteration of the shutdown timer
+                app.Services.GetRequiredService<ShutdownTimer>().ResetTimer();
+            }
+
             if (startOptions.Dump)
             {
                 var config = app.Services?.GetService<IConfiguration>();
@@ -227,6 +238,8 @@ public void ConfigureServices(IServiceCollection services)
             );
 
             services.AddSingleton<RecordingHandler>(singletonRecordingHandler);
+            services.AddSingleton<ShutdownConfiguration>();
+            services.AddSingleton<ShutdownTimer>();
         }
 
         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<HttpExceptionMiddleware>();
+            app.UseMiddleware<ShutdownTimerMiddleware>();
 
             DebugLogger.ConfigureLogger(loggerFactory);