Expose session state on ServiceBusReceivedMessageContext

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

@@ -4,6 +4,12 @@
 
 ### 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, 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
 
 ### Bugs Fixed

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();
+    }
 }

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

@@ -707,6 +707,19 @@ public Mono<byte[]> 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<byte[]> 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<Void> 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<Void> 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.
@@ -1843,6 +1871,9 @@ private Mono<Void> 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 || sessionId.isEmpty()) {
+            return monoError(LOGGER,
+                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);
@@ -1860,12 +1891,15 @@ private Mono<byte[]> 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 || sessionId.isEmpty()) {
+            return monoError(LOGGER,
+                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);
 
         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));

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<ServiceBusManagementNode> getSessionManagement() {
+            return sessionManagement;
+        }
+
         /**
          * Begin the recurring lock renewal for the session.
          *

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;
@@ -24,8 +25,10 @@
 
 final class ServiceBusSessionReactorReceiver implements AmqpReceiveLink {
     private final ClientLogger logger;
+    private final ServiceBusTracer tracer;
     private final String sessionId;
     private final AmqpReceiveLink sessionLink;
+    private final Mono<ServiceBusManagementNode> sessionManagement;
     private final boolean hasIdleTimeout;
     private final Sinks.Many<Boolean> nextItemIdleTimeoutSink = Sinks.many().multicast().onBackpressureBuffer();
     private final Sinks.Empty<Void> terminateEndpointStatesSink = Sinks.empty();
@@ -34,8 +37,10 @@ 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();
         this.hasIdleTimeout = sessionIdleTimeout != null;
         if (hasIdleTimeout) {
             this.disposables
@@ -53,6 +58,41 @@ 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<byte[]> getSessionState() {
+        return tracer
+            .traceMono("ServiceBus.getSessionState",
+                sessionManagement.flatMap(mgmt -> mgmt.getSessionState(sessionId, getLinkName())))
+            .onErrorMap(ServiceBusSessionReactorReceiver::mapSessionStateError);
+    }
+
+    /**
+     * 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<Void> setSessionState(byte[] sessionState) {
+        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
     public String getHostname() {
         return sessionLink.getHostname();

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<String, ServiceBusSessionReactorReceiver> 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,43 @@ Mono<Void> defer(ServiceBusReceivedMessage message, DeferOptions options) {
                 options.getTransactionContext());
         }
 
+        Mono<byte[]> getSessionState(ServiceBusReceivedMessage message) {
+            final String sessionId = message.getSessionId();
+            if (sessionId == null || sessionId.isEmpty()) {
+                return notSessionMessageError("getSessionState");
+            }
+            final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT));
+            if (receiver == null) {
+                return sessionNotActiveError(sessionId, "getSessionState");
+            }
+            return receiver.getSessionState();
+        }
+
+        Mono<Void> setSessionState(ServiceBusReceivedMessage message, byte[] sessionState) {
+            final String sessionId = message.getSessionId();
+            if (sessionId == null || sessionId.isEmpty()) {
+                return notSessionMessageError("setSessionState");
+            }
+            final ServiceBusSessionReactorReceiver receiver = receivers.get(sessionId.toLowerCase(Locale.ROOT));
+            if (receiver == null) {
+                return sessionNotActiveError(sessionId, "setSessionState");
+            }
+            return receiver.setSessionState(sessionState);
+        }
+
+        private <T> Mono<T> 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 <T> Mono<T> 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, sessionId);
+            return monoError(logger, new IllegalStateException(m));
+        }
+
         private Mono<Void> updateDisposition(ServiceBusReceivedMessage message, DispositionStatus dispositionStatus,
             Map<String, Object> propertiesToModify, String deadLetterReason, String deadLetterDescription,
             ServiceBusTransactionContext transactionContext) {