Add EventGrid distributed tracing

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/CloudEvent.java

@@ -20,6 +20,7 @@
 import java.nio.charset.StandardCharsets;
 import java.time.OffsetDateTime;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -336,6 +337,9 @@ public String getSpecVersion() {
      * @return the extension attributes as an unmodifiable map.
      */
     public Map<String, Object> getExtensionAttributes() {
+        if (this.cloudEvent.getAdditionalProperties() == null) {
+            return null;
+        }
         return Collections.unmodifiableMap(this.cloudEvent.getAdditionalProperties());
     }
 
@@ -348,6 +352,9 @@ public Map<String, Object> getExtensionAttributes() {
      * @return the cloud event itself.
      */
     public CloudEvent addExtensionAttribute(String name, Object value) {
+        if (this.cloudEvent.getAdditionalProperties() == null) {
+            this.cloudEvent.setAdditionalProperties(new HashMap<>());
+        }
         this.cloudEvent.getAdditionalProperties().put(name.toLowerCase(Locale.ENGLISH), value);
         return this;
     }

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/EventGridPublisherAsyncClient.java

@@ -10,12 +10,15 @@
 import com.azure.core.http.rest.Response;
 import com.azure.core.util.Context;
 import com.azure.core.util.serializer.SerializerAdapter;
+import com.azure.core.util.tracing.TracerProxy;
+import com.azure.messaging.eventgrid.implementation.Constants;
 import com.azure.messaging.eventgrid.implementation.EventGridPublisherClientImpl;
 import com.azure.messaging.eventgrid.implementation.EventGridPublisherClientImplBuilder;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import static com.azure.core.util.FluxUtil.withContext;
+import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY;
 
 /**
  * A service client that publishes events to an EventGrid topic or domain. Use {@link EventGridPublisherClientBuilder}
@@ -33,6 +36,9 @@ public final class EventGridPublisherAsyncClient {
 
     private final EventGridServiceVersion serviceVersion;
 
+    // Please see <a href=https://docs.microsoft.com/en-us/azure/azure-resource-manager/management/azure-services-resource-providers>here</a>
+    // for more information on Azure resource provider namespaces.
+
     EventGridPublisherAsyncClient(HttpPipeline pipeline, String hostname, SerializerAdapter serializerAdapter,
                                   EventGridServiceVersion serviceVersion) {
         this.impl = new EventGridPublisherClientImplBuilder()
@@ -70,7 +76,8 @@ Mono<Void> sendEvents(Iterable<EventGridEvent> events, Context context) {
         return Flux.fromIterable(events)
             .map(EventGridEvent::toImpl)
             .collectList()
-            .flatMap(list -> this.impl.publishEventsAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishEventsAsync(this.hostname,
+                list, context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
     }
 
     /**
@@ -85,10 +92,12 @@ public Mono<Void> sendCloudEvents(Iterable<CloudEvent> events) {
     }
 
     Mono<Void> sendCloudEvents(Iterable<CloudEvent> events, Context context) {
+        this.addCloudEventTracePlaceHolder(events);
         return Flux.fromIterable(events)
             .map(CloudEvent::toImpl)
             .collectList()
-            .flatMap(list -> this.impl.publishCloudEventEventsAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishCloudEventEventsAsync(this.hostname, list,
+                context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
     }
 
     /**
@@ -105,7 +114,8 @@ public Mono<Void> sendCustomEvents(Iterable<Object> events) {
     Mono<Void> sendCustomEvents(Iterable<Object> events, Context context) {
         return Flux.fromIterable(events)
             .collectList()
-            .flatMap(list -> this.impl.publishCustomEventEventsAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishCustomEventEventsAsync(this.hostname, list,
+                context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
     }
 
     /**
@@ -123,7 +133,8 @@ Mono<Response<Void>> sendEventsWithResponse(Iterable<EventGridEvent> events, Con
         return Flux.fromIterable(events)
             .map(EventGridEvent::toImpl)
             .collectList()
-            .flatMap(list -> this.impl.publishEventsWithResponseAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishEventsWithResponseAsync(this.hostname, list,
+                context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
     }
 
     /**
@@ -138,10 +149,12 @@ public Mono<Response<Void>> sendCloudEventsWithResponse(Iterable<CloudEvent> eve
     }
 
     Mono<Response<Void>> sendCloudEventsWithResponse(Iterable<CloudEvent> events, Context context) {
+        this.addCloudEventTracePlaceHolder(events);
         return Flux.fromIterable(events)
             .map(CloudEvent::toImpl)
             .collectList()
-            .flatMap(list -> this.impl.publishCloudEventEventsWithResponseAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishCloudEventEventsWithResponseAsync(this.hostname, list,
+                context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
     }
 
     /**
@@ -158,6 +171,21 @@ public Mono<Response<Void>> sendCustomEventsWithResponse(Iterable<Object> events
     Mono<Response<Void>> sendCustomEventsWithResponse(Iterable<Object> events, Context context) {
         return Flux.fromIterable(events)
             .collectList()
-            .flatMap(list -> this.impl.publishCustomEventEventsWithResponseAsync(this.hostname, list, context));
+            .flatMap(list -> this.impl.publishCustomEventEventsWithResponseAsync(this.hostname, list,
+                context.addData(AZ_TRACING_NAMESPACE_KEY, Constants.EVENT_GRID_TRACING_NAMESPACE_VALUE)));
+    }
+
+    private void addCloudEventTracePlaceHolder(Iterable<CloudEvent> events) {
+        if (TracerProxy.isTracingEnabled()) {
+            for (CloudEvent event : events) {
+                if (event.getExtensionAttributes() == null ||
+                    (event.getExtensionAttributes().get(Constants.TRACE_PARENT) == null &&
+                    event.getExtensionAttributes().get(Constants.TRACE_STATE) == null)) {
+
+                    event.addExtensionAttribute(Constants.TRACE_PARENT, Constants.TRACE_PARENT_PLACEHOLDER);
+                    event.addExtensionAttribute(Constants.TRACE_STATE, Constants.TRACE_STATE_PLACEHOLDER);
+                }
+            }
+        }
     }
 }

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/EventGridPublisherClientBuilder.java

@@ -22,6 +22,8 @@
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.serializer.JacksonAdapter;
 import com.azure.core.util.serializer.SerializerAdapter;
+import com.azure.core.util.tracing.TracerProxy;
+import com.azure.messaging.eventgrid.implementation.CloudEventTracingPipelinePolicy;
 
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -144,6 +146,9 @@ public EventGridPublisherAsyncClient buildAsyncClient() {
 
         HttpPolicyProviders.addAfterRetryPolicies(httpPipelinePolicies);
 
+        if (TracerProxy.isTracingEnabled()) {
+            httpPipelinePolicies.add(new CloudEventTracingPipelinePolicy());
+        }
         httpPipelinePolicies.add(new HttpLoggingPolicy(httpLogOptions));
 
         HttpPipeline buildPipeline = new HttpPipelineBuilder()

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/implementation/CloudEventTracingPipelinePolicy.java

@@ -0,0 +1,46 @@
+package com.azure.messaging.eventgrid.implementation;
+
+import com.azure.core.http.HttpHeader;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.util.tracing.TracerProxy;
+import reactor.core.publisher.Mono;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+
+public class CloudEventTracingPipelinePolicy implements HttpPipelinePolicy {
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        final HttpRequest request = context.getHttpRequest();
+        final HttpHeader contentType = request.getHeaders().get(Constants.CONTENT_TYPE);
+        if (TracerProxy.isTracingEnabled() && contentType != null &&
+            Constants.CLOUD_EVENT_CONTENT_TYPE.equals(contentType.getValue())) {
+            return request.getBody().map(byteBuffer ->
+                replaceTracingPlaceHolder(request, byteBuffer)).then(next.process());
+        }
+        else {
+            return next.process();
+        }
+    }
+
+    static String replaceTracingPlaceHolder(HttpRequest request, ByteBuffer byteBuffer) {
+        String bodyString = new String(byteBuffer.array(), StandardCharsets.UTF_8);
+        final HttpHeader traceparentHeader = request.getHeaders().get(Constants.TRACE_PARENT);
+        final HttpHeader tracestateHeader = request.getHeaders().get(Constants.TRACE_STATE);
+        bodyString = bodyString.replace(Constants.TRACE_PARENT_REPLACE,
+            traceparentHeader != null
+                ? String.format(",\"%s\":\"%s\"", Constants.TRACE_PARENT,
+                traceparentHeader.getValue()) : "");
+        bodyString = bodyString.replace(Constants.TRACE_STATE_REPLACE,
+            tracestateHeader != null
+                ? String.format(",\"%s\":\"%s\"", Constants.TRACE_STATE, tracestateHeader.getValue()) : "");
+        request.setHeader(Constants.CONTENT_LENGTH, String.valueOf(bodyString.length()));
+        request.setBody(bodyString);
+        return bodyString;
+    }
+}

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/implementation/Constants.java

@@ -0,0 +1,14 @@
+package com.azure.messaging.eventgrid.implementation;
+
+public class Constants {
+    public static final String CONTENT_TYPE = "Content-Type";
+    public static final String CONTENT_LENGTH = "Content-Length";
+    public static final String CLOUD_EVENT_CONTENT_TYPE = "application/cloudevents-batch+json; charset=utf-8";
+    public static final String TRACE_PARENT = "traceparent";
+    public static final String TRACE_STATE = "tracestate";
+    public static final String TRACE_PARENT_PLACEHOLDER = "TP-14b6b15b-74b6-4178-847e-d142aa2727b2";
+    public static final String TRACE_STATE_PLACEHOLDER = "TS-14b6b15b-74b6-4178-847e-d142aa2727b2";
+    public static final String TRACE_PARENT_REPLACE = ",\"" + TRACE_PARENT + "\":\"TP-14b6b15b-74b6-4178-847e-d142aa2727b2\"";
+    public static final String TRACE_STATE_REPLACE = ",\"" + TRACE_STATE + "\":\"TS-14b6b15b-74b6-4178-847e-d142aa2727b2\"";
+    public static final String EVENT_GRID_TRACING_NAMESPACE_VALUE = "Microsoft.EventGrid";
+}

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/systemevents/AcsChatMessageEventBaseProperties.java

@@ -45,7 +45,7 @@ public class AcsChatMessageEventBaseProperties extends AcsChatEventBasePropertie
      * The version of the message
      */
     @JsonProperty(value = "version")
-    private Integer version;
+    private Long version;
 
     /**
      * Get the messageId property: The chat message id.
@@ -152,7 +152,7 @@ public AcsChatMessageEventBaseProperties setType(String type) {
      *
      * @return the version value.
      */
-    public Integer getVersion() {
+    public Long getVersion() {
         return this.version;
     }
 
@@ -162,7 +162,7 @@ public Integer getVersion() {
      * @param version the version value to set.
      * @return the AcsChatMessageEventBaseProperties object itself.
      */
-    public AcsChatMessageEventBaseProperties setVersion(Integer version) {
+    public AcsChatMessageEventBaseProperties setVersion(Long version) {
         this.version = version;
         return this;
     }

sdk/eventgrid/azure-messaging-eventgrid/src/main/java/com/azure/messaging/eventgrid/systemevents/AcsChatThreadEventBaseProperties.java

@@ -21,7 +21,7 @@ public class AcsChatThreadEventBaseProperties extends AcsChatEventBaseProperties
      * The version of the thread
      */
     @JsonProperty(value = "version")
-    private Integer version;
+    private Long version;
 
     /**
      * Get the createTime property: The original creation time of the thread.
@@ -48,7 +48,7 @@ public AcsChatThreadEventBaseProperties setCreateTime(OffsetDateTime createTime)
      *
      * @return the version value.
      */
-    public Integer getVersion() {
+    public Long getVersion() {
         return this.version;
     }
 
@@ -58,7 +58,7 @@ public Integer getVersion() {
      * @param version the version value to set.
      * @return the AcsChatThreadEventBaseProperties object itself.
      */
-    public AcsChatThreadEventBaseProperties setVersion(Integer version) {
+    public AcsChatThreadEventBaseProperties setVersion(Long version) {
         this.version = version;
         return this;
     }

sdk/eventgrid/azure-messaging-eventgrid/src/test/java/com/azure/messaging/eventgrid/implementation/CloudEventTracingPipelinePolicyTests.java

@@ -0,0 +1,74 @@
+package com.azure.messaging.eventgrid.implementation;
+
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpRequest;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.nio.ByteBuffer;
+
+public class CloudEventTracingPipelinePolicyTests {
+    private HttpRequest httpRequest;
+
+    @BeforeEach
+    public void setup() throws MalformedURLException {
+        httpRequest = new HttpRequest(HttpMethod.POST, new URL("https://something.com"));
+    }
+
+    @Test
+    void processBodyWithNoHeader() {
+        String testBodyString = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\",\"tracestate\":\"TS-14b6b15b-74b6-4178-847e-d142aa2727b2\"," +
+            "\"traceparent\":\"TP-14b6b15b-74b6-4178-847e-d142aa2727b2\"}]";
+        String expectedNewBody = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\"}]";
+
+        httpRequest.setBody(testBodyString);
+        httpRequest.setHeader(Constants.CONTENT_LENGTH, testBodyString);
+        String newBody = CloudEventTracingPipelinePolicy.replaceTracingPlaceHolder(
+            httpRequest, ByteBuffer.wrap(testBodyString.getBytes()));
+        Assertions.assertEquals(expectedNewBody, newBody);
+        Assertions.assertEquals(httpRequest.getHeaders().get(Constants.CONTENT_LENGTH).getValue(),
+            String.valueOf(newBody.length()));
+    }
+
+    @Test
+    void processBodyWithTraceParentHeader() {
+        httpRequest.setHeader(Constants.TRACE_PARENT, "aTraceParent");
+        String testBodyString = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\",\"tracestate\":\"TS-14b6b15b-74b6-4178-847e-d142aa2727b2\"," +
+            "\"traceparent\":\"TP-14b6b15b-74b6-4178-847e-d142aa2727b2\"}]";
+        String expectedNewBody = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\"," +
+            "\"traceparent\":\"aTraceParent\"}]";
+        httpRequest.setBody(testBodyString);
+        httpRequest.setHeader(Constants.CONTENT_LENGTH, testBodyString);
+        String newBody = CloudEventTracingPipelinePolicy.replaceTracingPlaceHolder(
+            httpRequest, ByteBuffer.wrap(testBodyString.getBytes()));
+        Assertions.assertEquals(expectedNewBody, newBody);
+        Assertions.assertEquals(httpRequest.getHeaders().get(Constants.CONTENT_LENGTH).getValue(),
+            String.valueOf(newBody.length()));
+    }
+
+    @Test
+    void processBodyWithTraceStateHeader() {
+        httpRequest.setHeader(Constants.TRACE_STATE, "aTraceState");
+        String testBodyString = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\",\"tracestate\":\"TS-14b6b15b-74b6-4178-847e-d142aa2727b2\"," +
+            "\"traceparent\":\"TP-14b6b15b-74b6-4178-847e-d142aa2727b2\"}]";
+        String expectedNewBody = "[{\"id\":\"313ac785-2dca-467e-a6a7-623f1baa2890\",\"source\":\"source\"," +
+            "\"type\":\"json\",\"specversion\":\"1.0\"," +
+            "\"tracestate\":\"aTraceState\"}]";
+
+        httpRequest.setBody(testBodyString);
+        httpRequest.setHeader(Constants.CONTENT_LENGTH, testBodyString);
+        String newBody = CloudEventTracingPipelinePolicy.replaceTracingPlaceHolder(
+            httpRequest, ByteBuffer.wrap(testBodyString.getBytes()));
+        Assertions.assertEquals(expectedNewBody, newBody);
+        Assertions.assertEquals(httpRequest.getHeaders().get(Constants.CONTENT_LENGTH).getValue(),
+            String.valueOf(newBody.length()));
+    }
+}