Add more ServiceBus test cases for W3C enabled scenario (#1106)

* Add more ServiceBus test cases for W3C enabled scenario

* Fix broken tests after merge with #1107

Src/Common/W3C/W3COperationCorrelationTelemetryInitializer.cs

@@ -23,10 +23,11 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
     /// Telemetry Initializer that sets correlation ids for W3C.
     /// </summary>
     [Obsolete("Not ready for public consumption.")]
-    [SuppressMessage("Microsoft.Performance", "CA1812:AvoidUninstantiatedInternalClasses", Justification = "TelemetryInitializers are intended to be instatiated by the framework when added to a config.")]
+    [SuppressMessage("Microsoft.Performance", "CA1812:AvoidUninstantiatedInternalClasses", Justification =
+        "TelemetryInitializers are intended to be instantiated by the framework when added to a config.")]
     [EditorBrowsable(EditorBrowsableState.Never)]
 #if DEPENDENCY_COLLECTOR
-    public 
+    public
 #else
     internal
 #endif
@@ -45,7 +46,8 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
             UpdateTelemetry(telemetry, currentActivity, false);
         }
 
-        [SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", Justification = "This method has different code for Net45/NetCore")]
+        [SuppressMessage("Microsoft.Usage", "CA1801:ReviewUnusedParameters", Justification =
+            "This method has different code for Net45/NetCore")]
         internal static void UpdateTelemetry(ITelemetry telemetry, Activity activity, bool forceUpdate)
         {
             if (activity == null)
@@ -55,7 +57,7 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
 
             activity.UpdateContextOnActivity();
 
-            // Requests and dependnecies are initialized from the current Activity 
+            // Requests and dependencies are initialized from the current Activity 
             // (i.e. telemetry.Id = current.Id). Activity is created for such requests specifically
             // Traces, exceptions, events on the other side are children of current activity
             // There is one exception - SQL DiagnosticSource where current Activity is a parent
@@ -67,9 +69,9 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
             if (initializeFromCurrent)
             {
                 initializeFromCurrent &= !(opTelemetry is DependencyTelemetry dependency &&
-                                           dependency.Type == SqlRemoteDependencyType && 
+                                           dependency.Type == SqlRemoteDependencyType &&
                                            dependency.Context.GetInternalContext().SdkVersion
-                                               .StartsWith(RddDiagnosticSourcePrefix, StringComparison.Ordinal)); 
+                                               .StartsWith(RddDiagnosticSourcePrefix, StringComparison.Ordinal));
             }
 
             string spanId = null, parentSpanId = null;
@@ -78,15 +80,6 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
                 switch (tag.Key)
                 {
                     case W3CConstants.TraceIdTag:
-#if NET45
-                        // on .NET Fx Activities are not always reliable, this code prevents update
-                        // of the telemetry that was forcibly updated during Activity lifetime
-                        // ON .NET Core there is no such problem 
-                        if (telemetry.Context.Operation.Id == tag.Value && initializeFromCurrent && !forceUpdate)
-                        {
-                            return;
-                        }
-#endif
                         telemetry.Context.Operation.Id = tag.Value;
                         break;
                     case W3CConstants.SpanIdTag:
@@ -107,15 +100,27 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
 
             if (initializeFromCurrent)
             {
+#if NET45
+                // on .NET Fx Activities are not always reliable, this code prevents update
+                // of the telemetry that was forcibly updated during Activity lifetime
+                // ON .NET Core there is no such problem 
+                // if spanId is valid already and update is not forced, ignore it
+                if (!forceUpdate && IsValidTelemetryId(opTelemetry.Id, telemetry.Context.Operation.Id))
+                {
+                    return;
+                }
+#endif
                 opTelemetry.Id = StringUtilities.FormatRequestId(telemetry.Context.Operation.Id, spanId);
                 if (parentSpanId != null)
                 {
-                    telemetry.Context.Operation.ParentId = StringUtilities.FormatRequestId(telemetry.Context.Operation.Id, parentSpanId);
+                    telemetry.Context.Operation.ParentId =
+                        StringUtilities.FormatRequestId(telemetry.Context.Operation.Id, parentSpanId);
                 }
             }
             else
             {
-                telemetry.Context.Operation.ParentId = StringUtilities.FormatRequestId(telemetry.Context.Operation.Id, spanId);
+                telemetry.Context.Operation.ParentId =
+                    StringUtilities.FormatRequestId(telemetry.Context.Operation.Id, spanId);
             }
 
             if (opTelemetry != null)
@@ -131,5 +136,16 @@ namespace Microsoft.ApplicationInsights.W3C.Internal
                 }
             }
         }
+
+#if NET45
+        private static bool IsValidTelemetryId(string id, string operationId)
+        {
+            return id.Length == 51 &&
+                   id[0] == '|' &&
+                   id[33] == '.' &&
+                   id.IndexOf('.', 34) == 50 &&
+                   id.IndexOf(operationId, 1, 33, StringComparison.Ordinal) == 1;
+        }
+#endif
     }
 }

Src/DependencyCollector/Shared.Tests/Implementation/ServiceBusDiagnosticListenerTests.cs

@@ -28,6 +28,7 @@
             this.sentItems = new List<ITelemetry>();
             this.configuration.TelemetryChannel = new StubTelemetryChannel { OnSend = item => this.sentItems.Add(item), EndpointAddress = "https://dc.services.visualstudio.com/v2/track" };
             this.configuration.InstrumentationKey = Guid.NewGuid().ToString();
+            this.configuration.TelemetryInitializers.Add(new OperationCorrelationTelemetryInitializer());
         }
 
         [TestCleanup]
@@ -142,6 +143,9 @@
         [TestMethod]
         public void ServiceBusProcessHanding()
         {
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
             using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
             using (var module = new DependencyTrackingTelemetryModule())
             {
@@ -149,78 +153,236 @@
                 module.Initialize(this.configuration);
 
                 Activity parentActivity = new Activity("parent").AddBaggage("k1", "v1").Start();
-                var telemetry = this.TrackOperation<RequestTelemetry>(listener,
-                    "Microsoft.Azure.ServiceBus.Process", TaskStatus.RanToCompletion);
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process",
+                    TaskStatus.RanToCompletion,
+                    operation: TrackTraceDuringProcessing);
 
-                Assert.IsNotNull(telemetry);
-                Assert.AreEqual("Process", telemetry.Name);
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
                 Assert.AreEqual(
                     $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
-                    telemetry.Source);
-                Assert.IsTrue(telemetry.Success.Value);
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
 
-                Assert.AreEqual(parentActivity.Id, telemetry.Context.Operation.ParentId);
-                Assert.AreEqual(parentActivity.RootId, telemetry.Context.Operation.Id);
-                Assert.AreEqual("v1", telemetry.Properties["k1"]);
-                Assert.AreEqual("messageId", telemetry.Properties["MessageId"]);
+                Assert.AreEqual(parentActivity.Id, requestTelemetry.Context.Operation.ParentId);
+                Assert.AreEqual(parentActivity.RootId, requestTelemetry.Context.Operation.Id);
+                Assert.AreEqual("v1", requestTelemetry.Properties["k1"]);
+
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
             }
         }
 
         [TestMethod]
         public void ServiceBusProcessHandingExternalParent()
         {
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
             using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
             using (var module = new DependencyTrackingTelemetryModule())
             {
                 module.IncludeDiagnosticSourceActivities.Add("Microsoft.Azure.ServiceBus");
                 module.Initialize(this.configuration);
 
-                var telemetry = this.TrackOperation<RequestTelemetry>(listener,
-                    "Microsoft.Azure.ServiceBus.Process", TaskStatus.RanToCompletion, "parent");
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process", 
+                    TaskStatus.RanToCompletion, 
+                    "parent",
+                    TrackTraceDuringProcessing);
 
-                Assert.IsNotNull(telemetry);
-                Assert.AreEqual("Process", telemetry.Name);
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
                 Assert.AreEqual(
                     $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
-                    telemetry.Source);
-                Assert.IsTrue(telemetry.Success.Value);
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
 
-                Assert.AreEqual("parent", telemetry.Context.Operation.ParentId);
-                Assert.AreEqual("parent", telemetry.Context.Operation.Id);
-                Assert.AreEqual("messageId", telemetry.Properties["MessageId"]);
+                Assert.AreEqual("parent", requestTelemetry.Context.Operation.ParentId);
+                Assert.AreEqual("parent", requestTelemetry.Context.Operation.Id);
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
             }
         }
 
         [TestMethod]
         public void ServiceBusProcessHandingWithoutParent()
         {
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
             using (var module = new DependencyTrackingTelemetryModule())
             using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
             {
                 module.IncludeDiagnosticSourceActivities.Add("Microsoft.Azure.ServiceBus");
                 module.Initialize(this.configuration);
 
-                var telemetry = this.TrackOperation<RequestTelemetry>(listener,
-                    "Microsoft.Azure.ServiceBus.Process", TaskStatus.RanToCompletion);
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process",
+                    TaskStatus.RanToCompletion,
+                    operation: TrackTraceDuringProcessing);
 
-                Assert.IsNotNull(telemetry);
-                Assert.AreEqual("Process", telemetry.Name);
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
                 Assert.AreEqual(
                     $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
-                    telemetry.Source);
-                Assert.IsTrue(telemetry.Success.Value);
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
 
                 // W3C compatible-Id ( should go away when W3C is implemented in .NET https://github.com/dotnet/corefx/issues/30331 TODO)
-                Assert.AreEqual(32, telemetry.Context.Operation.Id.Length);
-                Assert.IsTrue(Regex.Match(telemetry.Context.Operation.Id, @"[a-z][0-9]").Success);
+                Assert.AreEqual(32, requestTelemetry.Context.Operation.Id.Length);
+                Assert.IsTrue(Regex.Match(requestTelemetry.Context.Operation.Id, @"[a-z][0-9]").Success);
                 // end of workaround test
 
-                Assert.AreEqual("messageId", telemetry.Properties["MessageId"]);
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
             }
         }
 
 #pragma warning disable 612, 618
         [TestMethod]
+        public void ServiceBusProcessHandingW3C()
+        {
+            this.configuration.TelemetryInitializers.Add(new W3COperationCorrelationTelemetryInitializer());
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
+            using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
+            using (var module = new DependencyTrackingTelemetryModule())
+            {
+                module.EnableW3CHeadersInjection = true;
+                module.IncludeDiagnosticSourceActivities.Add("Microsoft.Azure.ServiceBus");
+                module.Initialize(this.configuration);
+
+                Activity parentActivity = new Activity("parent").AddBaggage("k1", "v1").Start();
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process",
+                    TaskStatus.RanToCompletion,
+                    operation: TrackTraceDuringProcessing);
+
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
+                Assert.AreEqual(
+                    $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
+
+                Assert.AreEqual($"|{parentActivity.GetTraceId()}.{parentActivity.GetSpanId()}.", requestTelemetry.Context.Operation.ParentId);
+                Assert.AreEqual(parentActivity.GetTraceId(), requestTelemetry.Context.Operation.Id);
+                Assert.AreEqual("v1", requestTelemetry.Properties["k1"]);
+
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
+            }
+        }
+
+        [TestMethod]
+        public void ServiceBusProcessHandingExternalParentW3C()
+        {
+            this.configuration.TelemetryInitializers.Add(new W3COperationCorrelationTelemetryInitializer());
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
+            using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
+            using (var module = new DependencyTrackingTelemetryModule())
+            {
+                module.EnableW3CHeadersInjection = true;
+                module.IncludeDiagnosticSourceActivities.Add("Microsoft.Azure.ServiceBus");
+                module.Initialize(this.configuration);
+
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process",
+                    TaskStatus.RanToCompletion,
+                    "parent",
+                    TrackTraceDuringProcessing);
+
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
+                Assert.AreEqual(
+                    $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
+
+                Assert.AreEqual("parent", requestTelemetry.Context.Operation.ParentId);
+                Assert.AreEqual("parent", requestTelemetry.Properties[W3CConstants.LegacyRootIdProperty]);
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
+            }
+        }
+
+        [TestMethod]
+        public void ServiceBusProcessHandingWithoutParentW3C()
+        {
+            this.configuration.TelemetryInitializers.Add(new W3COperationCorrelationTelemetryInitializer());
+            var tc = new TelemetryClient(this.configuration);
+            void TrackTraceDuringProcessing() => tc.TrackTrace("trace");
+
+            using (var module = new DependencyTrackingTelemetryModule())
+            using (var listener = new DiagnosticListener("Microsoft.Azure.ServiceBus"))
+            {
+                module.EnableW3CHeadersInjection = true;
+                module.IncludeDiagnosticSourceActivities.Add("Microsoft.Azure.ServiceBus");
+                module.Initialize(this.configuration);
+
+                var requestTelemetry = this.TrackOperation<RequestTelemetry>(
+                    listener,
+                    "Microsoft.Azure.ServiceBus.Process",
+                    TaskStatus.RanToCompletion,
+                    operation: TrackTraceDuringProcessing);
+
+                Assert.IsNotNull(requestTelemetry);
+                Assert.AreEqual("Process", requestTelemetry.Name);
+                Assert.AreEqual(
+                    $"type:{RemoteDependencyConstants.AzureServiceBus} | name:queueName | endpoint:sb://queuename.myservicebus.com/",
+                    requestTelemetry.Source);
+                Assert.IsTrue(requestTelemetry.Success.Value);
+
+                // W3C compatible-Id ( should go away when W3C is implemented in .NET https://github.com/dotnet/corefx/issues/30331 TODO)
+                Assert.AreEqual(32, requestTelemetry.Context.Operation.Id.Length);
+                Assert.IsTrue(Regex.Match(requestTelemetry.Context.Operation.Id, @"[a-z][0-9]").Success);
+                // end of workaround test
+
+                Assert.AreEqual("messageId", requestTelemetry.Properties["MessageId"]);
+
+                var traceTelemetry = this.sentItems.OfType<TraceTelemetry>();
+                Assert.AreEqual(1, traceTelemetry.Count());
+
+                Assert.AreEqual(requestTelemetry.Context.Operation.Id, traceTelemetry.Single().Context.Operation.Id);
+                Assert.AreEqual(requestTelemetry.Id, traceTelemetry.Single().Context.Operation.ParentId);
+            }
+        }
+
         public void ServiceBusProcessHandingExternalParentW3CCompatibleRequestId()
         {
             this.configuration.TelemetryInitializers.Add(new W3COperationCorrelationTelemetryInitializer());
@@ -274,7 +436,12 @@
             }
         }
 
-        private T TrackOperation<T>(DiagnosticListener listener, string activityName, TaskStatus status, string parentId = null) where T : OperationTelemetry
+        private T TrackOperation<T>(
+            DiagnosticListener listener,
+            string activityName, 
+            TaskStatus status, 
+            string parentId = null,
+            Action operation = null) where T : OperationTelemetry
         {
             Activity activity = null;
 
@@ -297,13 +464,15 @@
                 }
             }
 
+            operation?.Invoke();
+
             if (activity != null)
             {
                 listener.StopActivity(activity, new { Entity = "queueName", Endpoint = new Uri("sb://queuename.myservicebus.com/"), Status = status });
-                return this.sentItems.Last() as T;
+                return this.sentItems.OfType<T>().Last();
             }
 
-            return null;
+            return default(T);
         }
     }
 }

Src/DependencyCollector/Shared/Implementation/ClientServerDependencyTracker.cs

@@ -59,7 +59,7 @@
 
                 // Every operation must have its own Activity
                 // if dependency is tracked with profiler of event source, we need to generate a proper Id for it
-                // in case of HTTP it will be propagated into the requert header.
+                // in case of HTTP it will be propagated into the request header.
                 // So, we will create a new Activity for the dependency, just to generate an Id.
                 activity = new Activity(DependencyActivityName);
 
@@ -80,21 +80,23 @@
                 activity.Start();
             }
 
-            // telemetry is initialized from current Activity (root and parent Id, but not the Id)
-            telemetry.Id = activity.Id;
-
-            // set operation root Id in case there was no parent activity (e.g. HttpRequest in background thread)
-            if (string.IsNullOrEmpty(telemetry.Context.Operation.Id))
-            {
-                telemetry.Context.Operation.Id = activity.RootId;
-            }
-
-#pragma warning disable 612, 618
             if (IsW3CEnabled)
             {
+#pragma warning disable 612, 618
                 W3COperationCorrelationTelemetryInitializer.UpdateTelemetry(telemetry, activity, true);
-            }
 #pragma warning restore 612, 618
+            }
+            else
+            {
+                // telemetry is initialized from current Activity (root and parent Id, but not the Id)
+                telemetry.Id = activity.Id;
+
+                // set operation root Id in case there was no parent activity (e.g. HttpRequest in background thread)
+                if (string.IsNullOrEmpty(telemetry.Context.Operation.Id))
+                {
+                    telemetry.Context.Operation.Id = activity.RootId;
+                }
+            }
 
             PretendProfilerIsAttached = false;
             return telemetry;

Src/DependencyCollector/Shared/Implementation/HttpProcessing.cs

@@ -281,10 +281,6 @@ namespace Microsoft.ApplicationInsights.DependencyCollector.Implementation
                         {
                             statusCode = (int)responseObj.StatusCode;
                             this.SetTarget(telemetry, responseObj.Headers);
-                            if (this.injectW3CHeaders && request is HttpWebRequest httpRequest)
-                            {
-                                // this.SetLegacyId(telemetry, httpRequest.Headers);
-                            }
 
                             // Set the operation details for the response
                             telemetry.SetOperationDetail(RemoteDependencyConstants.HttpResponseOperationDetailName, responseObj);