Diagnostics : Adds Continuation Token from PkRange cache

Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs

@@ -264,6 +264,10 @@ private async Task<CollectionRoutingMap> GetRoutingMapForCollectionAsync(
                 throw new NotFoundException($"{DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture)}: GetRoutingMapForCollectionAsync(collectionRid: {collectionRid}), Range information either doesn't exist or is not complete.");
             }
 
+            trace.AddDatum($"PKRangeCache Info({DateTime.UtcNow.ToString("o", CultureInfo.InvariantCulture)})",
+                                          new PartitionKeyRangeCacheTraceDatum(
+                                              previousContinuationToken: previousRoutingMap?.ChangeFeedNextIfNoneMatch,
+                                              continuationToken: routingMap.ChangeFeedNextIfNoneMatch));
             return routingMap;
         }
 

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

@@ -16,5 +16,6 @@ internal interface ITraceDatumVisitor
         void Visit(ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum);
         void Visit(CpuHistoryTraceDatum cpuHistoryTraceDatum);
         void Visit(ClientConfigurationTraceDatum clientConfigurationTraceDatum);
+        void Visit(PartitionKeyRangeCacheTraceDatum partitionKeyRangeCacheTraceDatum);
     }
 }

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

@@ -0,0 +1,23 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Tracing.TraceData
+{
+    internal class PartitionKeyRangeCacheTraceDatum : TraceDatum
+    {
+        public string PreviousContinuationToken { get; }
+        public string ContinuationToken { get; }
+
+        public PartitionKeyRangeCacheTraceDatum(string previousContinuationToken, string continuationToken)
+        {
+            this.PreviousContinuationToken = previousContinuationToken; 
+            this.ContinuationToken = continuationToken;
+        }
+
+        internal override void Accept(ITraceDatumVisitor traceDatumVisitor)
+        {
+            traceDatumVisitor.Visit(this);
+        }
+    }
+}

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

@@ -441,6 +441,19 @@ public void Visit(StoreResult storeResult)
                 this.jsonWriter.WriteObjectEnd();
             }
 
+            public void Visit(PartitionKeyRangeCacheTraceDatum partitionKeyRangeCacheTraceDatum)
+            {
+                this.jsonWriter.WriteObjectStart();
+
+                this.jsonWriter.WriteFieldName("Previous Continuation Token");
+                this.WriteStringValueOrNull(partitionKeyRangeCacheTraceDatum.PreviousContinuationToken);
+
+                this.jsonWriter.WriteFieldName("Continuation Token");
+                this.WriteStringValueOrNull(partitionKeyRangeCacheTraceDatum.ContinuationToken);
+
+                this.jsonWriter.WriteObjectEnd();
+            }
+
             private void WriteJsonUriArray(string propertyName, IEnumerable<TransportAddressUri> uris)
             {
                 this.jsonWriter.WriteFieldName(propertyName);
@@ -560,6 +573,7 @@ private void WriteDateTimeStringValue(DateTime value)
                     this.jsonWriter.WriteStringValue(value.ToString("o", CultureInfo.InvariantCulture));
                 }
             }
+
         }
     }
 }

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

@@ -496,6 +496,15 @@ public override string ToString()
                 {
                     return this.toStringValue;
                 }
+
+                public void Visit(PartitionKeyRangeCacheTraceDatum partitionKeyRangeCacheTraceDatum)
+                {
+                    StringBuilder stringBuilder = new StringBuilder();
+                    stringBuilder.AppendLine($"Previous Continuation Token: {partitionKeyRangeCacheTraceDatum.PreviousContinuationToken ?? "<null>"}");
+                    stringBuilder.AppendLine($"Continuation Token: {partitionKeyRangeCacheTraceDatum.ContinuationToken ?? "<null>"}");
+
+                    this.toStringValue = stringBuilder.ToString();
+                }
             }
 
             /// <summary>

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

@@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     using System.Threading;
     using Microsoft.Azure.Cosmos.Routing;
     using Newtonsoft.Json.Linq;
+    using System.Linq;
 
     [TestClass]
     public class CosmosItemSessionTokenTests : BaseCosmosClientHelper
@@ -371,6 +372,18 @@ public async Task InvalidSessionTokenAfterContainerRecreationAndCollectionCacheR
                 while (queryIteratorOldContainer.HasMoreResults)
                 {
                     FeedResponse<JObject> response = await queryIteratorOldContainer.ReadNextAsync();
+                    if(i == 0)
+                    {
+                        string diagnosticString = response.Diagnostics.ToString();
+                        Assert.IsTrue(diagnosticString.Contains("PKRangeCache Info("));
+                        JObject diagnosticJobject = JObject.Parse(diagnosticString);
+                        JToken actualToken = diagnosticJobject.SelectToken("$.children[0].children[?(@.name=='Get Partition Key Ranges')].children[?(@.name=='Try Get Overlapping Ranges')].data");
+                        JToken actualNode = actualToken.Children().First().First();
+
+                        Assert.IsTrue(actualNode["Previous Continuation Token"].ToString().Length == 0);
+                        Assert.IsTrue(actualNode["Continuation Token"].ToString().Length > 0);
+                    }
+
                     itemCountOldContainer += response.Count;
                 }