Skip to content

Commit

Permalink
Merge
Browse files Browse the repository at this point in the history
  • Loading branch information
schmittjoseph committed Nov 10, 2022
1 parent 67367ae commit f475bc5
Show file tree
Hide file tree
Showing 33 changed files with 893 additions and 223 deletions.
28 changes: 27 additions & 1 deletion documentation/openapi.json
Original file line number Diff line number Diff line change
Expand Up @@ -1207,6 +1207,14 @@
"type": "string",
"format": "uuid"
}
},
{
"name": "stop",
"in": "query",
"schema": {
"type": "boolean",
"default": false
}
}
],
"responses": {
Expand All @@ -1215,6 +1223,9 @@
},
"200": {
"description": "Success"
},
"202": {
"description": "Accepted"
}
}
}
Expand Down Expand Up @@ -1496,7 +1507,8 @@
"Running",
"Succeeded",
"Failed",
"Cancelled"
"Cancelled",
"Stopping"
],
"type": "string"
},
Expand All @@ -1517,6 +1529,13 @@
"process": {
"$ref": "#/components/schemas/OperationProcessInfo"
},
"egressProviderName": {
"type": "string",
"nullable": true
},
"isStoppable": {
"type": "boolean"
},
"resourceLocation": {
"type": "string",
"nullable": true
Expand Down Expand Up @@ -1544,6 +1563,13 @@
},
"process": {
"$ref": "#/components/schemas/OperationProcessInfo"
},
"egressProviderName": {
"type": "string",
"nullable": true
},
"isStoppable": {
"type": "boolean"
}
},
"additionalProperties": false,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ public partial class DiagController : ControllerBase
private readonly ICollectionRuleService _collectionRuleService;
private readonly ProfilerChannel _profilerChannel;
private readonly ILogsOperationFactory _logsOperationFactory;
private readonly ITraceOperationFactory _traceOperationFactory;

public DiagController(ILogger<DiagController> logger,
IServiceProvider serviceProvider)
Expand All @@ -67,6 +68,7 @@ public DiagController(ILogger<DiagController> logger,
_collectionRuleService = serviceProvider.GetRequiredService<ICollectionRuleService>();
_profilerChannel = serviceProvider.GetRequiredService<ProfilerChannel>();
_logsOperationFactory = serviceProvider.GetRequiredService<ILogsOperationFactory>();
_traceOperationFactory = serviceProvider.GetRequiredService<ITraceOperationFactory>();
}

/// <summary>
Expand Down Expand Up @@ -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]
Expand All @@ -221,6 +222,8 @@ public Task<ActionResult> CaptureDump(
[FromQuery]
string egressProvider = null)
{
const string artifactType = Utilities.ArtifactType_Dump;

ProcessKey? processKey = Utilities.GetProcessKey(pid, uid, name);

return InvokeForProcess(async processInfo =>
Expand All @@ -229,6 +232,7 @@ public Task<ActionResult> CaptureDump(

if (string.IsNullOrEmpty(egressProvider))
{
await RegisterCurrentHttpResponseAsOperation(processInfo, artifactType);
Stream dumpStream = await _dumpService.DumpAsync(processInfo.EndpointInfo, type, HttpContext.RequestAborted);

_logger.WrittenToHttpStream();
Expand All @@ -238,17 +242,17 @@ public Task<ActionResult> CaptureDump(
}
else
{
KeyValueLogScope scope = Utilities.CreateArtifactScope(Utilities.ArtifactType_Dump, processInfo.EndpointInfo);
KeyValueLogScope scope = Utilities.CreateArtifactScope(artifactType, processInfo.EndpointInfo);

return await SendToEgress(new EgressOperation(
token => _dumpService.DumpAsync(processInfo.EndpointInfo, type, token),
egressProvider,
dumpFileName,
processInfo,
ContentTypes.ApplicationOctetStream,
scope), limitKey: Utilities.ArtifactType_Dump);
scope), limitKey: artifactType);
}
}, processKey, Utilities.ArtifactType_Dump);
}, processKey, artifactType);
}

/// <summary>
Expand All @@ -266,7 +270,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]
Expand Down Expand Up @@ -325,7 +328,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]
Expand Down Expand Up @@ -369,7 +371,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]
Expand Down Expand Up @@ -420,7 +421,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]
Expand Down Expand Up @@ -476,7 +476,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]
Expand Down Expand Up @@ -593,7 +592,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]
Expand Down Expand Up @@ -636,29 +634,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);
}

Expand Down Expand Up @@ -723,75 +713,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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
Expand Down
Loading

0 comments on commit f475bc5

Please sign in to comment.