CosmosDiagnostics: Fixes a race condition causing InvalidOperationException

Microsoft.Azure.Cosmos/src/Tracing/TraceData/ClientSideRequestStatisticsTraceDatum.cs

@@ -7,16 +7,26 @@ namespace Microsoft.Azure.Cosmos.Tracing.TraceData
  using System;
  using System.Collections.Generic;
  using System.Diagnostics;
- using System.Net;
  using System.Net.Http;
  using System.Text;
  using Microsoft.Azure.Documents;
- using Microsoft.Azure.Documents.Collections;
 
  internal sealed class ClientSideRequestStatisticsTraceDatum : TraceDatum, IClientSideRequestStatistics
  {
+ private static readonly IEnumerable<KeyValuePair<string, AddressResolutionStatistics>> EmptyEndpointToAddressResolutionStatistics = new List<KeyValuePair<string, AddressResolutionStatistics>>();
+ private static readonly IReadOnlyList<StoreResponseStatistics> EmptyStoreResponseStatistics = new List<StoreResponseStatistics>();
+ private static readonly IReadOnlyList<HttpResponseStatistics> EmptyHttpResponseStatistics = new List<HttpResponseStatistics>();
+
  private readonly object lockObject = new object();
  private readonly long clientSideRequestStatisticsCreateTime;
+ private readonly Dictionary<string, AddressResolutionStatistics> endpointToAddressResolutionStats;
+ private readonly Dictionary<int, DateTime> recordRequestHashCodeToStartTime;
+ private readonly List<StoreResponseStatistics> storeResponseStatistics;
+ private readonly List<HttpResponseStatistics> httpResponseStatistics;
+
+ private IEnumerable<KeyValuePair<string, AddressResolutionStatistics>> shallowCopyOfEndpointToAddressResolutionStatistics = null;
+ private IReadOnlyList<StoreResponseStatistics> shallowCopyOfStoreResponseStatistics = null;
+ private IReadOnlyList<HttpResponseStatistics> shallowCopyOfHttpResponseStatistics = null;
 
  private long? firstStartRequestTimestamp;
  private long? lastStartRequestTimestamp;
@@ -27,36 +37,79 @@ public ClientSideRequestStatisticsTraceDatum(DateTime startTime)
  {
  this.RequestStartTimeUtc = startTime;
  this.RequestEndTimeUtc = null;
- this.EndpointToAddressResolutionStatistics = new Dictionary<string, AddressResolutionStatistics>();
- this.RecordRequestHashCodeToStartTime = new Dictionary<int, DateTime>();
+ this.endpointToAddressResolutionStats = new Dictionary<string, AddressResolutionStatistics>();
+ this.recordRequestHashCodeToStartTime = new Dictionary<int, DateTime>();
  this.ContactedReplicas = new List<Uri>();
- this.StoreResponseStatisticsList = new List<StoreResponseStatistics>();
+ this.storeResponseStatistics = new List<StoreResponseStatistics>();
  this.FailedReplicas = new HashSet<Uri>();
  this.RegionsContactedWithName = new HashSet<(string, Uri)>();
  this.clientSideRequestStatisticsCreateTime = Stopwatch.GetTimestamp();
- this.HttpResponseStatisticsList = new List<HttpResponseStatistics>();
+ this.httpResponseStatistics = new List<HttpResponseStatistics>();
  }
 
  public DateTime RequestStartTimeUtc { get; }
 
  public DateTime? RequestEndTimeUtc { get; set; }
 
- public Dictionary<string, AddressResolutionStatistics> EndpointToAddressResolutionStatistics { get; }
+ public IEnumerable<KeyValuePair<string, AddressResolutionStatistics>> EndpointToAddressResolutionStatistics
+ {
+ get
+ {
+ if (this.endpointToAddressResolutionStats.Count == 0)
+ {
+ return ClientSideRequestStatisticsTraceDatum.EmptyEndpointToAddressResolutionStatistics;
+ }
 
- private Dictionary<int, DateTime> RecordRequestHashCodeToStartTime { get; }
+ lock (this.endpointToAddressResolutionStats)
+ {
+ this.shallowCopyOfEndpointToAddressResolutionStatistics ??= CreateShallowCopy(this.endpointToAddressResolutionStats);
+ return this.shallowCopyOfEndpointToAddressResolutionStatistics;
+ }
+ }
+ }
 
  public List<Uri> ContactedReplicas { get; set; }
 
- public List<StoreResponseStatistics> StoreResponseStatisticsList { get; }
-
- public List<HttpResponseStatistics> HttpResponseStatisticsList { get; }
-
  public HashSet<Uri> FailedReplicas { get; }
 
  public HashSet<Uri> RegionsContacted { get; }
 
  public HashSet<(string, Uri)> RegionsContactedWithName { get; }
 
+ public IReadOnlyList<StoreResponseStatistics> StoreResponseStatisticsList
+ {
+ get
+ {
+ if (this.storeResponseStatistics.Count == 0)
+ {
+ return ClientSideRequestStatisticsTraceDatum.EmptyStoreResponseStatistics;
+ }
+
+ lock (this.storeResponseStatistics)
+ {
+ this.shallowCopyOfStoreResponseStatistics ??= ClientSideRequestStatisticsTraceDatum.CreateShallowCopy(this.storeResponseStatistics);
+ return this.shallowCopyOfStoreResponseStatistics;
+ }
+ }
+ }
+
+ public IReadOnlyList<HttpResponseStatistics> HttpResponseStatisticsList
+ {
+ get
+ {
+ if (this.httpResponseStatistics.Count == 0)
+ {
+ return ClientSideRequestStatisticsTraceDatum.EmptyHttpResponseStatistics;
+ }
+
+ lock (this.httpResponseStatistics)
+ {
+ this.shallowCopyOfHttpResponseStatistics ??= ClientSideRequestStatisticsTraceDatum.CreateShallowCopy(this.httpResponseStatistics);
+ return this.shallowCopyOfHttpResponseStatistics;
+ }
+ }
+ }
+
  public TimeSpan RequestLatency
  {
  get
@@ -91,7 +144,7 @@ public TimeSpan EstimatedClientDelayFromAllCauses
 
  public void RecordRequest(DocumentServiceRequest request)
  {
- lock (this.lockObject)
+ lock (this.storeResponseStatistics)
  {
  long timestamp = Stopwatch.GetTimestamp();
  if (this.received429ResponseSinceLastStartRequest)
@@ -109,14 +162,14 @@ public void RecordRequest(DocumentServiceRequest request)
  this.received429ResponseSinceLastStartRequest = false;
  }
 
- this.RecordRequestHashCodeToStartTime[request.GetHashCode()] = DateTime.UtcNow;
+ this.recordRequestHashCodeToStartTime[request.GetHashCode()] = DateTime.UtcNow;
  }
 
  public void RecordResponse(DocumentServiceRequest request, StoreResult storeResult)
  {
  // One DocumentServiceRequest can map to multiple store results
  DateTime? startDateTime = null;
- if (this.RecordRequestHashCodeToStartTime.TryGetValue(request.GetHashCode(), out DateTime startRequestTime))
+ if (this.recordRequestHashCodeToStartTime.TryGetValue(request.GetHashCode(), out DateTime startRequestTime))
  {
  startDateTime = startRequestTime;
  }
@@ -125,7 +178,7 @@ public void RecordResponse(DocumentServiceRequest request, StoreResult storeResu
  Debug.Fail("DocumentServiceRequest start time not recorded");
  }
 
- DateTime responseTime = DateTime.UtcNow;
+ DateTime responseTime = this.GetAndUpdateRequestEndTime();
  Uri locationEndpoint = request.RequestContext.LocationEndpointToRoute;
  string regionName = request.RequestContext.RegionName;
  StoreResponseStatistics responseStatistics = new StoreResponseStatistics(
@@ -141,19 +194,16 @@ public void RecordResponse(DocumentServiceRequest request, StoreResult storeResu
  this.IsCpuOverloaded = true;
  }
 
- lock (this.lockObject)
+ lock (this.storeResponseStatistics)
  {
- if (!this.RequestEndTimeUtc.HasValue || responseTime > this.RequestEndTimeUtc)
- {
- this.RequestEndTimeUtc = responseTime;
- }
-
  if (locationEndpoint != null)
  {
  this.RegionsContactedWithName.Add((regionName, locationEndpoint));
  }
 
- this.StoreResponseStatisticsList.Add(responseStatistics);
+ // Reset the shallow copy
+ this.shallowCopyOfStoreResponseStatistics = null;
+ this.storeResponseStatistics.Add(responseStatistics);
 
  if (!this.received429ResponseSinceLastStartRequest &&
  storeResult.StatusCode == StatusCodes.TooManyRequests)
@@ -171,9 +221,11 @@ public string RecordAddressResolutionStart(Uri targetEndpoint)
  endTime: DateTime.MaxValue,
  targetEndpoint: targetEndpoint == null ? "<NULL>" : targetEndpoint.ToString());
 
- lock (this.lockObject)
+ lock (this.endpointToAddressResolutionStats)
  {
- this.EndpointToAddressResolutionStatistics.Add(identifier, resolutionStats);
+ // Reset the shallow copy
+ this.shallowCopyOfEndpointToAddressResolutionStatistics = null;
+ this.endpointToAddressResolutionStats.Add(identifier, resolutionStats);
  }
 
  return identifier;
@@ -186,37 +238,36 @@ public void RecordAddressResolutionEnd(string identifier)
  return;
  }
 
- DateTime responseTime = DateTime.UtcNow;
- lock (this.lockObject)
+ DateTime responseTime = this.GetAndUpdateRequestEndTime();
+
+ lock (this.endpointToAddressResolutionStats)
  {
- if (!this.EndpointToAddressResolutionStatistics.ContainsKey(identifier))
+ if (!this.endpointToAddressResolutionStats.ContainsKey(identifier))
  {
  throw new ArgumentException("Identifier {0} does not exist. Please call start before calling end.", identifier);
  }
 
- if (!this.RequestEndTimeUtc.HasValue || responseTime > this.RequestEndTimeUtc)
- {
- this.RequestEndTimeUtc = responseTime;
- }
-
- AddressResolutionStatistics start = this.EndpointToAddressResolutionStatistics[identifier];
+ AddressResolutionStatistics start = this.endpointToAddressResolutionStats[identifier];
 
- this.EndpointToAddressResolutionStatistics[identifier] = new AddressResolutionStatistics(
+ // Reset the shallow copy
+ this.shallowCopyOfEndpointToAddressResolutionStatistics = null;
+ this.endpointToAddressResolutionStats[identifier] = new AddressResolutionStatistics(
  start.StartTime,
  responseTime,
  start.TargetEndpoint);
  }
  }
 
- public void RecordHttpResponse(HttpRequestMessage request, 
- HttpResponseMessage response, 
+ public void RecordHttpResponse(HttpRequestMessage request,
+ HttpResponseMessage response,
  ResourceType resourceType,
  DateTime requestStartTimeUtc)
  {
- lock (this.lockObject)
+ lock (this.httpResponseStatistics)
  {
- DateTime requestEndTimeUtc = this.RecordHttpResponseEndTime();
- this.HttpResponseStatisticsList.Add(new HttpResponseStatistics(requestStartTimeUtc,
+ this.shallowCopyOfHttpResponseStatistics = null;
+ DateTime requestEndTimeUtc = this.GetAndUpdateRequestEndTime();
+ this.httpResponseStatistics.Add(new HttpResponseStatistics(requestStartTimeUtc,
  requestEndTimeUtc,
  request.RequestUri,
  request.Method,
@@ -231,10 +282,11 @@ public void RecordHttpException(HttpRequestMessage request,
  ResourceType resourceType,
  DateTime requestStartTimeUtc)
  {
- lock (this.lockObject)
+ lock (this.httpResponseStatistics)
  {
- DateTime requestEndTimeUtc = this.RecordHttpResponseEndTime();
- this.HttpResponseStatisticsList.Add(new HttpResponseStatistics(requestStartTimeUtc,
+ this.shallowCopyOfHttpResponseStatistics = null;
+ DateTime requestEndTimeUtc = this.GetAndUpdateRequestEndTime();
+ this.httpResponseStatistics.Add(new HttpResponseStatistics(requestStartTimeUtc,
  requestEndTimeUtc,
  request.RequestUri,
  request.Method,
@@ -244,17 +296,42 @@ public void RecordHttpException(HttpRequestMessage request,
  }
  }
 
- private DateTime RecordHttpResponseEndTime()
+ private DateTime GetAndUpdateRequestEndTime()
  {
  DateTime requestEndTimeUtc = DateTime.UtcNow;
- if (!this.RequestEndTimeUtc.HasValue || requestEndTimeUtc > this.RequestEndTimeUtc)
+ lock (this.lockObject)
  {
- this.RequestEndTimeUtc = requestEndTimeUtc;
+ if (!this.RequestEndTimeUtc.HasValue || requestEndTimeUtc > this.RequestEndTimeUtc)
+ {
+ this.RequestEndTimeUtc = requestEndTimeUtc;
+ }
  }
 
  return requestEndTimeUtc;
  }
 
+ private static IReadOnlyList<T> CreateShallowCopy<T>(List<T> originalList)
+ {
+ List<T> shallowCopy = new List<T>(originalList.Count);
+ foreach (T item in originalList)
+ {
+ shallowCopy.Add(item);
+ }
+
+ return shallowCopy;
+ }
+
+ private static IEnumerable<KeyValuePair<string, T>> CreateShallowCopy<T>(Dictionary<string, T> originalDict)
+ {
+ List<KeyValuePair<string, T>> shallowCopy = new List<KeyValuePair<string, T>>(originalDict.Count);
+ foreach (KeyValuePair<string, T> item in originalDict)
+ {
+ shallowCopy.Add(item);
+ }
+
+ return shallowCopy;
+ }
+
  internal override void Accept(ITraceDatumVisitor traceDatumVisitor)
  {
  traceDatumVisitor.Visit(this);

Microsoft.Azure.Cosmos/src/Tracing/TraceWriter.TraceJsonWriter.cs

@@ -213,9 +213,9 @@ public void Visit(ClientSideRequestStatisticsTraceDatum clientSideRequestStatist
  this.jsonWriter.WriteFieldName("AddressResolutionStatistics");
  this.jsonWriter.WriteArrayStart();
 
- foreach (ClientSideRequestStatisticsTraceDatum.AddressResolutionStatistics stat in clientSideRequestStatisticsTraceDatum.EndpointToAddressResolutionStatistics.Values)
+ foreach (KeyValuePair<string, ClientSideRequestStatisticsTraceDatum.AddressResolutionStatistics> stat in clientSideRequestStatisticsTraceDatum.EndpointToAddressResolutionStatistics)
  {
- this.VisitAddressResolutionStatistics(stat);
+ this.VisitAddressResolutionStatistics(stat.Value);
  }
 
  this.jsonWriter.WriteArrayEnd();

Microsoft.Azure.Cosmos/src/Tracing/TraceWriter.TraceTextWriter.cs

@@ -382,12 +382,12 @@ public void Visit(ClientSideRequestStatisticsTraceDatum clientSideRequestStatist
  stringBuilder.AppendLine(AddressResolutionStatisticsTextTable.Singleton.TopLine);
  stringBuilder.AppendLine(AddressResolutionStatisticsTextTable.Singleton.Header);
  stringBuilder.AppendLine(AddressResolutionStatisticsTextTable.Singleton.MiddleLine);
- foreach (AddressResolutionStatistics stat in clientSideRequestStatisticsTraceDatum.EndpointToAddressResolutionStatistics.Values)
+ foreach (KeyValuePair<string, AddressResolutionStatistics> stat in clientSideRequestStatisticsTraceDatum.EndpointToAddressResolutionStatistics)
  {
  string row = AddressResolutionStatisticsTextTable.Singleton.GetRow(
- stat.StartTime.ToString("hh:mm:ss:fff", CultureInfo.InvariantCulture),
- stat.EndTime.HasValue ? stat.EndTime.Value.ToString("hh:mm:ss:fff", CultureInfo.InvariantCulture) : "NO END TIME",
- stat.TargetEndpoint);
+ stat.Value.StartTime.ToString("hh:mm:ss:fff", CultureInfo.InvariantCulture),
+ stat.Value.EndTime.HasValue ? stat.Value.EndTime.Value.ToString("hh:mm:ss:fff", CultureInfo.InvariantCulture) : "NO END TIME",
+ stat.Value.TargetEndpoint);
  stringBuilder.AppendLine(row);
  }
 
@@ -440,7 +440,7 @@ public void Visit(ClientSideRequestStatisticsTraceDatum clientSideRequestStatist
  }
  }
 
- if (clientSideRequestStatisticsTraceDatum.HttpResponseStatisticsList.Count > 0)
+ if (clientSideRequestStatisticsTraceDatum.HttpResponseStatisticsList.Any())
  {
  stringBuilder.AppendLine("Http Response Statistics");
  foreach (HttpResponseStatistics stat in clientSideRequestStatisticsTraceDatum.HttpResponseStatisticsList)