Servicebus track2 sync queue up multiple receive calls #10940
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,114 +5,203 @@ | |
|
|
||
|
Member
There was a problem hiding this comment. Only thing you need are unit tests. The logic in this subscriber is complex and I can see it being hard to debug. |
||
| import com.azure.core.util.logging.ClientLogger; | ||
| import org.reactivestreams.Subscription; | ||
| import reactor.core.Disposable; | ||
| import reactor.core.publisher.BaseSubscriber; | ||
| import reactor.core.publisher.Mono; | ||
|
|
||
| import java.util.Queue; | ||
| import java.util.concurrent.ConcurrentLinkedQueue; | ||
| import java.util.concurrent.atomic.AtomicBoolean; | ||
| import java.util.concurrent.atomic.AtomicInteger; | ||
| import java.util.concurrent.atomic.AtomicLong; | ||
|
|
||
| /** | ||
| * Subscriber that listens to events and publishes them downstream and publishes events to them in the order received. | ||
| */ | ||
| class SynchronousMessageSubscriber extends BaseSubscriber<ServiceBusReceivedMessageContext> { | ||
| private final ClientLogger logger = new ClientLogger(SynchronousMessageSubscriber.class); | ||
| private final AtomicBoolean isDisposed = new AtomicBoolean(); | ||
| private final long prefetch; | ||
| private final AtomicInteger wip = new AtomicInteger(); | ||
| private final Queue<SynchronousReceiveWork> workQueue = new ConcurrentLinkedQueue<>(); | ||
| private final Queue<ServiceBusReceivedMessageContext> bufferMessages = new ConcurrentLinkedQueue<>(); | ||
| private final SynchronousReceiveWork initialWork; | ||
| private final AtomicLong remaining = new AtomicLong(); | ||
|
|
||
| private volatile Subscription subscription; | ||
|
|
||
| private SynchronousReceiveWork currentWork; | ||
| private Disposable timeoutOperation; | ||
|
|
||
| SynchronousMessageSubscriber(long prefetch, SynchronousReceiveWork initialWork) { | ||
|
hemanttanwar marked this conversation as resolved.
|
||
| this.prefetch = prefetch; | ||
| this.initialWork = initialWork; | ||
| } | ||
|
|
||
| /** | ||
| * On an initial subscription, will take the first work item, and request that amount of work for it. | ||
| * @param subscription Subscription for upstream. | ||
| */ | ||
| @Override | ||
| protected void hookOnSubscribe(Subscription subscription) { | ||
| this.subscription = subscription; | ||
|
Member
There was a problem hiding this comment. Do you need this variable? if the subscription != null, I'll assume there is an upstream. Also, if we call hookonSubscribe twice, it'll stomp over the previous subscription. We should guard against this by only setting it when it is null and erroring when someone wants to set it again. You'll see a setOnce in some of the reactor operations. |
||
|
|
||
| logger.info("[{}] onSubscribe Pending: {}, Scheduling receive timeout task '{}'.", initialWork.getId(), | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| initialWork.getNumberOfEvents(), initialWork.getTimeout()); | ||
|
|
||
| // This will trigger subscription.request(N) and queue up the work | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| queueWork(initialWork); | ||
| } | ||
|
|
||
| /** | ||
| * Publishes the event to the current {@link SynchronousReceiveWork}. If that work item is complete, will dispose of | ||
| * the subscriber. | ||
| * @param message Event to publish. | ||
| */ | ||
| @Override | ||
| protected void hookOnNext(ServiceBusReceivedMessageContext message) { | ||
|
hemanttanwar marked this conversation as resolved.
|
||
| if (currentWork == null) { | ||
| //Boundary condition(timeout cases), buffer the received message for future requests. | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| bufferMessages.add(message); | ||
|
conniey marked this conversation as resolved.
Outdated
|
||
| return; | ||
| } | ||
| currentWork.next(message); | ||
| remaining.decrementAndGet(); | ||
|
|
||
| if (currentWork.isTerminal()) { | ||
| currentWork.complete(); | ||
| if (timeoutOperation != null && !timeoutOperation.isDisposed()) { | ||
| timeoutOperation.dispose(); | ||
| } | ||
|
|
||
| // Now see if there is more queued up work | ||
| currentWork = workQueue.poll(); | ||
| if (currentWork != null) { | ||
|
|
||
| logger.verbose("[{}] Picking up next receive request.", currentWork.getId()); | ||
|
|
||
| // timer to complete the current in case of timeout trigger | ||
| timeoutOperation = getTimeoutOperation(); | ||
|
|
||
| requestCredits(currentWork.getNumberOfEvents()); | ||
| } else { | ||
| if (wip.decrementAndGet() != 0) { | ||
| logger.warning("There is another worker in drainLoop. But there should only be 1 worker."); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * | ||
| * @param requested credits for current {@link SynchronousReceiveWork}. | ||
| */ | ||
| private void requestCredits(long requested) { | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| long creditToAdd = requested - (remaining.get() + bufferMessages.size()); | ||
| if (creditToAdd > 0) { | ||
| remaining.addAndGet(creditToAdd); | ||
| subscription.request(creditToAdd); | ||
| } else { | ||
| logger.verbose("[{}] No need to request credit. ", currentWork.getId()); | ||
| } | ||
|
|
||
| } | ||
|
|
||
| void queueWork(SynchronousReceiveWork work) { | ||
|
|
||
| logger.info("[{}] Pending: {}, Scheduling receive timeout task '{}'.", work.getId(), work.getNumberOfEvents(), | ||
| work.getTimeout()); | ||
| workQueue.add(work); | ||
| drain(); | ||
| } | ||
|
|
||
| private void drain() { | ||
| if (workQueue.size() == 0) { | ||
| return; | ||
| } | ||
|
|
||
| // If someone is already in this loop, then we are already clearing the queue. | ||
| if (!wip.compareAndSet(0, 1)) { | ||
| return; | ||
| } | ||
|
|
||
| // Drain queue.. | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| drainQueue(); | ||
|
|
||
| } | ||
|
|
||
|
|
||
| private void drainQueue() { | ||
|
Member
There was a problem hiding this comment.
Member
There was a problem hiding this comment. You can look at any of the operators for inspiration. |
||
| if (isTerminated()) { | ||
| return; | ||
| } | ||
| currentWork = workQueue.poll(); | ||
| if (currentWork == null) { | ||
| return; | ||
| } | ||
| long sentFromBuffer = 0; | ||
| if (bufferMessages.size() > 0) { | ||
| // If we already have messages in buffer, we should send it first | ||
|
|
||
| while (!bufferMessages.isEmpty() || sentFromBuffer < currentWork.getNumberOfEvents()) { | ||
| currentWork.next(bufferMessages.poll()); | ||
| remaining.decrementAndGet(); | ||
| ++sentFromBuffer; | ||
| } | ||
| if (sentFromBuffer == currentWork.getNumberOfEvents()) { | ||
| currentWork.complete(); | ||
| logger.verbose("[{}] Sent [{}] messages from buffer.", currentWork.getId(), sentFromBuffer); | ||
| drainQueue(); | ||
| } | ||
| } | ||
| // timer to complete the current in case of timeout trigger | ||
| timeoutOperation = getTimeoutOperation(); | ||
|
|
||
| requestCredits(currentWork.getNumberOfEvents() - sentFromBuffer); | ||
| } | ||
|
|
||
| private Disposable getTimeoutOperation() { | ||
| return Mono.delay(currentWork.getTimeout()) | ||
| .subscribe(l -> { | ||
| if (currentWork != null && !currentWork.isTerminal()) { | ||
| logger.verbose("[{}] Timeout triggered.", currentWork.getId()); | ||
| currentWork.complete(); | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| } | ||
| if (wip.decrementAndGet() != 0) { | ||
|
hemanttanwar marked this conversation as resolved.
Outdated
|
||
| logger.warning("There is another worker in drainLoop. But there should only be 1 worker."); | ||
| } | ||
| //any work which needs to be processed. | ||
| drain(); | ||
| }); | ||
| } | ||
| /** | ||
| * {@inheritDoc} | ||
| */ | ||
| @Override | ||
| protected void hookOnError(Throwable throwable) { | ||
| logger.error("[{}] Errors occurred upstream", currentWork.getId(), throwable); | ||
| currentWork.error(throwable); | ||
| dispose(); | ||
| } | ||
|
|
||
| @Override | ||
| protected void hookOnCancel() { | ||
| if (isDisposed.getAndSet(true)) { | ||
| return; | ||
| } | ||
|
|
||
| if (currentWork != null) { | ||
| currentWork.complete(); | ||
| } | ||
|
|
||
| subscription.cancel(); | ||
|
|
||
| if (timeoutOperation != null && !timeoutOperation.isDisposed()) { | ||
| timeoutOperation.dispose(); | ||
| } | ||
| } | ||
|
|
||
| private boolean isTerminated() { | ||
| return isDisposed.get(); | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.