OrleansRuntimeDiagnostics.cs

    public static class OrleansRuntimeDiagnosticsConfigurationExtensions
    {
        public static ClusterConfiguration UseDiagnostics(this ClusterConfiguration config)
        {
            config.Globals.RegisterBootstrapProvider<OrleansRuntimeDiagnosticsBootstrapper>(nameof(OrleansRuntimeDiagnosticsBootstrapper));

            return config;
        }
    }

    public class OrleansRuntimeDiagnosticsBootstrapper : IBootstrapProvider
    {
        public Task Init(string name, IProviderRuntime providerRuntime, IProviderConfiguration config)
        {
            this.Name = name;

            var interceptor = new OrleansRuntimeDiagnosticsInterceptor(providerRuntime.GetInvokeInterceptor());
            providerRuntime.SetInvokeInterceptor(interceptor.Handle);

            return Task.CompletedTask;
        }

        public Task Close()
        {
            return Task.CompletedTask;
        }

        public string Name { get; private set; }
    }

    public class OrleansRuntimeDiagnosticsInterceptor
    {
        private readonly InvokeInterceptor _inner;

        /// <summary>
        /// Creates a diagnostic interceptor
        /// </summary>
        /// <param name="inner">The decorated existing InvokeInterceptor, can be null</param>
        public OrleansRuntimeDiagnosticsInterceptor(InvokeInterceptor inner)
        {
            // default when there is no existing invoke interceptor
            _inner = inner ?? ((method, request, grain, invoker) => invoker.Invoke(grain, request));
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public async Task<object> Handle(MethodInfo targetMethod, InvokeMethodRequest request, IGrain target,
            IGrainMethodInvoker invoker)
        {

            if (!DiagnosticListener.IsEnabled())
                return await _inner(targetMethod, request, target, invoker).ConfigureAwait(false);


            Activity activity = null;
            IDictionary<string, object> requestContext = null;

            if (DiagnosticListener.IsEnabled(OrleansRuntimeDiagnosticsLoggingStrings.ActivityName))
            {
                // export requestcontext once
                requestContext = RequestContext.Export();
                activity = StartActivity(request, targetMethod, requestContext);
            }

            try
            {

                var result = await _inner(targetMethod, request, target, invoker).ConfigureAwait(false);

                if (activity != null)
                {
                    // is set when diagnosticslistener is enabled
                    // requestcontext dictionary is also not null
                    StopActivity(request, targetMethod, activity, requestContext);
                }

                return result;
            }
            catch (Exception ex)
            {
                //capture failed method...
                if (DiagnosticListener.IsEnabled(OrleansRuntimeDiagnosticsLoggingStrings.DiagnosticsUnhandledExceptionName))
                {
                    var timestamp = Stopwatch.GetTimestamp();
                    // Diagnostics is enabled for UnhandledException, but it may not be for BeginRequest
                    // so call GetTimestamp if currentTimestamp is zero (from above)
                    RecordUnhandledExceptionDiagnostics(targetMethod, request, timestamp, ex);
                }
                throw;
            }
        }


        [MethodImpl(MethodImplOptions.NoInlining)]
        private void RecordUnhandledExceptionDiagnostics(MethodInfo targetMethod, InvokeMethodRequest request, long timestamp, Exception exception)
        {
            //isenabled check is done in parent
            DiagnosticListener.Write(
                OrleansRuntimeDiagnosticsLoggingStrings.DiagnosticsUnhandledExceptionName,
                new
                {
                    RequestContext = RequestContext.Export(),
                    Request = request,
                    TargetMethod = targetMethod,
                    Exception = exception,
                    Timestamp = timestamp
                });
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private Activity StartActivity(InvokeMethodRequest request, MethodInfo targetMethod, IDictionary<string, object> requestContext)
        {
            var activity = new Activity(OrleansRuntimeDiagnosticsLoggingStrings.ActivityName);


            if (requestContext.TryGetValue(OrleansRuntimeDiagnosticsLoggingStrings.RequestIdHeaderName, out object requestId))
            {
                activity.SetParentId((string)requestId);

                // We expect baggage to be empty by default
                // Only very advanced users will be using it in near future, we encourage them to keep baggage small (few items)

                if (requestContext.TryGetValue(OrleansRuntimeDiagnosticsLoggingStrings.CorrelationContextHeaderName, out object baggage))
                {
                    KeyValuePair<string, string>[] values = (KeyValuePair<string, string>[])baggage;

                    foreach (var item in values)
                    {
                        activity.AddBaggage(item.Key, item.Value);
                    }
                }
            }

            if (DiagnosticListener.IsEnabled(OrleansRuntimeDiagnosticsLoggingStrings.ActivityStartName))
            {
                DiagnosticListener.StartActivity(activity, new { RequestContext = requestContext, Request = request, TargetMethod = targetMethod });
            }
            else
            {
                activity.Start();
            }

            return activity;
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        private void StopActivity(InvokeMethodRequest request, MethodInfo targetMethod, Activity activity, IDictionary<string, object> requestContext)
        {
            DiagnosticListener.StopActivity(activity, new { RequestContext = requestContext, Request = request, TargetMethod = targetMethod });
        }

        private static readonly DiagnosticListener DiagnosticListener =
            new DiagnosticListener(OrleansRuntimeDiagnosticsLoggingStrings.DiagnosticListenerName);
    }

    public static class OrleansRuntimeDiagnosticsLoggingStrings
    {
        public const string DiagnosticListenerName = "Orleans.Runtime";
        public const string ActivityName = "Orleans.Runtime.GrainRequestIn";
        public const string ActivityStartName = "Orleans.Runtime.GrainRequestIn.Start";


        public const string RequestIdHeaderName = "Request-Id";
        public const string CorrelationContextHeaderName = "Correlation-Context";
        public const string DiagnosticsUnhandledExceptionName = "Orleans.Runtime.UnhandledException";
    }