147 lines
4.3 KiB
C#
147 lines
4.3 KiB
C#
|
|
using Microsoft.Extensions.Logging;
|
|
using System.Collections.Concurrent;
|
|
|
|
namespace StellaOps.Microservice;
|
|
|
|
/// <summary>
|
|
/// Tracks in-flight requests and manages their cancellation tokens.
|
|
/// </summary>
|
|
public sealed class InflightRequestTracker : IDisposable
|
|
{
|
|
private readonly ConcurrentDictionary<Guid, InflightRequest> _inflight = new();
|
|
private readonly ILogger<InflightRequestTracker> _logger;
|
|
private bool _disposed;
|
|
|
|
/// <summary>
|
|
/// Initializes a new instance of the <see cref="InflightRequestTracker"/> class.
|
|
/// </summary>
|
|
public InflightRequestTracker(ILogger<InflightRequestTracker> logger)
|
|
{
|
|
_logger = logger;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Gets the count of in-flight requests.
|
|
/// </summary>
|
|
public int Count => _inflight.Count;
|
|
|
|
/// <summary>
|
|
/// Starts tracking a request and returns a cancellation token for it.
|
|
/// </summary>
|
|
/// <param name="correlationId">The correlation ID of the request.</param>
|
|
/// <returns>A cancellation token that will be triggered if the request is cancelled.</returns>
|
|
public CancellationToken Track(Guid correlationId)
|
|
{
|
|
ObjectDisposedException.ThrowIf(_disposed, this);
|
|
|
|
var cts = new CancellationTokenSource();
|
|
var request = new InflightRequest(cts);
|
|
|
|
if (!_inflight.TryAdd(correlationId, request))
|
|
{
|
|
cts.Dispose();
|
|
throw new InvalidOperationException($"Request {correlationId} is already being tracked");
|
|
}
|
|
|
|
_logger.LogDebug("Started tracking request {CorrelationId}", correlationId);
|
|
return cts.Token;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Cancels a specific request.
|
|
/// </summary>
|
|
/// <param name="correlationId">The correlation ID of the request to cancel.</param>
|
|
/// <param name="reason">The reason for cancellation.</param>
|
|
/// <returns>True if the request was found and cancelled; otherwise false.</returns>
|
|
public bool Cancel(Guid correlationId, string? reason)
|
|
{
|
|
if (_inflight.TryGetValue(correlationId, out var request))
|
|
{
|
|
try
|
|
{
|
|
request.Cts.Cancel();
|
|
_logger.LogInformation(
|
|
"Cancelled request {CorrelationId}: {Reason}",
|
|
correlationId,
|
|
reason ?? "Unknown");
|
|
return true;
|
|
}
|
|
catch (ObjectDisposedException)
|
|
{
|
|
// CTS was already disposed, request completed
|
|
return false;
|
|
}
|
|
}
|
|
|
|
_logger.LogDebug(
|
|
"Cannot cancel request {CorrelationId}: not found (may have already completed)",
|
|
correlationId);
|
|
return false;
|
|
}
|
|
|
|
/// <summary>
|
|
/// Marks a request as completed and removes it from tracking.
|
|
/// </summary>
|
|
/// <param name="correlationId">The correlation ID of the completed request.</param>
|
|
public void Complete(Guid correlationId)
|
|
{
|
|
if (_inflight.TryRemove(correlationId, out var request))
|
|
{
|
|
request.Cts.Dispose();
|
|
_logger.LogDebug("Completed request {CorrelationId}", correlationId);
|
|
}
|
|
}
|
|
|
|
/// <summary>
|
|
/// Cancels all in-flight requests.
|
|
/// </summary>
|
|
/// <param name="reason">The reason for cancellation.</param>
|
|
public void CancelAll(string reason)
|
|
{
|
|
var count = 0;
|
|
foreach (var kvp in _inflight)
|
|
{
|
|
try
|
|
{
|
|
kvp.Value.Cts.Cancel();
|
|
count++;
|
|
}
|
|
catch (ObjectDisposedException)
|
|
{
|
|
// Already disposed
|
|
}
|
|
}
|
|
|
|
_logger.LogInformation("Cancelled {Count} in-flight requests: {Reason}", count, reason);
|
|
|
|
// Clear and dispose all
|
|
foreach (var kvp in _inflight)
|
|
{
|
|
if (_inflight.TryRemove(kvp.Key, out var request))
|
|
{
|
|
request.Cts.Dispose();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// <inheritdoc />
|
|
public void Dispose()
|
|
{
|
|
if (_disposed) return;
|
|
_disposed = true;
|
|
|
|
CancelAll("Disposing tracker");
|
|
}
|
|
|
|
private sealed class InflightRequest
|
|
{
|
|
public CancellationTokenSource Cts { get; }
|
|
|
|
public InflightRequest(CancellationTokenSource cts)
|
|
{
|
|
Cts = cts;
|
|
}
|
|
}
|
|
}
|