Fix end tracing span for Service bus processor client (Azure#19436)

Author: samvaity

sdk/eventhubs/azure-messaging-eventhubs/CHANGELOG.md

@@ -1,7 +1,8 @@
 # Release History
 
 ## 5.6.0-beta.1 (Unreleased)
-
+### Bug Fixes
+- Update to end the trace span regardless of the scope instance type for process operation tracing spans.
 
 ## 5.5.0 (2020-02-15)
 ### New features

sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/PartitionPumpManager.java

@@ -33,8 +33,6 @@
 import com.azure.messaging.eventhubs.models.PartitionEvent;
 import com.azure.messaging.eventhubs.models.PartitionOwnership;
 import com.azure.messaging.eventhubs.models.ReceiveOptions;
-import java.io.Closeable;
-import java.io.IOException;
 import java.time.Duration;
 import java.util.List;
 import java.util.Locale;
@@ -376,12 +374,12 @@ private void endProcessTracingSpan(Context processSpanContext, Signal<Void> sign
         }
 
         Object spanObject = spanScope.get();
-        if (spanObject instanceof Closeable) {
-            Closeable close = (Closeable) spanObject;
+        if (spanObject instanceof AutoCloseable) {
+            AutoCloseable close = (AutoCloseable) spanObject;
             try {
                 close.close();
-            } catch (IOException ioException) {
-                logger.error(Messages.EVENT_PROCESSOR_RUN_END, ioException);
+            } catch (Exception exception) {
+                logger.error(Messages.EVENT_PROCESSOR_RUN_END, exception);
             }
 
         } else {

sdk/eventhubs/azure-messaging-eventhubs/src/main/resources/eventhubs-messages.properties

@@ -1,4 +1,4 @@
-PROCESS_SPAN_SCOPE_TYPE_ERROR=Process span scope type is not of type Closeable, but type: %s. Not closing the scope and span
+PROCESS_SPAN_SCOPE_TYPE_ERROR=Process span scope type is not of type AutoCloseable, but type: %s. Not closing the scope and span
 MESSAGE_NOT_OF_TYPE=Message body type is not of type Data, but type: %s. Not setting body contents.
 REQUEST_VALUE_NOT_VALID=Back pressure request value not valid. It must be between {} and {}.
 EVENT_DATA_DOES_NOT_FIT=EventData does not fit into maximum number of batches. '%s'

sdk/eventhubs/azure-messaging-eventhubs/src/test/java/com/azure/messaging/eventhubs/EventProcessorClientTest.java

@@ -29,7 +29,6 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
-import java.io.Closeable;
 import java.time.Instant;
 import java.util.Collections;
 import java.util.HashMap;
@@ -139,7 +138,7 @@ public void testWithSimplePartitionProcessor() throws Exception {
             invocation -> {
                 Context passed = invocation.getArgument(1, Context.class);
                 return passed.addData(SPAN_CONTEXT_KEY, "value1")
-                    .addData("scope", (Closeable) () -> {
+                    .addData("scope", (AutoCloseable) () -> {
                     })
                     .addData(PARENT_SPAN_KEY, "value2");
             }
@@ -229,7 +228,7 @@ public void testProcessSpans() throws Exception {
             invocation -> {
                 Context passed = invocation.getArgument(1, Context.class);
                 assertTrue(passed.getData(MESSAGE_ENQUEUED_TIME).isPresent());
-                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (Closeable) () -> {
+                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (AutoCloseable) () -> {
                     return;
                 }).addData(PARENT_SPAN_KEY, "value2");
             }
@@ -450,7 +449,7 @@ public void testSingleEventReceiveHeartBeat() throws InterruptedException {
         when(tracer.start(eq("EventHubs.process"), any(), eq(ProcessKind.PROCESS))).thenAnswer(
             invocation -> {
                 Context passed = invocation.getArgument(1, Context.class);
-                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (Closeable) () -> {
+                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (AutoCloseable) () -> {
                     return;
                 }).addData(PARENT_SPAN_KEY, "value2");
             }

sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md

@@ -1,7 +1,8 @@
 # Release History
 
 ## 7.1.0-beta.1 (Unreleased)
-
+### Bug Fixes
+- Update to end the trace span regardless of the scope instance type for process operation tracing spans.
 
 ## 7.0.2 (2021-02-10)
 ### Dependency Updates

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorClient.java

@@ -15,8 +15,6 @@
 import reactor.core.publisher.Signal;
 import reactor.core.scheduler.Schedulers;
 
-import java.io.Closeable;
-import java.io.IOException;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.Optional;
@@ -264,20 +262,20 @@ private void endProcessTracingSpan(Context processSpanContext, Signal<Void> sign
         if (!spanScope.isPresent() || !tracerProvider.isEnabled()) {
             return;
         }
-        if (spanScope.get() instanceof Closeable) {
-            Closeable close = (Closeable) processSpanContext.getData(SCOPE_KEY).get();
+        if (spanScope.get() instanceof AutoCloseable) {
+            AutoCloseable close = (AutoCloseable) processSpanContext.getData(SCOPE_KEY).get();
             try {
                 close.close();
-                tracerProvider.endSpan(processSpanContext, signal);
-            } catch (IOException ioException) {
-                logger.error("endTracingSpan().close() failed with an error %s", ioException);
+            } catch (Exception exception) {
+                logger.error("endTracingSpan().close() failed with an error %s", exception);
             }
 
         } else {
             logger.warning(String.format(Locale.US,
-                "Process span scope type is not of type Closeable, but type: %s. Not closing the scope and span",
-                spanScope.get() != null ? spanScope.getClass() : "null"));
+                "Process span scope type is not of type AutoCloseable, but type: %s. Not closing the scope"
+                    + " and span", spanScope.get() != null ? spanScope.getClass() : "null"));
         }
+        tracerProvider.endSpan(processSpanContext, signal);
     }
 
     private Context startProcessTracingSpan(ServiceBusReceivedMessage receivedMessage, String entityPath,

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorTest.java

@@ -15,7 +15,6 @@
 import reactor.core.publisher.FluxSink;
 import reactor.core.publisher.Mono;
 
-import java.io.Closeable;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -353,7 +352,7 @@ public void testProcessorWithTracingEnabled() throws InterruptedException {
             invocation -> {
                 Context passed = invocation.getArgument(1, Context.class);
                 assertTrue(passed.getData(MESSAGE_ENQUEUED_TIME).isPresent());
-                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (Closeable) () -> {
+                return passed.addData(SPAN_CONTEXT_KEY, "value1").addData("scope", (AutoCloseable) () -> {
                     return;
                 }).addData(PARENT_SPAN_KEY, "value2");
             }