[release/7.x] Support gracefully stopping traces on-demand

documentation/openapi.json

@@ -1207,6 +1207,14 @@
               "type": "string",
               "format": "uuid"
             }
+          },
+          {
+            "name": "stop",
+            "in": "query",
+            "schema": {
+              "type": "boolean",
+              "default": false
+            }
           }
         ],
         "responses": {
@@ -1215,6 +1223,9 @@
           },
           "200": {
             "description": "Success"
+          },
+          "202": {
+            "description": "Accepted"
           }
         }
       }
@@ -1496,7 +1507,8 @@
           "Running",
           "Succeeded",
           "Failed",
-          "Cancelled"
+          "Cancelled",
+          "Stopping"
         ],
         "type": "string"
       },
@@ -1517,6 +1529,13 @@
           "process": {
             "$ref": "#/components/schemas/OperationProcessInfo"
           },
+          "egressProviderName": {
+            "type": "string",
+            "nullable": true
+          },
+          "isStoppable": {
+            "type": "boolean"
+          },
           "resourceLocation": {
             "type": "string",
             "nullable": true
@@ -1544,6 +1563,13 @@
           },
           "process": {
             "$ref": "#/components/schemas/OperationProcessInfo"
+          },
+          "egressProviderName": {
+            "type": "string",
+            "nullable": true
+          },
+          "isStoppable": {
+            "type": "boolean"
           }
         },
         "additionalProperties": false,

src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.Metrics.cs

@@ -25,7 +25,6 @@ partial class DiagController
         [ProducesWithProblemDetails(ContentTypes.ApplicationJsonSequence)]
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Metrics)]
         [EgressValidation]
         public Task<ActionResult> CaptureMetrics(
             [FromQuery]
@@ -75,7 +74,6 @@ public Task<ActionResult> CaptureMetrics(
         [ProducesWithProblemDetails(ContentTypes.ApplicationJsonSequence)]
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Metrics)]
         [EgressValidation]
         public Task<ActionResult> CaptureMetricsCustom(
             [FromBody][Required]

src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/DiagController.cs

@@ -52,6 +52,7 @@ public partial class DiagController : ControllerBase
         private readonly ProfilerChannel _profilerChannel;
         private readonly IDumpOperationFactory _dumpOperationFactory;
         private readonly ILogsOperationFactory _logsOperationFactory;
+        private readonly ITraceOperationFactory _traceOperationFactory;
 
         public DiagController(ILogger<DiagController> logger,
             IServiceProvider serviceProvider)
@@ -67,6 +68,7 @@ public DiagController(ILogger<DiagController> logger,
             _profilerChannel = serviceProvider.GetRequiredService<ProfilerChannel>();
             _dumpOperationFactory = serviceProvider.GetRequiredService<IDumpOperationFactory>();
             _logsOperationFactory = serviceProvider.GetRequiredService<ILogsOperationFactory>();
+            _traceOperationFactory = serviceProvider.GetRequiredService<ITraceOperationFactory>();
         }
 
         /// <summary>
@@ -207,7 +209,6 @@ public Task<ActionResult<Dictionary<string, string>>> GetProcessEnvironment(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(FileResult), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Dump)]
         [EgressValidation]
         public Task<ActionResult> CaptureDump(
             [FromQuery]
@@ -248,7 +249,6 @@ public Task<ActionResult> CaptureDump(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(FileResult), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_GCDump)]
         [EgressValidation]
         public Task<ActionResult> CaptureGcDump(
             [FromQuery]
@@ -307,7 +307,6 @@ public Task<ActionResult> CaptureGcDump(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(FileResult), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Trace)]
         [EgressValidation]
         public Task<ActionResult> CaptureTrace(
             [FromQuery]
@@ -351,7 +350,6 @@ public Task<ActionResult> CaptureTrace(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(FileResult), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Trace)]
         [EgressValidation]
         public Task<ActionResult> CaptureTraceCustom(
             [FromBody][Required]
@@ -402,7 +400,6 @@ public Task<ActionResult> CaptureTraceCustom(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(string), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Logs)]
         [EgressValidation]
         public Task<ActionResult> CaptureLogs(
             [FromQuery]
@@ -458,7 +455,6 @@ public Task<ActionResult> CaptureLogs(
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(string), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Logs)]
         [EgressValidation]
         public Task<ActionResult> CaptureLogsCustom(
             [FromBody]
@@ -575,7 +571,6 @@ public Task<ActionResult<CollectionRuleDetailedDescription>> GetCollectionRuleDe
         [ProducesResponseType(typeof(ProblemDetails), StatusCodes.Status429TooManyRequests)]
         [ProducesResponseType(typeof(string), StatusCodes.Status200OK)]
         [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
-        [RequestLimit(LimitKey = Utilities.ArtifactType_Stacks)]
         [EgressValidation]
         public async Task<ActionResult> CaptureStacks(
             [FromQuery]
@@ -618,29 +613,21 @@ private Task<ActionResult> StartTrace(
             TimeSpan duration,
             string egressProvider)
         {
-            string fileName = TraceUtilities.GenerateTraceFileName(processInfo.EndpointInfo);
+            IArtifactOperation traceOperation = _traceOperationFactory.Create(
+                processInfo.EndpointInfo,
+                configuration,
+                duration);
+
+            if (_diagnosticPortOptions.Value.ConnectionMode == DiagnosticPortConnectionMode.Listen)
+            {
+                IDisposable operationRegistration = _operationTrackerService.Register(processInfo.EndpointInfo);
+                HttpContext.Response.RegisterForDispose(operationRegistration);
+            }
 
             return Result(
                 Utilities.ArtifactType_Trace,
                 egressProvider,
-                async (outputStream, token) =>
-                {
-                    IDisposable operationRegistration = null;
-                    try
-                    {
-                        if (_diagnosticPortOptions.Value.ConnectionMode == DiagnosticPortConnectionMode.Listen)
-                        {
-                            operationRegistration = _operationTrackerService.Register(processInfo.EndpointInfo);
-                        }
-                        await TraceUtilities.CaptureTraceAsync(null, processInfo.EndpointInfo, configuration, duration, outputStream, token);
-                    }
-                    finally
-                    {
-                        operationRegistration?.Dispose();
-                    }
-                },
-                fileName,
-                ContentTypes.ApplicationOctetStream,
+                traceOperation,
                 processInfo);
         }
 
@@ -705,75 +692,89 @@ private Task<ActionResult> StartLogs(
             return null;
         }
 
-        private Task<ActionResult> Result(
+        private async Task<ActionResult> Result(
             string artifactType,
             string providerName,
-            Func<Stream, CancellationToken, Task> action,
-            string fileName,
-            string contentType,
+            IArtifactOperation operation,
             IProcessInfo processInfo,
             bool asAttachment = true)
         {
             KeyValueLogScope scope = Utilities.CreateArtifactScope(artifactType, processInfo.EndpointInfo);
 
             if (string.IsNullOrEmpty(providerName))
             {
-                return Task.FromResult<ActionResult>(new OutputStreamResult(
-                    action,
-                    contentType,
-                    asAttachment ? fileName : null,
-                    scope));
+                await RegisterCurrentHttpResponseAsOperation(processInfo, artifactType, operation);
+                return new OutputStreamResult(
+                    operation,
+                    asAttachment ? operation.GenerateFileName() : null,
+                    scope);
             }
             else
             {
-                return SendToEgress(new EgressOperation(
-                    action,
+                return await SendToEgress(new EgressOperation(
+                    operation,
                     providerName,
-                    fileName,
                     processInfo,
-                    contentType,
                     scope),
                     limitKey: artifactType);
             }
         }
 
-        private Task<ActionResult> Result(
+        private async Task<ActionResult> Result(
             string artifactType,
             string providerName,
-            IArtifactOperation operation,
+            Func<Stream, CancellationToken, Task> action,
+            string fileName,
+            string contentType,
             IProcessInfo processInfo,
             bool asAttachment = true)
         {
             KeyValueLogScope scope = Utilities.CreateArtifactScope(artifactType, processInfo.EndpointInfo);
 
             if (string.IsNullOrEmpty(providerName))
             {
-                return Task.FromResult<ActionResult>(new OutputStreamResult(
-                    operation,
-                    asAttachment ? operation.GenerateFileName() : null,
-                    scope));
+                await RegisterCurrentHttpResponseAsOperation(processInfo, artifactType);
+                return new OutputStreamResult(
+                    action,
+                    contentType,
+                    asAttachment ? fileName : null,
+                    scope);
             }
             else
             {
-                return SendToEgress(new EgressOperation(
-                    operation,
+                return await SendToEgress(new EgressOperation(
+                    action,
                     providerName,
+                    fileName,
                     processInfo,
+                    contentType,
                     scope),
                     limitKey: artifactType);
             }
         }
 
-        private async Task<ActionResult> SendToEgress(EgressOperation egressStreamResult, string limitKey)
+        private async Task RegisterCurrentHttpResponseAsOperation(IProcessInfo processInfo, string artifactType, IArtifactOperation operation = null)
+        {
+            // While not strictly a Location redirect, use the same header as externally egressed operations for consistency.
+            HttpContext.Response.Headers["Location"] = await RegisterOperation(
+                new HttpResponseEgressOperation(HttpContext, processInfo, operation),
+                limitKey: artifactType);
+        }
+
+        private async Task<string> RegisterOperation(IEgressOperation egressOperation, string limitKey)
         {
             // Will throw TooManyRequestsException if there are too many concurrent operations.
-            Guid operationId = await _operationsStore.AddOperation(egressStreamResult, limitKey);
-            string newUrl = this.Url.Action(
+            Guid operationId = await _operationsStore.AddOperation(egressOperation, limitKey);
+            return this.Url.Action(
                 action: nameof(OperationsController.GetOperationStatus),
                 controller: OperationsController.ControllerName, new { operationId = operationId },
                 protocol: this.HttpContext.Request.Scheme, this.HttpContext.Request.Host.ToString());
+        }
 
-            return Accepted(newUrl);
+        private async Task<ActionResult> SendToEgress(IEgressOperation egressOperation, string limitKey)
+        {
+            string operationUrl = await RegisterOperation(egressOperation, limitKey);
+            return Accepted(operationUrl);
         }
 
         private Task<ActionResult> InvokeForProcess(Func<IProcessInfo, ActionResult> func, ProcessKey? processKey, string artifactType = null)

src/Microsoft.Diagnostics.Monitoring.WebApi/Controllers/OperationsController.cs

@@ -72,14 +72,31 @@ public IActionResult GetOperationStatus(Guid operationId)
         [HttpDelete("{operationId}")]
         [ProducesWithProblemDetails(ContentTypes.ApplicationJson)]
         [ProducesResponseType(typeof(void), StatusCodes.Status200OK)]
-        public IActionResult CancelOperation(Guid operationId)
+        [ProducesResponseType(typeof(void), StatusCodes.Status202Accepted)]
+        public IActionResult CancelOperation(
+            Guid operationId,
+            [FromQuery]
+            bool stop = false)
         {
             return this.InvokeService(() =>
             {
                 //Note that if the operation is not found, it will throw an InvalidOperationException and
                 //return an error code.
-                _operationsStore.CancelOperation(operationId);
-                return Ok();
+                if (stop)
+                {
+                    // If stopping an operation fails, it's undefined behavior.
+                    // Leave the operation in the "Stopping" state and it'll either complete on its own
+                    // or the user will cancel it.
+                    _operationsStore.StopOperation(operationId, (ex) => _logger.StopOperationFailed(operationId, ex));
+
+                    // Stop operations are not instant, they are instead queued and can take an indeterminate amount of time.
+                    return Accepted();
+                }
+                else
+                {
+                    _operationsStore.CancelOperation(operationId);
+                    return Ok();
+                }
             }, _logger);
         }
     }

src/Microsoft.Diagnostics.Monitoring.WebApi/ITraceOperationFactory.cs

@@ -0,0 +1,35 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.Diagnostics.Monitoring.EventPipe;
+using System;
+using System.Collections.Generic;
+
+namespace Microsoft.Diagnostics.Monitoring.WebApi
+{
+    /// <summary>
+    /// Factory for creating operations that produce trace artifacts.
+    /// </summary>
+    internal interface ITraceOperationFactory
+    {
+        /// <summary>
+        /// Creates an operation that produces a trace artifact.
+        /// </summary>
+        IArtifactOperation Create(
+            IEndpointInfo endpointInfo,
+            MonitoringSourceConfiguration configuration,
+            TimeSpan duration);
+
+        /// <summary>
+        /// Creates an operation that produces a trace artifact and supports a stopping event.
+        /// </summary>
+        IArtifactOperation Create(
+            IEndpointInfo endpointInfo,
+            MonitoringSourceConfiguration configuration,
+            TimeSpan duration,
+            string providerName,
+            string eventName,
+            IDictionary<string, string> payloadFilter);
+    }
+}