CosmosTraceDiagnostics.cs

// ------------------------------------------------------------
// Copyright (c) Microsoft Corporation.  All rights reserved.
// ------------------------------------------------------------

namespace Microsoft.Azure.Cosmos.Diagnostics
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Text;
    using Microsoft.Azure.Cosmos.Json;
    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
    using Microsoft.Azure.Cosmos.Tracing;
    using Microsoft.Azure.Cosmos.Tracing.TraceData;
    using static Microsoft.Azure.Cosmos.Tracing.TraceData.ClientSideRequestStatisticsTraceDatum;

    internal sealed class CosmosTraceDiagnostics : CosmosDiagnostics
    {
        private readonly Lazy<ServerSideCumulativeMetrics> accumulatedMetrics;
        private readonly Lazy<string> cachedSummaryJson;

        public CosmosTraceDiagnostics(ITrace trace)
            : this(trace, CosmosClientOptions.DefaultMaxDiagnosticsSummarySizeBytes)
        {
        }

        internal CosmosTraceDiagnostics(ITrace trace, int maxDiagnosticsSummarySizeBytes)
        {
            if (trace == null)
            {
                throw new ArgumentNullException(nameof(trace));
            }

            // Need to set to the root trace, since we don't know which layer of the stack the response message was returned from.
            ITrace rootTrace = trace;
            while (rootTrace.Parent != null)
            {
                rootTrace = rootTrace.Parent;
            }

            this.Value = rootTrace;
            this.accumulatedMetrics = new Lazy<ServerSideCumulativeMetrics>(() => PopulateServerSideCumulativeMetrics(this.Value));
            this.cachedSummaryJson = new Lazy<string>(() =>
                DiagnosticsSummaryWriter.WriteSummary(this.Value, maxDiagnosticsSummarySizeBytes));
        }

        public ITrace Value { get; }

        public override string ToString()
        {
            if (this.Value is Tracing.Trace rootConcreteTrace)
            {
                rootConcreteTrace.SetWalkingStateRecursively();
            }
            
            return this.ToJsonString();
        }

        public override string ToString(DiagnosticsVerbosity verbosity)
        {
            return verbosity switch
            {
                DiagnosticsVerbosity.Summary => this.cachedSummaryJson.Value,
                DiagnosticsVerbosity.Detailed => this.ToString(),
                _ => this.ToString(),
            };
        }

        public override TimeSpan GetClientElapsedTime()
        {
            return this.Value.Duration;
        }

        public override IReadOnlyList<(string regionName, Uri uri)> GetContactedRegions()
        {
            if (this.Value is Tracing.Trace rootConcreteTrace)
            {
                rootConcreteTrace.SetWalkingStateRecursively();
            }
            
            return this.Value?.Summary?.RegionsContacted;
        }

        public override ServerSideCumulativeMetrics GetQueryMetrics()
        {
            if (this.Value is Tracing.Trace rootConcreteTrace)
            {
                rootConcreteTrace.SetWalkingStateRecursively();
            }
            
            return this.accumulatedMetrics.Value;
        }

        internal bool IsGoneExceptionHit()
        {
            if (this.Value is Tracing.Trace rootConcreteTrace)
            {
                rootConcreteTrace.SetWalkingStateRecursively();
            }
            
            return this.WalkTraceTreeForGoneException(this.Value);
        }

        private bool WalkTraceTreeForGoneException(ITrace currentTrace)
        {
            if (currentTrace == null)
            {
                return false;
            }

            foreach (object datum in currentTrace.Data.Values)
            {
                if (datum is ClientSideRequestStatisticsTraceDatum clientSideRequestStatisticsTraceDatum)
                {
                    foreach (StoreResponseStatistics responseStatistics in clientSideRequestStatisticsTraceDatum.StoreResponseStatisticsList)
                    {
                        if (responseStatistics.StoreResult != null && responseStatistics.StoreResult.StatusCode == Documents.StatusCodes.Gone)
                        {
                            return true;
                        }
                    }
                }
            }

            foreach (ITrace childTrace in currentTrace.Children)
            {
                if (this.WalkTraceTreeForGoneException(childTrace))
                {
                    return true;
                }
            }

            return false;
        }

        private string ToJsonString()
        {
            ReadOnlyMemory<byte> utf8String = this.WriteTraceToJsonWriter(JsonSerializationFormat.Text);
            return Encoding.UTF8.GetString(utf8String.Span);
        }

        private ReadOnlyMemory<byte> WriteTraceToJsonWriter(JsonSerializationFormat jsonSerializationFormat)
        {
            IJsonWriter jsonTextWriter = JsonWriter.Create(jsonSerializationFormat);
            TraceWriter.WriteTrace(jsonTextWriter, this.Value);
            return jsonTextWriter.GetResult();
        }

        private static ServerSideCumulativeMetrics PopulateServerSideCumulativeMetrics(ITrace trace)
        {
            ServerSideMetricsInternalAccumulator accumulator = new ServerSideMetricsInternalAccumulator();
            ServerSideMetricsTraceExtractor.WalkTraceTreeForQueryMetrics(trace, accumulator);

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

            ServerSideCumulativeMetrics accumulatedMetrics = new ServerSideCumulativeMetricsInternal(serverSideMetricsList);
            return accumulatedMetrics.PartitionedMetrics.Count != 0 ? accumulatedMetrics : null;
        }

        public override DateTime? GetStartTimeUtc()
        {
            if (this.Value == null || this.Value.StartTime == null)
            {
                return null;
            }

            return this.Value.StartTime;
        }

        public override int GetFailedRequestCount()
        {
            if (this.Value == null || this.Value.Summary == null)
            {
                return 0;
            }

            return this.Value.Summary.GetFailedCount();
       }
    }
}