ServiceBus tracing: Fix more occurrences of NPE when message is null in tracing instrumentation (Azure#36800)

* Fix more occurences of NPE when message is null in tracing instr

Author: Liudmila Molkova

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

@@ -35,6 +35,9 @@
 
 - Fixed issue causing updates to TopicProperties with AuthorizationRules to return 400 Bad request. ([#34880](https://github.com/Azure/azure-sdk-for-java/issues/34880))
 
+- Fixed `NullPointerException` that happens when session processor or receiver encounters an error and distributed tracing is enabled.
+  ([#36800](https://github.com/Azure/azure-sdk-for-java/issues/36800))
+
 ### Other Changes
 
 #### Dependency Updates

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

@@ -55,6 +55,11 @@ protected void hookOnSubscribe(Subscription subscription) {
 
         @Override
         protected void hookOnNext(ServiceBusMessageContext message) {
+            if (message == null || message.getMessage() == null)  {
+                downstream.onNext(message);
+                return;
+            }
+
             Context span = instrumentation.instrumentProcess("ServiceBus.process", message.getMessage(), Context.NONE);
             message.getMessage().setContext(span);
             AutoCloseable scope = tracer.makeSpanCurrent(span);

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

@@ -374,7 +374,7 @@ public void onSubscribe(Subscription subscription) {
                 @SuppressWarnings("try")
                 @Override
                 public void onNext(ServiceBusMessageContext serviceBusMessageContext) {
-                    Context span = serviceBusMessageContext.getMessage().getContext();
+                    Context span = serviceBusMessageContext.getMessage() != null ? serviceBusMessageContext.getMessage().getContext() : Context.NONE;
                     Exception exception = null;
                     AutoCloseable scope = tracer.makeSpanCurrent(span);
                     try {

sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/implementation/instrumentation/ContextAccessor.java

@@ -21,6 +21,7 @@ public interface SendMessageContextAccessor {
     }
 
     public static ServiceBusReceivedMessage setContext(ServiceBusReceivedMessage message, Context context) {
+        assert message != null; // message is never null on this path.
         return receiveAccessor.setContext(message, context);
     }
 

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

@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.messaging.servicebus;
+
+import com.azure.core.util.BinaryData;
+import com.azure.core.util.Context;
+import com.azure.core.util.tracing.SpanKind;
+import com.azure.core.util.tracing.StartSpanOptions;
+import com.azure.core.util.tracing.Tracer;
+import com.azure.messaging.servicebus.implementation.instrumentation.ReceiverKind;
+import com.azure.messaging.servicebus.implementation.instrumentation.ServiceBusReceiverInstrumentation;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import reactor.test.StepVerifier;
+import reactor.test.publisher.TestPublisher;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class FluxTraceTest {
+    private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(30);
+    private final ServiceBusReceivedMessage receivedMessage = new ServiceBusReceivedMessage(BinaryData.fromString("Some Data"));
+    private final ServiceBusMessageContext message = new ServiceBusMessageContext(receivedMessage);
+    private final TestPublisher<ServiceBusMessageContext> messagesPublisher = TestPublisher.create();
+
+    @ParameterizedTest
+    @EnumSource(ReceiverKind.class)
+    public void testProcessSpans(ReceiverKind receiverKind) {
+        TestTracer tracer = new TestTracer();
+        ServiceBusReceiverInstrumentation instrumentation = new ServiceBusReceiverInstrumentation(tracer, null, "fqdn", "entityPath", null, receiverKind);
+        FluxTrace fluxTrace = new FluxTrace(messagesPublisher.flux(), instrumentation);
+
+        StepVerifier.create(fluxTrace)
+            .then(() -> messagesPublisher.next(message))
+            .assertNext(m -> {
+                switch (receiverKind) {
+                    case SYNC_RECEIVER:
+                        assertEquals(0, tracer.getStartedSpans().size());
+                        break;
+                    default:
+                        assertEquals(1, tracer.getStartedSpans().size());
+                        assertEquals("ServiceBus.process", tracer.getStartedSpans().get(0));
+                        break;
+                }
+            })
+            .thenCancel()
+            .verify(DEFAULT_TIMEOUT);
+    }
+
+    @ParameterizedTest
+    @EnumSource(ReceiverKind.class)
+    public void nullMessage(ReceiverKind receiverKind) {
+        TestTracer tracer = new TestTracer();
+        ServiceBusReceiverInstrumentation instrumentation = new ServiceBusReceiverInstrumentation(tracer, null, "fqdn", "entityPath", null, receiverKind);
+        FluxTrace fluxTrace = new FluxTrace(messagesPublisher.flux(), instrumentation);
+
+        StepVerifier.create(fluxTrace)
+            .then(() -> messagesPublisher.next(new ServiceBusMessageContext("sessionId", new RuntimeException("foo"))))
+            .assertNext(m -> assertEquals(0, tracer.getStartedSpans().size()))
+            .thenCancel()
+            .verify(DEFAULT_TIMEOUT);
+    }
+
+    private static class TestTracer implements Tracer {
+        private final List<String> startedSpans = new ArrayList<>();
+        @Override
+        public Context start(String methodName, StartSpanOptions options, Context context) {
+            startedSpans.add(methodName);
+            return context;
+        }
+
+        @Override
+        public Context start(String methodName, Context context) {
+            return start(methodName, new StartSpanOptions(SpanKind.INTERNAL), context);
+        }
+
+        @Override
+        public void end(String errorMessage, Throwable throwable, Context context) {
+        }
+
+        @Override
+        public void setAttribute(String key, String value, Context context) {
+        }
+
+        public List<String> getStartedSpans() {
+            return startedSpans;
+        }
+    }
+}

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

@@ -39,6 +39,7 @@
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
@@ -439,6 +440,37 @@ public void testProcessorWithTracingEnabled() throws InterruptedException {
         verify(tracer, atLeast(numberOfTimes - 1)).end(isNull(), isNull(), any());
     }
 
+    @Test
+    @SuppressWarnings("unchecked")
+    public void testProcessorWithTracingEnabledAndNullMessage() throws InterruptedException {
+        final Tracer tracer = mock(Tracer.class);
+        final int numberOfTimes = 1;
+
+        when(tracer.isEnabled()).thenReturn(true);
+        when(tracer.extractContext(any())).thenReturn(Context.NONE);
+
+        when(tracer.start(eq("ServiceBus.process"), any(StartSpanOptions.class), any())).thenReturn(new Context(PARENT_TRACE_CONTEXT_KEY, "span"));
+
+        Flux<ServiceBusMessageContext> messageFlux = Flux.just(new ServiceBusMessageContext("sessionId", new RuntimeException("foo")));
+        ServiceBusClientBuilder.ServiceBusReceiverClientBuilder receiverBuilder = getBuilder(messageFlux, tracer);
+
+        CountDownLatch countDownLatch = new CountDownLatch(numberOfTimes);
+        ServiceBusProcessorClient serviceBusProcessorClient = new ServiceBusProcessorClient(receiverBuilder, ENTITY_NAME,
+            null, null,
+            messageContext -> fail("Should not have received a message"),
+            error -> {
+                assertEquals("foo", error.getException().getMessage());
+                countDownLatch.countDown();
+            },
+            new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1));
+
+        serviceBusProcessorClient.start();
+        assertTrue(countDownLatch.await(20, TimeUnit.SECONDS));
+        serviceBusProcessorClient.close();
+
+        verify(tracer, never()).start(eq("ServiceBus.process"), any(StartSpanOptions.class), any(Context.class));
+    }
+
     @Test
     @SuppressWarnings("unchecked")
     public void testProcessorWithTracingDisabled() throws InterruptedException {