Store Traceparent in CloudEvent attributes (Azure#15515)

Author: JoshLove-msft

sdk/core/Azure.Core.TestFramework/src/RecordedTestSanitizer.cs

@@ -15,6 +15,13 @@ public class RecordedTestSanitizer
         public const string SanitizeValue = "Sanitized";
         public List<string> JsonPathSanitizers { get; } = new List<string>();
 
+        /// <summary>
+        /// This is just a temporary workaround to avoid breaking tests that need to be re-recorded
+        //  when updating the JsonPathSanitizer logic to avoid changing date formats when deserializing requests.
+        //  this property will be removed in the future.
+        /// </summary>
+        public bool DoNotConvertJsonDateTokens { get; set; }
+
         private static readonly string[] s_sanitizeValueArray = { SanitizeValue };
 
         public List<string> SanitizedHeaders { get; } = new List<string> { "Authorization" };
@@ -41,15 +48,31 @@ public virtual string SanitizeTextBody(string contentType, string body)
                 return body;
             try
             {
-                var jsonO = JToken.Parse(body);
+                var settings = new JsonSerializerSettings
+                {
+                    DateParseHandling = DateParseHandling.None
+                };
+
+                JToken jsonO;
+                // Prevent default behavior where JSON.NET will convert DateTimeOffset
+                // into a DateTime.
+                if (DoNotConvertJsonDateTokens)
+                {
+                    jsonO = JsonConvert.DeserializeObject<JToken>(body, settings);
+                }
+                else
+                {
+                    jsonO = JToken.Parse(body);
+                }
+
                 foreach (string jsonPath in JsonPathSanitizers)
                 {
                     foreach (JToken token in jsonO.SelectTokens(jsonPath))
                     {
                         token.Replace(JToken.FromObject(SanitizeValue));
                     }
                 }
-                return JsonConvert.SerializeObject(jsonO);
+                return JsonConvert.SerializeObject(jsonO, settings);
             }
             catch
             {

sdk/eventgrid/Azure.Messaging.EventGrid/api/Azure.Messaging.EventGrid.netstandard2.0.cs

@@ -59,7 +59,7 @@ public EventGridPublisherClient(System.Uri endpoint, Azure.Messaging.EventGrid.E
     public partial class EventGridPublisherClientOptions : Azure.Core.ClientOptions
     {
         public EventGridPublisherClientOptions(Azure.Messaging.EventGrid.EventGridPublisherClientOptions.ServiceVersion version = Azure.Messaging.EventGrid.EventGridPublisherClientOptions.ServiceVersion.V2018_01_01) { }
-        public Azure.Core.Serialization.ObjectSerializer DataSerializer { get { throw null; } set { } }
+        public Azure.Core.Serialization.ObjectSerializer Serializer { get { throw null; } set { } }
         public enum ServiceVersion
         {
             V2018_01_01 = 1,

sdk/eventgrid/Azure.Messaging.EventGrid/samples/Sample1_PublishEventsToTopic.md

@@ -16,7 +16,7 @@ EventGridPublisherClient client = new EventGridPublisherClient(
 ```C# Snippet:CreateClientWithOptions
 EventGridPublisherClientOptions clientOptions = new EventGridPublisherClientOptions()
 {
-    DataSerializer = myCustomDataSerializer
+    Serializer = myCustomDataSerializer
 };
 
 EventGridPublisherClient client = new EventGridPublisherClient(

sdk/eventgrid/Azure.Messaging.EventGrid/src/Customization/EventGridEvent.cs

@@ -117,7 +117,10 @@ public static EventGridEvent[] Parse(string requestContent)
                     egEventInternal.EventType,
                     egEventInternal.DataVersion,
                     egEventInternal.EventTime,
-                    egEventInternal.Id);
+                    egEventInternal.Id)
+                {
+                    Topic = egEventInternal.Topic
+                };
 
                 egEvents.Add(egEvent);
             }

sdk/eventgrid/Azure.Messaging.EventGrid/src/Customization/EventGridPublisherClient.cs

@@ -3,9 +3,9 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Globalization;
 using System.IO;
-using System.Linq;
 using System.Security.Cryptography;
 using System.Text;
 using System.Text.Json;
@@ -32,6 +32,9 @@ public class EventGridPublisherClient
         private readonly string _apiVersion;
         private readonly ObjectSerializer _dataSerializer;
 
+        private const string TraceParentHeaderName = "traceparent";
+        private const string TraceStateHeaderName = "tracestate";
+
         /// <summary>Initalizes an instance of EventGridClient.</summary>
         protected EventGridPublisherClient()
         {
@@ -62,7 +65,7 @@ public EventGridPublisherClient(Uri endpoint, AzureKeyCredential credential, Eve
             Argument.AssertNotNull(credential, nameof(credential));
             options ??= new EventGridPublisherClientOptions();
             _apiVersion = options.Version.GetVersionString();
-            _dataSerializer = options.DataSerializer ?? new JsonObjectSerializer();
+            _dataSerializer = options.Serializer ?? new JsonObjectSerializer();
             _endpoint = endpoint;
             _key = credential;
             HttpPipeline pipeline = HttpPipelineBuilder.Build(options, new AzureKeyCredentialPolicy(credential, Constants.SasKeyName));
@@ -80,7 +83,7 @@ public EventGridPublisherClient(Uri endpoint, EventGridSharedAccessSignatureCred
         {
             Argument.AssertNotNull(credential, nameof(credential));
             options ??= new EventGridPublisherClientOptions();
-            _dataSerializer = options.DataSerializer ?? new JsonObjectSerializer();
+            _dataSerializer = options.Serializer ?? new JsonObjectSerializer();
             _endpoint = endpoint;
             HttpPipeline pipeline = HttpPipelineBuilder.Build(options, new EventGridSharedAccessSignatureCredentialPolicy(credential));
             _serviceRestClient = new EventGridRestClient(new ClientDiagnostics(options), pipeline, options.Version.GetVersionString());
@@ -198,6 +201,15 @@ private async Task<Response> SendCloudEventsInternal(IEnumerable<CloudEvent> eve
                 // List of events cannot be null
                 Argument.AssertNotNull(events, nameof(events));
 
+                string activityId = null;
+                string traceState = null;
+                Activity currentActivity = Activity.Current;
+                if (currentActivity != null && currentActivity.IsW3CFormat())
+                {
+                    activityId = currentActivity.Id;
+                    currentActivity.TryGetTraceState(out traceState);
+                }
+
                 List<CloudEventInternal> eventsWithSerializedPayloads = new List<CloudEventInternal>();
                 foreach (CloudEvent cloudEvent in events)
                 {
@@ -222,6 +234,17 @@ private async Task<Response> SendCloudEventsInternal(IEnumerable<CloudEvent> eve
                         newCloudEvent.Add(kvp.Key, new CustomModelSerializer(kvp.Value, _dataSerializer, cancellationToken));
                     }
 
+                    if (activityId != null &&
+                        !cloudEvent.ExtensionAttributes.ContainsKey(TraceParentHeaderName) &&
+                        !cloudEvent.ExtensionAttributes.ContainsKey(TraceStateHeaderName))
+                    {
+                        newCloudEvent.Add(TraceParentHeaderName, activityId);
+                        if (traceState != null)
+                        {
+                            newCloudEvent.Add(TraceStateHeaderName, traceState);
+                        }
+                    }
+
                     // The 'Data' property is optional for CloudEvents
                     // Additionally, if the type of data is binary, 'Data' will not be populated (data will be stored in 'DataBase64' instead)
                     if (cloudEvent.Data != null)

sdk/eventgrid/Azure.Messaging.EventGrid/src/Customization/EventGridPublisherClientOptions.cs

@@ -40,7 +40,7 @@ public enum ServiceVersion
         /// <summary>
         /// Used to serialize the payloads of given events to UTF-8 encoded JSON.
         /// </summary>
-        public ObjectSerializer DataSerializer { get; set; }
+        public ObjectSerializer Serializer { get; set; }
 
         /// <summary>
         /// Initializes a new instance of the <see cref="EventGridPublisherClientOptions"/>

sdk/eventgrid/Azure.Messaging.EventGrid/tests/CloudEventTests.cs

@@ -0,0 +1,131 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
+using System.Text.Json;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core;
+using Azure.Core.Pipeline;
+using Azure.Core.TestFramework;
+using Azure.Messaging.EventGrid.Models;
+using NUnit.Framework;
+
+namespace Azure.Messaging.EventGrid.Tests
+{
+    public class CloudEventTests
+    {
+
+        private const string TraceParentHeaderName = "traceparent";
+        private const string TraceStateHeaderName = "tracestate";
+
+        [Test]
+        [TestCase(false, false)]
+        [TestCase(true, true)]
+        [TestCase(true, false)]
+        [TestCase(false, true)]
+        public async Task SetsTraceParentExtension(bool inclTraceparent, bool inclTracestate)
+        {
+            var mockTransport = new MockTransport(new MockResponse(200));
+            var options = new EventGridPublisherClientOptions
+            {
+                Transport = mockTransport
+            };
+            EventGridPublisherClient client =
+               new EventGridPublisherClient(
+                   new Uri("http://localHost"),
+                   new AzureKeyCredential("fakeKey"),
+                   options);
+            var activity = new Activity($"{nameof(EventGridPublisherClient)}.{nameof(EventGridPublisherClient.SendEvents)}");
+            activity.SetW3CFormat();
+            activity.Start();
+            List<CloudEvent> eventsList = new List<CloudEvent>();
+            for (int i = 0; i < 10; i++)
+            {
+                CloudEvent cloudEvent = new CloudEvent(
+                    "record",
+                    "Microsoft.MockPublisher.TestEvent",
+                    JsonDocument.Parse("{\"property1\": \"abc\",  \"property2\": 123}").RootElement)
+                {
+                    Id = "id",
+                    Subject = $"Subject-{i}",
+                    Time = DateTimeOffset.UtcNow
+                };
+                if (inclTraceparent && i % 2 == 0)
+                {
+                    cloudEvent.ExtensionAttributes.Add("traceparent", "traceparentValue");
+                }
+                if (inclTracestate && i % 2 == 0)
+                {
+                    cloudEvent.ExtensionAttributes.Add("tracestate", "param:value");
+                }
+                eventsList.Add(cloudEvent);
+            }
+            await client.SendEventsAsync(eventsList);
+
+            activity.Stop();
+            List<CloudEventInternal> cloudEvents = DeserializeRequest(mockTransport.SingleRequest);
+            IEnumerator<CloudEvent> cloudEnum = eventsList.GetEnumerator();
+            foreach (CloudEventInternal cloudInternal in cloudEvents)
+            {
+                cloudEnum.MoveNext();
+                Dictionary<string, object> cloudEventAttr = cloudEnum.Current.ExtensionAttributes;
+                if (cloudEventAttr.ContainsKey(TraceParentHeaderName) &&
+                    cloudEventAttr.ContainsKey(TraceStateHeaderName))
+                {
+                    Assert.AreEqual(
+                        cloudEventAttr[TraceParentHeaderName],
+                        cloudInternal[TraceParentHeaderName]);
+
+                    Assert.AreEqual(
+                        cloudEventAttr[TraceStateHeaderName],
+                        cloudInternal[TraceStateHeaderName]);
+                }
+                else if (cloudEventAttr.ContainsKey(TraceParentHeaderName))
+                {
+                    Assert.AreEqual(
+                        cloudEventAttr[TraceParentHeaderName],
+                        cloudInternal[TraceParentHeaderName]);
+                }
+                else if (cloudEventAttr.ContainsKey(TraceStateHeaderName))
+                {
+                    Assert.AreEqual(
+                       cloudEventAttr[TraceStateHeaderName],
+                       cloudInternal[TraceStateHeaderName]);
+                }
+                else
+                {
+                    Assert.AreEqual(
+                       activity.Id,
+                       cloudInternal[TraceParentHeaderName]);
+                }
+            }
+        }
+
+        private static List<CloudEventInternal> DeserializeRequest(Request request)
+        {
+            var content = request.Content as Utf8JsonRequestContent;
+            var stream = new MemoryStream();
+            content.WriteTo(stream, CancellationToken.None);
+            stream.Position = 0;
+            JsonDocument requestDocument = JsonDocument.Parse(stream);
+            var cloudEvents = new List<CloudEventInternal>();
+            // Parse JsonElement into separate events, deserialize event envelope properties
+            if (requestDocument.RootElement.ValueKind == JsonValueKind.Object)
+            {
+                cloudEvents.Add(CloudEventInternal.DeserializeCloudEventInternal(requestDocument.RootElement));
+            }
+            else if (requestDocument.RootElement.ValueKind == JsonValueKind.Array)
+            {
+                foreach (JsonElement property in requestDocument.RootElement.EnumerateArray())
+                {
+                    cloudEvents.Add(CloudEventInternal.DeserializeCloudEventInternal(property));
+                }
+            }
+            return cloudEvents;
+        }
+    }
+}