Diagnostics : Adds Continuation Token from PkRange cache (#3180)

Diagnostics : Adds Continuation Token from PkRange cache. Right now, NoOpTrace has been passed in most of the cases for PartitionKeyRangeCache call. In future PRs, we will be removing NoOpTrace from these calls to get the proper diagnostics.
Azure · May 13, 2022 · ba47e73 · ba47e73
1 parent 69d5ef4
commit ba47e73
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 0 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionKeyRangeCache.cs b/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;
         }
 

diff --git a/Microsoft.Azure.Cosmos/src/Tracing/ITraceDatumVisitor.cs b/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);
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Tracing/TraceData/PartitionKeyRangeCacheTraceDatum.cs b/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);
+        }
+    }
+}
diff --git a/Microsoft.Azure.Cosmos/src/Tracing/TraceWriter.TraceJsonWriter.cs b/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));
                 }
             }
+
         }
     }
 }
diff --git a/Microsoft.Azure.Cosmos/src/Tracing/TraceWriter.TraceTextWriter.cs b/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>

diff --git a/...ft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemSessionTokenTests.cs b/...ft.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;
                 }