Refactor registration to artifact operation (dotnet#3960)

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

@@ -642,12 +642,6 @@ private Task<ActionResult> StartTrace(
  configuration,
  duration);
 
- if (_diagnosticPortOptions.Value.ConnectionMode == DiagnosticPortConnectionMode.Listen)
- {
- IDisposable operationRegistration = _operationTrackerService.Register(processInfo.EndpointInfo);
- HttpContext.Response.RegisterForDispose(operationRegistration);
- }
-
  return Result(
  Utilities.ArtifactType_Trace,
  egressProvider,

src/Tools/dotnet-monitor/CollectionRules/Actions/CollectTraceAction.cs

@@ -45,14 +45,12 @@ private sealed class CollectTraceAction :
  {
  private readonly IServiceProvider _serviceProvider;
  private readonly IOptionsMonitor<GlobalCounterOptions> _counterOptions;
- private readonly OperationTrackerService _operationTrackerService;
 
  public CollectTraceAction(IServiceProvider serviceProvider, IEndpointInfo endpointInfo, CollectTraceOptions options)
  : base(endpointInfo, options)
  {
  _serviceProvider = serviceProvider ?? throw new ArgumentNullException(nameof(serviceProvider));
  _counterOptions = _serviceProvider.GetRequiredService<IOptionsMonitor<GlobalCounterOptions>>();
- _operationTrackerService = _serviceProvider.GetRequiredService<OperationTrackerService>();
  }
 
  protected override async Task<CollectionRuleActionResult> ExecuteCoreAsync(
@@ -109,7 +107,6 @@ protected override async Task<CollectionRuleActionResult> ExecuteCoreAsync(
  EgressOperation egressOperation = new EgressOperation(
  async (outputStream, token) =>
  {
- using IDisposable operationRegistration = _operationTrackerService.Register(EndpointInfo);
  await operation.ExecuteAsync(outputStream, startCompletionSource, token);
  },
  egressProvider,

src/Tools/dotnet-monitor/Logs/LogsOperation.cs

@@ -19,8 +19,8 @@ internal sealed class LogsOperation : PipelineArtifactOperation<EventLogsPipelin
  private readonly LogFormat _format;
  private readonly EventLogsPipelineSettings _settings;
 
- public LogsOperation(IEndpointInfo endpointInfo, EventLogsPipelineSettings settings, LogFormat format, ILogger logger)
- : base(logger, Utils.ArtifactType_Logs, endpointInfo)
+ public LogsOperation(IEndpointInfo endpointInfo, EventLogsPipelineSettings settings, LogFormat format, OperationTrackerService trackerService, ILogger logger)
+ : base(trackerService, logger, Utils.ArtifactType_Logs, endpointInfo)
  {
  _format = format;
  _settings = settings;

src/Tools/dotnet-monitor/Logs/LogsOperationFactory.cs

@@ -10,16 +10,18 @@ namespace Microsoft.Diagnostics.Tools.Monitor
 {
  internal sealed class LogsOperationFactory : ILogsOperationFactory
  {
+ private readonly OperationTrackerService _operationTrackerService;
  private readonly ILogger<LogsOperation> _logger;
 
- public LogsOperationFactory(ILogger<LogsOperation> logger)
+ public LogsOperationFactory(OperationTrackerService operationTrackerService, ILogger<LogsOperation> logger)
  {
+ _operationTrackerService = operationTrackerService;
  _logger = logger;
  }
 
  public IArtifactOperation Create(IEndpointInfo endpointInfo, EventLogsPipelineSettings settings, LogFormat format)
  {
- return new LogsOperation(endpointInfo, settings, format, _logger);
+ return new LogsOperation(endpointInfo, settings, format, _operationTrackerService, _logger);
  }
  }
 }

src/Tools/dotnet-monitor/Metrics/MetricsOperation.cs

@@ -17,8 +17,8 @@ internal sealed class MetricsOperation : PipelineArtifactOperation<MetricsPipeli
  {
  private readonly MetricsPipelineSettings _settings;
 
- public MetricsOperation(IEndpointInfo endpointInfo, MetricsPipelineSettings settings, ILogger logger)
- : base(logger, Utils.ArtifactType_Metrics, endpointInfo)
+ public MetricsOperation(IEndpointInfo endpointInfo, MetricsPipelineSettings settings, OperationTrackerService trackerService, ILogger logger)
+ : base(trackerService, logger, Utils.ArtifactType_Metrics, endpointInfo)
  {
  _settings = settings;
  }

src/Tools/dotnet-monitor/Metrics/MetricsOperationFactory.cs

@@ -9,16 +9,18 @@ namespace Microsoft.Diagnostics.Tools.Monitor
 {
  internal sealed class MetricsOperationFactory : IMetricsOperationFactory
  {
+ private readonly OperationTrackerService _operationTrackerService;
  private readonly ILogger<MetricsOperation> _logger;
 
- public MetricsOperationFactory(ILogger<MetricsOperation> logger)
+ public MetricsOperationFactory(OperationTrackerService operationTrackerService, ILogger<MetricsOperation> logger)
  {
+ _operationTrackerService = operationTrackerService;
  _logger = logger;
  }
 
  public IArtifactOperation Create(IEndpointInfo endpointInfo, MetricsPipelineSettings settings)
  {
- return new MetricsOperation(endpointInfo, settings, _logger);
+ return new MetricsOperation(endpointInfo, settings, _operationTrackerService, _logger);
  }
  }
 }

src/Tools/dotnet-monitor/PipelineArtifactOperation.cs

@@ -19,13 +19,17 @@ internal abstract class PipelineArtifactOperation<T> :
 
  private Func<CancellationToken, Task> _stopFunc;
 
- protected PipelineArtifactOperation(ILogger logger, string artifactType, IEndpointInfo endpointInfo, bool isStoppable = true)
+ protected OperationTrackerService OperationTrackerService { get; }
+
+ protected PipelineArtifactOperation(OperationTrackerService trackerService, ILogger logger, string artifactType, IEndpointInfo endpointInfo, bool isStoppable = true, bool register = false)
  {
  _artifactType = artifactType;
+ OperationTrackerService = trackerService;
 
  Logger = logger;
  EndpointInfo = endpointInfo;
  IsStoppable = isStoppable;
+ Register = register;
  }
 
  public async Task ExecuteAsync(Stream outputStream, TaskCompletionSource<object> startCompletionSource, CancellationToken token)
@@ -34,6 +38,8 @@ public async Task ExecuteAsync(Stream outputStream, TaskCompletionSource<object>
 
  _stopFunc = pipeline.StopAsync;
 
+ using IDisposable trackerRegistration = Register ? OperationTrackerService.Register(EndpointInfo) : null;
+
  Task runTask = await StartPipelineAsync(pipeline, token);
 
  Logger.StartCollectArtifact(_artifactType);
@@ -65,6 +71,8 @@ public async Task StopAsync(CancellationToken token)
 
  public bool IsStoppable { get; }
 
+ public bool Register { get; }
+
  protected abstract T CreatePipeline(Stream outputStream);
 
  /// <summary>

src/Tools/dotnet-monitor/Trace/AbstractTraceOperation.cs

@@ -16,8 +16,8 @@ internal abstract class AbstractTraceOperation : PipelineArtifactOperation<Event
 
  protected readonly EventTracePipelineSettings _settings;
 
- public AbstractTraceOperation(IEndpointInfo endpointInfo, EventTracePipelineSettings settings, ILogger logger)
- : base(logger, Utils.ArtifactType_Trace, endpointInfo)
+ public AbstractTraceOperation(IEndpointInfo endpointInfo, EventTracePipelineSettings settings, OperationTrackerService trackerService, ILogger logger)
+ : base(trackerService, logger, Utils.ArtifactType_Trace, endpointInfo, register: true)
  {
  _settings = settings;
  }

src/Tools/dotnet-monitor/Trace/TraceOperation.cs

@@ -15,8 +15,8 @@ internal sealed class TraceOperation : AbstractTraceOperation
  {
  private readonly TaskCompletionSource<object> _eventStreamAvailableCompletionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
- public TraceOperation(IEndpointInfo endpointInfo, EventTracePipelineSettings settings, ILogger logger)
- : base(endpointInfo, settings, logger) { }
+ public TraceOperation(IEndpointInfo endpointInfo, EventTracePipelineSettings settings, OperationTrackerService trackerService, ILogger logger)
+ : base(endpointInfo, settings, trackerService, logger) { }
 
  protected override EventTracePipeline CreatePipeline(Stream outputStream)
  {

src/Tools/dotnet-monitor/Trace/TraceOperationFactory.cs

@@ -11,10 +11,12 @@ namespace Microsoft.Diagnostics.Tools.Monitor
 {
  internal sealed class TraceOperationFactory : ITraceOperationFactory
  {
+ private readonly OperationTrackerService _operationTrackerService;
  private readonly ILogger<TraceOperation> _logger;
 
- public TraceOperationFactory(ILogger<TraceOperation> logger)
+ public TraceOperationFactory(OperationTrackerService operationTrackerService, ILogger<TraceOperation> logger)
  {
+ _operationTrackerService = operationTrackerService;
  _logger = logger;
  }
 
@@ -26,7 +28,7 @@ public IArtifactOperation Create(IEndpointInfo endpointInfo, MonitoringSourceCon
  Duration = duration
  };
 
- return new TraceOperation(endpointInfo, settings, _logger);
+ return new TraceOperation(endpointInfo, settings, _operationTrackerService, _logger);
  }
 
  public IArtifactOperation Create(IEndpointInfo endpointInfo, MonitoringSourceConfiguration configuration, TimeSpan duration, string providerName, string eventName, IDictionary<string, string> payloadFilter)
@@ -37,7 +39,7 @@ public IArtifactOperation Create(IEndpointInfo endpointInfo, MonitoringSourceCon
  Duration = duration
  };
 
- return new TraceUntilEventOperation(endpointInfo, settings, providerName, eventName, payloadFilter, _logger);
+ return new TraceUntilEventOperation(endpointInfo, settings, providerName, eventName, payloadFilter, _operationTrackerService, _logger);
  }
  }
 }

src/Tools/dotnet-monitor/Trace/TraceUntilEventOperation.cs

@@ -27,8 +27,9 @@ public TraceUntilEventOperation(
  string providerName,
  string eventName,
  IDictionary<string, string> payloadFilter,
+ OperationTrackerService trackerService,
  ILogger logger)
- : base(endpointInfo, settings, logger)
+ : base(endpointInfo, settings, trackerService, logger)
  {
  _providerName = providerName;
  _eventName = eventName;