Include session token to diagnostics

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

@@ -27,6 +27,8 @@ internal class PointOperationStatistics : CosmosDiagnostics
         public string ErrorMessage { get; }
         public HttpMethod Method { get; }
         public Uri RequestUri { get; }
+        public string RequestSessionToken { get; }
+        public string ResponseSessionToken { get; }
 
         public DateTime requestStartTime { get; private set; }
 
@@ -54,6 +56,8 @@ internal PointOperationStatistics(
             string errorMessage,
             HttpMethod method,
             Uri requestUri,
+            string requestSessionToken,
+            string responseSessionToken,
             CosmosClientSideRequestStatistics clientSideRequestStatistics)
         {
             this.ActivityId = activityId;
@@ -63,6 +67,8 @@ internal PointOperationStatistics(
             this.ErrorMessage = errorMessage;
             this.Method = method;
             this.RequestUri = requestUri;
+            this.RequestSessionToken = requestSessionToken;
+            this.ResponseSessionToken = responseSessionToken;
             if (clientSideRequestStatistics != null)
             {
                 this.requestStartTime = clientSideRequestStatistics.requestStartTime;

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

@@ -130,7 +130,7 @@ public static QueryIterator Create(
             }
             catch (Documents.DocumentClientException exception)
             {
-                response = exception.ToCosmosResponseMessage(request: null);
+                response = exception.ToCosmosResponseMessage(requestMessage: null);
             }
             catch (CosmosException exception)
             {

Microsoft.Azure.Cosmos/src/Util/Extensions.cs

@@ -18,163 +18,167 @@ internal static class Extensions
     {
         private static readonly char[] NewLineCharacters = new[] { '\r', '\n' };
 
-        internal static ResponseMessage ToCosmosResponseMessage(this DocumentServiceResponse response, RequestMessage requestMessage)
+        internal static ResponseMessage ToCosmosResponseMessage(this DocumentServiceResponse documentServiceResponse, RequestMessage requestMessage)
         {
             Debug.Assert(requestMessage != null, nameof(requestMessage));
 
-            ResponseMessage cosmosResponse = new ResponseMessage(response.StatusCode, requestMessage);
-            if (response.ResponseBody != null)
+            ResponseMessage responseMessage = new ResponseMessage(documentServiceResponse.StatusCode, requestMessage);
+            if (documentServiceResponse.ResponseBody != null)
             {
-                cosmosResponse.Content = response.ResponseBody;
+                responseMessage.Content = documentServiceResponse.ResponseBody;
             }
 
-            if (response.Headers != null)
+            if (documentServiceResponse.Headers != null)
             {
-                foreach (string key in response.Headers)
+                foreach (string key in documentServiceResponse.Headers)
                 {
-                    cosmosResponse.Headers.Add(key, response.Headers[key]);
+                    responseMessage.Headers.Add(key, documentServiceResponse.Headers[key]);
                 }
             }
 
-            CosmosClientSideRequestStatistics cosmosClientSideRequestStatistics = response.RequestStats as CosmosClientSideRequestStatistics;
-            cosmosResponse.Diagnostics = new PointOperationStatistics(
-                activityId: cosmosResponse.Headers.ActivityId,
-                statusCode: response.StatusCode,
-                subStatusCode: response.SubStatusCode,
-                requestCharge: cosmosResponse.Headers.RequestCharge,
-                errorMessage: cosmosResponse.ErrorMessage,
+            CosmosClientSideRequestStatistics cosmosClientSideRequestStatistics = documentServiceResponse.RequestStats as CosmosClientSideRequestStatistics;
+            responseMessage.Diagnostics = new PointOperationStatistics(
+                activityId: responseMessage.Headers.ActivityId,
+                statusCode: documentServiceResponse.StatusCode,
+                subStatusCode: documentServiceResponse.SubStatusCode,
+                requestCharge: responseMessage.Headers.RequestCharge,
+                errorMessage: responseMessage.ErrorMessage,
                 method: requestMessage?.Method,
                 requestUri: requestMessage?.RequestUri,
+                requestSessionToken: requestMessage?.Headers?.Session,
+                responseSessionToken: responseMessage.Headers.Session,
                 clientSideRequestStatistics: cosmosClientSideRequestStatistics);
 
-            return cosmosResponse;
+            return responseMessage;
         }
 
-        internal static ResponseMessage ToCosmosResponseMessage(this DocumentClientException dce, RequestMessage request)
+        internal static ResponseMessage ToCosmosResponseMessage(this DocumentClientException documentClientException, RequestMessage requestMessage)
         {
             // if StatusCode is null it is a client business logic error and it never hit the backend, so throw
-            if (dce.StatusCode == null)
+            if (documentClientException.StatusCode == null)
             {
-                throw dce;
+                throw documentClientException;
             }
 
             // if there is a status code then it came from the backend, return error as http error instead of throwing the exception
-            ResponseMessage cosmosResponse = new ResponseMessage(dce.StatusCode ?? HttpStatusCode.InternalServerError, request);
+            ResponseMessage responseMessage = new ResponseMessage(documentClientException.StatusCode ?? HttpStatusCode.InternalServerError, requestMessage);
             string reasonPhraseString = string.Empty;
-            if (!string.IsNullOrEmpty(dce.Message))
+            if (!string.IsNullOrEmpty(documentClientException.Message))
             {
-                if (dce.Message.IndexOfAny(Extensions.NewLineCharacters) >= 0)
+                if (documentClientException.Message.IndexOfAny(Extensions.NewLineCharacters) >= 0)
                 {
-                    StringBuilder sb = new StringBuilder(dce.Message);
+                    StringBuilder sb = new StringBuilder(documentClientException.Message);
                     sb = sb.Replace("\r", string.Empty);
                     sb = sb.Replace("\n", string.Empty);
                     reasonPhraseString = sb.ToString();
                 }
                 else
                 {
-                    reasonPhraseString = dce.Message;
+                    reasonPhraseString = documentClientException.Message;
                 }
             }
 
-            cosmosResponse.ErrorMessage = reasonPhraseString;
-            cosmosResponse.Error = dce.Error;
+            responseMessage.ErrorMessage = reasonPhraseString;
+            responseMessage.Error = documentClientException.Error;
 
-            if (dce.Headers != null)
+            if (documentClientException.Headers != null)
             {
-                foreach (string header in dce.Headers.AllKeys())
+                foreach (string header in documentClientException.Headers.AllKeys())
                 {
-                    cosmosResponse.Headers.Add(header, dce.Headers[header]);
+                    responseMessage.Headers.Add(header, documentClientException.Headers[header]);
                 }
             }
 
-            cosmosResponse.Diagnostics = new PointOperationStatistics(
-                activityId: cosmosResponse.Headers.ActivityId,
-                statusCode: dce.StatusCode.Value,
+            responseMessage.Diagnostics = new PointOperationStatistics(
+                activityId: responseMessage.Headers.ActivityId,
+                statusCode: documentClientException.StatusCode.Value,
                 subStatusCode: SubStatusCodes.Unknown,
-                requestCharge: cosmosResponse.Headers.RequestCharge,
-                errorMessage: cosmosResponse.ErrorMessage,
-                method: request?.Method,
-                requestUri: request?.RequestUri,
-                clientSideRequestStatistics: dce.RequestStatistics as CosmosClientSideRequestStatistics);
+                requestCharge: responseMessage.Headers.RequestCharge,
+                errorMessage: responseMessage.ErrorMessage,
+                method: requestMessage?.Method,
+                requestUri: requestMessage?.RequestUri,
+                requestSessionToken: requestMessage?.Headers?.Session,
+                responseSessionToken: responseMessage.Headers.Session,
+                clientSideRequestStatistics: documentClientException.RequestStatistics as CosmosClientSideRequestStatistics);
 
-            if (request != null)
+            if (requestMessage != null)
             {
-                request.Properties.Remove(nameof(DocumentClientException));
-                request.Properties.Add(nameof(DocumentClientException), dce);
+                requestMessage.Properties.Remove(nameof(DocumentClientException));
+                requestMessage.Properties.Add(nameof(DocumentClientException), documentClientException);
             }
 
-            return cosmosResponse;
+            return responseMessage;
         }
 
-        internal static ResponseMessage ToCosmosResponseMessage(this StoreResponse response, RequestMessage request)
+        internal static ResponseMessage ToCosmosResponseMessage(this StoreResponse storeResponse, RequestMessage requestMessage)
         {
             // Is status code conversion lossy? 
-            ResponseMessage httpResponse = new ResponseMessage((HttpStatusCode)response.Status, request);
-            if (response.ResponseBody != null)
+            ResponseMessage responseMessage = new ResponseMessage((HttpStatusCode)storeResponse.Status, requestMessage);
+            if (storeResponse.ResponseBody != null)
             {
-                httpResponse.Content = response.ResponseBody;
+                responseMessage.Content = storeResponse.ResponseBody;
             }
 
-            for (int i = 0; i < response.ResponseHeaderNames.Length; i++)
+            for (int i = 0; i < storeResponse.ResponseHeaderNames.Length; i++)
             {
-                httpResponse.Headers.Add(response.ResponseHeaderNames[i], response.ResponseHeaderValues[i]);
+                responseMessage.Headers.Add(storeResponse.ResponseHeaderNames[i], storeResponse.ResponseHeaderValues[i]);
             }
 
-            return httpResponse;
+            return responseMessage;
         }
 
-        internal static void TraceException(Exception e)
+        internal static void TraceException(Exception exception)
         {
-            AggregateException aggregateException = e as AggregateException;
+            AggregateException aggregateException = exception as AggregateException;
             if (aggregateException != null)
             {
-                foreach (Exception exception in aggregateException.InnerExceptions)
+                foreach (Exception tempException in aggregateException.InnerExceptions)
                 {
-                    Extensions.TraceExceptionInternal(exception);
+                    Extensions.TraceExceptionInternal(tempException);
                 }
             }
             else
             {
-                Extensions.TraceExceptionInternal(e);
+                Extensions.TraceExceptionInternal(exception);
             }
         }
 
         public static async Task<IDisposable> UsingWaitAsync(
             this SemaphoreSlim semaphoreSlim,
-            CancellationToken cancellationToken = default(CancellationToken))
+            CancellationToken cancellationToken)
         {
             await semaphoreSlim.WaitAsync(cancellationToken).ConfigureAwait(false);
             return new UsableSemaphoreWrapper(semaphoreSlim);
         }
 
-        private static void TraceExceptionInternal(Exception e)
+        private static void TraceExceptionInternal(Exception exception)
         {
-            while (e != null)
+            while (exception != null)
             {
                 Uri requestUri = null;
 
-                SocketException socketException = e as SocketException;
+                SocketException socketException = exception as SocketException;
                 if (socketException != null)
                 {
                     DefaultTrace.TraceWarning(
                         "Exception {0}: RequesteUri: {1}, SocketErrorCode: {2}, {3}, {4}",
-                        e.GetType(),
+                        exception.GetType(),
                         requestUri,
                         socketException.SocketErrorCode,
-                        e.Message,
-                        e.StackTrace);
+                        exception.Message,
+                        exception.StackTrace);
                 }
                 else
                 {
                     DefaultTrace.TraceWarning(
                         "Exception {0}: RequestUri: {1}, {2}, {3}",
-                        e.GetType(),
+                        exception.GetType(),
                         requestUri,
-                        e.Message,
-                        e.StackTrace);
+                        exception.Message,
+                        exception.StackTrace);
                 }
 
-                e = e.InnerException;
+                exception = exception.InnerException;
             }
         }
     }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosDiagnosticsTests.cs

@@ -45,49 +45,49 @@ public async Task PointOperationDiagnostic()
             //Checking point operation diagnostics on typed operations
             ToDoActivity testItem = ToDoActivity.CreateRandomToDoActivity();
             ItemResponse<ToDoActivity> createResponse = await this.Container.CreateItemAsync<ToDoActivity>(item: testItem);
-            Assert.IsNotNull(createResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(createResponse.Diagnostics);
 
             ItemResponse<ToDoActivity> readResponse = await this.Container.ReadItemAsync<ToDoActivity>(id: testItem.id, partitionKey: new PartitionKey(testItem.status));
             Assert.IsNotNull(readResponse.Diagnostics);
 
             testItem.description = "NewDescription";
             ItemResponse<ToDoActivity> replaceResponse = await this.Container.ReplaceItemAsync<ToDoActivity>(item: testItem, id: testItem.id, partitionKey: new PartitionKey(testItem.status));
             Assert.AreEqual(replaceResponse.Resource.description, "NewDescription");
-            Assert.IsNotNull(replaceResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(replaceResponse.Diagnostics);
 
             ItemResponse<ToDoActivity> deleteResponse = await this.Container.DeleteItemAsync<ToDoActivity>(partitionKey: new Cosmos.PartitionKey(testItem.status), id: testItem.id);
             Assert.IsNotNull(deleteResponse);
-            Assert.IsNotNull(deleteResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(deleteResponse.Diagnostics);
 
             //Checking point operation diagnostics on stream operations
             ResponseMessage createStreamResponse = await this.Container.CreateItemStreamAsync(
                 partitionKey: new PartitionKey(testItem.status),
                 streamPayload: TestCommon.Serializer.ToStream<ToDoActivity>(testItem));
-            Assert.IsNotNull(createStreamResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(createStreamResponse.Diagnostics);
 
             ResponseMessage readStreamResponse = await this.Container.ReadItemStreamAsync(
                 id: testItem.id,
                 partitionKey: new PartitionKey(testItem.status));
-            Assert.IsNotNull(readStreamResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(readStreamResponse.Diagnostics);
 
             ResponseMessage replaceStreamResponse = await this.Container.ReplaceItemStreamAsync(
                streamPayload: TestCommon.Serializer.ToStream<ToDoActivity>(testItem),
                id: testItem.id,
                partitionKey: new PartitionKey(testItem.status));
-            Assert.IsNotNull(replaceStreamResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(replaceStreamResponse.Diagnostics);
 
             ResponseMessage deleteStreamResponse = await this.Container.DeleteItemStreamAsync(
                id: testItem.id,
                partitionKey: new PartitionKey(testItem.status));
-            Assert.IsNotNull(deleteStreamResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(deleteStreamResponse.Diagnostics);
 
             // Ensure diagnostics are set even on failed operations
             testItem.description = new string('x', Microsoft.Azure.Documents.Constants.MaxResourceSizeInBytes + 1);
             ResponseMessage createTooBigStreamResponse = await this.Container.CreateItemStreamAsync(
                 partitionKey: new PartitionKey(testItem.status),
                 streamPayload: TestCommon.Serializer.ToStream<ToDoActivity>(testItem));
             Assert.IsFalse(createTooBigStreamResponse.IsSuccessStatusCode);
-            Assert.IsNotNull(createTooBigStreamResponse.Diagnostics);
+            CosmosDiagnosticsTests.VerifyPointDiagnostics(createTooBigStreamResponse.Diagnostics);
         }
 
         [TestMethod]
@@ -156,6 +156,34 @@ public static void VerifyQueryDiagnostics(CosmosDiagnostics diagnostics)
             }
         }
 
+        public static void VerifyPointDiagnostics(CosmosDiagnostics diagnostics)
+        {
+            string info = diagnostics.ToString();
+            Assert.IsNotNull(info);
+            JObject jObject = JObject.Parse(info);
+            Assert.IsNotNull(jObject["ActivityId"].ToString());
+            Assert.IsNotNull(jObject["StatusCode"].ToString());
+            Assert.IsNotNull(jObject["RequestCharge"].ToString());
+            Assert.IsNotNull(jObject["RequestUri"].ToString());
+            Assert.IsNotNull(jObject["requestStartTime"].ToString());
+            Assert.IsNotNull(jObject["requestEndTime"].ToString());
+            Assert.IsNotNull(jObject["responseStatisticsList"].ToString());
+            Assert.IsNotNull(jObject["supplementalResponseStatisticsList"].ToString());
+            Assert.IsNotNull(jObject["addressResolutionStatistics"].ToString());
+            Assert.IsNotNull(jObject["contactedReplicas"].ToString());
+            Assert.IsNotNull(jObject["failedReplicas"].ToString());
+            Assert.IsNotNull(jObject["regionsContacted"].ToString());
+            Assert.IsNotNull(jObject["requestLatency"].ToString());
+
+            int statusCode = jObject["StatusCode"].ToObject<int>();
+
+            // Session token only expected on success
+            if (statusCode >= 200 && statusCode < 300)
+            {
+                Assert.IsNotNull(jObject["ResponseSessionToken"].ToString());
+            }
+        }
+
         private async Task<long> ExecuteQueryAndReturnOutputDocumentCount(string queryText, int expectedItemCount)
         {
             QueryDefinition sql = new QueryDefinition(queryText);