Skip to content

Commit 001766a

Browse files
Address reviewer feedback: centralize walking state management in CosmosTraceDiagnostics
Co-authored-by: kirankumarkolli <6880899+kirankumarkolli@users.noreply.github.com>
1 parent 4e93c98 commit 001766a

5 files changed

Lines changed: 107 additions & 102 deletions

File tree

Microsoft.Azure.Cosmos/src/Diagnostics/CosmosTraceDiagnostics.cs

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,19 @@ public CosmosTraceDiagnostics(ITrace trace)
4040

4141
public override string ToString()
4242
{
43+
if (this.Value is Tracing.Trace rootConcreteTrace)
44+
{
45+
rootConcreteTrace.SetWalkingStateRecursively(true);
46+
try
47+
{
48+
return this.ToJsonString();
49+
}
50+
finally
51+
{
52+
rootConcreteTrace.SetWalkingStateRecursively(false);
53+
}
54+
}
55+
4356
return this.ToJsonString();
4457
}
4558

@@ -129,7 +142,23 @@ private ReadOnlyMemory<byte> WriteTraceToJsonWriter(JsonSerializationFormat json
129142
private static ServerSideCumulativeMetrics PopulateServerSideCumulativeMetrics(ITrace trace)
130143
{
131144
ServerSideMetricsInternalAccumulator accumulator = new ServerSideMetricsInternalAccumulator();
132-
ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(trace, accumulator);
145+
146+
if (trace is Tracing.Trace rootConcreteTrace)
147+
{
148+
rootConcreteTrace.SetWalkingStateRecursively(true);
149+
try
150+
{
151+
ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(trace, accumulator);
152+
}
153+
finally
154+
{
155+
rootConcreteTrace.SetWalkingStateRecursively(false);
156+
}
157+
}
158+
else
159+
{
160+
ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(trace, accumulator);
161+
}
133162

134163
IReadOnlyList<ServerSidePartitionedMetricsInternal> serverSideMetricsList = accumulator.GetPartitionedServerSideMetrics().Select(metrics => new ServerSidePartitionedMetricsInternal(metrics)).ToList();
135164

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/ServerSideMetricsTraceExtractor.cs

Lines changed: 17 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -18,37 +18,26 @@ public static void WalkTraceTreeForQueryMetrics(ITrace currentTrace, ServerSideM
1818
return;
1919
}
2020

21-
// Set walking state to prevent modifications during the walk
21+
// Assert that walking state is set
2222
if (currentTrace is Tracing.Trace concreteTrace)
2323
{
24-
concreteTrace.SetWalkingStateRecursively(true);
24+
System.Diagnostics.Debug.Assert(concreteTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
2525
}
2626

27-
try
27+
foreach (object datum in currentTrace.Data.Values)
2828
{
29-
foreach (object datum in currentTrace.Data.Values)
29+
if (datum is QueryMetricsTraceDatum queryMetricsTraceDatum)
3030
{
31-
if (datum is QueryMetricsTraceDatum queryMetricsTraceDatum)
32-
{
33-
ServerSideMetricsInternal serverSideMetrics = queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics;
34-
serverSideMetrics.FeedRange = currentTrace.Name;
35-
ServerSideMetricsTraceExtractor.WalkTraceTreeForPartitionInfo(currentTrace, serverSideMetrics);
36-
accumulator.Accumulate(serverSideMetrics);
37-
}
38-
}
39-
40-
foreach (ITrace childTrace in currentTrace.Children)
41-
{
42-
ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(childTrace, accumulator);
31+
ServerSideMetricsInternal serverSideMetrics = queryMetricsTraceDatum.QueryMetrics.ServerSideMetrics;
32+
serverSideMetrics.FeedRange = currentTrace.Name;
33+
ServerSideMetricsTraceExtractor.WalkTraceTreeForPartitionInfo(currentTrace, serverSideMetrics);
34+
accumulator.Accumulate(serverSideMetrics);
4335
}
4436
}
45-
finally
37+
38+
foreach (ITrace childTrace in currentTrace.Children)
4639
{
47-
// Clear walking state when done
48-
if (currentTrace is Tracing.Trace concreteTraceFinally)
49-
{
50-
concreteTraceFinally.SetWalkingStateRecursively(false);
51-
}
40+
ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(childTrace, accumulator);
5241
}
5342
}
5443

@@ -59,6 +48,12 @@ private static void WalkTraceTreeForPartitionInfo(ITrace currentTrace, ServerSid
5948
return;
6049
}
6150

51+
// Assert that walking state is set
52+
if (currentTrace is Tracing.Trace concreteTrace)
53+
{
54+
System.Diagnostics.Debug.Assert(concreteTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
55+
}
56+
6257
foreach (Object datum in currentTrace.Data.Values)
6358
{
6459
if (datum is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)

Microsoft.Azure.Cosmos/src/Tracing/Trace.cs

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ private Trace(
5959

6060
public IReadOnlyDictionary<string, object> Data => this.data.IsValueCreated ? this.data.Value : Trace.EmptyDictionary;
6161

62+
internal bool IsBeingWalked => this.isBeingWalked;
63+
6264
public void Dispose()
6365
{
6466
this.stopwatch.Stop();
@@ -67,6 +69,11 @@ public void Dispose()
6769
public ITrace StartChild(
6870
string name)
6971
{
72+
if (this.isBeingWalked)
73+
{
74+
return this; // Return self when being walked to avoid modifications
75+
}
76+
7077
return this.StartChild(
7178
name,
7279
level: TraceLevel.Verbose,
@@ -78,6 +85,11 @@ public ITrace StartChild(
7885
TraceComponent component,
7986
TraceLevel level)
8087
{
88+
if (this.isBeingWalked)
89+
{
90+
return this; // Return self when being walked to avoid modifications
91+
}
92+
8193
if (this.Parent != null && !this.stopwatch.IsRunning)
8294
{
8395
return this.Parent.StartChild(name, component, level);
@@ -168,14 +180,9 @@ public void AddOrUpdateDatum(string key, object value)
168180
}
169181
}
170182

171-
internal void SetWalkingState(bool isWalking)
172-
{
173-
this.isBeingWalked = isWalking;
174-
}
175-
176183
internal void SetWalkingStateRecursively(bool isWalking)
177184
{
178-
this.SetWalkingState(isWalking);
185+
this.isBeingWalked = isWalking;
179186

180187
lock (this.children)
181188
{

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

Lines changed: 10 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -38,36 +38,25 @@ public SummaryDiagnostics(ITrace trace)
3838

3939
private void CollectSummaryFromTraceTree(ITrace currentTrace)
4040
{
41-
// Set walking state to prevent modifications during the walk
41+
// Assert that walking state is set
4242
if (currentTrace is Trace concreteTrace)
4343
{
44-
concreteTrace.SetWalkingStateRecursively(true);
44+
System.Diagnostics.Debug.Assert(concreteTrace.IsBeingWalked, "SetWalkingStateRecursively should be set to true");
4545
}
4646

47-
try
47+
foreach (var datum in currentTrace.Data)
4848
{
49-
foreach (var datum in currentTrace.Data)
49+
if (datum.Value is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)
5050
{
51-
if (datum.Value is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)
52-
{
53-
this.AggregateStatsFromStoreResults(clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList);
54-
this.AggregateGatewayStatistics(clientSideRequestStatisticsTraceDatum.HttpResponseStatisticsList);
55-
this.AggregateRegionsContacted(clientSideRequestStatisticsTraceDatum.RegionsContacted);
56-
}
57-
}
58-
59-
foreach (ITrace childTrace in currentTrace.Children)
60-
{
61-
this.CollectSummaryFromTraceTree(childTrace);
51+
this.AggregateStatsFromStoreResults(clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList);
52+
this.AggregateGatewayStatistics(clientSideRequestStatisticsTraceDatum.HttpResponseStatisticsList);
53+
this.AggregateRegionsContacted(clientSideRequestStatisticsTraceDatum.RegionsContacted);
6254
}
6355
}
64-
finally
56+
57+
foreach (ITrace childTrace in currentTrace.Children)
6558
{
66-
// Clear walking state when done
67-
if (currentTrace is Trace concreteTraceFinally)
68-
{
69-
concreteTraceFinally.SetWalkingStateRecursively(false);
70-
}
59+
this.CollectSummaryFromTraceTree(childTrace);
7160
}
7261
}
7362

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

Lines changed: 37 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -34,74 +34,59 @@ public static void WriteTrace(
3434
throw new ArgumentNullException(nameof(trace));
3535
}
3636

37-
// Set walking state to prevent modifications during the walk
38-
if (isRootTrace && trace is Trace concreteTrace)
37+
writer.WriteObjectStart();
38+
39+
if (isRootTrace)
3940
{
40-
concreteTrace.SetWalkingStateRecursively(true);
41+
writer.WriteFieldName("Summary");
42+
SummaryDiagnostics summaryDiagnostics = new SummaryDiagnostics(trace);
43+
summaryDiagnostics.WriteSummaryDiagnostics(writer);
4144
}
45+
writer.WriteFieldName("name");
46+
writer.WriteStringValue(trace.Name);
4247

43-
try
48+
if (isRootTrace)
4449
{
50+
writer.WriteFieldName("start datetime");
51+
writer.WriteStringValue(trace.StartTime.ToString(TraceWriter.DateTimeFormatString));
52+
}
53+
writer.WriteFieldName("duration in milliseconds");
54+
writer.WriteNumberValue(trace.Duration.TotalMilliseconds);
55+
56+
// Use direct access since walking state prevents modifications
57+
if (trace.Data.Any())
58+
{
59+
writer.WriteFieldName("data");
4560
writer.WriteObjectStart();
4661

47-
if (isRootTrace)
62+
foreach (KeyValuePair<string, object> kvp in trace.Data)
4863
{
49-
writer.WriteFieldName("Summary");
50-
SummaryDiagnostics summaryDiagnostics = new SummaryDiagnostics(trace);
51-
summaryDiagnostics.WriteSummaryDiagnostics(writer);
52-
}
53-
writer.WriteFieldName("name");
54-
writer.WriteStringValue(trace.Name);
64+
string key = kvp.Key;
65+
object value = kvp.Value;
5566

56-
if (isRootTrace)
57-
{
58-
writer.WriteFieldName("start datetime");
59-
writer.WriteStringValue(trace.StartTime.ToString(TraceWriter.DateTimeFormatString));
67+
writer.WriteFieldName(key);
68+
WriteTraceDatum(writer, value);
6069
}
61-
writer.WriteFieldName("duration in milliseconds");
62-
writer.WriteNumberValue(trace.Duration.TotalMilliseconds);
63-
64-
if (trace.Data.Any())
65-
{
66-
writer.WriteFieldName("data");
67-
writer.WriteObjectStart();
68-
69-
foreach (KeyValuePair<string, object> kvp in trace.Data)
70-
{
71-
string key = kvp.Key;
72-
object value = kvp.Value;
73-
74-
writer.WriteFieldName(key);
75-
WriteTraceDatum(writer, value);
76-
}
7770

78-
writer.WriteObjectEnd();
79-
}
80-
81-
if (trace.Children.Any())
82-
{
83-
writer.WriteFieldName("children");
84-
writer.WriteArrayStart();
85-
86-
foreach (ITrace child in trace.Children)
87-
{
88-
WriteTrace(writer,
89-
child,
90-
isRootTrace: false);
91-
}
92-
93-
writer.WriteArrayEnd();
94-
}
9571
writer.WriteObjectEnd();
9672
}
97-
finally
73+
74+
// Use direct access since walking state prevents modifications
75+
if (trace.Children.Any())
9876
{
99-
// Clear walking state when done
100-
if (isRootTrace && trace is Trace concreteTraceFinally)
77+
writer.WriteFieldName("children");
78+
writer.WriteArrayStart();
79+
80+
foreach (ITrace child in trace.Children)
10181
{
102-
concreteTraceFinally.SetWalkingStateRecursively(false);
82+
WriteTrace(writer,
83+
child,
84+
isRootTrace: false);
10385
}
86+
87+
writer.WriteArrayEnd();
10488
}
89+
writer.WriteObjectEnd();
10590
}
10691
}
10792

0 commit comments

Comments
 (0)