diff --git a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md index 7034ec765efe..c48dd251f024 100644 --- a/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md +++ b/sdk/servicebus/azure-messaging-servicebus/CHANGELOG.md @@ -4,6 +4,12 @@ ### Features Added +- Added `ServiceBusProcessorAsyncClient`, an asynchronous processor whose message and error handlers return a + `reactor.core.publisher.Mono` instead of running as blocking `Consumer` callbacks. It provides the same + push-based, auto-recovering, concurrency-managed message pump as `ServiceBusProcessorClient` for applications that + perform I/O-bound work reactively in the handler. Build it via `ServiceBusClientBuilder.processorAsync()` and + `ServiceBusClientBuilder.sessionProcessorAsync()`. ([#46564](https://github.com/Azure/azure-sdk-for-java/issues/46564)) + ### Breaking Changes ### Bugs Fixed diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java index f9ccc0f52bf2..dfa98e45d3b6 100644 --- a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusClientBuilder.java @@ -75,6 +75,7 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; +import java.util.function.Function; import java.util.function.Supplier; import static com.azure.core.amqp.implementation.ClientConstants.ENTITY_PATH_KEY; @@ -917,6 +918,26 @@ public ServiceBusSessionProcessorClientBuilder sessionProcessor() { return new ServiceBusSessionProcessorClientBuilder(); } + /** + * A new instance of {@link ServiceBusProcessorAsyncClientBuilder} used to configure a + * {@link ServiceBusProcessorAsyncClient} instance that processes messages with reactive, non-blocking handlers. + * + * @return A new instance of {@link ServiceBusProcessorAsyncClientBuilder}. + */ + public ServiceBusProcessorAsyncClientBuilder processorAsync() { + return new ServiceBusProcessorAsyncClientBuilder(); + } + + /** + * A new instance of {@link ServiceBusSessionProcessorAsyncClientBuilder} used to configure a + * {@link ServiceBusProcessorAsyncClient} instance that processes sessions with reactive, non-blocking handlers. + * + * @return A new instance of {@link ServiceBusSessionProcessorAsyncClientBuilder}. + */ + public ServiceBusSessionProcessorAsyncClientBuilder sessionProcessorAsync() { + return new ServiceBusSessionProcessorAsyncClientBuilder(); + } + /** * A new instance of {@link ServiceBusRuleManagerBuilder} used to configure a Service Bus rule manager instance. * @@ -1895,6 +1916,263 @@ private void validateInputs() { } } + /** + * Builder for creating {@link ServiceBusProcessorAsyncClient} to process sessions with reactive, non-blocking + * message handlers from a session-enabled Service Bus entity. + * + * @see ServiceBusProcessorAsyncClient + */ + @ServiceClientBuilder(serviceClients = { ServiceBusProcessorAsyncClient.class }) + public final class ServiceBusSessionProcessorAsyncClientBuilder { + private final ServiceBusProcessorClientOptions processorClientOptions; + private final ServiceBusSessionReceiverClientBuilder sessionReceiverClientBuilder; + private Function> processMessage; + private Function> processError; + + private ServiceBusSessionProcessorAsyncClientBuilder() { + sessionReceiverClientBuilder = new ServiceBusSessionReceiverClientBuilder(); + processorClientOptions = new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1); + sessionReceiverClientBuilder.maxConcurrentSessions(1); + } + + /** + * Sets the amount of time to continue auto-renewing the lock. Setting {@link Duration#ZERO} or {@code null} + * disables auto-renewal. For {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} mode, + * auto-renewal is disabled. + * + * @param maxAutoLockRenewDuration the amount of time to continue auto-renewing the lock. {@link Duration#ZERO} + * or {@code null} indicates that auto-renewal is disabled. + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException If {@code maxAutoLockRenewDuration} is negative. + */ + public ServiceBusSessionProcessorAsyncClientBuilder + maxAutoLockRenewDuration(Duration maxAutoLockRenewDuration) { + validateAndThrow(maxAutoLockRenewDuration, "maxAutoLockRenewDuration"); + sessionReceiverClientBuilder.maxAutoLockRenewDuration(maxAutoLockRenewDuration); + return this; + } + + /** + * Sets the maximum amount of time to wait for a message to be received for the currently active session. + * After this time has elapsed, the processor will close the session and attempt to process another session. + * + *

If not specified, the {@link AmqpRetryOptions#getTryTimeout()} will be used.

+ * @param sessionIdleTimeout Session idle timeout. + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException If {@code sessionIdleTimeout} is negative. + */ + public ServiceBusSessionProcessorAsyncClientBuilder sessionIdleTimeout(Duration sessionIdleTimeout) { + validateAndThrow(sessionIdleTimeout, "sessionIdleTimeout"); + sessionReceiverClientBuilder.sessionIdleTimeout(sessionIdleTimeout); + return this; + } + + /** + * Enables session processing roll-over by processing at most {@code maxConcurrentSessions}. + * + * @param maxConcurrentSessions Maximum number of concurrent sessions to process at any given time. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException if {@code maxConcurrentSessions} is less than 1. + */ + public ServiceBusSessionProcessorAsyncClientBuilder maxConcurrentSessions(int maxConcurrentSessions) { + if (maxConcurrentSessions < 1) { + throw LOGGER + .logExceptionAsError(new IllegalArgumentException("'maxConcurrentSessions' cannot be less than 1")); + } + sessionReceiverClientBuilder.maxConcurrentSessions(maxConcurrentSessions); + return this; + } + + /** + * Sets the prefetch count of the processor. For both {@link ServiceBusReceiveMode#PEEK_LOCK PEEK_LOCK} and + * {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} modes the default value is 0. + * + * @param prefetchCount The prefetch count. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder prefetchCount(int prefetchCount) { + sessionReceiverClientBuilder.prefetchCount(prefetchCount); + return this; + } + + /** + * Sets the name of the queue to create a processor for. + * @param queueName Name of the queue. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder queueName(String queueName) { + sessionReceiverClientBuilder.queueName(queueName); + return this; + } + + /** + * Sets the receive mode for the processor. + * @param receiveMode Mode for receiving messages. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder receiveMode(ServiceBusReceiveMode receiveMode) { + sessionReceiverClientBuilder.receiveMode(receiveMode); + return this; + } + + /** + * Sets the type of the {@link SubQueue} to connect to. + * + * @param subQueue The type of the sub queue. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @see #queueName A queue name or #topicName A topic name should be set as well. + * @see SubQueue + */ + public ServiceBusSessionProcessorAsyncClientBuilder subQueue(SubQueue subQueue) { + this.sessionReceiverClientBuilder.subQueue(subQueue); + return this; + } + + /** + * Sets the name of the subscription in the topic to listen to. {@link #topicName(String)} must also be set. + * + * @param subscriptionName Name of the subscription. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @see #topicName A topic name should be set as well. + */ + public ServiceBusSessionProcessorAsyncClientBuilder subscriptionName(String subscriptionName) { + sessionReceiverClientBuilder.subscriptionName(subscriptionName); + return this; + } + + /** + * Sets the name of the topic. {@link #subscriptionName(String)} must also be set. + * @param topicName Name of the topic. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @see #subscriptionName A subscription name should be set as well. + */ + public ServiceBusSessionProcessorAsyncClientBuilder topicName(String topicName) { + sessionReceiverClientBuilder.topicName(topicName); + return this; + } + + /** + * The async message processing callback for the processor that will be executed when a message is received. + * The returned {@link Mono} represents the asynchronous processing of the message. + * + * @param processMessage The async message processing function invoked when a message is received. + * + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder + processMessage(Function> processMessage) { + this.processMessage = processMessage; + return this; + } + + /** + * The async error handler for the processor which will be invoked in the event of an error while receiving + * messages. + * + * @param processError The async error handler invoked when an error occurs. + * + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder + processError(Function> processError) { + this.processError = processError; + return this; + } + + /** + * The maximum number of messages processed concurrently. For the async session processor this bounds the + * total number of in-flight handlers across all active sessions (the processor rolls through up to + * {@code maxConcurrentSessions} sessions); it is not a per-session limit. + * + * @param maxConcurrentCalls the maximum number of messages processed concurrently across all sessions. + * + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException if {@code maxConcurrentCalls} is less than 1. + */ + public ServiceBusSessionProcessorAsyncClientBuilder maxConcurrentCalls(int maxConcurrentCalls) { + if (maxConcurrentCalls < 1) { + throw LOGGER + .logExceptionAsError(new IllegalArgumentException("'maxConcurrentCalls' cannot be less than 1")); + } + processorClientOptions.setMaxConcurrentCalls(maxConcurrentCalls); + return this; + } + + /** + * Sets the maximum time to wait for in-flight message handlers to complete when the processor is closed. + * + *

If not specified, defaults to 30 seconds.

+ * + * @param drainTimeout The maximum time to wait for in-flight handlers. Must be positive. + * @return The updated {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + * @throws NullPointerException if {@code drainTimeout} is null. + * @throws IllegalArgumentException if {@code drainTimeout} is zero or negative. + */ + public ServiceBusSessionProcessorAsyncClientBuilder drainTimeout(Duration drainTimeout) { + processorClientOptions.setDrainTimeout(drainTimeout); + return this; + } + + /** + * Disables auto-complete and auto-abandon of received messages. By default, a message whose handler + * {@link Mono} completes successfully is {@link ServiceBusReceivedMessageContext#complete() completed}, and a + * message whose handler {@link Mono} signals an error is {@link ServiceBusReceivedMessageContext#abandon() + * abandoned}. + * + * @return The modified {@link ServiceBusSessionProcessorAsyncClientBuilder} object. + */ + public ServiceBusSessionProcessorAsyncClientBuilder disableAutoComplete() { + processorClientOptions.setDisableAutoComplete(true); + return this; + } + + /** + * Creates a session-aware async Service Bus processor responsible for reading + * {@link ServiceBusReceivedMessage messages} from a specific queue or subscription with reactive, + * non-blocking handlers. + * + * @return A new {@link ServiceBusProcessorAsyncClient} that receives messages from a queue or subscription. + * @throws IllegalStateException if {@link #queueName(String) queueName} or {@link #topicName(String) + * topicName} are not set or, both of these fields are set. It is also thrown if the Service Bus {@link + * #connectionString(String) connectionString} contains an {@code EntityPath} that does not match one set in + * {@link #queueName(String) queueName} or {@link #topicName(String) topicName}. Lastly, if a {@link + * #topicName(String) topicName} is set, but {@link #subscriptionName(String) subscriptionName} is not. + * @throws IllegalArgumentException Queue or topic name are not set via {@link #queueName(String) + * queueName()} or {@link #topicName(String) topicName()}, respectively. + * @throws NullPointerException if the {@link #processMessage(Function)} or {@link #processError(Function)} + * callbacks are not set. + */ + public ServiceBusProcessorAsyncClient buildProcessorAsyncClient() { + // Validate entity/connection inputs eagerly so a misconfiguration fails at build time (matching the + // documented exceptions) rather than being deferred to start(). + validateInputs(); + // The async processor manages message settlement explicitly in its reactive pipeline, so the underlying + // receiver must never auto-settle. The user's auto-complete preference is carried in + // processorClientOptions and applied by the processor itself. + sessionReceiverClientBuilder.disableAutoComplete(); + return new ServiceBusProcessorAsyncClient(sessionReceiverClientBuilder, + sessionReceiverClientBuilder.queueName, sessionReceiverClientBuilder.topicName, + sessionReceiverClientBuilder.subscriptionName, + Objects.requireNonNull(processMessage, "'processMessage' cannot be null"), + Objects.requireNonNull(processError, "'processError' cannot be null"), processorClientOptions); + } + + private void validateInputs() { + final ServiceBusSessionReceiverClientBuilder builder = sessionReceiverClientBuilder; + final MessagingEntityType entityType + = validateEntityPaths(connectionStringEntityName, builder.topicName, builder.queueName); + getEntityPath(entityType, builder.queueName, builder.topicName, builder.subscriptionName, builder.subQueue); + getConnectionOptions(); + } + } + /** * Builder for creating {@link ServiceBusReceiverClient} and {@link ServiceBusReceiverAsyncClient} to consume * messages from a session aware Service Bus entity. @@ -2652,6 +2930,232 @@ private void validateInputs() { } } + /** + * Builder for creating {@link ServiceBusProcessorAsyncClient} to process messages with reactive, non-blocking + * message handlers from a non-session-enabled Service Bus entity. + * + * @see ServiceBusProcessorAsyncClient + */ + @ServiceClientBuilder(serviceClients = { ServiceBusProcessorAsyncClient.class }) + public final class ServiceBusProcessorAsyncClientBuilder { + private final ServiceBusReceiverClientBuilder serviceBusReceiverClientBuilder; + private final ServiceBusProcessorClientOptions processorClientOptions; + private Function> processMessage; + private Function> processError; + + private ServiceBusProcessorAsyncClientBuilder() { + serviceBusReceiverClientBuilder = new ServiceBusReceiverClientBuilder(); + processorClientOptions = new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(1); + } + + /** + * Sets the prefetch count of the processor. For both {@link ServiceBusReceiveMode#PEEK_LOCK PEEK_LOCK} and + * {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} modes the default value is 0. + * + * Prefetch speeds up the message flow by aiming to have a message readily available for local retrieval when + * and before the application starts the processor. + * Setting a non-zero value will prefetch that number of messages. Setting the value to zero turns prefetch off. + * + * @param prefetchCount The prefetch count. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder prefetchCount(int prefetchCount) { + serviceBusReceiverClientBuilder.prefetchCount(prefetchCount); + return this; + } + + /** + * Sets the name of the queue to create a processor for. + * @param queueName Name of the queue. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder queueName(String queueName) { + serviceBusReceiverClientBuilder.queueName(queueName); + return this; + } + + /** + * Sets the receive mode for the processor. + * @param receiveMode Mode for receiving messages. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder receiveMode(ServiceBusReceiveMode receiveMode) { + serviceBusReceiverClientBuilder.receiveMode(receiveMode); + return this; + } + + /** + * Sets the type of the {@link SubQueue} to connect to. + * + * @param subQueue The type of the sub queue. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + * @see #queueName A queue name or #topicName A topic name should be set as well. + * @see SubQueue + */ + public ServiceBusProcessorAsyncClientBuilder subQueue(SubQueue subQueue) { + serviceBusReceiverClientBuilder.subQueue(subQueue); + return this; + } + + /** + * Sets the name of the subscription in the topic to listen to. {@link #topicName(String)} must also be set. + * + * @param subscriptionName Name of the subscription. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + * @see #topicName A topic name should be set as well. + */ + public ServiceBusProcessorAsyncClientBuilder subscriptionName(String subscriptionName) { + serviceBusReceiverClientBuilder.subscriptionName(subscriptionName); + return this; + } + + /** + * Sets the name of the topic. {@link #subscriptionName(String)} must also be set. + * @param topicName Name of the topic. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + * @see #subscriptionName A subscription name should be set as well. + */ + public ServiceBusProcessorAsyncClientBuilder topicName(String topicName) { + serviceBusReceiverClientBuilder.topicName(topicName); + return this; + } + + /** + * The async message processing callback for the processor which will be executed when a message is received. + * The returned {@link Mono} represents the asynchronous processing of the message. Up to + * {@code maxConcurrentCalls} messages are processed concurrently; a new message is requested from the broker as + * each in-flight handler's {@link Mono} terminates. + * + * @param processMessage The async message processing function invoked when a message is received. + * + * @return The updated {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder + processMessage(Function> processMessage) { + this.processMessage = processMessage; + return this; + } + + /** + * The async error handler for the processor which will be invoked in the event of an error while receiving + * messages. + * + * @param processError The async error handler invoked when an error occurs. + * + * @return The updated {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder + processError(Function> processError) { + this.processError = processError; + return this; + } + + /** + * Sets the amount of time to continue auto-renewing the lock. Setting {@link Duration#ZERO} or {@code null} + * disables auto-renewal. For {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} mode, + * auto-renewal is disabled. + * + * @param maxAutoLockRenewDuration the amount of time to continue auto-renewing the lock. {@link Duration#ZERO} + * or {@code null} indicates that auto-renewal is disabled. + * + * @return The updated {@link ServiceBusProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException If {@code maxAutoLockRenewDuration} is negative. + */ + public ServiceBusProcessorAsyncClientBuilder maxAutoLockRenewDuration(Duration maxAutoLockRenewDuration) { + validateAndThrow(maxAutoLockRenewDuration, "maxAutoLockRenewDuration"); + serviceBusReceiverClientBuilder.maxAutoLockRenewDuration(maxAutoLockRenewDuration); + return this; + } + + /** + * Max concurrent messages that this processor should process. By default, this is set to 1. + * + * @param maxConcurrentCalls max concurrent messages that this processor should process. + * @return The updated {@link ServiceBusProcessorAsyncClientBuilder} object. + * @throws IllegalArgumentException if the {@code maxConcurrentCalls} is set to a value less than 1. + */ + public ServiceBusProcessorAsyncClientBuilder maxConcurrentCalls(int maxConcurrentCalls) { + if (maxConcurrentCalls < 1) { + throw LOGGER + .logExceptionAsError(new IllegalArgumentException("'maxConcurrentCalls' cannot be less than 1")); + } + processorClientOptions.setMaxConcurrentCalls(maxConcurrentCalls); + return this; + } + + /** + * Sets the maximum time to wait for in-flight message handlers to complete when the processor is closed. + * + *

If not specified, defaults to 30 seconds.

+ * + * @param drainTimeout The maximum time to wait for in-flight handlers. Must be positive. + * @return The updated {@link ServiceBusProcessorAsyncClientBuilder} object. + * @throws NullPointerException if {@code drainTimeout} is null. + * @throws IllegalArgumentException if {@code drainTimeout} is zero or negative. + */ + public ServiceBusProcessorAsyncClientBuilder drainTimeout(Duration drainTimeout) { + processorClientOptions.setDrainTimeout(drainTimeout); + return this; + } + + /** + * Disables auto-complete and auto-abandon of received messages. By default, a message whose handler + * {@link Mono} completes successfully is {@link ServiceBusReceivedMessageContext#complete() completed}, and a + * message whose handler {@link Mono} signals an error is {@link ServiceBusReceivedMessageContext#abandon() + * abandoned}. + * + * @return The modified {@link ServiceBusProcessorAsyncClientBuilder} object. + */ + public ServiceBusProcessorAsyncClientBuilder disableAutoComplete() { + processorClientOptions.setDisableAutoComplete(true); + return this; + } + + /** + * Creates an async Service Bus message processor responsible for reading {@link ServiceBusReceivedMessage + * messages} from a specific queue or subscription with reactive, non-blocking handlers. + * + * @return A new {@link ServiceBusProcessorAsyncClient} that processes messages from a queue or subscription. + * @throws IllegalStateException if {@link #queueName(String) queueName} or {@link #topicName(String) + * topicName} are not set or, both of these fields are set. It is also thrown if the Service Bus {@link + * #connectionString(String) connectionString} contains an {@code EntityPath} that does not match one set in + * {@link #queueName(String) queueName} or {@link #topicName(String) topicName}. Lastly, if a {@link + * #topicName(String) topicName} is set, but {@link #subscriptionName(String) subscriptionName} is not. + * @throws IllegalArgumentException Queue or topic name are not set via {@link #queueName(String) + * queueName()} or {@link #topicName(String) topicName()}, respectively. + * @throws NullPointerException if the {@link #processMessage(Function)} or {@link #processError(Function)} + * callbacks are not set. + */ + public ServiceBusProcessorAsyncClient buildProcessorAsyncClient() { + // Validate entity/connection inputs eagerly so a misconfiguration fails at build time (matching the + // documented exceptions) rather than being deferred to start(). + validateInputs(); + // The async processor manages message settlement explicitly in its reactive pipeline, so the underlying + // receiver must never auto-settle. The user's auto-complete preference is carried in + // processorClientOptions and applied by the processor itself. + serviceBusReceiverClientBuilder.disableAutoComplete(); + return new ServiceBusProcessorAsyncClient(serviceBusReceiverClientBuilder, + serviceBusReceiverClientBuilder.queueName, serviceBusReceiverClientBuilder.topicName, + serviceBusReceiverClientBuilder.subscriptionName, + Objects.requireNonNull(processMessage, "'processMessage' cannot be null"), + Objects.requireNonNull(processError, "'processError' cannot be null"), processorClientOptions); + } + + private void validateInputs() { + final ServiceBusReceiverClientBuilder builder = serviceBusReceiverClientBuilder; + final MessagingEntityType entityType + = validateEntityPaths(connectionStringEntityName, builder.topicName, builder.queueName); + getEntityPath(entityType, builder.queueName, builder.topicName, builder.subscriptionName, builder.subQueue); + getConnectionOptions(); + } + } + /** * Builder for creating {@link ServiceBusReceiverClient} and {@link ServiceBusReceiverAsyncClient} to consume * messages from Service Bus. diff --git a/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClient.java b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClient.java new file mode 100644 index 000000000000..d13c262e661f --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/main/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClient.java @@ -0,0 +1,507 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.messaging.servicebus; + +import com.azure.core.util.logging.ClientLogger; +import com.azure.messaging.servicebus.ServiceBusClientBuilder.ServiceBusReceiverClientBuilder; +import com.azure.messaging.servicebus.ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder; +import com.azure.messaging.servicebus.implementation.ServiceBusProcessorClientOptions; +import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; +import reactor.core.Disposable; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.time.Duration; +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +/** + * The asynchronous processor client for processing Service Bus messages with reactive, non-blocking message handlers. + * {@link ServiceBusProcessorAsyncClient} provides the same push-based, auto-recovering, concurrency-managed message + * pump as {@link ServiceBusProcessorClient}, but the message and error handlers return a {@link Mono} instead of + * running as blocking {@link java.util.function.Consumer Consumer} callbacks. + * + *

This is the async counterpart to {@link ServiceBusProcessorClient}, completing the family of asynchronous + * clients ({@link ServiceBusSenderAsyncClient}, {@link ServiceBusReceiverAsyncClient}). It targets applications + * that perform I/O-bound work in the handler (for example, reactive HTTP calls, non-blocking database writes, or + * further messaging) and want to compose that work reactively without blocking a processing thread for the duration + * of each message, and without re-implementing auto-recovery, concurrency management, and lifecycle handling on top + * of {@link ServiceBusReceiverAsyncClient}.

+ * + *

Messages are dispatched to the handler with up to {@code maxConcurrentCalls} concurrent invocations in flight; + * a new message is requested from the broker only as an in-flight handler completes. By default a message whose + * handler {@link Mono} completes successfully is {@link ServiceBusReceivedMessageContext#complete() completed}, and a + * message whose handler {@link Mono} signals an error is {@link ServiceBusReceivedMessageContext#abandon() abandoned}; + * this auto-settlement can be disabled via + * {@link ServiceBusClientBuilder.ServiceBusProcessorAsyncClientBuilder#disableAutoComplete() + * disableAutoComplete()}.

+ * + *

A {@link ServiceBusProcessorAsyncClient} can be created for a session-enabled or a non-session-enabled Service + * Bus entity through {@link ServiceBusClientBuilder#processorAsync()} and + * {@link ServiceBusClientBuilder#sessionProcessorAsync()} respectively.

+ * + *

Auto-settlement and manual settlement

+ *

Auto-settlement (the default) is fully non-blocking - the completion or abandonment is issued reactively when the + * handler {@link Mono} terminates. When auto-settlement is disabled, the handler is responsible for settling each + * message through the {@link ServiceBusReceivedMessageContext} passed to it. Note that the settlement methods on + * {@link ServiceBusReceivedMessageContext} ({@link ServiceBusReceivedMessageContext#complete() complete()}, + * {@link ServiceBusReceivedMessageContext#abandon() abandon()}, etc.) are blocking; a handler that + * settles manually should perform that call on a scheduler that permits blocking (for example by wrapping it in + * {@link Mono#fromRunnable(Runnable)} subscribed on {@code Schedulers.boundedElastic()}), or + * rely on auto-settlement, which is non-blocking. A future revision may add reactive settlement methods.

+ * + *

Lifecycle

+ *

The lifecycle operations {@link #start()}, {@link #stop()}, and {@link #close()} are not designed to be invoked + * concurrently from multiple threads. Drive the processor's lifecycle from a single controlling thread; interleaving + * these calls from different threads leads to undefined behavior.

+ * + * @see ServiceBusProcessorClient + * @see ServiceBusClientBuilder.ServiceBusProcessorAsyncClientBuilder + * @see ServiceBusClientBuilder.ServiceBusSessionProcessorAsyncClientBuilder + */ +public final class ServiceBusProcessorAsyncClient implements AutoCloseable { + + private static final int SCHEDULER_INTERVAL_IN_SECONDS = 10; + private static final ClientLogger LOGGER = new ClientLogger(ServiceBusProcessorAsyncClient.class); + + private final ServiceBusReceiverClientBuilder receiverBuilder; + private final ServiceBusSessionReceiverClientBuilder sessionReceiverBuilder; + private final Function> processMessage; + private final Function> processError; + private final ServiceBusProcessorClientOptions processorOptions; + private final boolean autoComplete; + private final int maxConcurrentCalls; + + private final String queueName; + private final String topicName; + private final String subscriptionName; + + private final AtomicReference asyncClient = new AtomicReference<>(); + private final AtomicReference receiveDisposable = new AtomicReference<>(); + private final AtomicBoolean isRunning = new AtomicBoolean(); + private final AtomicInteger activeHandlerCount = new AtomicInteger(0); + private final Object drainLock = new Object(); + private volatile String cachedFullyQualifiedNamespace; + private volatile String cachedEntityPath; + // True while close() is draining. New PEEK_LOCK dispatches are skipped (without counting toward the drain) so the + // in-flight count can reach zero; RECEIVE_AND_DELETE dispatches are never skipped because the broker has already + // removed the message and dropping it here would lose it. Reset by start() so the processor can restart. + private volatile boolean closing; + // PEEK_LOCK is safe to skip during drain (the broker still owns the lock and will redeliver). Cached per receive + // cycle; defaults to false (no-skip) when the receive mode cannot be determined, so messages are never dropped. + private volatile boolean skipDuringDrain; + // Auto-settlement (complete/abandon) is applied only for PEEK_LOCK: in RECEIVE_AND_DELETE the broker already + // removed the message on delivery, so a settlement call is meaningless and would fail. Cached per receive cycle. + private volatile boolean settleMessages; + private Disposable monitorDisposable; + + /** + * Constructor to create a non-session async processor. + * + * @param receiverBuilder The receiver builder used to create new receivers for the processor. + * @param queueName The name of the queue this processor is associated with. + * @param topicName The name of the topic this processor is associated with. + * @param subscriptionName The name of the subscription this processor is associated with. + * @param processMessage The async message processing callback. + * @param processError The async error handler. + * @param processorOptions Options to configure this instance of the processor. + */ + ServiceBusProcessorAsyncClient(ServiceBusReceiverClientBuilder receiverBuilder, String queueName, String topicName, + String subscriptionName, Function> processMessage, + Function> processError, ServiceBusProcessorClientOptions processorOptions) { + this.receiverBuilder = Objects.requireNonNull(receiverBuilder, "'receiverBuilder' cannot be null"); + this.sessionReceiverBuilder = null; + this.processMessage = Objects.requireNonNull(processMessage, "'processMessage' cannot be null"); + this.processError = Objects.requireNonNull(processError, "'processError' cannot be null"); + this.processorOptions = Objects.requireNonNull(processorOptions, "'processorOptions' cannot be null"); + this.autoComplete = !processorOptions.isDisableAutoComplete(); + this.maxConcurrentCalls = validateMaxConcurrentCalls(processorOptions.getMaxConcurrentCalls()); + this.queueName = queueName; + this.topicName = topicName; + this.subscriptionName = subscriptionName; + } + + /** + * Constructor to create a session-enabled async processor. + * + * @param sessionReceiverBuilder The session receiver builder used to create new receivers for the processor. + * @param queueName The name of the queue this processor is associated with. + * @param topicName The name of the topic this processor is associated with. + * @param subscriptionName The name of the subscription this processor is associated with. + * @param processMessage The async message processing callback. + * @param processError The async error handler. + * @param processorOptions Options to configure this instance of the processor. + */ + ServiceBusProcessorAsyncClient(ServiceBusSessionReceiverClientBuilder sessionReceiverBuilder, String queueName, + String topicName, String subscriptionName, + Function> processMessage, + Function> processError, ServiceBusProcessorClientOptions processorOptions) { + this.sessionReceiverBuilder + = Objects.requireNonNull(sessionReceiverBuilder, "'sessionReceiverBuilder' cannot be null"); + this.receiverBuilder = null; + this.processMessage = Objects.requireNonNull(processMessage, "'processMessage' cannot be null"); + this.processError = Objects.requireNonNull(processError, "'processError' cannot be null"); + this.processorOptions = Objects.requireNonNull(processorOptions, "'processorOptions' cannot be null"); + this.autoComplete = !processorOptions.isDisableAutoComplete(); + this.maxConcurrentCalls = validateMaxConcurrentCalls(processorOptions.getMaxConcurrentCalls()); + this.queueName = queueName; + this.topicName = topicName; + this.subscriptionName = subscriptionName; + } + + /** + * Starts the processor in the background. When the returned {@link Mono} completes, the processor has wired a + * message receiver that invokes the async message handler as messages become available, and the async error + * handler when an error occurs. Control returns immediately - the returned {@link Mono} does not wait for messages + * to be processed. + *

The returned {@link Mono} is cold: the processor does not start until you subscribe to (or + * {@code block()} on) it. Calling {@code start()} without subscribing is a no-op.

+ *

+ * This method is idempotent - subscribing to the {@link Mono} returned when the processor is already running is a + * no-op. Calling {@code start()} after {@link #stop() stop()} resumes processing using the same underlying + * connection; calling {@code start()} after {@link #close() close()} starts the processor with a new connection. + *

+ * + * @return A {@link Mono} that completes when the processor has started. + */ + public Mono start() { + return Mono.fromRunnable(() -> { + synchronized (this) { + if (isRunning.getAndSet(true)) { + LOGGER.info("Processor is already running"); + return; + } + closing = false; + if (asyncClient.get() == null) { + asyncClient.set(createNewReceiver()); + } + subscribeToReceiver(asyncClient.get()); + startMonitor(); + } + }); + } + + /** + * Stops message processing for this processor. The receiving links and sessions are kept active and processing can + * be resumed by subscribing to {@link #start()} again. This does not wait for in-flight handlers to finish; + * disposing the receive subscription cancels any handlers that are still running. Use {@link #close()} to give + * in-flight handlers a best-effort chance to drain. + *

The returned {@link Mono} is cold: it takes effect only when subscribed to (or + * blocked on).

+ * + * @return A {@link Mono} that completes when the processor has stopped requesting new messages. + */ + public Mono stop() { + return Mono.fromRunnable(() -> { + synchronized (this) { + isRunning.set(false); + final Disposable disposable = receiveDisposable.getAndSet(null); + if (disposable != null) { + disposable.dispose(); + } + } + }); + } + + /** + * Stops message processing and closes the processor. The receiving links and sessions are closed; subscribing to + * {@link #start()} afterwards creates a new processing cycle with a new connection. + * + *

This method blocks while waiting for in-flight message handlers to complete (up to the configured drain + * timeout, default 30 seconds) before cancelling the subscription and closing the underlying client, so handlers + * that are already running can finish settlement against a live receiver. Draining is best-effort: + * a message the pump had already accepted at the instant {@code close()} begins may occasionally start after the + * drain observes an empty in-flight set, in which case it may settle against a closing receiver - for + * {@link ServiceBusReceiveMode#PEEK_LOCK PEEK_LOCK} the broker simply redelivers such a message, and for + * {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} no settlement is required. Callers should + * avoid invoking {@code close()} on latency-sensitive threads.

+ * + *

RECEIVE_AND_DELETE and shutdown: the broker removes + * {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE RECEIVE_AND_DELETE} messages on delivery, and the receiver keeps + * delivering while the drain waits. Under sustained traffic the in-flight set may not reach zero before the drain + * timeout; handlers still running when it elapses are cancelled and those messages are lost (they cannot be + * redelivered). A future revision may decouple handler execution from the receiver subscription to bound this.

+ * + *

Do not call {@code close()} from within a message or error handler: it blocks until in-flight handlers + * drain, so invoking it from a handler that has not yet returned would deadlock. Calling {@link #stop()} from a + * handler is also unsafe: it disposes the receive subscription and cancels the very handler that invoked it. + * Schedule shutdown from a separate controlling thread.

+ */ + @Override + public void close() { + final Duration drainTimeout; + synchronized (this) { + if (!isRunning.getAndSet(false) && asyncClient.get() == null) { + return; + } + // Signal the pump to stop accepting new PEEK_LOCK dispatches so the in-flight count can drain to zero. + closing = true; + drainTimeout = processorOptions.getDrainTimeout(); + } + + // Drain in-flight handlers BEFORE disposing the subscription and closing the receiver, so handlers that are + // already running can settle against a live receiver. Best-effort: a dispatch accepted by the pump at the + // instant close() begins may still slip past the drain (see the close() Javadoc). + drainHandlers(drainTimeout); + + synchronized (this) { + final Disposable disposable = receiveDisposable.getAndSet(null); + if (disposable != null) { + disposable.dispose(); + } + if (monitorDisposable != null) { + monitorDisposable.dispose(); + monitorDisposable = null; + } + final ServiceBusReceiverAsyncClient client = asyncClient.getAndSet(null); + if (client != null) { + client.close(); + } + } + } + + /** + * Returns {@code true} if the processor is running. If the processor is stopped or closed, this returns + * {@code false}. + * + * @return {@code true} if the processor is running; {@code false} otherwise. + */ + public boolean isRunning() { + return isRunning.get(); + } + + /** + * Returns the queue name associated with this instance of {@link ServiceBusProcessorAsyncClient}. + * + * @return the queue name, or {@code null} if the processor instance is for a topic and subscription. + */ + public String getQueueName() { + return this.queueName; + } + + /** + * Returns the topic name associated with this instance of {@link ServiceBusProcessorAsyncClient}. + * + * @return the topic name, or {@code null} if the processor instance is for a queue. + */ + public String getTopicName() { + return this.topicName; + } + + /** + * Returns the subscription name associated with this instance of {@link ServiceBusProcessorAsyncClient}. + * + * @return the subscription name, or {@code null} if the processor instance is for a queue. + */ + public String getSubscriptionName() { + return this.subscriptionName; + } + + /** + * Gets the identifier of the instance of {@link ServiceBusProcessorAsyncClient}. + * + * @return The identifier that can identify the instance of {@link ServiceBusProcessorAsyncClient}, or {@code null} + * if no receiver has been created yet (before the first {@link #start() start()}). + */ + public String getIdentifier() { + final ServiceBusReceiverAsyncClient client = asyncClient.get(); + return client == null ? null : client.getIdentifier(); + } + + private void subscribeToReceiver(ServiceBusReceiverAsyncClient receiverClient) { + cachedFullyQualifiedNamespace = receiverClient.getFullyQualifiedNamespace(); + cachedEntityPath = receiverClient.getEntityPath(); + // Only PEEK_LOCK is safe to skip during drain, and PEEK_LOCK is the only mode the processor settles. Default + // to no-skip and no-settle when the mode is unavailable. + final ReceiverOptions receiverOptions = receiverClient.getReceiverOptions(); + final boolean isPeekLock + = receiverOptions != null && receiverOptions.getReceiveMode() == ServiceBusReceiveMode.PEEK_LOCK; + this.skipDuringDrain = isPeekLock; + this.settleMessages = autoComplete && isPeekLock; + + final Disposable disposable = receiverClient.receiveMessagesWithContext() + .flatMap(messageContext -> dispatchMessage(messageContext, receiverClient), maxConcurrentCalls, 1) + .subscribe(ignored -> { + }, throwable -> { + LOGGER.info("Error receiving messages.", throwable); + handleError(throwable).subscribe(); + scheduleRestart(); + }, () -> { + LOGGER.info("Completed receiving messages."); + scheduleRestart(); + }); + receiveDisposable.set(disposable); + } + + private Mono dispatchMessage(ServiceBusMessageContext messageContext, + ServiceBusReceiverAsyncClient receiverClient) { + if (messageContext.hasError()) { + return handleError(messageContext.getThrowable()); + } + + return Mono.defer(() -> { + // Fast path: a redeliverable PEEK_LOCK message that arrives while close() is draining is skipped WITHOUT + // touching the in-flight counter, so churn from skipped messages under sustained load cannot keep the + // drain above zero and push shutdown toward the drain timeout. (The broker redelivers a skipped message.) + if (closing && skipDuringDrain) { + LOGGER.verbose("Skipping dispatch, processor is closing."); + return Mono.empty(); + } + // Publish intent-to-run, then re-check: this handles the check-then-act race where closing flips between + // the fast-path check and the increment, and narrows (best-effort, per the close() Javadoc) the window in + // which a dispatch could start after the drain has observed an empty in-flight set. + activeHandlerCount.incrementAndGet(); + if (closing && skipDuringDrain) { + LOGGER.verbose("Skipping dispatch, processor is closing."); + decrementAndNotifyDrain(); + return Mono.empty(); + } + final ServiceBusReceivedMessageContext receivedMessageContext; + try { + receivedMessageContext = new ServiceBusReceivedMessageContext(receiverClient, messageContext); + } catch (RuntimeException e) { + // Constructor threw before the doFinally decrement could attach; release the count so a synchronous + // failure here cannot leave the drain permanently above zero. + decrementAndNotifyDrain(); + throw LOGGER.logExceptionAsError(e); + } + + return Mono.defer(() -> processMessage.apply(receivedMessageContext)).then(Mono.defer(() -> { + // The handler succeeded. Auto-settle only in PEEK_LOCK; a completion failure is reported to the error + // handler with its OWN error source (e.g. COMPLETE) and must NOT trigger an abandon - the handler + // did its job. + if (!settleMessages) { + return Mono.empty(); + } + return receiverClient.complete(messageContext.getMessage()) + .onErrorResume(completeError -> handleError(completeError)); + })) + // Only a handler error reaches here (completion errors were handled above). Report it as USER_CALLBACK + // and abandon the message (when settling, i.e. PEEK_LOCK with auto-complete enabled). + .onErrorResume(handlerError -> handleError( + new ServiceBusException(handlerError, ServiceBusErrorSource.USER_CALLBACK)).then(Mono.defer(() -> { + if (settleMessages) { + LOGGER.warning("Error when processing message. Abandoning message.", handlerError); + return receiverClient.abandon(messageContext.getMessage()).onErrorResume(abandonError -> { + LOGGER.verbose("Failed to abandon message", abandonError); + return Mono.empty(); + }); + } + return Mono.empty(); + }))) + .then() + .doFinally(signalType -> decrementAndNotifyDrain()); + }); + } + + /** + * Decrements the in-flight handler count and, when it reaches zero, wakes any thread blocked in + * {@link #drainHandlers(Duration)} waiting for the drain to complete. + */ + private void decrementAndNotifyDrain() { + if (activeHandlerCount.decrementAndGet() <= 0) { + synchronized (drainLock) { + drainLock.notifyAll(); + } + } + } + + private static int validateMaxConcurrentCalls(int maxConcurrentCalls) { + if (maxConcurrentCalls < 1) { + throw LOGGER + .logExceptionAsError(new IllegalArgumentException("'maxConcurrentCalls' cannot be less than 1")); + } + return maxConcurrentCalls; + } + + private Mono handleError(Throwable throwable) { + return Mono.defer(() -> { + final ServiceBusErrorContext errorContext + = new ServiceBusErrorContext(throwable, cachedFullyQualifiedNamespace, cachedEntityPath); + // Wrap apply() in its own defer so a synchronously-throwing error handler (one that throws instead of + // returning Mono.error) is converted to an error signal and caught by onErrorResume - never propagated to + // the receive subscription, where it would otherwise trigger a spurious restart. + return Mono.defer(() -> processError.apply(errorContext)).onErrorResume(errorHandlerError -> { + LOGGER.verbose("Error from error handler. Ignoring error.", errorHandlerError); + return Mono.empty(); + }); + }); + } + + // Restart off the reactive receive thread: restartMessageReceiver() calls the blocking client close(), which must + // not run on the subscription's onError/onComplete callback thread. restartMessageReceiver() guards on + // isRunning/closing, so a stale schedule after stop()/close() is a no-op. + private void scheduleRestart() { + Mono.fromRunnable(this::restartMessageReceiver).subscribeOn(Schedulers.boundedElastic()).subscribe(); + } + + // Auto-recovery is one restart per receive-stream terminal (error/completion). There is deliberately no backoff + // or attempt cap here: the underlying ServiceBusReceiverAsyncClient applies the configured retry policy (with + // backoff) before surfacing a terminal error, so transient failures are absorbed there. This matches the sync + // ServiceBusProcessorClient - a fast persistent failure (entity deleted, credentials revoked) relies on that + // receiver-level policy rather than a second backoff layer here. Unlike close(), a restart does NOT drain in-flight + // handlers before closing the stale receiver: a handler still running settles against a closing receiver, so a + // PEEK_LOCK message is redelivered (no loss) and RECEIVE_AND_DELETE needs no settlement. + private synchronized void restartMessageReceiver() { + if (!isRunning.get() || closing) { + return; + } + final Disposable disposable = receiveDisposable.getAndSet(null); + if (disposable != null) { + disposable.dispose(); + } + final ServiceBusReceiverAsyncClient staleClient = asyncClient.get(); + if (staleClient != null) { + staleClient.close(); + } + final ServiceBusReceiverAsyncClient freshClient = createNewReceiver(); + asyncClient.set(freshClient); + subscribeToReceiver(freshClient); + } + + private synchronized void startMonitor() { + if (monitorDisposable == null) { + monitorDisposable = Schedulers.boundedElastic().schedulePeriodically(() -> { + final ServiceBusReceiverAsyncClient currentClient = this.asyncClient.get(); + if (currentClient == null) { + return; + } + if (currentClient.isConnectionClosed()) { + restartMessageReceiver(); + } + }, SCHEDULER_INTERVAL_IN_SECONDS, SCHEDULER_INTERVAL_IN_SECONDS, TimeUnit.SECONDS); + } + } + + private ServiceBusReceiverAsyncClient createNewReceiver() { + return this.receiverBuilder == null + ? this.sessionReceiverBuilder.buildAsyncClientForProcessor() + : this.receiverBuilder.buildAsyncClientForProcessor(); + } + + private void drainHandlers(Duration timeout) { + final long deadlineNanos = System.nanoTime() + timeout.toNanos(); + synchronized (drainLock) { + while (activeHandlerCount.get() > 0) { + final long remainingNanos = deadlineNanos - System.nanoTime(); + if (remainingNanos <= 0) { + LOGGER.info("Drain timeout elapsed with {} handler(s) still in flight; proceeding with shutdown.", + activeHandlerCount.get()); + break; + } + try { + final long remainingMillis = Math.max(1, remainingNanos / 1_000_000); + drainLock.wait(remainingMillis); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + } + } +} diff --git a/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClientTest.java b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClientTest.java new file mode 100644 index 000000000000..02a8bb141748 --- /dev/null +++ b/sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusProcessorAsyncClientTest.java @@ -0,0 +1,697 @@ +// 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.messaging.servicebus.implementation.ServiceBusProcessorClientOptions; +import com.azure.messaging.servicebus.models.ServiceBusReceiveMode; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Flux; +import reactor.core.publisher.FluxSink; +import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link ServiceBusProcessorAsyncClient}. + */ +public class ServiceBusProcessorAsyncClientTest { + + private static final String NAMESPACE = "namespace"; + private static final String ENTITY_NAME = "entity"; + + private ServiceBusReceiverAsyncClient mockReceiver; + + /** + * Messages flow to the async handler and, by default, each is completed after the handler {@link Mono} finishes. + */ + @Test + public void receivesMessagesAndAutoCompletes() throws InterruptedException { + final Flux messages = messageContexts(5); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final CountDownLatch latch = new CountDownLatch(5); + final AtomicInteger nextId = new AtomicInteger(); + final AtomicReference unexpectedError = new AtomicReference<>(); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + assertEquals(String.valueOf(nextId.getAndIncrement()), context.getMessage().getMessageId()); + latch.countDown(); + return Mono.empty(); + }, error -> { + unexpectedError.set(error.getException()); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + final boolean success = latch.await(5, TimeUnit.SECONDS); + processor.close(); + + assertTrue(success, "Failed to receive all expected messages"); + assertNull(unexpectedError.get(), "Unexpected error-handler invocation: " + unexpectedError.get()); + verify(mockReceiver, times(5)).complete(any()); + verify(mockReceiver, never()).abandon(any()); + } + + /** + * Session messages flow to the async handler through the session receiver builder. + */ + @Test + public void sessionMessagesAreDispatched() throws InterruptedException { + final Flux messages = messageContexts(6); + final ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder builder = sessionBuilder(messages); + final CountDownLatch latch = new CountDownLatch(6); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + latch.countDown(); + return Mono.empty(); + }, error -> Mono.empty(), options(1, false)); + + processor.start().block(); + final boolean success = latch.await(5, TimeUnit.SECONDS); + + assertTrue(success, "Failed to receive all expected session messages"); + // Settlement parity with the non-session path: each dispatched session message is auto-completed. + verify(mockReceiver, timeout(5000).times(6)).complete(any()); + processor.close(); + } + + /** + * When the handler {@link Mono} signals an error and auto-complete is enabled, the error handler is invoked and the + * message is abandoned (not completed). + */ + @Test + public void handlerErrorInvokesErrorHandlerAndAbandons() throws InterruptedException { + final Flux messages = messageContexts(1); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final CountDownLatch errorLatch = new CountDownLatch(1); + final AtomicReference source = new AtomicReference<>(); + + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.error(new IllegalStateException("boom")), error -> { + if (error.getException() instanceof ServiceBusException) { + source.set(((ServiceBusException) error.getException()).getErrorSource()); + } + errorLatch.countDown(); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + final boolean errored = errorLatch.await(5, TimeUnit.SECONDS); + processor.close(); + + assertTrue(errored, "Error handler was not invoked"); + assertEquals(ServiceBusErrorSource.USER_CALLBACK, source.get()); + verify(mockReceiver, times(1)).abandon(any()); + verify(mockReceiver, never()).complete(any()); + } + + /** + * When auto-complete is disabled, the processor settles nothing itself - neither complete nor abandon is called. + */ + @Test + public void disableAutoCompleteSkipsSettlement() throws InterruptedException { + final Flux messages = messageContexts(3); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final CountDownLatch latch = new CountDownLatch(3); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + latch.countDown(); + return Mono.empty(); + }, error -> Mono.empty(), options(1, true)); + + processor.start().block(); + final boolean success = latch.await(5, TimeUnit.SECONDS); + processor.close(); + + assertTrue(success, "Failed to receive all expected messages"); + verify(mockReceiver, never()).complete(any()); + verify(mockReceiver, never()).abandon(any()); + } + + /** + * A terminal error on the receive stream is surfaced to the error handler. + */ + @Test + public void receiverErrorInvokesErrorHandler() throws InterruptedException { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + when(mockReceiver.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(mockReceiver.getEntityPath()).thenReturn(ENTITY_NAME); + when(mockReceiver.isConnectionClosed()).thenReturn(false); + doNothing().when(mockReceiver).close(); + // First subscription errors; subsequent restart subscribes to a stream that never emits, bounding the retry. + when(mockReceiver.receiveMessagesWithContext()) + .thenReturn(Flux.error(new IllegalStateException("receive failed"))) + .thenReturn(Flux.never()); + + final CountDownLatch errorLatch = new CountDownLatch(1); + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> Mono.empty(), error -> { + errorLatch.countDown(); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + final boolean errored = errorLatch.await(5, TimeUnit.SECONDS); + + assertTrue(errored, "Receive error was not surfaced to the error handler"); + // The terminal receive error triggers a restart - a second receiver is built. + verify(builder, timeout(5000).times(2)).buildAsyncClientForProcessor(); + processor.close(); + // close() must not trigger a further restart - the count stays at two. + verify(builder, times(2)).buildAsyncClientForProcessor(); + } + + /** + * {@code start()} is idempotent and {@code isRunning()} reflects the lifecycle state. + */ + @Test + public void startIsIdempotentAndIsRunningReflectsState() { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(Flux.never()); + + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.empty(), error -> Mono.empty(), options(1, false)); + + assertFalse(processor.isRunning()); + processor.start().block(); + assertTrue(processor.isRunning()); + // Second start is a no-op - the receiver is built only once. + processor.start().block(); + assertTrue(processor.isRunning()); + verify(builder, times(1)).buildAsyncClientForProcessor(); + + processor.stop().block(); + assertFalse(processor.isRunning()); + processor.close(); + } + + /** + * Closing a processor that was never started is a no-op and does not throw. + */ + @Test + public void closeWithoutStartIsNoOp() { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(Flux.never()); + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.empty(), error -> Mono.empty(), options(1, false)); + + processor.close(); + assertFalse(processor.isRunning()); + } + + /** + * Concurrency is bounded by {@code maxConcurrentCalls} - no more than that many handlers run at once. + */ + @Test + public void concurrencyBoundedByMaxConcurrentCalls() throws InterruptedException { + final int maxConcurrentCalls = 2; + final int messageCount = 8; + final Flux messages = messageContexts(messageCount); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + + final CountDownLatch latch = new CountDownLatch(messageCount); + final AtomicInteger inFlight = new AtomicInteger(); + final AtomicInteger maxObserved = new AtomicInteger(); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> Mono.fromRunnable(() -> { + final int current = inFlight.incrementAndGet(); + maxObserved.accumulateAndGet(current, Math::max); + }).then(Mono.delay(Duration.ofMillis(40))).then(Mono.fromRunnable(() -> { + inFlight.decrementAndGet(); + latch.countDown(); + })), error -> Mono.empty(), options(maxConcurrentCalls, false)); + + processor.start().block(); + final boolean success = latch.await(10, TimeUnit.SECONDS); + processor.close(); + + assertTrue(success, "Failed to process all messages"); + assertTrue(maxObserved.get() <= maxConcurrentCalls, + "Observed " + maxObserved.get() + " concurrent handlers, expected at most " + maxConcurrentCalls); + // Lower bound: the pump must actually reach the concurrency limit, otherwise a regression to serial dispatch + // (concurrency 1) would still satisfy the upper bound and pass silently. + assertTrue(maxObserved.get() >= maxConcurrentCalls, "Observed max " + maxObserved.get() + + " concurrent handlers, expected the bound " + maxConcurrentCalls + " to be reached"); + } + + /** + * {@code stop()} followed by {@code start()} resumes processing on the same receiver (no new receiver is built). + */ + @Test + public void stopThenStartResumesWithSameReceiver() { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(Flux.never()); + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.empty(), error -> Mono.empty(), options(1, false)); + + processor.start().block(); + processor.stop().block(); + assertFalse(processor.isRunning()); + processor.start().block(); + assertTrue(processor.isRunning()); + // The receiver is reused across stop/start - it is built exactly once. + verify(builder, times(1)).buildAsyncClientForProcessor(); + processor.close(); + } + + /** + * {@code close()} blocks until an in-flight handler finishes settling against a still-open receiver. + */ + @Test + public void closeDrainsInFlightHandlerBeforeClosingReceiver() throws InterruptedException { + final Flux messages = messageContexts(1); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final CountDownLatch handlerStarted = new CountDownLatch(1); + final CountDownLatch allowFinish = new CountDownLatch(1); + + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.fromRunnable(handlerStarted::countDown).then(Mono.fromRunnable(() -> { + try { + allowFinish.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }).subscribeOn(Schedulers.boundedElastic())).then(), error -> Mono.empty(), options(1, false)); + + processor.start().block(); + assertTrue(handlerStarted.await(5, TimeUnit.SECONDS), "Handler never started"); + + final AtomicBoolean closeReturned = new AtomicBoolean(); + final Thread closeThread = new Thread(() -> { + processor.close(); + closeReturned.set(true); + }); + closeThread.setDaemon(true); + closeThread.start(); + + // close() must still be draining while the handler is held. + Thread.sleep(300); + assertFalse(closeReturned.get(), "close() returned before the in-flight handler finished"); + + allowFinish.countDown(); + closeThread.join(TimeUnit.SECONDS.toMillis(5)); + assertTrue(closeReturned.get(), "close() did not return after the handler finished"); + // The message was settled against the still-open receiver during the drain. + verify(mockReceiver, times(1)).complete(any()); + } + + /** + * A receive stream that completes triggers exactly one restart (not an unbounded loop). + */ + @Test + public void completingStreamTriggersExactlyOneRestart() throws InterruptedException { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + stubReceiverBasics(); + final ServiceBusReceivedMessage message = new ServiceBusReceivedMessage(BinaryData.fromString("hi")); + message.setMessageId("0"); + // First stream emits one message then completes (triggers a restart); the second stream stays open. + when(mockReceiver.receiveMessagesWithContext()).thenReturn(Flux.just(new ServiceBusMessageContext(message))) + .thenReturn(Flux.never()); + + final CountDownLatch processed = new CountDownLatch(1); + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + processed.countDown(); + return Mono.empty(); + }, error -> Mono.empty(), options(1, false)); + + processor.start().block(); + assertTrue(processed.await(5, TimeUnit.SECONDS), "Message was not processed"); + // Exactly one restart: the original receiver plus one rebuild. The second (open) stream does not restart. + verify(builder, timeout(5000).times(2)).buildAsyncClientForProcessor(); + processor.close(); + verify(builder, times(2)).buildAsyncClientForProcessor(); + } + + /** + * A failure to abandon a message (after a handler error) is swallowed and does not break the pump: a subsequent + * message is still processed, and no spurious restart is triggered. + */ + @Test + public void abandonFailureIsSwallowed() throws InterruptedException { + final AtomicReference> sinkRef = new AtomicReference<>(); + final Flux messages = Flux.create(sinkRef::set); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + when(mockReceiver.abandon(any())).thenReturn(Mono.error(new IllegalStateException("abandon failed"))); + final CountDownLatch errorLatch = new CountDownLatch(1); + final CountDownLatch secondProcessed = new CountDownLatch(1); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + if ("boom".equals(context.getMessage().getMessageId())) { + return Mono.error(new IllegalStateException("boom")); + } + return Mono.fromRunnable(secondProcessed::countDown); + }, error -> { + errorLatch.countDown(); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + sinkRef.get().next(messageContext("boom")); + assertTrue(errorLatch.await(5, TimeUnit.SECONDS), "Error handler was not invoked"); + // The abandon failure must be swallowed - the pump survives, so a subsequent message is still processed... + sinkRef.get().next(messageContext("ok")); + assertTrue(secondProcessed.await(5, TimeUnit.SECONDS), "Pump did not survive the abandon failure"); + processor.close(); + + verify(mockReceiver, times(1)).abandon(any()); + // ...and no spurious restart was triggered by the swallowed error (only the initial receiver was built). + verify(builder, times(1)).buildAsyncClientForProcessor(); + } + + /** + * The constructor rejects a {@code maxConcurrentCalls} value below 1 (fail-fast defense in depth beyond the + * builder validation). + */ + @Test + public void constructorRejectsInvalidMaxConcurrentCalls() { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(Flux.never()); + assertThrows(IllegalArgumentException.class, () -> new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, + null, null, context -> Mono.empty(), error -> Mono.empty(), options(0, false))); + } + + /** + * A completion failure (after the handler succeeds) is reported to the error handler with its OWN error source + * ({@code COMPLETE}, not {@code USER_CALLBACK}) and does NOT trigger an abandon - the handler did its job. + */ + @Test + public void completeFailureIsReportedAndNotAbandoned() throws InterruptedException { + final AtomicReference> sinkRef = new AtomicReference<>(); + final Flux messages = Flux.create(sinkRef::set); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final ServiceBusException completeError + = new ServiceBusException(new IllegalStateException("complete failed"), ServiceBusErrorSource.COMPLETE); + when(mockReceiver.complete(any())).thenReturn(Mono.error(completeError)); + final CountDownLatch errorLatch = new CountDownLatch(1); + final AtomicReference source = new AtomicReference<>(); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> Mono.empty(), error -> { + source.set(error.getErrorSource()); + errorLatch.countDown(); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + sinkRef.get().next(messageContext("m1")); + assertTrue(errorLatch.await(5, TimeUnit.SECONDS), "Completion failure was not surfaced to the error handler"); + processor.close(); + + // The completion failure keeps its COMPLETE source (not misclassified as USER_CALLBACK), and the message is + // NOT abandoned - the handler succeeded. + assertEquals(ServiceBusErrorSource.COMPLETE, source.get(), "error source should be COMPLETE"); + verify(mockReceiver, times(1)).complete(any()); + verify(mockReceiver, never()).abandon(any()); + } + + /** + * {@code getIdentifier()} returns {@code null} before the first {@code start()} and delegates to the underlying + * receiver afterwards. + */ + @Test + public void getIdentifierNullBeforeStartThenDelegates() { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(Flux.never()); + when(mockReceiver.getIdentifier()).thenReturn("receiver-id-123"); + final ServiceBusProcessorAsyncClient processor = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, + null, context -> Mono.empty(), error -> Mono.empty(), options(1, false)); + + assertNull(processor.getIdentifier(), "identifier must be null before start()"); + processor.start().block(); + assertEquals("receiver-id-123", processor.getIdentifier(), "identifier must delegate to the receiver"); + processor.close(); + } + + /** + * An in-band receive error context (a message context carrying a throwable) is surfaced to the error handler and is + * not settled. + */ + @Test + public void inBandErrorContextInvokesErrorHandler() throws InterruptedException { + final ServiceBusMessageContext errorContext + = new ServiceBusMessageContext("session", new IllegalStateException("receive error")); + final Flux messages = Flux.just(errorContext).concatWith(Flux.never()); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder = nonSessionBuilder(messages); + final CountDownLatch errorLatch = new CountDownLatch(1); + final AtomicBoolean handlerRan = new AtomicBoolean(); + + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + handlerRan.set(true); + return Mono.empty(); + }, error -> { + errorLatch.countDown(); + return Mono.empty(); + }, options(1, false)); + + processor.start().block(); + final boolean errored = errorLatch.await(5, TimeUnit.SECONDS); + processor.close(); + + assertTrue(errored, "Error context was not surfaced to the error handler"); + assertFalse(handlerRan.get(), "Message handler should not run for an in-band error context"); + verify(mockReceiver, never()).complete(any()); + verify(mockReceiver, never()).abandon(any()); + } + + /** + * During the close drain window, new PEEK_LOCK dispatches are skipped (not processed or settled) so the in-flight + * count can reach zero, while the already in-flight handler still settles. + */ + @Test + public void skipsNewPeekLockDispatchesDuringDrain() throws InterruptedException { + final AtomicReference> sinkRef = new AtomicReference<>(); + final Flux messages = Flux.create(sinkRef::set); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + stubReceiverBasics(); + when(mockReceiver.receiveMessagesWithContext()).thenReturn(messages); + final ReceiverOptions receiverOptions = mock(ReceiverOptions.class); + when(receiverOptions.getReceiveMode()).thenReturn(ServiceBusReceiveMode.PEEK_LOCK); + when(mockReceiver.getReceiverOptions()).thenReturn(receiverOptions); + + final CountDownLatch slowStarted = new CountDownLatch(1); + final CountDownLatch allowFinish = new CountDownLatch(1); + final Set processedIds = Collections.synchronizedSet(new HashSet<>()); + // maxConcurrentCalls = 3 so the two messages that arrive while the slow handler is held enter the dispatch + // pipeline concurrently (rather than being buffered behind a single slot). This is what makes the skip path + // actually exercised: at concurrency 1 they would never enter dispatch and the test could pass even if the + // skip logic were removed. + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + processedIds.add(context.getMessage().getMessageId()); + if ("slow".equals(context.getMessage().getMessageId())) { + return Mono.fromRunnable(slowStarted::countDown).then(Mono.fromRunnable(() -> { + try { + allowFinish.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }).subscribeOn(Schedulers.boundedElastic())).then(); + } + return Mono.empty(); + }, error -> Mono.empty(), options(3, false)); + + processor.start().block(); + sinkRef.get().next(messageContext("slow")); + assertTrue(slowStarted.await(5, TimeUnit.SECONDS), "Slow handler never started"); + + final AtomicBoolean closeReturned = new AtomicBoolean(); + final Thread closeThread = new Thread(() -> { + processor.close(); + closeReturned.set(true); + }); + closeThread.setDaemon(true); + closeThread.start(); + // close() is now draining (closing == true), waiting for the slow handler. + Thread.sleep(300); + + // Messages that arrive during the drain window must be skipped, not processed/settled. + sinkRef.get().next(messageContext("fast1")); + sinkRef.get().next(messageContext("fast2")); + Thread.sleep(300); + + allowFinish.countDown(); + closeThread.join(TimeUnit.SECONDS.toMillis(5)); + assertTrue(closeReturned.get(), "close() did not return"); + // Only the slow message was completed; the two PEEK_LOCK messages that arrived during drain were skipped. + verify(mockReceiver, times(1)).complete(any()); + // Assert on the skip path directly: the slow handler ran, but the two PEEK_LOCK messages that entered the + // dispatch pipeline during the drain window were skipped before reaching the handler. + assertTrue(processedIds.contains("slow"), "Slow handler should have run"); + assertFalse(processedIds.contains("fast1"), "fast1 arrived during drain and must be skipped (PEEK_LOCK)"); + assertFalse(processedIds.contains("fast2"), "fast2 arrived during drain and must be skipped (PEEK_LOCK)"); + } + + /** + * During the close drain window, {@link ServiceBusReceiveMode#RECEIVE_AND_DELETE} dispatches are NOT skipped - the + * broker has already removed the message, so dropping it here would lose it. This is the asymmetric counterpart to + * {@link #skipsNewPeekLockDispatchesDuringDrain()}: identical setup, only the receive mode differs, and the + * messages that arrive during drain must still be handled. + */ + @Test + public void processesReceiveAndDeleteDispatchesDuringDrain() throws InterruptedException { + final AtomicReference> sinkRef = new AtomicReference<>(); + final Flux messages = Flux.create(sinkRef::set); + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + stubReceiverBasics(); + when(mockReceiver.receiveMessagesWithContext()).thenReturn(messages); + final ReceiverOptions receiverOptions = mock(ReceiverOptions.class); + when(receiverOptions.getReceiveMode()).thenReturn(ServiceBusReceiveMode.RECEIVE_AND_DELETE); + when(mockReceiver.getReceiverOptions()).thenReturn(receiverOptions); + + final Set processedIds = Collections.synchronizedSet(new HashSet<>()); + final CountDownLatch slowStarted = new CountDownLatch(1); + final CountDownLatch allowFinish = new CountDownLatch(1); + final CountDownLatch fastProcessed = new CountDownLatch(2); + // maxConcurrentCalls = 3 so the two drain-window messages run concurrently with the still-held slow handler. + final ServiceBusProcessorAsyncClient processor + = new ServiceBusProcessorAsyncClient(builder, ENTITY_NAME, null, null, context -> { + final String id = context.getMessage().getMessageId(); + if ("slow".equals(id)) { + return Mono.fromRunnable(slowStarted::countDown).then(Mono.fromRunnable(() -> { + try { + allowFinish.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }).subscribeOn(Schedulers.boundedElastic())).then(); + } + return Mono.fromRunnable(() -> { + processedIds.add(id); + fastProcessed.countDown(); + }); + }, error -> Mono.empty(), options(3, false)); + + processor.start().block(); + sinkRef.get().next(messageContext("slow")); + assertTrue(slowStarted.await(5, TimeUnit.SECONDS), "Slow handler never started"); + + final AtomicBoolean closeReturned = new AtomicBoolean(); + final Thread closeThread = new Thread(() -> { + processor.close(); + closeReturned.set(true); + }); + closeThread.setDaemon(true); + closeThread.start(); + // close() is now draining (closing == true), waiting for the slow handler. + Thread.sleep(300); + + // RECEIVE_AND_DELETE messages that arrive during the drain window must still be handled, not dropped. They are + // processed while the slow handler is still held (close() still blocked), proving the not-skip branch. + sinkRef.get().next(messageContext("fast1")); + sinkRef.get().next(messageContext("fast2")); + final boolean bothHandled = fastProcessed.await(5, TimeUnit.SECONDS); + + allowFinish.countDown(); + closeThread.join(TimeUnit.SECONDS.toMillis(5)); + assertTrue(closeReturned.get(), "close() did not return"); + assertTrue(bothHandled, "RECEIVE_AND_DELETE messages during drain must be processed, not skipped"); + assertTrue(processedIds.contains("fast1") && processedIds.contains("fast2"), + "Both RECEIVE_AND_DELETE messages should have been handled during drain"); + // The broker already removed RECEIVE_AND_DELETE messages on delivery, so the processor must not settle them. + verify(mockReceiver, never()).complete(any()); + verify(mockReceiver, never()).abandon(any()); + } + + private ServiceBusProcessorClientOptions options(int maxConcurrentCalls, boolean disableAutoComplete) { + return new ServiceBusProcessorClientOptions().setMaxConcurrentCalls(maxConcurrentCalls) + .setDisableAutoComplete(disableAutoComplete); + } + + private static ServiceBusMessageContext messageContext(String messageId) { + final ServiceBusReceivedMessage message = new ServiceBusReceivedMessage(BinaryData.fromString("hello")); + message.setMessageId(messageId); + return new ServiceBusMessageContext(message); + } + + private static Flux messageContexts(int count) { + final List contexts = Collections.synchronizedList(new ArrayList<>()); + for (int i = 0; i < count; i++) { + final ServiceBusReceivedMessage message = new ServiceBusReceivedMessage(BinaryData.fromString("hello")); + message.setMessageId(String.valueOf(i)); + message.setSessionId(String.valueOf(i % 3)); + contexts.add(new ServiceBusMessageContext(message)); + } + // Emit the messages then stay open (like a live receiver stream). A completing stream would trigger the + // processor's restart-on-complete path and re-subscribe this cold flux, re-emitting the same messages. + return Flux.fromIterable(contexts).concatWith(Flux.never()); + } + + private ServiceBusClientBuilder.ServiceBusReceiverClientBuilder + nonSessionBuilder(Flux messages) { + final ServiceBusClientBuilder.ServiceBusReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + stubReceiver(messages); + return builder; + } + + private ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder + sessionBuilder(Flux messages) { + final ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder builder + = mock(ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder.class); + mockReceiver = mock(ServiceBusReceiverAsyncClient.class); + when(builder.buildAsyncClientForProcessor()).thenReturn(mockReceiver); + stubReceiver(messages); + return builder; + } + + private void stubReceiver(Flux messages) { + stubReceiverBasics(); + when(mockReceiver.receiveMessagesWithContext()).thenReturn(messages); + } + + private void stubReceiverBasics() { + when(mockReceiver.getFullyQualifiedNamespace()).thenReturn(NAMESPACE); + when(mockReceiver.getEntityPath()).thenReturn(ENTITY_NAME); + when(mockReceiver.isConnectionClosed()).thenReturn(false); + when(mockReceiver.complete(any())).thenReturn(Mono.empty()); + when(mockReceiver.abandon(any())).thenReturn(Mono.empty()); + // Default the receiver to PEEK_LOCK so auto-settlement is exercised; tests that need RECEIVE_AND_DELETE + // override getReceiverOptions() after calling this. + final ReceiverOptions peekLockOptions = mock(ReceiverOptions.class); + when(peekLockOptions.getReceiveMode()).thenReturn(ServiceBusReceiveMode.PEEK_LOCK); + when(mockReceiver.getReceiverOptions()).thenReturn(peekLockOptions); + doNothing().when(mockReceiver).close(); + } +}