/// <summary>
/// Wraps the consumer ensuring that the consumer task continues to run
/// in the event of unforeseen failures.
/// </summary>
/// <param name="collection"></param>
private void ExecuteConsumer(PartitionedBlockingCollection <TRequest> collection)
{
var cancellationToken = CancellationToken;
while (!cancellationToken.IsCancellationRequested)
{
var shouldRetryImmediately = false;
try
{
shouldRetryImmediately = StreamRequests(collection, cancellationToken);
}
catch (OperationCanceledException)
{
return;
}
catch (Exception ex)
{
LogMessage(LogLevel.Finest, "Received an exception while attempting to StreamRequests.", ex);
}
if (!cancellationToken.IsCancellationRequested)
{
var delayInMs = shouldRetryImmediately ? 0 : _delayBetweenRpcCallsMs;
_delayer.Delay(delayInMs, cancellationToken);
}
}
}
private bool DequeueItems(PartitionedBlockingCollection <TRequest> collection, int maxBatchSize, CancellationToken cancellationToken, out IList <TRequest> items)
{
items = null;
if (!collection.Take(out var firstItem, cancellationToken))
{
return(false);
}
if (cancellationToken.IsCancellationRequested)
{
return(false);
}
items = new List <TRequest>();
items.Add(firstItem);
for (var i = 0; i < maxBatchSize - 1 && !cancellationToken.IsCancellationRequested && collection.TryTake(out var item); i++)
{
items.Add(item);
}
if (cancellationToken.IsCancellationRequested)
{
ProcessFailedItems(items, collection);
return(false);
}
return(true);
}
private void ProcessFailedItems(IList <TRequest> items, PartitionedBlockingCollection <TRequest> collection)
{
foreach (var item in items)
{
if (!collection.TryAdd(item))
{
_agentHealthReporter.ReportInfiniteTracingSpanEventsDropped(1);
}
}
}
/// <summary>
/// This executes every time a local configuration change is made. It is more convenient
/// that OnConfigurationChanged
/// </summary>
/// <param name="_"></param>
private void AgentConnected(AgentConnectedEvent _)
{
_spanStreamingService.Shutdown(false);
_schedulerSvc.StopExecuting(ReportSupportabilityMetrics);
var oldCapacity = (_spanEvents?.Capacity).GetValueOrDefault(0);
var oldPartitionCount = (_spanEvents?.PartitionCount).GetValueOrDefault(0);
var oldCount = (_spanEvents?.Count).GetValueOrDefault(0);
var oldCollection = _spanEvents != null
? Interlocked.Exchange(ref _spanEvents, null)
: null;
var newCapacity = _configuration.InfiniteTracingQueueSizeSpans;
var newPartitionCount = Math.Min(_configuration.InfiniteTracingQueueSizeSpans, _configuration.InfiniteTracingPartitionCountSpans);
if (!IsServiceEnabled || newCapacity <= 0 || newPartitionCount <= 0 || newPartitionCount > 62)
{
if (oldCount > 0)
{
RecordDroppedSpans(oldCount);
}
if (IsServiceEnabled)
{
Log.Info($"SpanEventAggregatorInfiniteTracing: Configuration is invalid - Infinite Tracing will NOT be enabled.");
LogConfiguration();
}
return;
}
if (oldCapacity == newCapacity && oldPartitionCount == newPartitionCount && oldCollection != null)
{
_spanEvents = oldCollection;
}
else
{
var overflowCount = oldCount - newCapacity;
if (overflowCount > 0)
{
RecordDroppedSpans(overflowCount);
}
_spanEvents = oldCollection != null
? new PartitionedBlockingCollection <Span>(newCapacity, newPartitionCount, oldCollection)
: new PartitionedBlockingCollection <Span>(newCapacity, newPartitionCount);
}
LogConfiguration();
_schedulerSvc.ExecuteEvery(ReportSupportabilityMetrics, TimeSpan.FromMinutes(1));
_spanStreamingService.StartConsumingCollection(_spanEvents);
}
/// <summary>
/// Designed to be called by the aggregator.
/// </summary>
/// <param name="collection"></param>
public void StartConsumingCollection(PartitionedBlockingCollection <TRequest> collection)
{
if (collection == null)
{
Log.Debug("Unable to start Data Streaming Service because queue was NULL.");
return;
}
Task.Run(() =>
{
Restart(collection);
});
}
private bool StreamRequests(PartitionedBlockingCollection <TRequest> collection, CancellationToken serviceCancellationToken)
{
var consumerId = _consumerId.Increment();
using (var streamCancellationTokenSource = new CancellationTokenSource())
using (var serviceAndStreamCancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(serviceCancellationToken, streamCancellationTokenSource.Token))
{
if (!GetRequestStreamWithRetry(consumerId, serviceAndStreamCancellationTokenSource.Token, out var requestStream, out var responseStream))
{
LogMessage(LogLevel.Debug, consumerId, "Unable to obtain Stream, exiting consumer");
streamCancellationTokenSource.Cancel();
return(false);
}
_hasAnyStreamStarted = true;
_responseStreamsDic[consumerId] = new ResponseStreamWrapper <TResponse>(consumerId, responseStream, serviceAndStreamCancellationTokenSource.Token);
while (!serviceCancellationToken.IsCancellationRequested && _grpcWrapper.IsConnected)
{
if (!DequeueItems(collection, BatchSizeConfigValue, serviceCancellationToken, out var items))
{
return(false);
}
if (items == null || items.Count == 0)
{
LogMessage(LogLevel.Debug, consumerId, $"Expected a {_modelType} from the collection, but it was null");
continue;
}
var trySendStatus = TrySend(consumerId, requestStream, items, serviceCancellationToken);
if (trySendStatus != TrySendStatus.Success)
{
ProcessFailedItems(items, collection);
_grpcWrapper.TryCloseRequestStream(requestStream);
streamCancellationTokenSource.Cancel();
return(trySendStatus == TrySendStatus.ErrorWithImmediateRetry);
}
IsStreaming = true;
}
streamCancellationTokenSource.Cancel();
}
return(false);
}
private void Restart(PartitionedBlockingCollection <TRequest> collection)
{
_grpcWrapper.Shutdown();
_collection = collection;
do
{
_shouldRestart = false;
if (StartService())
{
StartConsumers(); //This blocks until the cancellation token is triggered
}
} while (_shouldRestart);
}
public void QueueSize_PartitionedCorrectly(int configQueueSize, int configPartitionCount)
{
Mock.Arrange(() => _currentConfiguration.InfiniteTracingPartitionCountSpans).Returns(configPartitionCount);
Mock.Arrange(() => _currentConfiguration.InfiniteTracingQueueSizeSpans).Returns(configQueueSize);
var streamingSvc = GetMockStreamingService(true, true);
PartitionedBlockingCollection <Span> actualCollection = default;
Mock.Arrange(() => streamingSvc.StartConsumingCollection(Arg.IsAny <PartitionedBlockingCollection <Span> >()))
.DoInstead <PartitionedBlockingCollection <Span> >((c) => { actualCollection = c; });
var aggregator = CreateAggregator(streamingSvc);
FireAgentConnectedEvent();
NrAssert.Multiple
(
() => Assert.AreEqual(configQueueSize, aggregator.Capacity),
() => Assert.AreEqual(configQueueSize, actualCollection.Capacity)
);
}
public int PartitionCount_IsApplied(int configQueueSize, int?configPartitionCount)
{
if (configPartitionCount.HasValue)
{
Mock.Arrange(() => _currentConfiguration.InfiniteTracingPartitionCountSpans).Returns(configPartitionCount.Value);
}
Mock.Arrange(() => _currentConfiguration.InfiniteTracingQueueSizeSpans).Returns(configQueueSize);
var streamingSvc = GetMockStreamingService(true, true);
PartitionedBlockingCollection <Span> actualCollection = default;
Mock.Arrange(() => streamingSvc.StartConsumingCollection(Arg.IsAny <PartitionedBlockingCollection <Span> >()))
.DoInstead <PartitionedBlockingCollection <Span> >((c) => { actualCollection = c; });
var aggregator = CreateAggregator(streamingSvc);
FireAgentConnectedEvent();
return((actualCollection?.PartitionCount).GetValueOrDefault(0));
}
