Restore Activity with OnExecuteRequestStep where available

src/Microsoft.AspNet.TelemetryCorrelation/ActivityHelper.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System.Collections;
 using System.Diagnostics;
 using System.Web;
 
@@ -27,53 +28,35 @@ internal static class ActivityHelper
         public const string AspNetActivityStartName = "Microsoft.AspNet.HttpReqIn.Start";
 
         /// <summary>
-        /// Event name for the activity stop event.
+        /// Event name for the lost activity stop event.
         /// </summary>
         public const string AspNetActivityLostStopName = "Microsoft.AspNet.HttpReqIn.ActivityLost.Stop";
 
+        /// <summary>
+        /// Event name for the restored activity stop event.
+        /// </summary>
+        public const string AspNetActivityRestoredStopName = "Microsoft.AspNet.HttpReqIn.ActivityRestored.Stop";
+
         /// <summary>
         /// Key to store the activity in HttpContext.
         /// </summary>
         public const string ActivityKey = "__AspnetActivity__";
 
-        private const int MaxActivityStackSize = 128;
-        private static readonly DiagnosticListener AspNetListener = new DiagnosticListener(AspNetListenerName);
-
         /// <summary>
-        /// It's possible that a request is executed in both native threads and managed threads,
-        /// in such case Activity.Current will be lost during native thread and managed thread switch.
-        /// This method is intended to restore the current activity in order to correlate the child
-        /// activities with the root activity of the request.
+        /// Key to store the restored activity in HttpContext.
         /// </summary>
-        /// <param name="root">Root activity id for the current request.</param>
-        /// <returns>If it returns an activity, it will be silently stopped with the parent activity</returns>
-        public static Activity RestoreCurrentActivity(Activity root)
-        {
-            Debug.Assert(root != null);
-
-            // workaround to restore the root activity, because we don't
-            // have a way to change the Activity.Current
-            var childActivity = new Activity(root.OperationName);
-            childActivity.SetParentId(root.Id);
-            childActivity.SetStartTime(root.StartTimeUtc);
-            foreach (var item in root.Baggage)
-            {
-                childActivity.AddBaggage(item.Key, item.Value);
-            }
+        public const string RestoredActivityKey = "__AspnetActivityRestored__";
 
-            childActivity.Start();
-
-            AspNetTelemetryCorrelationEventSource.Log.ActivityStarted(childActivity.Id);
-            return childActivity;
-        }
+        private const int MaxActivityStackSize = 128;
+        private static readonly DiagnosticListener AspNetListener = new DiagnosticListener(AspNetListenerName);
 
         /// <summary>
         /// Stops the activity and notifies listeners about it.
         /// </summary>
         /// <param name="activity">Activity to stop.</param>
-        /// <param name="context">Current HttpContext.</param>
+        /// <param name="contextItems">HttpContext.Items.</param>
         /// <returns>True if activity was found in the stack, false otherwise.</returns>
-        public static bool StopAspNetActivity(Activity activity, HttpContext context)
+        public static bool StopAspNetActivity(Activity activity, IDictionary contextItems)
         {
             var currentActivity = Activity.Current;
             if (activity != null && currentActivity != null)
@@ -87,11 +70,11 @@ public static bool StopAspNetActivity(Activity activity, HttpContext context)
 
                     if (newCurrentActivity == null)
                     {
-                        break;
+                        return false;
                     }
 
                     // there could be a case when request or any child activity is stopped
-                    // from the child execution context. In this case, Activity is present in the Current Stack, 
+                    // from the child execution context. In this case, Activity is present in the Current Stack,
                     // but is finished, i.e. stopping it has no effect on the Current.
                     if (newCurrentActivity == currentActivity)
                     {
@@ -116,13 +99,11 @@ public static bool StopAspNetActivity(Activity activity, HttpContext context)
                 }
 
                 // if activity is in the stack, stop it with Stop event
-                if (Activity.Current != null)
-                {
-                    AspNetListener.StopActivity(Activity.Current, new { });
-                    RemoveCurrentActivity(context);
-                    AspNetTelemetryCorrelationEventSource.Log.ActivityStopped(activity.Id);
-                    return true;
-                }
+                AspNetListener.StopActivity(currentActivity, new { });
+                contextItems[ActivityKey] = null;
+
+                AspNetTelemetryCorrelationEventSource.Log.ActivityStopped(currentActivity.Id, currentActivity.OperationName);
+                return true;
             }
 
             return false;
@@ -135,12 +116,21 @@ public static bool StopAspNetActivity(Activity activity, HttpContext context)
         /// <param name="context">Current HttpContext.</param>
         public static void StopLostActivity(Activity activity, HttpContext context)
         {
-            if (activity != null)
-            {
-                AspNetListener.Write(AspNetActivityLostStopName, new { activity });
-                RemoveCurrentActivity(context);
-                AspNetTelemetryCorrelationEventSource.Log.ActivityStopped(activity.Id, true);
-            }
+            context.Items[ActivityKey] = null;
+            AspNetListener.Write(AspNetActivityLostStopName, new { activity });
+            AspNetTelemetryCorrelationEventSource.Log.ActivityStopped(activity.Id, AspNetActivityLostStopName);
+        }
+
+        /// <summary>
+        /// Notifies listeners that there the lost activity was lost during execution and there was an intermediate activity.
+        /// </summary>
+        /// <param name="activity">Activity to notify about.</param>
+        /// <param name="context">Current HttpContext.</param>
+        public static void StopRestoredActivity(Activity activity, HttpContext context)
+        {
+            context.Items[RestoredActivityKey] = null;
+            AspNetListener.Write(AspNetActivityRestoredStopName, new { Activity = activity });
+            AspNetTelemetryCorrelationEventSource.Log.ActivityStopped(activity.Id, AspNetActivityRestoredStopName);
         }
 
         /// <summary>
@@ -152,12 +142,12 @@ public static Activity CreateRootActivity(HttpContext context)
         {
             if (AspNetListener.IsEnabled() && AspNetListener.IsEnabled(AspNetActivityName))
             {
-                var rootActivity = new Activity(ActivityHelper.AspNetActivityName);
+                var rootActivity = new Activity(AspNetActivityName);
 
                 rootActivity.Extract(context.Request.Unvalidated.Headers);
                 if (StartAspNetActivity(rootActivity))
                 {
-                    SaveCurrentActivity(context, rootActivity);
+                    context.Items[ActivityKey] = rootActivity;
                     AspNetTelemetryCorrelationEventSource.Log.ActivityStarted(rootActivity.Id);
                     return rootActivity;
                 }
@@ -167,16 +157,56 @@ public static Activity CreateRootActivity(HttpContext context)
         }
 
         /// <summary>
-        /// This should be called after the Activity starts and only for root activity of a request.
+        /// Saves activity in the HttpContext.Items.
         /// </summary>
-        /// <param name="context">Context to save context to.</param>
+        /// <param name="contextItems">Context to save context to.</param>
+        /// <param name="key">Slot name.</param>
         /// <param name="activity">Activity to save.</param>
-        internal static void SaveCurrentActivity(HttpContext context, Activity activity)
+        internal static void SaveCurrentActivity(IDictionary contextItems, string key, Activity activity)
         {
-            Debug.Assert(context != null);
+            Debug.Assert(contextItems != null);
             Debug.Assert(activity != null);
 
-            context.Items[ActivityKey] = activity;
+            contextItems[key] = activity;
+        }
+
+        /// <summary>
+        /// It's possible that a request is executed in both native threads and managed threads,
+        /// in such case Activity.Current will be lost during native thread and managed thread switch.
+        /// This method is intended to restore the current activity in order to correlate the child
+        /// activities with the root activity of the request.
+        /// </summary>
+        /// <param name="contextItems">HttpContext.Items dictionary.</param>
+        internal static void RestoreActivityIfNeeded(IDictionary contextItems)
+        {
+            if (Activity.Current == null)
+            {
+                var rootActivity = (Activity)contextItems[ActivityKey];
+                if (rootActivity != null && !contextItems.Contains(RestoredActivityKey))
+                {
+                    contextItems[RestoredActivityKey] = RestoreActivity(rootActivity);
+                }
+            }
+        }
+
+        private static Activity RestoreActivity(Activity root)
+        {
+            Debug.Assert(root != null);
+
+            // workaround to restore the root activity, because we don't
+            // have a way to change the Activity.Current
+            var childActivity = new Activity(root.OperationName);
+            childActivity.SetParentId(root.Id);
+            childActivity.SetStartTime(root.StartTimeUtc);
+            foreach (var item in root.Baggage)
+            {
+                childActivity.AddBaggage(item.Key, item.Value);
+            }
+
+            childActivity.Start();
+
+            AspNetTelemetryCorrelationEventSource.Log.ActivityRestored(childActivity.Id);
+            return childActivity;
         }
 
         private static bool StartAspNetActivity(Activity activity)
@@ -197,11 +227,5 @@ private static bool StartAspNetActivity(Activity activity)
 
             return false;
         }
-
-        private static void RemoveCurrentActivity(HttpContext context)
-        {
-            Debug.Assert(context != null);
-            context.Items[ActivityKey] = null;
-        }
     }
 }

src/Microsoft.AspNet.TelemetryCorrelation/AspNetTelemetryCorrelationEventSource.cs

@@ -29,10 +29,10 @@ public void ActivityStarted(string id)
             WriteEvent(2, id);
         }
 
-        [Event(3, Message = "Activity stopped, Id='{0}', lost {1}", Level = EventLevel.Verbose)]
-        public void ActivityStopped(string id, bool lost = false)
+        [Event(3, Message = "Activity stopped, Id='{0}', Name='{1}'", Level = EventLevel.Verbose)]
+        public void ActivityStopped(string id, string eventName)
         {
-            WriteEvent(3, id, lost);
+            WriteEvent(3, id, eventName);
         }
 
         [Event(4, Message = "Failed to parse header '{0}', value: '{1}'", Level = EventLevel.Informational)]
@@ -58,6 +58,12 @@ public void ActivityStackIsTooDeep(string id, string name)
         {
             WriteEvent(7, id, name);
         }
+
+        [Event(8, Message = "Activity restored, Id='{0}'", Level = EventLevel.Informational)]
+        public void ActivityRestored(string id)
+        {
+            WriteEvent(8, id);
+        }
     }
 }
 #pragma warning restore SA1600 // Elements must be documented

src/Microsoft.AspNet.TelemetryCorrelation/TelemetryCorrelationHttpModule.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Diagnostics;
+using System.Reflection;
 using System.Web;
 
 namespace Microsoft.AspNet.TelemetryCorrelation
@@ -13,6 +14,12 @@ namespace Microsoft.AspNet.TelemetryCorrelation
     public class TelemetryCorrelationHttpModule : IHttpModule
     {
         private const string BeginCalledFlag = "Microsoft.AspNet.TelemetryCorrelation.BeginCalled";
+        private static MethodInfo onStepMethodInfo = null;
+
+        static TelemetryCorrelationHttpModule()
+        {
+            onStepMethodInfo = typeof(HttpApplication).GetMethod("OnExecuteRequestStep");
+        }
 
         /// <inheritdoc />
         public void Dispose()
@@ -24,7 +31,40 @@ public void Init(HttpApplication context)
         {
             context.BeginRequest += Application_BeginRequest;
             context.EndRequest += Application_EndRequest;
-            context.PreRequestHandlerExecute += Application_PreRequestHandlerExecute;
+
+            // OnExecuteRequestStep is availabile starting with 4.7.1
+            // If this is executed in 4.7.1 runtime (regardless of targeted .NET version),
+            // we will use it to restore lost activity, otherwise keep PreRequestHandlerExecute
+            if (onStepMethodInfo != null)
+            {
+                onStepMethodInfo.Invoke(context, new object[] { (Action<HttpContextBase, Action>)OnExecuteRequestStep });
+            }
+            else
+            {
+                context.PreRequestHandlerExecute += Application_PreRequestHandlerExecute;
+            }
+        }
+
+        /// <summary>
+        /// Restores Activity before each pipeline step if it was lost.
+        /// </summary>
+        /// <param name="context">HttpContext instance.</param>
+        /// <param name="step">Step to be executed.</param>
+        internal void OnExecuteRequestStep(HttpContextBase context, Action step)
+        {
+            // Once we have public Activity.Current setter (https://github.com/dotnet/corefx/issues/29207) this method will be
+            // simplified to just assign Current if is was lost.
+            // In the mean time, we are creating child Activity to restore the context. We have to send
+            // event with this Activity to tracing system. It created a lot of issues for listeners as
+            // we may potentially have a lot of them for different stages.
+            // To reduce amount of events, we only care about ExecuteRequestHandler stage - restore activity here and
+            // stop/report it to tracing system in EndRequest.
+            if (context.CurrentNotification == RequestNotification.ExecuteRequestHandler && !context.IsPostNotification)
+            {
+                ActivityHelper.RestoreActivityIfNeeded(context.Items);
+            }
+
+            step();
         }
 
         private void Application_BeginRequest(object sender, EventArgs e)
@@ -38,13 +78,7 @@ private void Application_BeginRequest(object sender, EventArgs e)
         private void Application_PreRequestHandlerExecute(object sender, EventArgs e)
         {
             AspNetTelemetryCorrelationEventSource.Log.TraceCallback("Application_PreRequestHandlerExecute");
-            var context = ((HttpApplication)sender).Context;
-
-            var rootActivity = (Activity)context.Items[ActivityHelper.ActivityKey];
-            if (Activity.Current == null && rootActivity != null)
-            {
-                ActivityHelper.RestoreCurrentActivity(rootActivity);
-            }
+            ActivityHelper.RestoreActivityIfNeeded(((HttpApplication)sender).Context.Items);
         }
 
         private void Application_EndRequest(object sender, EventArgs e)
@@ -59,16 +93,25 @@ private void Application_EndRequest(object sender, EventArgs e)
             {
                 // Activity has never been started
                 var activity = ActivityHelper.CreateRootActivity(context);
-                ActivityHelper.StopAspNetActivity(activity, context);
+                ActivityHelper.StopAspNetActivity(activity, context.Items);
             }
             else
             {
                 var activity = (Activity)context.Items[ActivityHelper.ActivityKey];
 
                 // try to stop activity if it's in the Current stack
-                if (!ActivityHelper.StopAspNetActivity(activity, context))
+                // stop all running Activities on the way
+                if (!ActivityHelper.StopAspNetActivity(activity, context.Items))
                 {
-                    // Activity we created was lost, let's report it
+                    // perhaps we attempted to restore the Activity before
+                    var restoredActivity = (Activity)context.Items[ActivityHelper.RestoredActivityKey];
+                    if (restoredActivity != null)
+                    {
+                        // if so, report it
+                        ActivityHelper.StopRestoredActivity(restoredActivity, context);
+                    }
+
+                    // Activity we created was lost let's report it
                     if (activity != null)
                     {
                         ActivityHelper.StopLostActivity(activity, context);