From 0a38b0612d070faf7f15239c3a75d8159c39bbdb Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 19:35:35 -0700 Subject: [PATCH 1/6] feat: expose session state on ServiceBusReceivedMessageContext Add getSessionState() and setSessionState(byte[]) to ServiceBusReceivedMessageContext so session state can be read and written from a ServiceBusSessionProcessorClient message handler without opening a separate session receiver. The methods throw IllegalStateException when the message was not received from a session-enabled entity. Resolves #49207 --- .../azure-messaging-servicebus/CHANGELOG.md | 5 + .../ServiceBusReceivedMessageContext.java | 33 ++++ .../ServiceBusReceiverAsyncClient.java | 28 +++ .../servicebus/ServiceBusSessionAcquirer.java | 10 + .../ServiceBusSessionReactorReceiver.java | 22 +++ .../servicebus/SessionsMessagePump.java | 41 ++++- .../ServiceBusReceivedMessageContextTest.java | 174 ++++++++++++++++++ .../SessionReceiversTrackerTest.java | 97 ++++++++++ 8 files changed, 407 insertions(+), 3 deletions(-) create mode 100644 sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java create mode 100644 sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java diff --git a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md index ec97b08279fc..820c1f69fcfc 100644 --- a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md @@ -4,6 +4,11 @@ ### Features Added +- Added `getSessionState()` and `setSessionState(byte[])` to `ServiceBusReceivedMessageContext`, allowing + session state to be read and written directly from a `ServiceBusSessionProcessorClient` message handler + without opening a separate session receiver. The methods throw `IllegalStateException` when the message + was not received from a session-enabled entity. ([#49207](https://github.com/Azure/azure-sdk-for-java/issues/49207)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java index f89e58154ef8..3438bf44497b 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContext.java @@ -167,4 +167,37 @@ public void deadLetter(DeadLetterOptions options) { } receiverClient.deadLetter(receivedMessageContext.getMessage(), options).block(); } + + /** + * Gets the state of the session that the {@link #getMessage() message} in this context belongs to. This operation + * is only supported when the {@link ServiceBusProcessorClient} is consuming from a session-enabled entity. + * + * @return The session state, or {@code null} if no state is set for the session. + * @throws IllegalStateException if the message was not received from a session-enabled entity, or if the session + * that delivered the message is no longer held by this processor. + * @throws ServiceBusException if the session state could not be acquired. + */ + public byte[] getSessionState() { + if (sessionReceivers != null) { + return sessionReceivers.getSessionState(receivedMessageContext.getMessage()).block(); + } + return receiverClient.getSessionState(receivedMessageContext.getMessage()).block(); + } + + /** + * Sets the state of the session that the {@link #getMessage() message} in this context belongs to. This operation + * is only supported when the {@link ServiceBusProcessorClient} is consuming from a session-enabled entity. + * + * @param sessionState State to set on the session, or {@code null} to clear the session state. + * @throws IllegalStateException if the message was not received from a session-enabled entity, or if the session + * that delivered the message is no longer held by this processor. + * @throws ServiceBusException if the session state could not be set. + */ + public void setSessionState(byte[] sessionState) { + if (sessionReceivers != null) { + sessionReceivers.setSessionState(receivedMessageContext.getMessage(), sessionState).block(); + return; + } + receiverClient.setSessionState(receivedMessageContext.getMessage(), sessionState).block(); + } } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 0b3154ec33ef..954e707770df 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -707,6 +707,19 @@ public Mono getSessionState() { return getSessionState(receiverOptions.getSessionId()); } + /** + * Gets the state of the session that delivered the given {@code message}. Unlike {@link #getSessionState()}, this + * targets the message's own session id rather than the receiver's bound session id, so it is correct for a session + * processor whose receiver is not bound to a single session. For a non-session receiver the underlying operation + * fails with an {@link IllegalStateException}. Package-private; invoked by {@link ServiceBusReceivedMessageContext}. + * + * @param message the message whose session state to read. + * @return the session state, or an empty Mono if no state is set for the session. + */ + Mono getSessionState(ServiceBusReceivedMessage message) { + return getSessionState(message.getSessionId()); + } + /** * Reads the next active message without changing the state of the receiver or the message source. The first call to * {@code peek()} fetches the first active message for this receiver. Each subsequent call fetches the subsequent @@ -1278,6 +1291,21 @@ public Mono setSessionState(byte[] sessionState) { return this.setSessionState(receiverOptions.getSessionId(), sessionState); } + /** + * Sets the state of the session that delivered the given {@code message}. Unlike {@link #setSessionState(byte[])}, + * this targets the message's own session id rather than the receiver's bound session id, so it is correct for a + * session processor whose receiver is not bound to a single session. For a non-session receiver the underlying + * operation fails with an {@link IllegalStateException}. Package-private; invoked by + * {@link ServiceBusReceivedMessageContext}. + * + * @param message the message whose session state to set. + * @param sessionState the state to set on the session, or {@code null} to clear it. + * @return a Mono that completes when the session state is set. + */ + Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { + return this.setSessionState(message.getSessionId(), sessionState); + } + /** * Starts a new service side transaction. The {@link ServiceBusTransactionContext transaction context} should be * passed to all operations that needs to be in this transaction. diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java index ee9a0a9e0399..ee3e8dc3a0ac 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionAcquirer.java @@ -285,6 +285,16 @@ ServiceBusReceiveLink getLink() { return link; } + /** + * Gets the Mono that resolves the management node used for session-scoped operations such as session state + * get/set and session lock renewal. + * + * @return the Mono to obtain the management node for this session. + */ + Mono getSessionManagement() { + return sessionManagement; + } + /** * Begin the recurring lock renewal for the session. * diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java index 42cb50c79f8b..6366be5f0d49 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java @@ -6,6 +6,7 @@ import com.azure.core.amqp.implementation.AmqpReceiveLink; import com.azure.core.util.logging.ClientLogger; import com.azure.core.util.logging.LoggingEventBuilder; +import com.azure.messaging.servicebus.implementation.ServiceBusManagementNode; import com.azure.messaging.servicebus.implementation.instrumentation.ServiceBusTracer; import org.apache.qpid.proton.amqp.transport.DeliveryState; import org.apache.qpid.proton.message.Message; @@ -26,6 +27,7 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { private final ClientLogger logger; private final String sessionId; private final AmqpReceiveLink sessionLink; + private final Mono sessionManagement; private final boolean hasIdleTimeout; private final Sinks.Many nextItemIdleTimeoutSink = Sinks.many().multicast().onBackpressureBuffer(); private final Sinks.Empty terminateEndpointStatesSink = Sinks.empty(); @@ -36,6 +38,7 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { this.logger = logger; this.sessionId = session.getId(); this.sessionLink = session.getLink(); + this.sessionManagement = session.getSessionManagement(); this.hasIdleTimeout = sessionIdleTimeout != null; if (hasIdleTimeout) { this.disposables @@ -53,6 +56,25 @@ public String getSessionId() { return sessionId; } + /** + * Gets the state of the session this receiver is streaming messages from. + * + * @return A Mono that completes with the session state, or an empty Mono if there is no state set for the session. + */ + Mono getSessionState() { + return sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName())); + } + + /** + * Sets the state of the session this receiver is streaming messages from. + * + * @param sessionState State to set on the session, or {@code null} to clear the session state. + * @return A Mono that completes when the session state is set. + */ + Mono setSessionState(byte[] sessionState) { + return sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName())); + } + @Override public String getHostname() { return sessionLink.getHostname(); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java index 091f81ebee2c..e2c718a0d4a8 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java @@ -789,8 +789,10 @@ static final class SessionReceiversTracker { private final ConcurrentHashMap receivers; private final ServiceBusReceiverInstrumentation instrumentation; - private SessionReceiversTracker(ClientLogger logger, int size, String fullyQualifiedNamespace, - String entityPath, ServiceBusReceiveMode receiveMode, ServiceBusReceiverInstrumentation instrumentation) { + // Visible for testing (package-private) so SessionReceiversTrackerTest can construct the tracker and + // exercise the session-state paths directly; in production it is only constructed by SessionsMessagePump. + SessionReceiversTracker(ClientLogger logger, int size, String fullyQualifiedNamespace, String entityPath, + ServiceBusReceiveMode receiveMode, ServiceBusReceiverInstrumentation instrumentation) { this.logger = logger; this.fullyQualifiedNamespace = fullyQualifiedNamespace; this.entityPath = entityPath; @@ -799,7 +801,8 @@ private SessionReceiversTracker(ClientLogger logger, int size, String fullyQuali this.instrumentation = instrumentation; } - private void track(ServiceBusSessionReactorReceiver receiver) { + // Visible for testing (package-private) so SessionReceiversTrackerTest can populate the receivers map. + void track(ServiceBusSessionReactorReceiver receiver) { receivers.put(receiver.getSessionId().toLowerCase(Locale.ROOT), receiver); } @@ -872,6 +875,38 @@ Mono defer(ServiceBusReceivedMessage message, DeferOptions options) { options.getTransactionContext()); } + Mono getSessionState(ServiceBusReceivedMessage message) { + final ServiceBusSessionReactorReceiver receiver = sessionReceiver(message); + if (receiver == null) { + return sessionNotActiveError(message, "getSessionState"); + } + return receiver.getSessionState(); + } + + Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { + final ServiceBusSessionReactorReceiver receiver = sessionReceiver(message); + if (receiver == null) { + return sessionNotActiveError(message, "setSessionState"); + } + return receiver.setSessionState(sessionState); + } + + private ServiceBusSessionReactorReceiver sessionReceiver(ServiceBusReceivedMessage message) { + final String sessionId = message.getSessionId(); + if (sessionId == null) { + return null; + } + return receivers.get(sessionId.toLowerCase(Locale.ROOT)); + } + + private Mono sessionNotActiveError(ServiceBusReceivedMessage message, String operation) { + final String m = String.format( + "Cannot perform '%s'; the session '%s' is no longer active for this processor. Session state " + + "operations are only valid while the session that delivered the message is still held.", + operation, message.getSessionId()); + return monoError(logger, new IllegalStateException(m)); + } + private Mono updateDisposition(ServiceBusReceivedMessage message, DispositionStatus dispositionStatus, Map propertiesToModify, String deadLetterReason, String deadLetterDescription, ServiceBusTransactionContext transactionContext) { diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java new file mode 100644 index 000000000000..49d21b8969ef --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.BinaryData; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import reactor.core.publisher.Mono; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link ServiceBusReceivedMessageContext}, focused on the session-state operations exposed to a + * {@link ServiceBusProcessorClient} message handler. Covers both delegation branches: the receiver-client path + * (non-session and V1 processors) and the V2 session-processor path ({@code SessionsMessagePump}). + */ +class ServiceBusReceivedMessageContextTest { + private static final byte[] STATE = new byte[] { 10, 111, 23 }; + + @Mock + private ServiceBusReceiverAsyncClient receiverClient; + + private AutoCloseable mocksCloseable; + private ServiceBusReceivedMessage message; + private ServiceBusMessageContext messageContext; + + @BeforeEach + void setup() { + mocksCloseable = MockitoAnnotations.openMocks(this); + message = new ServiceBusReceivedMessage(BinaryData.fromString("some-data")); + messageContext = new ServiceBusMessageContext(message); + } + + @AfterEach + void teardown() throws Exception { + if (mocksCloseable != null) { + mocksCloseable.close(); + } + } + + // --- Receiver-client path: non-session processor and V1 processor --- + + @Test + void getSessionStateDelegatesToReceiver() { + when(receiverClient.getSessionState(message)).thenReturn(Mono.just(STATE)); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertEquals(STATE, context.getSessionState()); + verify(receiverClient).getSessionState(message); + } + + @Test + void getSessionStateReturnsNullWhenNoState() { + when(receiverClient.getSessionState(message)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertNull(context.getSessionState()); + } + + @Test + void setSessionStateDelegatesToReceiver() { + when(receiverClient.setSessionState(message, STATE)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + context.setSessionState(STATE); + + verify(receiverClient).setSessionState(message, STATE); + } + + /** + * Asymmetry test: on a non-session entity the underlying receiver rejects session-state reads, and the context + * must surface that {@link IllegalStateException} rather than silently succeeding. + */ + @Test + void getSessionStateOnNonSessionEntityThrows() { + when(receiverClient.getSessionState(message)) + .thenReturn(Mono.error(new IllegalStateException("Cannot get session state on a non-session receiver."))); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertThrows(IllegalStateException.class, context::getSessionState); + } + + /** + * Asymmetry test: on a non-session entity the underlying receiver rejects session-state writes, and the context + * must surface that {@link IllegalStateException} rather than silently succeeding. + */ + @Test + void setSessionStateOnNonSessionEntityThrows() { + when(receiverClient.setSessionState(message, STATE)) + .thenReturn(Mono.error(new IllegalStateException("Cannot set session state on a non-session receiver."))); + final ServiceBusReceivedMessageContext context + = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + + assertThrows(IllegalStateException.class, () -> context.setSessionState(STATE)); + } + + // --- Session-processor path: V2 SessionsMessagePump (the #49207 scenario) --- + + @Test + void getSessionStateDelegatesToSessionTracker() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)).thenReturn(Mono.just(STATE)); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertEquals(STATE, context.getSessionState()); + verify(tracker).getSessionState(message); + } + + @Test + void getSessionStateViaSessionTrackerReturnsNullWhenNoState() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertNull(context.getSessionState()); + } + + @Test + void setSessionStateDelegatesToSessionTracker() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.setSessionState(message, STATE)).thenReturn(Mono.empty()); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + context.setSessionState(STATE); + + verify(tracker).setSessionState(message, STATE); + } + + /** + * Asymmetry test: when the session that delivered the message is no longer held by the processor, the tracker + * errors and the context must surface that {@link IllegalStateException} for a read. + */ + @Test + void getSessionStateThrowsWhenSessionNoLongerActive() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.getSessionState(message)) + .thenReturn(Mono.error(new IllegalStateException("the session is no longer active for this processor."))); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertThrows(IllegalStateException.class, context::getSessionState); + } + + /** + * Asymmetry test: when the session that delivered the message is no longer held by the processor, the tracker + * errors and the context must surface that {@link IllegalStateException} for a write. + */ + @Test + void setSessionStateThrowsWhenSessionNoLongerActive() { + final SessionsMessagePump.SessionReceiversTracker tracker + = mock(SessionsMessagePump.SessionReceiversTracker.class); + when(tracker.setSessionState(message, STATE)) + .thenReturn(Mono.error(new IllegalStateException("the session is no longer active for this processor."))); + final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); + + assertThrows(IllegalStateException.class, () -> context.setSessionState(STATE)); + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java new file mode 100644 index 000000000000..3670e73a5d2a --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java @@ -0,0 +1,97 @@ +// 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.logging.ClientLogger; +import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link SessionsMessagePump.SessionReceiversTracker} session-state routing (#49207): looking up the + * session receiver by the message's session id, erroring when the session is not (or no longer) tracked, and + * delegating to the tracked {@link ServiceBusSessionReactorReceiver} otherwise. + */ +class SessionReceiversTrackerTest { + private static final byte[] STATE = new byte[] { 5, 6, 7 }; + + private static SessionsMessagePump.SessionReceiversTracker newTracker() { + return new SessionsMessagePump.SessionReceiversTracker(new ClientLogger(SessionReceiversTrackerTest.class), 4, + "fqdn.servicebus.windows.net", "entity", ServiceBusReceiveMode.PEEK_LOCK, null); + } + + private static ServiceBusReceivedMessage messageForSession(String sessionId) { + final ServiceBusReceivedMessage message = new ServiceBusReceivedMessage(BinaryData.fromString("data")); + if (sessionId != null) { + message.setSessionId(sessionId); + } + return message; + } + + @Test + void getSessionStateErrorsWhenMessageHasNoSession() { + StepVerifier.create(newTracker().getSessionState(messageForSession(null))) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void getSessionStateErrorsWhenSessionNotTracked() { + StepVerifier.create(newTracker().getSessionState(messageForSession("absent-session"))) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void setSessionStateErrorsWhenSessionNotTracked() { + StepVerifier.create(newTracker().setSessionState(messageForSession("absent-session"), STATE)) + .expectError(IllegalStateException.class) + .verify(); + } + + @Test + void getSessionStateDelegatesToTrackedReceiver() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("s1"); + when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); + tracker.track(receiver); + + StepVerifier.create(tracker.getSessionState(messageForSession("s1"))).expectNext(STATE).verifyComplete(); + verify(receiver).getSessionState(); + } + + @Test + void setSessionStateDelegatesToTrackedReceiver() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("s1"); + when(receiver.setSessionState(STATE)).thenReturn(Mono.empty()); + tracker.track(receiver); + + StepVerifier.create(tracker.setSessionState(messageForSession("s1"), STATE)).verifyComplete(); + verify(receiver).setSessionState(STATE); + } + + /** + * Session-id lookup is case-insensitive: a receiver tracked under one casing is found by a message whose session + * id differs only in case (mirrors how {@code track()} and the settlement path key the receivers map). + */ + @Test + void getSessionStateLookupIsCaseInsensitive() { + final SessionsMessagePump.SessionReceiversTracker tracker = newTracker(); + final ServiceBusSessionReactorReceiver receiver = mock(ServiceBusSessionReactorReceiver.class); + when(receiver.getSessionId()).thenReturn("Session-A"); + when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); + tracker.track(receiver); + + StepVerifier.create(tracker.getSessionState(messageForSession("SESSION-A"))).expectNext(STATE).verifyComplete(); + } +} From 03af6d9db701be7572546929480f546366846c7f Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 19:48:47 -0700 Subject: [PATCH 2/6] fix: predictable errors for null session id in session-state APIs Addresses PR review: emit a distinct IllegalStateException when a message has no session id (vs. an untracked session) in the session-processor path, guard the receiver-client session-state methods against a null session id (avoids an NPE from getLinkName), and document in the CHANGELOG that the context APIs also throw when the delivering session is no longer held. --- .../azure-messaging-servicebus/CHANGELOG.md | 3 +- .../ServiceBusReceiverAsyncClient.java | 6 ++++ .../servicebus/SessionsMessagePump.java | 29 +++++++++++-------- 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md index 363084dbb66a..217622bdbf03 100644 --- a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md @@ -7,7 +7,8 @@ - Added `getSessionState()` and `setSessionState(byte[])` to `ServiceBusReceivedMessageContext`, allowing session state to be read and written directly from a `ServiceBusSessionProcessorClient` message handler without opening a separate session receiver. The methods throw `IllegalStateException` when the message - was not received from a session-enabled entity. ([#49207](https://github.com/Azure/azure-sdk-for-java/issues/49207)) + was not received from a session-enabled entity, or when the session that delivered the message is no longer + held by the processor. ([#49207](https://github.com/Azure/azure-sdk-for-java/issues/49207)) ### Breaking Changes diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 954e707770df..528bc0fc4a92 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -1871,6 +1871,9 @@ private Mono setSessionState(String sessionId, byte[] sessionState) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "setSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot set session state on a non-session receiver.")); + } else if (sessionId == null) { + return monoError(LOGGER, + new IllegalStateException("Cannot set session state because the session id is null.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); @@ -1888,6 +1891,9 @@ private Mono getSessionState(String sessionId) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "getSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot get session state on a non-session receiver.")); + } else if (sessionId == null) { + return monoError(LOGGER, + new IllegalStateException("Cannot get session state because the session id is null.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java index e2c718a0d4a8..0a347fd2867b 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java @@ -876,34 +876,39 @@ Mono defer(ServiceBusReceivedMessage message, DeferOptions options) { } Mono getSessionState(ServiceBusReceivedMessage message) { - final ServiceBusSessionReactorReceiver receiver = sessionReceiver(message); + final String sessionId = message.getSessionId(); + if (sessionId == null) { + return notSessionMessageError("getSessionState"); + } + final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); if (receiver == null) { - return sessionNotActiveError(message, "getSessionState"); + return sessionNotActiveError(sessionId, "getSessionState"); } return receiver.getSessionState(); } Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { - final ServiceBusSessionReactorReceiver receiver = sessionReceiver(message); + final String sessionId = message.getSessionId(); + if (sessionId == null) { + return notSessionMessageError("setSessionState"); + } + final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); if (receiver == null) { - return sessionNotActiveError(message, "setSessionState"); + return sessionNotActiveError(sessionId, "setSessionState"); } return receiver.setSessionState(sessionState); } - private ServiceBusSessionReactorReceiver sessionReceiver(ServiceBusReceivedMessage message) { - final String sessionId = message.getSessionId(); - if (sessionId == null) { - return null; - } - return receivers.get(sessionId.toLowerCase(Locale.ROOT)); + private Mono notSessionMessageError(String operation) { + return monoError(logger, new IllegalStateException(String.format( + "Cannot perform '%s'; the message was not received from a session-enabled entity.", operation))); } - private Mono sessionNotActiveError(ServiceBusReceivedMessage message, String operation) { + private Mono sessionNotActiveError(String sessionId, String operation) { final String m = String.format( "Cannot perform '%s'; the session '%s' is no longer active for this processor. Session state " + "operations are only valid while the session that delivered the message is still held.", - operation, message.getSessionId()); + operation, sessionId); return monoError(logger, new IllegalStateException(m)); } From 3abd2c477107716b3f69439ce58df08294772178 Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 20:00:39 -0700 Subject: [PATCH 3/6] fix: treat empty session id like null in session-state routing Addresses PR re-review: ServiceBusReceivedMessage.getSessionId() may return an empty string; the session-processor path now routes an empty session id to the 'not a session-enabled entity' error instead of the misleading 'session is no longer active' error. Adds empty-session-id test coverage. --- .../servicebus/SessionsMessagePump.java | 4 ++-- .../servicebus/SessionReceiversTrackerTest.java | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java index 0a347fd2867b..80cd0aaa9af6 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/SessionsMessagePump.java @@ -877,7 +877,7 @@ Mono defer(ServiceBusReceivedMessage message, DeferOptions options) { Mono getSessionState(ServiceBusReceivedMessage message) { final String sessionId = message.getSessionId(); - if (sessionId == null) { + if (sessionId == null || sessionId.isEmpty()) { return notSessionMessageError("getSessionState"); } final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); @@ -889,7 +889,7 @@ Mono getSessionState(ServiceBusReceivedMessage message) { Mono setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) { final String sessionId = message.getSessionId(); - if (sessionId == null) { + if (sessionId == null || sessionId.isEmpty()) { return notSessionMessageError("setSessionState"); } final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT)); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java index 3670e73a5d2a..2a6c897ee509 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java @@ -42,6 +42,22 @@ void getSessionStateErrorsWhenMessageHasNoSession() { .verify(); } + @Test + void getSessionStateErrorsWhenMessageHasEmptySession() { + StepVerifier.create(newTracker().getSessionState(messageForSession(""))) + .expectErrorMatches(e -> e instanceof IllegalStateException + && e.getMessage().contains("was not received from a session-enabled entity")) + .verify(); + } + + @Test + void setSessionStateErrorsWhenMessageHasEmptySession() { + StepVerifier.create(newTracker().setSessionState(messageForSession(""), STATE)) + .expectErrorMatches(e -> e instanceof IllegalStateException + && e.getMessage().contains("was not received from a session-enabled entity")) + .verify(); + } + @Test void getSessionStateErrorsWhenSessionNotTracked() { StepVerifier.create(newTracker().getSessionState(messageForSession("absent-session"))) From 835bf072ed10b3c4698d068ef2dc416d1606af06 Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 20:10:56 -0700 Subject: [PATCH 4/6] fix: empty session-id guard on receiver client + content-based byte[] test assertions Addresses PR re-review: getSessionState/setSessionState on the receiver client now reject an empty session id (not just null) with a clear IllegalStateException; tests compare session-state bytes by content (assertArrayEquals / Arrays.equals) instead of reference identity. --- .../servicebus/ServiceBusReceiverAsyncClient.java | 8 ++++---- .../ServiceBusReceivedMessageContextTest.java | 6 +++--- .../servicebus/SessionReceiversTrackerTest.java | 10 ++++++++-- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 528bc0fc4a92..0154a08036df 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -1871,9 +1871,9 @@ private Mono setSessionState(String sessionId, byte[] sessionState) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "setSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot set session state on a non-session receiver.")); - } else if (sessionId == null) { + } else if (sessionId == null || sessionId.isEmpty()) { return monoError(LOGGER, - new IllegalStateException("Cannot set session state because the session id is null.")); + new IllegalStateException("Cannot set session state because the session id is null or empty.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); @@ -1891,9 +1891,9 @@ private Mono getSessionState(String sessionId) { new IllegalStateException(String.format(INVALID_OPERATION_DISPOSED_RECEIVER, "getSessionState"))); } else if (!isSessionEnabled) { return monoError(LOGGER, new IllegalStateException("Cannot get session state on a non-session receiver.")); - } else if (sessionId == null) { + } else if (sessionId == null || sessionId.isEmpty()) { return monoError(LOGGER, - new IllegalStateException("Cannot get session state because the session id is null.")); + new IllegalStateException("Cannot get session state because the session id is null or empty.")); } assert sessionManager != null; // guaranteed to be non-null when isSessionEnabled is true. final String linkName = sessionManager.getLinkName(sessionId); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java index 49d21b8969ef..857063280a98 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusReceivedMessageContextTest.java @@ -11,7 +11,7 @@ import org.mockito.MockitoAnnotations; import reactor.core.publisher.Mono; -import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; @@ -55,7 +55,7 @@ void getSessionStateDelegatesToReceiver() { final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(receiverClient, messageContext); - assertEquals(STATE, context.getSessionState()); + assertArrayEquals(STATE, context.getSessionState()); verify(receiverClient).getSessionState(message); } @@ -116,7 +116,7 @@ void getSessionStateDelegatesToSessionTracker() { when(tracker.getSessionState(message)).thenReturn(Mono.just(STATE)); final ServiceBusReceivedMessageContext context = new ServiceBusReceivedMessageContext(tracker, messageContext); - assertEquals(STATE, context.getSessionState()); + assertArrayEquals(STATE, context.getSessionState()); verify(tracker).getSessionState(message); } diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java index 2a6c897ee509..104558c9f25b 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/SessionReceiversTrackerTest.java @@ -10,6 +10,8 @@ import reactor.core.publisher.Mono; import reactor.test.StepVerifier; +import java.util.Arrays; + import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -80,7 +82,9 @@ void getSessionStateDelegatesToTrackedReceiver() { when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); tracker.track(receiver); - StepVerifier.create(tracker.getSessionState(messageForSession("s1"))).expectNext(STATE).verifyComplete(); + StepVerifier.create(tracker.getSessionState(messageForSession("s1"))) + .expectNextMatches(actual -> Arrays.equals(STATE, actual)) + .verifyComplete(); verify(receiver).getSessionState(); } @@ -108,6 +112,8 @@ void getSessionStateLookupIsCaseInsensitive() { when(receiver.getSessionState()).thenReturn(Mono.just(STATE)); tracker.track(receiver); - StepVerifier.create(tracker.getSessionState(messageForSession("SESSION-A"))).expectNext(STATE).verifyComplete(); + StepVerifier.create(tracker.getSessionState(messageForSession("SESSION-A"))) + .expectNextMatches(actual -> Arrays.equals(STATE, actual)) + .verifyComplete(); } } From 50f5bbf350de7048abe93f9ba88652febf8b0ddb Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 20:20:58 -0700 Subject: [PATCH 5/6] fix: consistent tracing for session-state operations Addresses PR re-review: the V2 session-processor session-state path now emits ServiceBus.getSessionState/setSessionState spans (the reactor receiver retains the tracer and wraps the management-node calls with traceMono), matching the receiver-client path; and corrects the getSessionState span name in ServiceBusReceiverAsyncClient (was mislabeled ServiceBus.setSessionState). --- .../servicebus/ServiceBusReceiverAsyncClient.java | 2 +- .../servicebus/ServiceBusSessionReactorReceiver.java | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java index 0154a08036df..38e71d22e60f 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusReceiverAsyncClient.java @@ -1899,7 +1899,7 @@ private Mono getSessionState(String sessionId) { final String linkName = sessionManager.getLinkName(sessionId); return tracer - .traceMono("ServiceBus.setSessionState", + .traceMono("ServiceBus.getSessionState", connectionProcessor.flatMap(connection -> connection.getManagementNode(entityPath, entityType)) .flatMap(channel -> channel.getSessionState(sessionId, linkName))) .onErrorMap((err) -> mapError(err, ServiceBusErrorSource.RECEIVE)); diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java index 6366be5f0d49..68ff1c34767c 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java @@ -25,6 +25,7 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { private final ClientLogger logger; + private final ServiceBusTracer tracer; private final String sessionId; private final AmqpReceiveLink sessionLink; private final Mono sessionManagement; @@ -36,6 +37,7 @@ final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink { ServiceBusSessionReactorReceiver(ClientLogger logger, ServiceBusTracer tracer, ServiceBusSessionAcquirer.Session session, Duration sessionIdleTimeout, Duration maxSessionLockRenew) { this.logger = logger; + this.tracer = tracer; this.sessionId = session.getId(); this.sessionLink = session.getLink(); this.sessionManagement = session.getSessionManagement(); @@ -62,7 +64,8 @@ public String getSessionId() { * @return A Mono that completes with the session state, or an empty Mono if there is no state set for the session. */ Mono getSessionState() { - return sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName())); + return tracer.traceMono("ServiceBus.getSessionState", + sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName()))); } /** @@ -72,7 +75,8 @@ Mono getSessionState() { * @return A Mono that completes when the session state is set. */ Mono setSessionState(byte[] sessionState) { - return sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName())); + return tracer.traceMono("ServiceBus.setSessionState", + sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName()))); } @Override From 6f6148c9e4ca96b4da57ef957bd0945313b52b02 Mon Sep 17 00:00:00 2001 From: Eldert Grootenboer Date: Thu, 16 Jul 2026 20:29:49 -0700 Subject: [PATCH 6/6] fix: map V2 session-state errors to ServiceBusException Addresses PR re-review: the session-processor session-state path now maps management-link failures to ServiceBusException (ServiceBusErrorSource.RECEIVE), matching the receiver-client path and the ServiceBusReceivedMessageContext Javadoc contract, so callers no longer observe a raw AmqpException. --- .../ServiceBusSessionReactorReceiver.java | 22 +++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java index 68ff1c34767c..1b35f4a54337 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusSessionReactorReceiver.java @@ -64,8 +64,10 @@ public String getSessionId() { * @return A Mono that completes with the session state, or an empty Mono if there is no state set for the session. */ Mono getSessionState() { - return tracer.traceMono("ServiceBus.getSessionState", - sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName()))); + return tracer + .traceMono("ServiceBus.getSessionState", + sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName()))) + .onErrorMap(ServiceBusSessionReactorReceiver::mapSessionStateError); } /** @@ -75,8 +77,20 @@ Mono getSessionState() { * @return A Mono that completes when the session state is set. */ Mono setSessionState(byte[] sessionState) { - return tracer.traceMono("ServiceBus.setSessionState", - sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName()))); + return tracer + .traceMono("ServiceBus.setSessionState", + sessionManagement.flatMap(mgmt -> mgmt.setSessionState(sessionId, sessionState, getLinkName()))) + .onErrorMap(ServiceBusSessionReactorReceiver::mapSessionStateError); + } + + // Mirror ServiceBusReceiverAsyncClient's session-state error handling: surface a ServiceBusException to the + // ServiceBusReceivedMessageContext caller (its Javadoc documents ServiceBusException), rather than a raw + // AmqpException from the management link. + private static Throwable mapSessionStateError(Throwable throwable) { + if (throwable instanceof ServiceBusException) { + return throwable; + } + return new ServiceBusException(throwable, ServiceBusErrorSource.RECEIVE); } @Override