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,250 @@ | |
|
|
||
| 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 reactor.core.publisher.Operators; | ||
|
|
||
| import java.time.Duration; | ||
| 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; | ||
| import java.util.concurrent.atomic.AtomicReferenceFieldUpdater; | ||
|
|
||
| /** | ||
| * 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 AtomicInteger wip = new AtomicInteger(); | ||
| private final Queue<SynchronousReceiveWork> workQueue = new ConcurrentLinkedQueue<>(); | ||
| private final Queue<ServiceBusReceivedMessageContext> bufferMessages = new ConcurrentLinkedQueue<>(); | ||
| private final AtomicLong remaining = new AtomicLong(); | ||
|
|
||
| private final long requested; | ||
| private final Object currentWorkLock = new Object(); | ||
|
|
||
| private Disposable currentTimeoutOperation; | ||
| private SynchronousReceiveWork currentWork; | ||
| private boolean subscriberInitialized; | ||
|
|
||
| private volatile Subscription subscription; | ||
|
|
||
| private static final AtomicReferenceFieldUpdater<SynchronousMessageSubscriber, Subscription> UPSTREAM = | ||
| AtomicReferenceFieldUpdater.newUpdater(SynchronousMessageSubscriber.class, Subscription.class, | ||
| "subscription"); | ||
|
|
||
|
|
||
| SynchronousMessageSubscriber(long prefetch, SynchronousReceiveWork initialWork) { | ||
hemanttanwar marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| this.workQueue.add(initialWork); | ||
| requested = initialWork.getNumberOfEvents() > prefetch ? initialWork.getNumberOfEvents() : prefetch; | ||
| } | ||
|
|
||
| /** | ||
| * 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) { | ||
|
|
||
| if (Operators.setOnce(UPSTREAM, this, subscription)) { | ||
| this.subscription = subscription; | ||
| remaining.addAndGet(requested); | ||
| subscription.request(requested); | ||
| subscriberInitialized = true; | ||
| drain(); | ||
| } else { | ||
| logger.error("Already subscribed once."); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * 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.
Show resolved
Hide resolved
|
||
| bufferMessages.add(message); | ||
| drain(); | ||
| } | ||
|
|
||
| /** | ||
| * Queue the work to be picked up by drain loop. | ||
| * @param work to be queued. | ||
| */ | ||
| void queueWork(SynchronousReceiveWork work) { | ||
|
|
||
| logger.info("[{}] Pending: {}, Scheduling receive timeout task '{}'.", work.getId(), work.getNumberOfEvents(), | ||
| work.getTimeout()); | ||
| workQueue.add(work); | ||
|
|
||
| // Do not drain if another thread want to queue the work before we have subscriber | ||
| if (subscriberInitialized) { | ||
| drain(); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Drain the work, only one thread can be in this loop at a time. | ||
| */ | ||
| private void drain() { | ||
| // If someone is already in this loop, then we are already clearing the queue. | ||
| if (!wip.compareAndSet(0, 1)) { | ||
| return; | ||
| } | ||
|
|
||
| try { | ||
| drainQueue(); | ||
| } finally { | ||
| final int decremented = wip.decrementAndGet(); | ||
| if (decremented != 0) { | ||
| logger.warning("There should be 0, but was: {}", decremented); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /*** | ||
| * Drain the queue using a lock on current work in progress. | ||
| */ | ||
| 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; | ||
| } | ||
|
|
||
| // Acquiring the lock | ||
| synchronized (currentWorkLock) { | ||
|
|
||
| // Making sure current work not become terminal since last drain queue cycle | ||
| if (currentWork != null && currentWork.isTerminal()) { | ||
| workQueue.remove(currentWork); | ||
| if (currentTimeoutOperation != null && !currentTimeoutOperation.isDisposed()) { | ||
| currentTimeoutOperation.dispose(); | ||
| } | ||
| currentTimeoutOperation = null; | ||
|
Member
There was a problem hiding this comment. Is there any benefit to setting this to null again (and in a few places)? Once a subscription is disposed calling dispose again is a no-op.
Member
There was a problem hiding this comment. Having a current work is always tied to a timeout operation. Checking currentWork != null should be enough.
Contributor
Author
There was a problem hiding this comment. We might pick currentWork more than one time from workQueue. We do not need to process currentWork if is picked up second time and no bufferMessages to send to it.
The timeout Operation is not removing currentWork from he queue, thus currentWork needs be picked up again, thus we need to threat this different and go in the loop and remove this timeout work. setting |
||
| } | ||
|
|
||
| // We should process a work when | ||
| // 1. it is first time getting picked up | ||
| // 2. or more messages have arrived while we were in drain loop. | ||
| // We might not have all the message in bufferMessages needed for workQueue, Thus we will only remove work | ||
| // from queue when we have delivered all the messages to currentWork. | ||
|
|
||
| while ((currentWork = workQueue.peek()) != null | ||
| && (!currentWork.isProcessingStarted() || bufferMessages.size() > 0)) { | ||
|
|
||
| // Additional check for safety, but normally this work should never be terminal | ||
| if (currentWork.isTerminal()) { | ||
| // This work already finished by either timeout or no more messages to send, process next work. | ||
| workQueue.remove(currentWork); | ||
| if (currentTimeoutOperation != null && !currentTimeoutOperation.isDisposed()) { | ||
| currentTimeoutOperation.dispose(); | ||
| } | ||
| continue; | ||
| } | ||
|
|
||
| if (!currentWork.isProcessingStarted()) { | ||
| // timer to complete the currentWork in case of timeout trigger | ||
| currentTimeoutOperation = getTimeoutOperation(currentWork); | ||
| currentWork.startedProcessing(); | ||
| } | ||
|
|
||
| // Send messages to currentWork from buffer | ||
| while (bufferMessages.size() > 0 && !currentWork.isTerminal()) { | ||
| currentWork.next(bufferMessages.poll()); | ||
| remaining.decrementAndGet(); | ||
| } | ||
|
|
||
| // if we have delivered all the messages to currentWork, we will complete it. | ||
| if (currentWork.isTerminal()) { | ||
| if (currentWork.getError() == null) { | ||
| currentWork.complete(); | ||
| } | ||
| // Now remove from queue since it is complete | ||
| workQueue.remove(currentWork); | ||
| if (currentTimeoutOperation != null && !currentTimeoutOperation.isDisposed()) { | ||
| currentTimeoutOperation.dispose(); | ||
| } | ||
| logger.verbose("The work [{}] is complete.", currentWork.getId()); | ||
| } else { | ||
| // Since this work is not complete, find out how much we should request from upstream | ||
| long creditToAdd = currentWork.getRemaining() - (remaining.get() + bufferMessages.size()); | ||
| if (creditToAdd > 0) { | ||
| remaining.addAndGet(creditToAdd); | ||
| subscription.request(creditToAdd); | ||
| logger.verbose("Requesting [{}] from upstream for work [{}].", creditToAdd, | ||
| currentWork.getId()); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * @param work on which timeout thread need to start. | ||
| * | ||
| * @return {@link Disposable} for the timeout operation. | ||
| */ | ||
| private Disposable getTimeoutOperation(SynchronousReceiveWork work) { | ||
| Duration timeout = work.getTimeout(); | ||
| return Mono.delay(timeout).thenReturn(work) | ||
| .subscribe(l -> { | ||
| synchronized (currentWorkLock) { | ||
| if (currentWork == work) { | ||
| work.timeout(); | ||
| } | ||
| } | ||
| }); | ||
| } | ||
|
|
||
| /** | ||
| * {@inheritDoc} | ||
| */ | ||
| @Override | ||
| protected void hookOnError(Throwable throwable) { | ||
| logger.error("[{}] Errors occurred upstream", currentWork.getId(), throwable); | ||
| synchronized (currentWorkLock) { | ||
| currentWork.error(throwable); | ||
| } | ||
| dispose(); | ||
| } | ||
|
|
||
| @Override | ||
| protected void hookOnCancel() { | ||
| if (isDisposed.getAndSet(true)) { | ||
| return; | ||
| } | ||
|
|
||
| synchronized (currentWorkLock) { | ||
| if (currentWork != null) { | ||
| currentWork.complete(); | ||
| } | ||
| if (currentTimeoutOperation != null && !currentTimeoutOperation.isDisposed()) { | ||
| currentTimeoutOperation.dispose(); | ||
| } | ||
| currentTimeoutOperation = null; | ||
| } | ||
|
|
||
| subscription.cancel(); | ||
| } | ||
|
|
||
| private boolean isTerminated() { | ||
| return isDisposed.get(); | ||
| } | ||
|
|
||
| int getWorkQueueSize() { | ||
| return this.workQueue.size(); | ||
| } | ||
|
|
||
| long getRequested() { | ||
| return this.requested; | ||
| } | ||
|
|
||
| boolean isSubscriberInitialized() { | ||
| return this.subscriberInitialized; | ||
| } | ||
| } | ||
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Only thing you need are unit tests. The logic in this subscriber is complex and I can see it being hard to debug.