Skip to content

Commit 59b70a6

Browse files
sourabh1007sourabh1007
authored
[Internal] AI Integration: Adds CorrelationId and Activity Id Attributes for query operation (#3630)
* add activityid in Otel attributes * added correlation id * operation type fix * remove test changes * test fix * fix baseline test * rename correlationId * fix tests again * include only not null attributes in test * fixed tests * changefeedxml * test fix * ordering activity in operationname oerder * fix test * review comments * refator header getter setter * clean up Co-authored-by: Sourabh Jain <[email protected]>
1 parent 416b154 commit 59b70a6

21 files changed

+312
-228
lines changed

‎Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,8 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties(
7373
SubStatusCodeLiteral = this.SubStatusCodeLiteral,
7474
ContentType = this.ContentType,
7575
QueryMetricsText = QueryMetricsText,
76-
IndexUtilizationText = IndexUtilizationText
76+
IndexUtilizationText = IndexUtilizationText,
77+
CorrelatedActivityId = this.CorrelatedActivityId
7778
};
7879
}
7980

@@ -108,7 +109,8 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
108109
SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral ?? (substatusCode.HasValue ? substatusCode.Value.ToString() : null),
109110
ContentType = sourceHeaders.ContentType,
110111
QueryMetricsText = sourceHeaders.QueryMetricsText,
111-
IndexUtilizationText = sourceHeaders.IndexUtilizationText
112+
IndexUtilizationText = sourceHeaders.IndexUtilizationText,
113+
CorrelatedActivityId = sourceHeaders.CorrelatedActivityId
112114
};
113115
}
114116
}

‎Microsoft.Azure.Cosmos/src/Headers/Headers.cs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,15 @@ public virtual string ContinuationToken
3636
internal set => this.CosmosMessageHeaders.Continuation = value;
3737
}
3838

39+
/// <summary>
40+
/// Gets or Set the CoorelatedActivityId in the current <see cref="ResponseMessage"/>.
41+
/// </summary>
42+
internal virtual string CorrelatedActivityId
43+
{
44+
get => this.CosmosMessageHeaders.Get(HttpConstants.HttpHeaders.CorrelatedActivityId);
45+
set => this.CosmosMessageHeaders.Set(HttpConstants.HttpHeaders.CorrelatedActivityId, value);
46+
}
47+
3948
/// <summary>
4049
/// Gets the request charge for this request from the Azure Cosmos DB service.
4150
/// </summary>

‎Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ public override async Task<ResponseMessage> ReadNextAsync(ITrace trace, Cancella
235235
{
236236
RequestCharge = tryGetQueryPage.Result.RequestCharge,
237237
ActivityId = tryGetQueryPage.Result.ActivityId,
238+
CorrelatedActivityId = this.correlatedActivityId.ToString(),
238239
SubStatusCode = Documents.SubStatusCodes.Unknown
239240
};
240241

‎Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributeKeys.cs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ internal sealed class OpenTelemetryAttributeKeys
2424
public const string ConnectionMode = "db.cosmosdb.connection_mode";
2525
public const string OperationType = "db.cosmosdb.operation_type";
2626

27-
// Request Specifics
27+
// Request/Response Specifics
2828
public const string ContainerName = "db.cosmosdb.container";
2929
public const string RequestContentLength = "db.cosmosdb.request_content_length_bytes";
3030
public const string ResponseContentLength = "db.cosmosdb.response_content_length_bytes";
@@ -35,6 +35,8 @@ internal sealed class OpenTelemetryAttributeKeys
3535
public const string RetryCount = "db.cosmosdb.retry_count";
3636
public const string ItemCount = "db.cosmosdb.item_count";
3737
public const string RequestDiagnostics = "db.cosmosdb.request_diagnostics";
38+
public const string ActivityId = "db.cosmosdb.activity_id";
39+
public const string CorrelatedActivityId = "db.cosmosdb.correlated_activity_id";
3840

3941
// Exceptions
4042
public const string ExceptionType = "exception.type";

‎Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributes.cs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,5 +54,20 @@ internal OpenTelemetryAttributes(RequestMessage requestMessage)
5454
/// SubStatusCode
5555
/// </summary>
5656
internal int SubStatusCode { get; set; }
57+
58+
/// <summary>
59+
/// ActivityId
60+
/// </summary>
61+
internal string ActivityId { get; set; }
62+
63+
/// <summary>
64+
/// CorrelatedActivityId
65+
/// </summary>
66+
internal string CorrelatedActivityId { get; set; }
67+
68+
/// <summary>
69+
/// OperationType
70+
/// </summary>
71+
internal Documents.OperationType OperationType { get; set; }
5772
}
5873
}

‎Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryCoreRecorder.cs‎

Lines changed: 28 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,13 @@ namespace Microsoft.Azure.Cosmos.Telemetry
1414
internal struct OpenTelemetryCoreRecorder : IDisposable
1515
{
1616
private const string CosmosDb = "cosmosdb";
17-
17+
1818
private readonly DiagnosticScope scope;
1919
private readonly DistributedTracingOptions config;
2020

2121
private readonly Documents.OperationType operationType;
22-
22+
private OpenTelemetryAttributes response = null;
23+
2324
internal static IDictionary<Type, Action<Exception, DiagnosticScope>> OTelCompatibleExceptions = new Dictionary<Type, Action<Exception, DiagnosticScope>>()
2425
{
2526
{ typeof(CosmosNullReferenceException), (exception, scope) => CosmosNullReferenceException.RecordOtelAttributes((CosmosNullReferenceException)exception, scope)},
@@ -41,15 +42,14 @@ public OpenTelemetryCoreRecorder(
4142
this.config = config;
4243
this.operationType = operationType;
4344

44-
if (this.IsEnabled)
45+
if (scope.IsEnabled)
4546
{
4647
this.scope.Start();
4748

4849
this.Record(
4950
operationName: operationName,
5051
containerName: containerName,
5152
databaseName: databaseName,
52-
operationType: operationType,
5353
clientContext: clientContext);
5454
}
5555
}
@@ -70,21 +70,18 @@ public void Record(string key, string value)
7070
/// <param name="operationName"></param>
7171
/// <param name="containerName"></param>
7272
/// <param name="databaseName"></param>
73-
/// <param name="operationType"></param>
7473
/// <param name="clientContext"></param>
7574
public void Record(
7675
string operationName,
7776
string containerName,
7877
string databaseName,
79-
Documents.OperationType operationType,
8078
CosmosClientContext clientContext)
8179
{
8280
if (this.IsEnabled)
8381
{
8482
this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbOperation, operationName);
8583
this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbName, databaseName);
8684
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ContainerName, containerName);
87-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, operationType);
8885

8986
// Other information
9087
this.scope.AddAttribute(OpenTelemetryAttributeKeys.DbSystemName, OpenTelemetryCoreRecorder.CosmosDb);
@@ -106,18 +103,7 @@ public void Record(OpenTelemetryAttributes response)
106103
{
107104
if (this.IsEnabled)
108105
{
109-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestContentLength, response.RequestContentLength);
110-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ResponseContentLength, response.ResponseContentLength);
111-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.StatusCode, (int)response.StatusCode);
112-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.SubStatusCode, (int)response.SubStatusCode);
113-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestCharge, response.RequestCharge);
114-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ItemCount, response.ItemCount);
115-
116-
if (response.Diagnostics != null)
117-
{
118-
this.scope.AddAttribute(OpenTelemetryAttributeKeys.Region, ClientTelemetryHelper.GetContactedRegions(response.Diagnostics.GetContactedRegions()));
119-
CosmosDbEventSource.RecordDiagnosticsForRequests(this.config, this.operationType, response);
120-
}
106+
this.response = response;
121107
}
122108
}
123109

@@ -168,6 +154,29 @@ public void Dispose()
168154
{
169155
if (this.scope.IsEnabled)
170156
{
157+
Documents.OperationType operationType
158+
= (this.response == null || this.response?.OperationType == Documents.OperationType.Invalid) ? this.operationType : this.response.OperationType;
159+
160+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.OperationType, operationType);
161+
162+
if (this.response != null)
163+
{
164+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestContentLength, this.response.RequestContentLength);
165+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ResponseContentLength, this.response.ResponseContentLength);
166+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.StatusCode, (int)this.response.StatusCode);
167+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.SubStatusCode, (int)this.response.SubStatusCode);
168+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.RequestCharge, this.response.RequestCharge);
169+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ItemCount, this.response.ItemCount);
170+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.ActivityId, this.response.ActivityId);
171+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.CorrelatedActivityId, this.response.CorrelatedActivityId);
172+
173+
if (this.response.Diagnostics != null)
174+
{
175+
this.scope.AddAttribute(OpenTelemetryAttributeKeys.Region, ClientTelemetryHelper.GetContactedRegions(this.response.Diagnostics.GetContactedRegions()));
176+
CosmosDbEventSource.RecordDiagnosticsForRequests(this.config, operationType, this.response);
177+
}
178+
}
179+
171180
this.scope.Dispose();
172181
}
173182
}

‎Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ internal OpenTelemetryResponse(TransactionalBatchResponse responseMessage)
1818
diagnostics: responseMessage.Diagnostics,
1919
itemCount: responseMessage.Headers?.ItemCount,
2020
requestMessage: null,
21-
subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
21+
subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
22+
activityId: responseMessage.Headers?.ActivityId,
23+
correlationId: responseMessage.Headers?.CorrelatedActivityId)
2224
{
2325
}
2426

@@ -30,7 +32,11 @@ internal OpenTelemetryResponse(ResponseMessage responseMessage)
3032
diagnostics: responseMessage.Diagnostics,
3133
itemCount: responseMessage.Headers?.ItemCount,
3234
requestMessage: responseMessage.RequestMessage,
33-
subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
35+
subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
36+
activityId: responseMessage.Headers?.ActivityId,
37+
correlationId: responseMessage.Headers?.CorrelatedActivityId,
38+
operationType: responseMessage is QueryResponse ? Documents.OperationType.Query : Documents.OperationType.Invalid
39+
)
3440
{
3541
}
3642

@@ -41,7 +47,10 @@ private OpenTelemetryResponse(
4147
CosmosDiagnostics diagnostics,
4248
string itemCount,
4349
RequestMessage requestMessage,
44-
int subStatusCode)
50+
int subStatusCode,
51+
string activityId,
52+
string correlationId,
53+
Documents.OperationType operationType = Documents.OperationType.Invalid)
4554
: base(requestMessage)
4655
{
4756
this.StatusCode = statusCode;
@@ -50,6 +59,9 @@ private OpenTelemetryResponse(
5059
this.Diagnostics = diagnostics;
5160
this.ItemCount = itemCount;
5261
this.SubStatusCode = subStatusCode;
62+
this.ActivityId = activityId;
63+
this.CorrelatedActivityId = correlationId;
64+
this.OperationType = operationType;
5365
}
5466

5567
private static string GetPayloadSize(ResponseMessage response)

‎Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse{T}.cs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@ internal OpenTelemetryResponse(FeedResponse<T> responseMessage)
1717
diagnostics: responseMessage.Diagnostics,
1818
itemCount: responseMessage.Headers?.ItemCount,
1919
requestMessage: responseMessage.RequestMessage,
20-
subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
20+
subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
21+
activityId: responseMessage.Headers?.ActivityId,
22+
correlatedActivityId: responseMessage.Headers?.CorrelatedActivityId,
23+
operationType: responseMessage is QueryResponse<T> ? Documents.OperationType.Query : Documents.OperationType.Invalid)
2124
{
2225
}
2326

@@ -29,7 +32,10 @@ internal OpenTelemetryResponse(Response<T> responseMessage)
2932
diagnostics: responseMessage.Diagnostics,
3033
itemCount: responseMessage.Headers?.ItemCount,
3134
requestMessage: responseMessage.RequestMessage,
32-
subStatusCode: (int)responseMessage.Headers?.SubStatusCode)
35+
subStatusCode: (int)responseMessage.Headers?.SubStatusCode,
36+
activityId: responseMessage.Headers?.ActivityId,
37+
correlatedActivityId: responseMessage.Headers?.CorrelatedActivityId,
38+
operationType: responseMessage is QueryResponse ? Documents.OperationType.Query : Documents.OperationType.Invalid)
3339
{
3440
}
3541

@@ -40,7 +46,10 @@ private OpenTelemetryResponse(
4046
CosmosDiagnostics diagnostics,
4147
string itemCount,
4248
RequestMessage requestMessage,
43-
int subStatusCode)
49+
int subStatusCode,
50+
string activityId,
51+
string correlatedActivityId,
52+
Documents.OperationType operationType)
4453
: base(requestMessage)
4554
{
4655
this.StatusCode = statusCode;
@@ -49,6 +58,9 @@ private OpenTelemetryResponse(
4958
this.Diagnostics = diagnostics;
5059
this.ItemCount = itemCount;
5160
this.SubStatusCode = subStatusCode;
61+
this.ActivityId = activityId;
62+
this.CorrelatedActivityId = correlatedActivityId;
63+
this.OperationType = operationType;
5264
}
5365
}
5466
}

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/BaselineTests.cs‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,17 +116,13 @@ public void ExecuteTestSuite(IEnumerable<TInput> inputs, [CallerMemberName] stri
116116
string outputText =
117117
Regex.Replace(
118118
Regex.Replace(
119-
Regex.Replace(
120119
File.ReadAllText(outputPath), @"\s+", string.Empty),
121-
@"<ATTRIBUTE-VALUE>[\w\W]*?</ATTRIBUTE-VALUE>", string.Empty),
122120
@"<OPERATION>[\w\W]*?</OPERATION>", string.Empty); // in changefeed test in was changing
123121

124122
string baselineText =
125123
Regex.Replace(
126124
Regex.Replace(
127-
Regex.Replace(
128125
File.ReadAllText(baselinePath), @"\s+", string.Empty),
129-
@"<ATTRIBUTE-VALUE>[\w\W]*?</ATTRIBUTE-VALUE>", string.Empty),
130126
@"<OPERATION>[\w\W]*?</OPERATION>", string.Empty);
131127

132128
int commonPrefixLength = 0;

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/BaselineTest/TestBaseline/EndToEndTraceWriterBaselineTests.BatchOperationsAsync.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@
129129
}
130130
]
131131
}]]></Json>
132-
<OTelActivities><ACTIVITY><OPERATION>Operation.ExecuteAsync</OPERATION><ATTRIBUTE-KEY>kind</ATTRIBUTE-KEY><ATTRIBUTE-KEY>az.namespace</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.operation</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.container</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.operation_type</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.system</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.machine_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>net.peer.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.client_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.user_agent</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.connection_mode</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.response_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.sub_status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_charge</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.item_count</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.regions_contacted</ATTRIBUTE-KEY></ACTIVITY>
132+
<OTelActivities><ACTIVITY><OPERATION>Operation.ExecuteAsync</OPERATION><ATTRIBUTE-KEY>kind</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>client</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>az.namespace</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Microsoft.DocumentDB</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.operation</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>ExecuteAsync</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.name</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.container</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.system</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>cosmosdb</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.machine_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>net.peer.name</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>127.0.0.1</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.client_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.user_agent</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.connection_mode</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Direct</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.operation_type</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>Replace</ATTRIBUTE-VALUE><ATTRIBUTE-KEY>db.cosmosdb.request_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.response_content_length_bytes</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.sub_status_code</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.request_charge</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.item_count</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.activity_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.correlated_activity_id</ATTRIBUTE-KEY><ATTRIBUTE-KEY>db.cosmosdb.regions_contacted</ATTRIBUTE-KEY><ATTRIBUTE-VALUE>South Central US</ATTRIBUTE-VALUE></ACTIVITY>
133133
<EVENT>Ideally, this should contain request diagnostics but request diagnostics is subject to change with each request as it contains few unique id. So just putting this tag with this static text to make sure event is getting generated for each test.</EVENT>
134134
</OTelActivities>
135135
</Output>

0 commit comments

Comments
 (0)